Run Nemotron 5 Training with WorkloadRun

1
apiVersion: cre.nvidia.com/v1alpha1
2
kind: WorkloadRun
3
metadata:
4
  name: nemotron5-56b
5
  namespace: default
6
spec:
7
  image: nvcr.io/nvidia/pytorch:25.08-py3
8
  framework:
9
    exec:
10
      command: ["/bin/bash", "/config/train.sh"]
11
  numNodes: 16
12
  config:
13
    inline:
14
      train.sh: |
15
        #!/bin/bash
16
        set -e
17
18
        WORKSPACE_DIR=${WORKSPACE_DIR:-/mnt/workspace}
19
        MEGATRON_PATH=${MEGATRON_PATH:-${WORKSPACE_DIR}/megatron-lm}
20
        CHECKPOINT_DIR=${CHECKPOINT_DIR:-${WORKSPACE_DIR}/checkpoints}
21
        TENSORBOARD_DIR=${TENSORBOARD_DIR:-${WORKSPACE_DIR}/tensorboard}
22
        mkdir -p ${CHECKPOINT_DIR} ${TENSORBOARD_DIR}
23
24
        if ! ls ${MEGATRON_PATH}/megatron/core/datasets/helpers_cpp*.so 1>/dev/null 2>&1; then
25
          echo "Building Megatron helpers_cpp..."
26
          cd ${MEGATRON_PATH}
27
          pip install -e . --no-deps --no-build-isolation 2>&1 | tail -5
28
          cd ${WORKSPACE_DIR}
29
        fi
30
31
        CHECKPOINT_ARGS=""
32
        DATA_CACHE_ARGS=""
33
        if [ "$ENABLE_CHECKPOINT" = "true" ]; then
34
          CHECKPOINT_ARGS="--save ${CHECKPOINT_DIR}"
35
          if [ -d "${CHECKPOINT_DIR}" ] && [ "$(ls -A ${CHECKPOINT_DIR} 2>/dev/null)" ]; then
36
            CHECKPOINT_ARGS="${CHECKPOINT_ARGS} --load ${CHECKPOINT_DIR}"
37
          fi
38
          DATA_CACHE_ARGS="--data-cache-path ${WORKSPACE_DIR}/data_cache"
39
          mkdir -p ${WORKSPACE_DIR}/data_cache
40
        fi
41
42
        TOTAL_GPUS=$((PET_NNODES * PET_NPROC_PER_NODE))
43
        TP=${TENSOR_PARALLELISM:-4}
44
        PP=${PIPELINE_PARALLELISM:-1}
45
        MBS=${MICRO_BATCH_SIZE:-1}
46
        GBS=${GLOBAL_BATCH_SIZE:-$TOTAL_GPUS}
47
48
        echo "TOTAL_GPUS=${TOTAL_GPUS} GBS=${GBS} MBS=${MBS} TP=${TP} PP=${PP}"
49
50
        exec torchrun \
51
          --nnodes $PET_NNODES \
52
          --nproc-per-node $PET_NPROC_PER_NODE \
53
          ${MEGATRON_PATH}/pretrain_gpt.py \
54
          --attention-backend flash \
55
          --distributed-timeout-minutes 230 \
56
          --use-mcore-models \
57
          --no-mmap-bin-files \
58
          --sequence-parallel \
59
          --untie-embeddings-and-output-weights \
60
          --disable-bias-linear \
61
          --init-method-std 0.014 \
62
          --position-embedding-type rope \
63
          --rotary-base 1000000 \
64
          --rotary-percent 1.0 \
65
          --squared-relu \
66
          --group-query-attention \
67
          --kv-channels 128 \
68
          --normalization RMSNorm \
69
          --attention-dropout 0.0 \
70
          --hidden-dropout 0.0 \
71
          --exit-duration-in-mins 30 \
72
          --train-iters 50 \
73
          --lr-decay-iters 1830030 \
74
          --lr 6e-4 \
75
          --min-lr 6e-6 \
76
          --weight-decay 0.1 \
77
          --clip-grad 1.0 \
78
          --lr-decay-style cosine \
79
          --lr-warmup-iters 5 \
80
          --eval-iters 1 \
81
          --eval-interval 50 \
82
          --log-interval 10 \
83
          --tokenizer-type NullTokenizer \
84
          --vocab-size 131072 \
85
          --mock-data \
86
          --num-workers 1 \
87
          --no-create-attention-mask-in-dataloader \
88
          --log-progress \
