Add A4X MAX Kimi-K2 FP8mx 256 GPUs recipe

[ngu3]

Add A4X MAX Kimi-K2 256 GPUs FP8mx recipe

[gemini-code-assist]

Code Review

This pull request introduces a Helm chart recipe to run a kimi-k2 pretraining workload on a4x-max GKE Node pools using the Nvidia Megatron-Bridge framework. The review feedback highlights several critical areas for improvement: hardcoded namespaces and Hugging Face tokens should be parameterized; runtime package installations and git cloning should be moved to a pre-built container image to avoid fragility and rate-limiting; Helm templates should handle null GCS mounts gracefully; and shell/Python scripting in the launcher should be optimized for robustness and memory efficiency.

[gemini-code-assist]

The script does not enable set -eo pipefail. Without this, any failures in the torchrun pipeline or other commands will be silently ignored, and the script will exit with status 0 (success), masking critical training crashes from Kubernetes. Enable set -eo pipefail at the start of the script.

Suggested change

	config_overrides=()
	set -eo pipefail
	config_overrides=()

[gemini-code-assist]

The hardcoded export HF_TOKEN=<YOUR_HF_TOKEN> in the launcher script overrides any Hugging Face token passed via container environment variables with the literal string <YOUR_HF_TOKEN>. This will cause authentication failures when downloading models or tokenizers. Use the existing environment variable if available.

Suggested change

	export HF_TOKEN=<YOUR_HF_TOKEN>
	export HF_TOKEN="${HF_TOKEN:-}"

[gemini-code-assist]

The namespace of the JobSet is hardcoded to default. If the Helm chart is installed in a different namespace, the JobSet will be created in the default namespace while other resources (like ConfigMaps and Services) are created in the target namespace, causing the deployment to fail. Use {{ .Release.Namespace }} instead.

  namespace: "{{ .Release.Namespace }}"

[gemini-code-assist]

The namespace default is hardcoded in RANK_0_FQDN, DOMAIN_NAME, and MASTER_ADDR. If this Helm chart is installed in any namespace other than default, the pods will fail to resolve the master address and the training job will fail to initialize. Use {{ .Release.Namespace }} instead of hardcoding default.

              - name: RANK_0_FQDN
                value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.{{.Release.Namespace}}.svc.cluster.local"
              - name: HOSTNAME_PREFIX
                value: "{{.Release.Name}}-workload-"
              - name: DOMAIN_NAME
                value: "{{.Release.Name}}.{{.Release.Namespace}}.svc.cluster.local"
              - name: MASTER_ADDR
                value: "{{.Release.Name}}-workload-0-0.{{.Release.Name}}.{{.Release.Namespace}}.svc.cluster.local"
              - name: MASTER_PORT

[gemini-code-assist]

Performing system-level package installations (apt update, doca-ofed-userspace, libnccl2, nccl-gib-plugins) at runtime on all 64 nodes simultaneously is highly inefficient and fragile. It significantly delays job startup, requires public internet access on all nodes (which fails in private clusters), and is highly susceptible to transient network failures or rate-limiting from external package repositories. These dependencies should be pre-installed in a custom Docker image instead.

[gemini-code-assist]

Cloning the Megatron-Bridge repository from GitHub at runtime on all 64 nodes simultaneously can lead to GitHub rate-limiting, slow startup times, and will fail entirely in private GKE clusters without internet access. Consider pre-building these dependencies into a custom container image or using an init container with a shared volume to clone it once.

[gemini-code-assist]

Using a list comprehension [sys.stdout.write(...) for line in ...] for its side effects causes Python to accumulate a list of return values (integers representing bytes written) in memory for the entire duration of the training run. For long-running jobs with large log volumes, this will lead to unnecessary memory growth. Use sys.stdout.writelines with a generator expression instead to process the stream lazily without memory overhead.

Suggested change

	--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, '')]"
	--no-python bash worker_command.sh 2>&1 | python3 -u -c "import sys, time; sys.stdout.writelines('[{}] {}'.format(time.strftime('%Y-%m-%d %H:%M:%S'), line) for line in iter(sys.stdin.readline, ''))"

[gemini-code-assist]

If volumes.gcsMounts contains null placeholders (as defined by default in values.yaml), the template will render invalid volume definitions with name "null" and bucketName "null". Wrap the volume rendering in a check for if $gcs.bucketName to avoid rendering empty/invalid mounts.

            {{- range $gcs := $root.Values.volumes.gcsMounts }}
            {{- if $gcs.bucketName }}
            - name: "{{ $gcs.bucketName }}"
              csi:
                driver: gcsfuse.csi.storage.gke.io
                volumeAttributes:
                  bucketName: "{{ $gcs.bucketName }}"
                  {{- if $gcs.mountOptions }}
                  mountOptions: "{{ $gcs.mountOptions }}"
                  {{- end }}
            {{- end }}
            {{- end}}

[gemini-code-assist]

The workload container does not specify any CPU or memory requests/limits. Since the pod uses a large shared-memory tmpfs volume of up to 250Gi (which consumes host memory), not specifying memory requests can lead to severe node overcommit and unpredictable Out-Of-Memory (OOM) kills by the kernel. It is highly recommended to set explicit CPU and memory requests/limits.

[gemini-code-assist]

If volumes.gcsMounts contains null placeholders (as defined by default in values.yaml), the template will render invalid volumeMount definitions with name "null" and mountPath "null". Wrap the volumeMount rendering in a check for if $gcs.bucketName to avoid rendering empty/invalid mounts.

                {{- range $gcs := $root.Values.volumes.gcsMounts }}
                {{- if $gcs.bucketName }}
                - name: "{{ $gcs.bucketName }}"
                  mountPath: "{{ $gcs.mountPath }}"
                {{- end }}
                {{- end }}