diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/Chart.yaml b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/Chart.yaml new file mode 100644 index 00000000..bd8f2122 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/Chart.yaml @@ -0,0 +1,20 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +apiVersion: v2 +name: a4x_max_jobset_workload +description: a4x_max_jobset_workload +type: application +version: 0.1.0 +appVersion: "1.16.0" \ No newline at end of file diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/README.md b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/README.md new file mode 100644 index 00000000..0e572721 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/README.md @@ -0,0 +1,154 @@ + +# Pretrain kimi-k2 workloads on a4x-max GKE Node pools with Nvidia Megatron-Bridge Framework + +This recipe outlines the steps for running a kimi-k2 pretraining +workload on [a4x-max GKE Node pools](https://cloud.google.com/kubernetes-engine) by using the +[Megatron-Bridge pretraining workload](https://github.com/NVIDIA-NeMo/Megatron-Bridge). + +## Orchestration and deployment tools + +For this recipe, the following setup is used: + +- Orchestration - [Google Kubernetes Engine (GKE)](https://cloud.google.com/kubernetes-engine) +- Pretraining job configuration and deployment - A Helm chart is used to + configure and deploy the [Kubernetes Jobset](https://kubernetes.io/blog/2025/03/23/introducing-jobset) resource which manages the execution of the + [Megatron-Bridge pretraining workload](https://github.com/NVIDIA-NeMo/Megatron-Bridge). + +## Test environment + +This recipe has been optimized for and tested with the following configuration: + +- GKE cluster +Please follow Cluster Toolkit [instructions](https://github.com/GoogleCloudPlatform/cluster-toolkit/) +to create your a4x-max GKE cluster. + +## Training dataset + +This recipe uses a mock pretraining dataset provided by the Megatron-Bridge framework. + +## Docker container image + +This recipe uses the following docker images: + +- `nvcr.io/nvidia/nemo:26.04.01` +**Installed Plugins:** +- `nccl-gib-plugins` version: 1.1.2-1 + +## Run the recipe + +From your client workstation, complete the following steps: + +### Configure environment settings + +Set the environment variables to match your environment: + + ```bash + export PROJECT_ID= + export CLUSTER_REGION= + export CLUSTER_NAME= + export GCS_BUCKET= # Note: path should not be prefixed with gs:// + export KUEUE_NAME= + export HF_TOKEN= + ``` + +Replace the following values: + + - ``: your Google Cloud project ID. + - ``: the region where your cluster is located. + - ``: the name of your GKE cluster. + - ``: the name of your Cloud Storage bucket. Don't include the `gs://` prefix. + - ``: the name of the Kueue local queue. The default queue created by the cluster toolkit is `a4x-max`. Make sure to verify the name of the local queue in your cluster. + - ``: Your HuggingFace token. + +Set the default project: + + ```bash + gcloud config set project $PROJECT_ID + ``` + +### Get the recipe + +Clone the `gpu-recipes` repository and set a reference to the recipe folder. + +``` +git clone https://github.com/ai-hypercomputer/gpu-recipes.git +cd gpu-recipes +export REPO_ROOT=`git rev-parse --show-toplevel` +export RECIPE_ROOT=$REPO_ROOT/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe +cd $RECIPE_ROOT +``` + +### Get cluster credentials + +``` +gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION +``` + +### Configure and submit a pretraining job + +#### Using 64 node (256 gpus) fp8mx precision +To execute the job with the default settings, run the following command from +your client: + +```bash +cd $RECIPE_ROOT +export WORKLOAD_NAME=$USER-a4x-max-kimi-k2-256gpus +helm install $WORKLOAD_NAME . -f values.yaml \ +--set-file workload_launcher=launcher.sh \ +--set workload.image=nvcr.io/nvidia/nemo:26.04.01 \ +--set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \ +--set volumes.gcsMounts[0].mountPath=/job-logs \ +--set workload.envs[0].value=/job-logs/$WORKLOAD_NAME \ +--set queue=${KUEUE_NAME} +``` + +**Examples** + +- To set the number of training steps to 100, run the following command from + your client: + + ```bash + cd $RECIPE_ROOT + export WORKLOAD_NAME=$USER-a4x-max-kimi-k2-256gpus + helm install $WORKLOAD_NAME . -f values.yaml \ + --set-file workload_launcher=launcher.sh \ + --set workload.image=nvcr.io/nvidia/nemo:26.04.01 \ + --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \ + --set volumes.gcsMounts[0].mountPath=/job-logs \ + --set workload.envs[0].value=/job-logs/$WORKLOAD_NAME \ + --set queue=${KUEUE_NAME} \ + --set workload.arguments[0]="trainer.max_steps=100" + ``` + +### Monitor the job + +To check the status of pods in your job, run the following command: + +``` +kubectl get pods | grep $USER-a4x-max-kimi-k2-64node +``` + +Replace the following: + +- JOB_NAME_PREFIX - your job name prefix. For example $USER-a4x-max-kimi-k2-64node. + +To get the logs for one of the pods, run the following command: + +``` +kubectl logs POD_NAME +``` + +Information about the training job's progress, including crucial details such as +loss, step count, and step time, is generated by the rank 0 process. +This process runs on the pod whose name begins with +`JOB_NAME_PREFIX-workload-0-0`. +For example: `$USER-a4x-max-kimi-k2-64node-workload-0-0-s9zrv`. + +### Uninstall the Helm release + +You can delete the job and other resources created by the Helm chart. To +uninstall Helm, run the following command from your client: + +```bash +helm uninstall $USER-a4x-max-kimi-k2-64node +``` diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/launcher.sh b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/launcher.sh new file mode 100644 index 00000000..e275af71 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/launcher.sh @@ -0,0 +1,231 @@ +usage() +{ +cat << EOF +usage: bash ./launcher.sh [config-override [config-override ...]] +config-override (Optional) A NeMo configuration override. E.g. trainer.max_steps=10000. +EOF +} + +parse_args() { + while [[ "$1" != "" ]]; do + case $(grep -o "=" <<< "$1" | wc -l) in + 1 ) + config_overrides+=("$1") + ;; + * ) + echo "Invalid config override: $1" + usage + exit 1 + esac + shift + done + config_overrides="${config_overrides[*]}" +} + +config_overrides=() +parse_args "$@" + +if [[ -z "${config_overrides[*]}" ]]; then + echo "No NeMo config overrides specified" +else + echo "NeMo config overrides:" + echo " ${config_overrides}" +fi + +export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH" +ldconfig "$LD_LIBRARY_PATH" +echo "Added $LD_LIBRARY_PATH to ldconfig:" +ldconfig -p | grep libcuda | sed 's/^/ /' +echo "" + +if [[ -n "${EXPLICIT_LOG_DIR}" ]]; then + explicit_log_dir="${EXPLICIT_LOG_DIR}" +else + explicit_log_dir="workload_logs" +fi + +# Ensure explicit_log_dir is an absolute path before any cd commands +if [[ "$explicit_log_dir" != /* ]]; then + explicit_log_dir="${PWD}/${explicit_log_dir}" +fi +echo "Logging to ${explicit_log_dir}" + +if [[ -n "${TOKENIZER_PATH}" ]]; then + echo "Getting tokenizer files" + cp "${TOKENIZER_PATH}"/* . + echo "" +fi + +echo "Launching Torch distributed on the node rank $JOB_COMPLETION_INDEX out of $NNODES nodes" + +# Create the nsys directory. +mkdir -p "${explicit_log_dir}/nsys" + +# Collect diagnostics +linux_kv="$(uname --kernel-release)" +cuda_driver_v="" +driver_v="" +vbios_v="" +if command -v nvidia-smi &> /dev/null; then + cuda_driver_v=$(nvidia-smi -q -x | grep -Po '(?<=).*(?=)' || true) + driver_v=$(nvidia-smi -q -x | grep -Po '(?<=).*(?=)' || true) + vbios_v=$(nvidia-smi -q -x | grep -Po '(?<=).*(?=)' | head -n1 || true) +fi +nccl_v=$(python3 -c "import torch; v=torch.cuda.nccl.version() if hasattr(torch.cuda, 'nccl') else 'unknown'; print('.'.join(map(str, v)) if isinstance(v, tuple) else v)" || echo "unknown") +cuda_container_v=$(python3 -c "import torch; print(torch.version.cuda)" || echo "unknown") + +kv="{\"linux_kernel_version\": \"${linux_kv}\"" +kv="${kv}, \"cuda_driver_version\": \"${cuda_driver_v}\"" +kv="${kv}, \"cuda_container_version\": \"${cuda_container_v}\"" +kv="${kv}, \"gpu_driver_version\": \"${driver_v}\"" +kv="${kv}, \"vbios_version\": \"${vbios_v}\"" +kv="${kv}, \"nccl_version\": \"${nccl_v}\"}" + +echo "VERSION_DIAGNOSTICS: ${kv}" + + +export HF_TOKEN= +export NCCL_IB_SPLIT_DATA_ON_QPS=1 +export NCCL_RAS_ENABLE=0 + +cd /opt +rm -rf Megatron-Bridge +git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git +cd Megatron-Bridge +git checkout c810129341a84e58f4cbed3093f70668a088c028 +git submodule update --init --recursive && sed -i 's/timeout=60/timeout=600/g' src/megatron/bridge/models/hf_pretrained/safe_config_loader.py +sed -i -e '/pretrain(config=recipe/i \ recipe.dist.distributed_timeout_minutes = 10' scripts/performance/run_script.py +ls + + + +worker_command=$(cat <<- EOM + if [ "\$RANK" -eq "0" ]; then + echo "Worker 0 is stalling for a few seconds.." ; + sleep 3 ; + echo "The detected environment within worker rank 0 is:" ; + env | sed 's/^/ /' ; + else + echo "Worker \$RANK is running" ; + fi ; + + echo "Collect nvidia-smi telemetry" + gpu_uuid=\$(nvidia-smi --id \$LOCAL_RANK --query-gpu uuid --format noheader) + + echo "NODE_NAME=${NODE_NAME}" + echo "Rank \$RANK starting on GPU \$gpu_uuid" + + echo "Runtime logging: " + echo "LOCAL_RANK: \$LOCAL_RANK" + echo "RANK: \$RANK" + + telemetry_dir="/runtime-logs/${JOBSET_NAME}/nvidia-smi" + mkdir -p \$telemetry_dir + touch "\$telemetry_dir/\$RANK.csv" + #touch "\$telemetry_dir/\$RANK_custom.csv" + echo "Telemetry dir: \$telemetry_dir" + ls "\$telemetry_dir" + + nv_smi_pid=\$! + + # power_dir="/runtime-logs/${JOBSET_NAME}/power" + # mkdir -p \$power_dir + # touch "\$power_dir/\$RANK.csv" + # nvidia-smi --id \$gpu_uuid --query-gpu=index,timestamp,power.draw.average,power.draw.instant,module.power.draw.average,module.power.draw.instant --format=csv,noheader,nounits -lms 100 -f \$power_dir/\$RANK.csv & + # nv_power_pid=\$! + + #cp /runtime-logs/custom_log_telemetry.sh . + #chmod +x ./custom_log_telemetry.sh + #./custom_log_telemetry.sh >> "\$telemetry_dir/\$RANK_custom.csv" & + #custom_log_telemetry_pid=\$! + + if [ "\$LOCAL_RANK" -eq "0" ] || [ "\$LOCAL_RANK" -eq "2" ]; then + # get CPU usage + echo "Starting CPU usage capture" + cpu_usage_dir="/runtime-logs/${JOBSET_NAME}/cpu-usage" + mkdir -p \$cpu_usage_dir + touch "\$cpu_usage_dir/cpu_usage_\$RANK.out" + + top -b -d 5 >> \$cpu_usage_dir/cpu_usage_\$RANK.out & + export cpu_usage_pid=\$! + + # get NIC usage + echo "Starting NIC usage capture" + cp /runtime-logs/custom_log_nic.sh . + chmod +x ./custom_log_nic.sh + + nic_usage_dir="/runtime-logs/${JOBSET_NAME}/nic-usage" + mkdir -p "\$nic_usage_dir" + touch "\$nic_usage_dir/nic_log_eth0_\$RANK.log" + + ./custom_log_nic.sh "eth0" >> "\$nic_usage_dir/nic_log_eth0_\$RANK.log" & + nic_usage_pid_eth0=\$! + echo "NIC RUN ID: \$nic_usage_pid_eth0" + fi + + cd /opt/Megatron-Bridge ; + + numactl \ + --cpunodebind=\$((LOCAL_RANK/2)) \ + --membind=\$((LOCAL_RANK/2)) \ + nice -10 \ + python scripts/performance/run_script.py \ + --model_family_name kimi \ + --model_recipe_name kimi_k2 \ + --gpu gb300 \ + --num_gpus 256 \ + --gpus_per_node 4 \ + --compute_dtype fp8_mx \ + --seq_length 4096 \ + --global_batch_size 4096 \ + --micro_batch_size 2 \ + --tensor_model_parallel_size 1 \ + --pipeline_model_parallel_size 4 \ + --virtual_pipeline_model_parallel_size 4 \ + --context_parallel_size 1 \ + --expert_model_parallel_size 64 \ + --expert_tensor_parallel_size 1 \ + --moe_flex_dispatcher_backend hybridep \ + --recompute_modules mla_up_proj \ + --max_step 50 \ + logger.log_throughput=True \ + train.manual_gc_interval=100 + + + + kill -9 \$nv_smi_pid + # kill -9 \$nv_power_pid + #kill -9 \$custom_log_telemetry_pid + + if [ "\$LOCAL_RANK" -eq "0" ] || [ "\$LOCAL_RANK" -eq "2" ]; then + kill -9 \$cpu_usage_pid + kill -9 \$nic_usage_pid_eth0 + fi + echo "Telemetry processes killed" + ps aux +EOM +) + +echo "$worker_command" > worker_command.sh +chmod 777 worker_command.sh + +mkdir -p "/runtime-logs/${JOBSET_NAME}/logs" + +torchrun \ +--nproc-per-node="4" \ +--nnodes="64" \ +--node_rank="${JOB_COMPLETION_INDEX}" \ +--rdzv_id="${JOB_IDENTIFIER}" \ +--master_addr="${MASTER_ADDR}" \ +--master_port="${MASTER_PORT}" \ +--no-python bash worker_command.sh 2>&1 | python3 -u -c "import sys, time; [sys.stdout.write('[{}] {}'.format(time.strftime('%Y-%m-%d %H:%M:%S'), line)) for line in iter(sys.stdin.readline, '')]" | tee "/runtime-logs/${JOBSET_NAME}/logs/nemo_mb_RANK_${JOB_COMPLETION_INDEX}.log" + + +if [[ "$JOB_COMPLETION_INDEX" == "0" ]]; then + mkdir -p "${ARTIFACT_DIR}" + cp -r "${explicit_log_dir}"/* "${ARTIFACT_DIR}/" + env > "${ARTIFACT_DIR}/environ.txt" + ls "${ARTIFACT_DIR}" +fi +echo "Training completed" +echo "Pod on $(hostname --fqdn) is exiting" diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-config-configmap.yaml b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-config-configmap.yaml new file mode 100644 index 00000000..a5900071 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-config-configmap.yaml @@ -0,0 +1,28 @@ +# yamllint disable +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +{{- if .Values.workload.configFile }} +apiVersion: v1 +kind: ConfigMap +metadata: + name: "{{ .Release.Name }}-config" +data: + workload-configuration: |- +{{- if .Values.workload_config }} +{{ .Values.workload_config | nindent 4 }} +{{- else }} +{{ "config: null" | nindent 4 }} +{{- end }} +{{- end }} \ No newline at end of file diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-job.yaml b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-job.yaml new file mode 100644 index 00000000..8e2583a1 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-job.yaml @@ -0,0 +1,390 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +{{$timestamp := now | date "2006-01-02-15-04-05"}} +{{$jobSuffix := randAlphaNum 4 | lower}} +{{$jobuuid := uuidv4}} +{{$nodes := div .Values.workload.gpus 4 | max 1}} +{{$gpusPerNode := min .Values.workload.gpus 4}} +{{- $root := . -}} +apiVersion: resource.nvidia.com/v1beta1 +kind: ComputeDomain +metadata: + name: "{{ .Release.Name }}-{{ $jobSuffix }}" +spec: + numNodes: {{ $nodes }} + channel: + resourceClaimTemplate: + name: "{{ .Release.Name }}-{{ $jobSuffix }}-channel" +--- +apiVersion: resource.k8s.io/v1 +kind: ResourceClaimTemplate +metadata: + name: "{{ .Release.Name }}-{{ $jobSuffix }}-all-mrdma" +spec: + spec: + devices: + requests: + - name: req-mrdma + exactly: + deviceClassName: mrdma.google.com + allocationMode: ExactCount + count: 8 +--- +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: "{{ .Release.Name }}" + namespace: default + labels: + {{- if $root.Values.queue }} + kueue.x-k8s.io/queue-name: "{{ $root.Values.queue }}" + {{- end }} +spec: + {{- if $root.Values.queue }} + suspend: true + {{- end }} + failurePolicy: + maxRestarts: {{ default 0 $root.Values.workload.max_workload_restarts }} + replicatedJobs: + - name: workload + replicas: 1 + template: + spec: + parallelism: {{ $nodes }} + completions: {{ $nodes }} + backoffLimit: 0 + completionMode: Indexed + activeDeadlineSeconds: 28800 # 4 hours (4 * 60 * 60) + ttlSecondsAfterFinished: 43200 # 12 hours (12 * 60 * 60) + template: + metadata: + annotations: + kubectl.kubernetes.io/default-container: workload + gke-gcsfuse/volumes: "true" + {{- if $root.Values.volumes.gcsVolumes }} + gke-gcsfuse/cpu-limit: "500m" + gke-gcsfuse/memory-limit: "1Ti" + gke-gcsfuse/ephemeral-storage-limit: "2Ti" + {{- end }} + {{- if $root.Values.volumes.psVolumes }} + gke-parallelstore/volumes: "true" + gke-parallelstore/cpu-limit: "0" + gke-parallelstore/memory-limit: "0" + {{- end }} + {{- if and $root.Values.queue $root.Values.tasSettings.topologyRequest }} + {{- toYaml .Values.tasSettings.topologyRequest | nindent 14 }} + {{- end }} + {{- if and $root.Values.queue $root.Values.dwsSettings.maxRunDurationSeconds }} + provreq.kueue.x-k8s.io/maxRunDurationSeconds: "{{ $root.Values.dwsSettings.maxRunDurationSeconds }}" + {{- end }} + + spec: + hostNetwork: {{ .Values.network.hostNetwork }} + hostIPC: {{ .Values.network.hostIPC }} + + subdomain: "{{.Release.Name}}" + restartPolicy: Never + affinity: + + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: kubernetes.io/arch + operator: In + values: + - arm64 + {{- if $root.Values.targetNodepools }} + - key: cloud.google.com/gke-nodepool + operator: In + values: + {{- range $nodepool := $root.Values.targetNodepools }} + - {{ $nodepool }} + {{- end }} + {{- end }} + + + tolerations: + - operator: "Exists" + key: nvidia.com/gpu + - operator: "Exists" + key: cloud.google.com/impending-node-termination + - key: "kubernetes.io/arch" + operator: "Equal" + value: "arm64" + effect: "NoSchedule" + + volumes: + - name: library-dir-host + hostPath: + path: /home/kubernetes/bin/nvidia + + + {{- if $root.Values.workload.configFile }} + - name: workload-configuration + configMap: + name: "{{.Release.Name}}-config" + items: + - key: workload-configuration + path: {{ $root.Values.workload.configFile | default "workload-configuration" }} + {{- end }} + + - name: workload-launcher + configMap: + name: "{{.Release.Name}}-launcher" + + - name: shared-memory + emptyDir: + medium: "Memory" + sizeLimit: 250Gi + + {{- range $pvc := $root.Values.volumes.pvcMounts }} + - name: "{{ $pvc.claimName }}" + persistentVolumeClaim: + claimName: "{{ $pvc.claimName }}" + {{- end }} + + {{- if $root.Values.volumes.gcsVolumes }} + {{- range $gcs := $root.Values.volumes.gcsMounts }} + - name: "{{ $gcs.bucketName }}" + csi: + driver: gcsfuse.csi.storage.gke.io + volumeAttributes: + bucketName: "{{ $gcs.bucketName }}" + {{- if $gcs.mountOptions }} + mountOptions: "{{ $gcs.mountOptions }}" + {{- end }} + {{- end}} + {{- else }} + {{- range $gcs := $root.Values.volumes.gcsMounts }} + - name: "{{ $gcs.bucketName }}" + emptyDir: {} + {{- end }} + {{- end }} + + {{- if $root.Values.volumes.ssdMountPath }} + - name: local-ssd + hostPath: + path: /mnt/stateful_partition/kube-ephemeral-ssd + {{- end }} + + + + resourceClaims: + - name: compute-domain-channel + resourceClaimTemplateName: "{{ .Release.Name }}-{{ $jobSuffix }}-channel" + - name: rdma + resourceClaimTemplateName: "{{ .Release.Name }}-{{ $jobSuffix }}-all-mrdma" + + containers: + {{- if $root.Values.workload.gcsSidecarImage }} + - name: gke-gcsfuse-sidecar + image: {{ $root.Values.workload.gcsSidecarImage }} + - name: gke-gcsfuse-metadata-prefetch + image: {{ $root.Values.workload.gcsSidecarImage }} + {{- end }} + {{- if $root.Values.workload.psSidecarImage }} + - name: gke-parallelstore-sidecar + image: {{ $root.Values.workload.psSidecarImage }} + {{- end }} + + - name: workload + image: "{{ $root.Values.workload.image }}" + imagePullPolicy: Always + resources: + requests: + nvidia.com/gpu: 4 + limits: + nvidia.com/gpu: 4 + claims: + - name: compute-domain-channel + - name: rdma + securityContext: + privileged: true + capabilities: + add: ["IPC_LOCK"] + env: + - name: JOB_IDENTIFIER + value: "{{ .Release.Name }}-{{ $timestamp }}" + - name: JOB_TIMESTAMP + value: "{{ $timestamp }}" + - name: JOBSET_NAME + value: "{{ .Release.Name }}" + - name: JOB_UUID + value: "{{ $jobuuid }}" + - name: JOB_ORCHESTRATOR + value: "gke" + # Add RANK based on the pod's index provided by the Indexed Job + # This is crucial for torch.distributed initialization. + - name: JOB_COMPLETION_INDEX + valueFrom: + fieldRef: + + fieldPath: metadata.annotations['batch.kubernetes.io/job-completion-index'] + - name: RANK_0_FQDN + value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.default.svc.cluster.local" + - name: HOSTNAME_PREFIX + value: "{{.Release.Name}}-workload-" + - name: DOMAIN_NAME + value: "{{.Release.Name}}.default.svc.cluster.local" + - name: MASTER_ADDR + value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.default.svc.cluster.local" + - name: MASTER_PORT + value: "6002" + - name: WORLD_SIZE + value: "{{ $root.Values.workload.gpus }}" + - name: NNODES + value: "{{ $nodes }}" + - name: GPUS_PER_NODE + value: "{{ $gpusPerNode }}" + - name: NODE_NAME + valueFrom: + fieldRef: + fieldPath: spec.nodeName + + - name: LD_LIBRARY_PATH + value: /usr/local/nvidia/lib64 + + {{ if $root.Values.network.ncclSettings }} + {{- toYaml .Values.network.ncclSettings | nindent 14 }} + {{ end }} + + {{ if $root.Values.workload.envs }} + {{- toYaml .Values.workload.envs | nindent 14 }} + {{ end }} + + command: + - bash + - -c + - | + echo "Pod on $(hostname --fqdn) is running" + echo "Pod is assigned job index of $JOB_COMPLETION_INDEX" + echo "Pod is using node $NODE_NAME" + + set -eo pipefail + + # Install DOCA-OFED + apt update -y + apt install -y curl + export DOCA_URL="https://linux.mellanox.com/public/repo/doca/3.1.0/ubuntu22.04/arm64-sbsa/" + BASE_URL=$([ "${DOCA_PREPUBLISH:-false}" = "true" ] && echo https://doca-repo-prod.nvidia.com/public/repo/doca || echo https://linux.mellanox.com/public/repo/doca) + DOCA_SUFFIX=${DOCA_URL#*public/repo/doca/}; DOCA_URL="$BASE_URL/$DOCA_SUFFIX" + curl $BASE_URL/GPG-KEY-Mellanox.pub | gpg --dearmor > /etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub + echo "deb [signed-by=/etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub] $DOCA_URL ./" > /etc/apt/sources.list.d/doca.list + apt update + apt install -y --allow-downgrades --allow-change-held-packages -o Dpkg::Options::="--force-overwrite" doca-ofed-userspace || apt --fix-broken install -y + + # Install NCCL and nccl-gib-plugins package + apt install --only-upgrade --allow-change-held-packages -y libnccl2 libnccl-dev + + # If image not from Google, trust the GCP signing key + curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /etc/apt/trusted.gpg.d/cloud.google.gpg + + # Add gpudirect-gib-apt repo + echo 'deb https://packages.cloud.google.com/apt gpudirect-gib-apt main' | tee /etc/apt/sources.list.d/nccl-gib.list + + apt update + apt install nccl-gib-plugins + + set +e + set +o pipefail + + export LD_LIBRARY_PATH="/usr/lib/aarch64-linux-gnu:/usr/lib/aarch64-linux-gnu/libibverbs:/usr/lib/libibverbs:/usr/local/nccl-plugin/lib64:/usr/local/nvidia/lib64/:${LD_LIBRARY_PATH}" + + echo "Launching workload with the following arguments:" + {{- range $root.Values.workload.defaultArguments }} + echo " {{ . }}" + {{- end }} + {{- range $root.Values.workload.arguments }} + echo " {{ . }}" + {{- end }} + echo "" + + echo "=== 🔦 GKE SYSTEM PROBE ===" + echo "--- env ---" + env | grep -iE "NCCL|CUDA|NVIDIA|GPU|TORCH|NEMO|MEGATRON|LD_LIBRARY" | sort + echo "--- ulimit -a ---" + ulimit -a + echo "--- limits ---" + cat /proc/self/limits + echo "--- df /dev/shm ---" + df -h /dev/shm + echo "--- libnccl files ---" + find / -name "libnccl*" 2>/dev/null + echo "--- ldconfig ---" + ldconfig -p | grep nccl || true + echo "--- framework versions ---" + python -c "import torch; print('PyTorch:', torch.__version__); print('CUDA:', torch.version.cuda); print('NCCL:', torch.cuda.nccl.version())" || true + echo "--- nvidia topo ---" + nvidia-smi topo -m || true + echo "--- nvidia devices ---" + nvidia-smi -L || true + echo "--- ip addr ---" + ip addr show + echo "--- cgroup ---" + cat /proc/1/cgroup + echo "--- caps ---" + cat /proc/self/status | grep Cap + echo "===========================" + + sleep 10 + + # Stream dmesg logs to stdout. + dmesg -W & + DMESG_PID=$! + trap "kill -9 $DMESG_PID" EXIT + + set -e + set -o pipefail + bash /workload/launcher/launch-workload.sh \ +{{- range $root.Values.workload.defaultArguments }} + {{ . }} \ + {{- end }} + {{- range $root.Values.workload.arguments }} + {{ . }} \ + {{- end }} + + + + volumeMounts: + - name: library-dir-host + mountPath: /usr/local/nvidia + + + {{- if $root.Values.workload.configFile }} + - name: workload-configuration + mountPath: {{ $root.Values.workload.configPath | default "/workload/configs" }} + {{- end }} + + - name: workload-launcher + mountPath: /workload/launcher + + - name: shared-memory + mountPath: /dev/shm + + {{- range $pvc := $root.Values.volumes.pvcMounts }} + - name: "{{ $pvc.claimName }}" + mountPath: "{{ $pvc.mountPath }}" + {{- end }} + + {{- range $gcs := $root.Values.volumes.gcsMounts }} + - name: "{{ $gcs.bucketName }}" + mountPath: "{{ $gcs.mountPath }}" + {{- end }} + + {{- if $root.Values.volumes.ssdMountPath }} + - name: local-ssd + mountPath: "{{ $root.Values.volumes.ssdMountPath }}" + {{- end }} \ No newline at end of file diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-launcher-configmap.yaml b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-launcher-configmap.yaml new file mode 100644 index 00000000..a40baaf3 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-launcher-configmap.yaml @@ -0,0 +1,28 @@ +# yamllint disable +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +apiVersion: v1 +kind: ConfigMap +metadata: + name: "{{ .Release.Name }}-launcher" +data: + launch-workload.sh: |- +{{- if .Values.workload_launcher }} +{{ .Values.workload_launcher | nindent 4 }} +{{- else }} + #!/bin/bash + echo "No workload launcher specified" + exit 1 +{{- end }} \ No newline at end of file diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-svc.yaml b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-svc.yaml new file mode 100644 index 00000000..fb0d4689 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/templates/workload-svc.yaml @@ -0,0 +1,22 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +apiVersion: v1 +kind: Service +metadata: + name: "{{ .Release.Name }}" +spec: + clusterIP: None + selector: + jobset.sigs.k8s.io/jobset-name: "{{ .Release.Name }}" \ No newline at end of file diff --git a/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/values.yaml b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/values.yaml new file mode 100644 index 00000000..a36c2914 --- /dev/null +++ b/training/a4x-max/kimi-k2/megatron-bridge-gke/nemo260401/256gpus-fp8mx-seq4096-gbs4096/recipe/values.yaml @@ -0,0 +1,33 @@ +queue: null +dwsSettings: + maxRunDurationSeconds: null +tasSettings: + topologyRequest: + kueue.x-k8s.io/podset-required-topology: cloud.google.com/gce-topology-block + kueue.x-k8s.io/podset-slice-required-topology: cloud.google.com/gce-topology-subblock + kueue.x-k8s.io/podset-slice-size: '16' +volumes: + gcsVolumes: false + psVolumes: false + gcsMounts: + - bucketName: null + mountPath: null +workload: + gpus: 256 + image: nvcr.io/nvidia/nemo:26.04.01 + defaultArguments: [] + arguments: [] + configFile: null + configPath: null + envs: + - name: ARTIFACT_DIR + value: null + - name: GLOO_SOCKET_IFNAME + value: eth0 +network: + hostNetwork: false + hostIPC: true + subnetworks[]: null + ncclSettings: + - name: NCCL_DEBUG + value: WARN