89
          --timing-log-option minmax \
90
          --log-params-norm \
91
          --log-num-zeros-in-grad \
92
          --log-throughput \
93
          --bf16 \
94
          --adam-beta1 0.9 \
95
          --adam-beta2 0.95 \
96
          --use-distributed-optimizer \
97
          --overlap-grad-reduce \
98
          --overlap-param-gather \
99
          --manual-gc \
100
          --log-straggler \
101
          --disable-straggler-on-startup \
102
          --straggler-minmax-count 16 \
103
          --check-weight-hash-across-dp-replicas-interval 20000 \
104
          --ckpt-fully-parallel-save \
105
          --ckpt-fully-parallel-load \
106
          --async-save \
107
          --ckpt-assume-constant-structure \
108
          --ckpt-format torch_dist \
109
          --num-layers 79 \
110
          --hidden-size 8192 \
111
          --ffn-hidden-size 32768 \
112
          --num-attention-heads 64 \
113
          --seq-length 8192 \
114
          --max-position-embeddings 8192 \
115
          --num-query-groups 8 \
116
          --tensor-model-parallel-size $TP \
117
          --pipeline-model-parallel-size $PP \
118
          --micro-batch-size $MBS \
119
          --global-batch-size $GBS \
120
          $CHECKPOINT_ARGS \
121
          $DATA_CACHE_ARGS \
122
          --save-interval ${SAVE_INTERVAL:-250} \
123
          --save-retain-interval ${SAVE_RETAIN_INTERVAL:-1000} \
124
          --tensorboard-dir ${TENSORBOARD_DIR}
125
  initContainers:
126
    - name: megatron-clone
127
      image: nvcr.io/nvidia/pytorch:25.08-py3
128
      command: ["/bin/bash", "-c"]
129
      args:
130
        - |
131
          set -ex
132
          if [ ! -d "/mnt/workspace/megatron-lm/.git" ]; then
133
            git clone --depth 1 -b core_v0.15.2 \
134
              https://github.com/NVIDIA/Megatron-LM.git /mnt/workspace/megatron-lm
135
          fi
136
          echo "Megatron-LM cloned successfully"
137
      volumeMounts:
138
        - mountPath: /mnt/workspace
139
          name: workspace
140
  volumes:
141
    - name: workspace
142
      emptyDir:
143
        medium: Memory
144
  volumeMounts:
145
    - mountPath: /mnt/workspace
146
      name: workspace
147
  env:
148
    - name: NVTE_FWD_LAYERNORM_SM_MARGIN
149
      value: "16"
150
    - name: NVTE_BWD_LAYERNORM_SM_MARGIN
151
      value: "16"
152
    - name: NVTE_FUSED_ATTN
153
      value: "0"
154
    - name: TORCHINDUCTOR_WORKER_START
155
      value: fork
156
    - name: CUDA_DEVICE_MAX_CONNECTIONS
157
      value: "1"
158
    - name: TORCH_NCCL_AVOID_RECORD_STREAMS
159
      value: "1"
160
    - name: TORCH_NCCL_HIGH_PRIORITY
161
      value: "1"
162
    - name: UCX_MEM_MMAP_HOOK_MODE
163
      value: none
164
    - name: UCX_MEM_CUDA_HOOK_MODE
165
      value: none
166
    - name: UCX_MEM_MALLOC_HOOKS
167
      value: "n"
168
    - name: UCX_ERROR_SIGNALS
169
      value: ""
170
    - name: EXIT_DURATION_MINS
171
      value: "30"
172
    - name: ENABLE_CHECKPOINT
173
      value: "false"
174
    - name: TRAIN_ITERS
175
      value: "50"
176
    - name: SAVE_INTERVAL
177
      value: "250"
178
    - name: SAVE_RETAIN_INTERVAL
179
      value: "1000"
180
    - name: PYTHONPATH
181
      value: /mnt/workspace/megatron-lm
182
  resources:
183
    limits:
184
      nvidia.com/gpu: "4"
185
      memory: 800Gi
186
      cpu: "128"
187
    requests:
188
      nvidia.com/gpu: "4"
189
      memory: 500Gi
190
      cpu: "64"
191
  goodputMeasurement:
192
    logProfileRef: megatron-training
193
    sampleInterval: 10s

1
spec:
2
  gangScheduler:
3
    schedulerName: kai-scheduler   # required; injected as schedulerName into every workload pod
4
    queue: high-priority           # optional; defaults to "default-queue"