diff --git a/megatron/core/tokenizers/utils/build_tokenizer.py b/megatron/core/tokenizers/utils/build_tokenizer.py index 7b9c0abe911..707a63775d5 100644 --- a/megatron/core/tokenizers/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/utils/build_tokenizer.py @@ -80,6 +80,14 @@ def build_tokenizer(args, **kwargs): _set_padded_vocab_size(args, tokenizer) return tokenizer + elif args.tokenizer_type == 'SFTTokenizer': + # SFTTokenizer uses the legacy tokenizer system + from megatron.core.tokenizers.text.libraries.sft_tokenizer import SFTTokenizer + tokenizer = SFTTokenizer( + args.tokenizer_model, + args.sft_tokenizer_prompt_format, + ) + return tokenizer if args.tokenizer_metadata: metadata = args.tokenizer_metadata diff --git a/megatron/training/datasets/sft_dataset.py b/megatron/training/datasets/sft_dataset.py index fd9d1fe7c14..9de5d2a52fe 100644 --- a/megatron/training/datasets/sft_dataset.py +++ b/megatron/training/datasets/sft_dataset.py @@ -117,10 +117,6 @@ def extend_with_padding(tokens, targets, positions, pad_len): tokens_list = tokens.tolist() targets_list = targets.tolist() - # Add EOD, unless it's already present - if tokens_list[-1] != eod: - tokens_list.append(eod) - targets_list.append(eod) pack_tokens.extend(tokens_list) pack_targets.extend(targets_list) diff --git a/scripts/bindpcie.sh b/scripts/bindpcie.sh new file mode 100755 index 00000000000..2a7356475a4 --- /dev/null +++ b/scripts/bindpcie.sh @@ -0,0 +1,212 @@ +#! /bin/bash +# Copyright (c) 2018-2022, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +set -euo pipefail + +print_usage() { + cat << EOF +${0} [options] [--] COMMAND [ARG...] + +Control binding policy for each task. Assumes one rank will be launched for each GPU. + +Options: + --cpu=MODE + * exclusive -- bind each rank to an exclusive set of cores near its GPU + * exclusive,nosmt -- bind each rank to an exclusive set of cores near its GPU, without hyperthreading + * node -- bind each rank to all cores in the NUMA node nearest its GPU [default] + * *.sh -- bind each rank using the bash associative array bind_cpu_cores or bind_cpu_nodes from a file + * off -- don't bind + --mem=MODE + * node -- bind each rank to the nearest NUMA node [default] + * *.sh -- bind each rank using the bash associative array bind_mem from a file + * off -- don't bind + --ib=MODE + * single -- bind each rank to a single IB device near its GPU + * off -- don't bind [default] +EOF +} + +################################################################################ +# Argument parsing +################################################################################ + +cpu_mode='node' +mem_mode='node' +ib_mode='off' +while [[ "$#" -gt "0" ]]; do + case "$1" in + -h|--help) print_usage ; exit 0 ;; + --cpu=*) cpu_mode="${1/*=/}"; shift ;; + --cpu) cpu_mode="$2"; shift 2 ;; + --mem=*) mem_mode="${1/*=/}"; shift ;; + --mem) mem_mode="$2"; shift 2 ;; + --ib=*) ib_mode="${1/*=/}"; shift ;; + --ib) ib_mode="$2"; shift 2 ;; + --) shift; break ;; + *) break ;; + esac +done +if [[ $# -lt 1 ]]; then + echo 'ERROR: no command given' 2>&1 + print_usage + exit 1 +fi + +################################################################################ +# Get system params +################################################################################ + +# LOCAL_RANK is set with an enroot hook for Pytorch containers +# SLURM_LOCALID is set by Slurm +# OMPI_COMM_WORLD_LOCAL_RANK is set by mpirun +readonly local_rank="${LOCAL_RANK:=${SLURM_LOCALID:=${OMPI_COMM_WORLD_LOCAL_RANK:-}}}" +if [[ ! "${local_rank}" ]]; then + echo 'ERROR: cannot read LOCAL_RANK from env' >&2 + exit 1 +fi + +num_gpus=$(nvidia-smi -i 0 --query-gpu=count --format=csv,noheader,nounits) +if [[ "${local_rank}" -ge "${num_gpus}" ]]; then + echo "ERROR: local rank is ${local_rank}, but there are only ${num_gpus} gpus available" >&2 + exit 1 +fi + +get_lscpu_value() { + awk -F: "(\$1 == \"${1}\"){gsub(/ /, \"\", \$2); print \$2; found=1} END{exit found!=1}" +} +lscpu_out=$(lscpu) +num_sockets=$(get_lscpu_value 'Socket(s)' <<< "${lscpu_out}") +num_nodes=$(lscpu --parse | grep -v '^#' | cut -f4 -d, | sort --unique | wc --lines) +cores_per_socket=$(get_lscpu_value 'Core(s) per socket' <<< "${lscpu_out}") + +echo "num_gpus=${num_gpus} num_sockets = ${num_sockets} num_nodes=${num_nodes} cores_per_socket=${cores_per_socket}" + +readonly cores_per_node=$(( (num_sockets * cores_per_socket) / num_nodes )) +if [[ "${num_gpus}" -gt "1" ]] && [[ "${num_gpus}" -ge "${num_nodes}" ]]; then + readonly gpus_per_node=$(( num_gpus / num_nodes )) +else + readonly gpus_per_node=1 + num_nodes="${num_gpus}" +fi +readonly cores_per_gpu=$(( cores_per_node / gpus_per_node )) +readonly local_node=$(( local_rank / gpus_per_node )) + + +declare -a ibdevs=() +if ibstat_out="$(ibv_devinfo --list | tail -n+2 | cut -f2 | grep -v '^$')"; then + mapfile -t ibdevs <<< "${ibstat_out}" +fi +readonly num_ibdevs="${#ibdevs[@]}" + +################################################################################ +# Setup for exec +################################################################################ + +declare -a numactl_args=() + +case "${cpu_mode}" in + exclusive) + numactl_args+=( "$(printf -- "--physcpubind=%u-%u,%u-%u" \ + $(( local_rank * cores_per_gpu )) \ + $(( (local_rank + 1) * cores_per_gpu - 1 )) \ + $(( local_rank * cores_per_gpu + (cores_per_gpu * gpus_per_node * num_nodes) )) \ + $(( (local_rank + 1) * cores_per_gpu + (cores_per_gpu * gpus_per_node * num_nodes) - 1 )) \ + )" ) + ;; + exclusive,nosmt) + numactl_args+=( "$(printf -- "--physcpubind=%u-%u" \ + $(( local_rank * cores_per_gpu )) \ + $(( (local_rank + 1) * cores_per_gpu - 1 )) \ + )" ) + ;; + node) + numactl_args+=( "--cpunodebind=${local_node}" ) + ;; + *.sh) + # shellcheck source=/dev/null + source "${cpu_mode}" + if [[ "${bind_cpu_cores:-}" ]]; then + numactl_args+=( "--physcpubind=${bind_cpu_cores[${local_rank}]}" ) + elif [[ "${bind_cpu_nodes:-}" ]]; then + numactl_args+=( "--cpunodebind=${bind_cpu_nodes[${local_rank}]}" ) + else + echo "ERROR: invalid CPU affinity file ${cpu_mode}." >&2 + exit 1 + fi + ;; + off|'') + ;; + *) + echo "ERROR: invalid cpu mode '${cpu_mode}'" 2>&1 + print_usage + exit 1 + ;; +esac + +case "${mem_mode}" in + node) + numactl_args+=( "--membind=${local_node}" ) + ;; + *.sh) + # shellcheck source=/dev/null + source "${mem_mode}" + if [[ ! "${bind_mem:-}" ]]; then + echo "ERROR: invalid memory affinity file ${mem_mode}." >&2 + exit 1 + fi + numactl_args+=( "--membind=${bind_mem[${local_rank}]}" ) + ;; + off|'') + ;; + *) + echo "ERROR: invalid mem mode '${mem_mode}'" 2>&1 + print_usage + exit 1 + ;; +esac + +case "${ib_mode}" in + single) + if [[ "${num_ibdevs}" -eq 0 ]]; then + echo "WARNING: used '$0 --ib=single', but there are 0 IB devices available; skipping IB binding." 2>&1 + elif (( num_ibdevs > num_gpus)) || (( num_gpus % num_ibdevs != 0 )) ; then + echo "ERROR: can't evenly map ${num_gpus} gpus to ${num_ibdevs} ibdevs" 2>&1 + echo "set MELLANOX_VISIBLE_DEVICES correctly or use --ib=off" 2>&1 + exit 1 + else + readonly ibdev="${ibdevs[$(( local_rank * num_ibdevs / num_gpus ))]}" + export OMPI_MCA_btl_openib_if_include="${OMPI_MCA_btl_openib_if_include-$ibdev}" + export UCX_NET_DEVICES="${UCX_NET_DEVICES-$ibdev:1}" + fi + ;; + off|'') + ;; + *) + echo "ERROR: invalid ib mode '${ib_mode}'" 2>&1 + print_usage + exit 1 + ;; +esac + +################################################################################ +# Exec +################################################################################ + +if [[ "${#numactl_args[@]}" -gt 0 ]] ; then + [[ "${DEBUG:-0}" = "1" ]] && set -x + exec numactl "${numactl_args[@]}" -- "${@}" +else + exec "${@}" +fi diff --git a/scripts/ultra_sft_64node_bindpcie.sh b/scripts/ultra_sft_64node_bindpcie.sh new file mode 100644 index 00000000000..cd7d560cd12 --- /dev/null +++ b/scripts/ultra_sft_64node_bindpcie.sh @@ -0,0 +1,362 @@ +#!/bin/bash +# +# ultra_sft_64node_bindpcie.sh +# ============================ +# 64-node Ultra SFT run with segment=16, HybridEP/UCX env, and bindpcie (NUMA binding). +# Based on ultra-v3-sft-hsg-mainfeb5merge-mxfp4_newbase; paths and model options unchanged. +# +# --- Summary of changes (for teammates) --- +# +# 1) SBATCH +# - Replaced --dependency=singleton with --segment=16 +# so all 64 nodes sit in the same NVLink domain. +# +# 2) Environment variables +# - NVTE_CPU_OFFLOAD_V1=1 +# - NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=64, USE_MNNVL=1 +# (segment=16, 4 GPUs/node => 16*4=64) +# - UCX (4 vars): UCX_MEM_MMAP_HOOK_MODE, UCX_MEM_CUDA_HOOK_MODE, +# UCX_MEM_MALLOC_HOOKS, UCX_ERROR_SIGNALS +# to avoid memory hook conflicts in multi-node. +# +# 3) bindpcie (CPU/memory NUMA binding) +# - BINDPCIE_SCRIPT="${MEGATRON_LM_DIR}/scripts/bindpcie.sh" +# - Launch: bindpcie --cpu=node --mem=node -- python ... (LAUNCH_CMD) +# - SLURM_LOCALID is passed into the container so bindpcie can use it as local rank. +# - Requires: scripts/bindpcie.sh in megatron-lm-ultra repo, and numactl in the container. +# +# 4) srun +# - Added --mpi=none +# - --container-env lists the above env vars + SLURM_LOCALID +# so they are visible inside the container. +# +# --- End of summary --- + +#SBATCH -p batch +#SBATCH -q normal +#SBATCH --account=llmservice_nemotron_ultra +#SBATCH --ntasks-per-node=4 +#SBATCH --nodes=64 +#SBATCH --time=3:45:00 +#SBATCH --exclusive +#SBATCH --gpus-per-node=4 +#SBATCH --mem=0 +#SBATCH --segment=16 +#SBATCH --job-name=ultra-v3-sft-hsg-mainfeb5merge-mxfp4_newbase + +################################################################ +### TransformerEngine +################################################################ +export NVTE_FWD_LAYERNORM_SM_MARGIN=16 +export NVTE_BWD_LAYERNORM_SM_MARGIN=16 +export NVTE_CPU_OFFLOAD_V1=1 +export TORCHINDUCTOR_WORKER_START=fork + +################################################################ +### HybridEP / MNNVL (segment=16 => 16*4=64 ranks per NVLink domain) +### EP % NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN == 0 for hybridep optimization. +################################################################ +export NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=64 +export USE_MNNVL=1 + +################################################################ +### UCX (prevents memory hook conflicts in multi-node) +################################################################ +export UCX_MEM_MMAP_HOOK_MODE=none +export UCX_MEM_CUDA_HOOK_MODE=none +export UCX_MEM_MALLOC_HOOKS=n +export UCX_ERROR_SIGNALS= + +################################################################ +### General +################################################################ +export QUANTIZATION_TYPE_DEBUG=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export OMP_NUM_THREADS=16 + +# Debug: See NCCL operations during checkpoint load +# export NCCL_DEBUG=INFO +# export TORCH_DISTRIBUTED_DEBUG=DETAIL + +export HF_HOME="/lustre/fsw/portfolios/llmservice/users/adithyare/.cache/huggingface/" + +NAME=${SLURM_JOB_NAME} + +OUTPUT_ROOT="/lustre/fs1/portfolios/llmservice/projects/llmservice_modelalignment_ppo/users/adithyare/nemotron_ultra/sft-runs" +MEGATRON_LM_DIR="/lustre/fs1/portfolios/llmservice/projects/llmservice_modelalignment_ppo/users/adithyare/code/megatron-lm-ultra" +IMAGE="/lustre/fs1/portfolios/llmservice/projects/llmservice_modelalignment_ppo/users/adithyare/containers/pt_ultra_mamba_ssmv230_23jan28.sqsh" + +# bindpcie: CPU/mem binding (--cpu=node, --mem=node). Requires numactl in container. +BINDPCIE_SCRIPT="${MEGATRON_LM_DIR}/scripts/bindpcie.sh" + +# WANDB_API_KEY: loaded from ~/.bashrc (do NOT hardcode here) +WANDB_PROJECT="ultra-v3-sft-hsg" + +RUN_DIR="${OUTPUT_ROOT}" +LOGS_DIR="${RUN_DIR}/${NAME}/logs/" +CHECKPOINT_DIR="${RUN_DIR}/${NAME}/checkpoints/" +DATACACHE_DIR="${RUN_DIR}/${NAME}/data_cache/" +TENSORBOARD_DIR="${RUN_DIR}/${NAME}/tensorboard/" + +mkdir -p ${LOGS_DIR} +mkdir -p ${CHECKPOINT_DIR} +mkdir -p ${DATACACHE_DIR} +mkdir -p ${TENSORBOARD_DIR} + +export TRITON_CACHE_DIR="/tmp/triton-cache" + + +DATETIME=`date +'date_%y-%m-%d_time_%H-%M-%S'` +if [ -n "${SLURM_JOB_ID:-}" ] ; then + SCRIPT_PATH=$(scontrol show job "$SLURM_JOB_ID" | awk -F= '/Command=/{print $2}') + ENV_LOG_FILENAME=${NAME}_${SLURM_JOB_ID}_${DATETIME}.env.log +else + SCRIPT_PATH=$(realpath "$0") + ENV_LOG_FILENAME=${NAME}_${DATETIME}.env.log +fi + +SCRIPT_DIR=$(dirname ${SCRIPT_PATH}) + +################################################################ +### Log environment +################################################################ +echo "<< START PATHS >>" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "IMAGE=${IMAGE}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "BINDPCIE_SCRIPT=${BINDPCIE_SCRIPT}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "MEGATRON_LM_DIR=${MEGATRON_LM_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "RUN_DIR=${RUN_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "LOGS_DIR=${LOGS_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "CHECKPOINT_DIR=${CHECKPOINT_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "DATACACHE_DIR=${DATACACHE_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "TENSORBOARD_DIR=${TENSORBOARD_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "SCRIPT_DIR=${SCRIPT_DIR}" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "<< END PATHS >>" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo -e "\n\n" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} + +echo "<< START GIT >>" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "GIT LOG" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +git -C ${MEGATRON_LM_DIR} log --oneline -1 |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo -e "\n\n" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "GIT STATUS" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +git -C ${MEGATRON_LM_DIR} status --porcelain --branch |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo -e "\n\n" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "GIT DIFF" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +git -C ${MEGATRON_LM_DIR} diff |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "<< END GIT >>" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo -e "\n\n" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} + +echo "<< START ENV >>" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +env |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} +echo "<< END ENV >>" |& tee -a ${LOGS_DIR}/${ENV_LOG_FILENAME} + + +#--result-rejected-tracker-filename ${RESULT_REJECTED_TRACKER_FILENAME} \ +#--iterations-to-skip ${ITERATIONS_TO_SKIP} \ +#--rerun-mode validate_results \ +# +#--enable-experimental \ +#--moe-shared-expert-overlap \ + + # Can not use with FP4 + + # MXFP8 + #--moe-router-padding-for-fp8 \ + #--fp8-format e4m3 \ + #--fp8-recipe mxfp8 \ + #--fp8-param-gather \ + #--reuse-grad-buf-for-mxfp8-param-ag \ + + # Additional options + #--recompute-modules layernorm moe_act \ + # + #--recompute-granularity selective \ + #--recompute-modules moe \ + # + #--tp-comm-overlap \ + + # Short context, use + # --enable-cuda-graph \ + # Long context, use + # --recompute-granularity selective \ + # --recompute-modules moe \ + + # NVFP4 args + # --keep-mtp-spec-in-bf16 \ + # --keep-mamba-stack-attention-linear-in-bf16 \ + # --keep-mamba-out-proj-in-mxfp8 \ + # --keep-moe-latent-projections-in-bf16 \ + # --first-last-layers-bf16 \ + # --num-layers-at-start-in-bf16 0 \ + # --num-layers-at-end-in-bf16 14 \ + # --fp4-format e2m1 \ + # --fp4-recipe nvfp4 \ + + # checkpoint load fix + # --cuda-graph-scope mamba attn moe_router \ + # --ckpt-fully-parallel-load \ + # --async-save \ + # --use-persistent-ckpt-worker \ + + +SEQ_LEN=262144 +TRAIN_SAMPLES=10000 +LR_WARMUP_SAMPLES=100 +LR_DECAY_SAMPLES=$((TRAIN_SAMPLES-LR_WARMUP_SAMPLES)) +LOG_INTERVAL=1 +SAVE_INTERVAL=20 +SAVE_RETAIN_INTERVAL=100 +GBS=64 +LR=1e-5 +MIN_LR=2e-6 + +TOKENIZER_MODEL_PATH="/lustre/fs1/portfolios/llmservice/projects/llmservice_modelalignment_ppo/users/adithyare/nemotron_super/tokenizer" +BASE_MODEL_PATH="/lustre/fs1/portfolios/llmservice/projects/llmservice_modelalignment_ppo/users/adithyare/nemotron_ultra/reinit-embeddings-ckpts/phase1_fp32rs_continued_iter_0600000/checkpoints" +BLEND_PATH="/lustre/fsw/portfolios/llmservice/users/adithyare/nemotron_ultra/blend_jan21.json" + + +OPTIONS=" \ + --sft \ + --sft-tokenizer-prompt-format identity \ + --distributed-timeout-minutes 5 \ + --num-dataset-builder-threads 32 \ + --tokenizer-type SFTTokenizer \ + --tokenizer-model ${TOKENIZER_MODEL_PATH} \ + \ + --recompute-granularity selective \ + --recompute-modules moe \ + --mtp-use-repeated-layer \ + \ + --context-parallel-size 16 \ + --tensor-model-parallel-size 8 \ + --expert-model-parallel-size 128 \ + --expert-tensor-parallel-size 1 \ + --pipeline-model-parallel-size 1 \ + --hybrid-override-pattern MEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME \ + --mtp-hybrid-override-pattern \"*E\" \ + \ + --pretrained-checkpoint ${BASE_MODEL_PATH} \ + --save-interval ${SAVE_INTERVAL} \ + --save-retain-interval ${SAVE_RETAIN_INTERVAL} \ + --lr $LR \ + --min-lr $MIN_LR \ + --lr-decay-style constant \ + --train-samples ${TRAIN_SAMPLES} \ + --lr-warmup-samples ${LR_WARMUP_SAMPLES} \ + --lr-decay-samples ${LR_DECAY_SAMPLES} \ + --seq-length ${SEQ_LEN} \ + --max-position-embeddings ${SEQ_LEN} \ + --log-interval ${LOG_INTERVAL} \ + --micro-batch-size 1 \ + --global-batch-size ${GBS} \ + --overlap-grad-reduce \ + --overlap-param-gather \ + \ + --mtp-num-layers 2 \ + --calculate-per-token-loss \ + --mtp-loss-scaling-factor 0.3 \ + \ + --cuda-graph-scope mamba attn moe_router \ + --te-rng-tracker \ + --high-priority-stream-groups ep \ + --manual-gc-interval 10 \ + --ddp-num-buckets 10 \ + --manual-gc \ + \ + --moe-latent-size 2048 \ + --moe-permute-fusion \ + --cross-entropy-loss-fusion \ + --cross-entropy-fusion-impl native \ + --use-fused-weighted-squared-relu \ + \ + --moe-token-dispatcher-type alltoall \ + --moe-router-score-function sigmoid \ + --moe-grouped-gemm \ + --num-experts 512 \ + --moe-router-topk 22 \ + --moe-aux-loss-coeff 1e-4 \ + --moe-router-topk-scaling-factor 5.0 \ + --moe-router-enable-expert-bias \ + --moe-router-dtype fp32 \ + --moe-router-load-balancing-type seq_aux_loss \ + --moe-shared-expert-intermediate-size 10240 \ + \ + --attention-backend flash \ + --num-workers 1 \ + --disable-gloo-process-groups \ + --ckpt-format torch_dist \ + --ckpt-fully-parallel-save \ + --ckpt-fully-parallel-load \ + --ckpt-assume-constant-structure \ + --use-persistent-ckpt-worker \ + \ + --squared-relu \ + --no-mmap-bin-files \ + --exit-duration-in-mins 5750 \ + --no-create-attention-mask-in-dataloader \ + \ + --sequence-parallel \ + --use-distributed-optimizer \ + --override-opt-param-scheduler \ + \ + --mamba-num-heads 256 \ + --is-hybrid-model \ + --untie-embeddings-and-output-weights \ + --init-method-std 0.014 \ + --position-embedding-type none \ + --num-layers 108 \ + --hidden-size 8192 \ + --num-attention-heads 64 \ + --group-query-attention \ + --num-query-groups 2 \ + --ffn-hidden-size 5120 \ + --kv-channels 128 \ + --save ${CHECKPOINT_DIR} \ + --load ${CHECKPOINT_DIR} \ + --per-split-data-args-path ${BLEND_PATH} \ + --data-cache-path ${DATACACHE_DIR} \ + --weight-decay 0.1 \ + --clip-grad 1.0 \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --disable-bias-linear \ + --normalization RMSNorm \ + --no-load-optim \ + --adam-beta1 0.9 \ + --adam-beta2 0.95 \ + --log-params-norm \ + --log-num-zeros-in-grad \ + --log-throughput \ + --log-progress \ + --log-energy \ + --log-memory-interval 200 \ + --logging-level 20 \ + --log-straggler \ + --disable-straggler-on-startup \ + --straggler-minmax-count 16 \ + --check-weight-hash-across-dp-replicas-interval 20000 \ + --ddp-pad-buckets-for-high-nccl-busbw \ + --timing-log-option minmax \ + --eval-interval 1000 \ + --eval-iters 14 \ + --te-precision-config-file /lustre/fs1/portfolios/llmservice/projects/llmservice_nlp_fm/nemotron6/code_ultra/te_quant.cfg \ + --bf16 \ + --use-mcore-models \ + --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \ + --wandb-project ${WANDB_PROJECT} \ + --wandb-exp-name ${NAME} \ + --dist-ckpt-strictness log_unexpected \ + --tensorboard-dir ${TENSORBOARD_DIR}" + +RUN_CMD="python -u ${MEGATRON_LM_DIR}/pretrain_mamba.py ${OPTIONS}" + +# Launch via bindpcie: CPU=node, mem=node (NUMA binding per rank). Script uses LOCAL_RANK or SLURM_LOCALID. +LAUNCH_CMD="${BINDPCIE_SCRIPT} --cpu=node --mem=node -- ${RUN_CMD}" + +srun -l \ + --mpi=none \ + --no-container-mount-home \ + --container-image=${IMAGE} \ + --container-mounts="/lustre:/lustre" \ + --container-env=NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN,USE_MNNVL,UCX_MEM_MMAP_HOOK_MODE,UCX_MEM_CUDA_HOOK_MODE,UCX_MEM_MALLOC_HOOKS,UCX_ERROR_SIGNALS,NVTE_CPU_OFFLOAD_V1,SLURM_LOCALID \ + --output="${LOGS_DIR}/%x_%j_${DATETIME}.log" \ + sh -c "${LAUNCH_CMD}"