|
| 1 | +usage() |
| 2 | +{ |
| 3 | +cat << EOF |
| 4 | +usage: bash ./launcher.sh [config-override [config-override ...]] |
| 5 | +config-override (Optional) A NeMo configuration override. E.g. trainer.max_steps=10000. |
| 6 | +EOF |
| 7 | +} |
| 8 | + |
| 9 | +parse_args() { |
| 10 | + while [[ "$1" != "" ]]; do |
| 11 | + case $(grep -o "=" <<< "$1" | wc -l) in |
| 12 | + 1 ) |
| 13 | + config_overrides+=("$1") |
| 14 | + ;; |
| 15 | + * ) |
| 16 | + echo "Invalid config override: $1" |
| 17 | + usage |
| 18 | + exit 1 |
| 19 | + esac |
| 20 | + shift |
| 21 | + done |
| 22 | + config_overrides="${config_overrides[*]}" |
| 23 | +} |
| 24 | + |
| 25 | +config_overrides=() |
| 26 | +parse_args "$@" |
| 27 | + |
| 28 | +if [[ -z "${config_overrides[*]}" ]]; then |
| 29 | + echo "No NeMo config overrides specified" |
| 30 | +else |
| 31 | + echo "NeMo config overrides:" |
| 32 | + echo " ${config_overrides}" |
| 33 | +fi |
| 34 | + |
| 35 | +export LD_LIBRARY_PATH="/usr/local/cuda/compat/lib:$NCCL_PLUGIN_PATH:$LD_LIBRARY_PATH" |
| 36 | +ldconfig "$LD_LIBRARY_PATH" |
| 37 | +echo "Added $LD_LIBRARY_PATH to ldconfig:" |
| 38 | +ldconfig -p | grep libcuda | sed 's/^/ /' |
| 39 | +echo "" |
| 40 | + |
| 41 | +if [[ -n "${EXPLICIT_LOG_DIR}" ]]; then |
| 42 | + explicit_log_dir="${EXPLICIT_LOG_DIR}" |
| 43 | +else |
| 44 | + explicit_log_dir="workload_logs" |
| 45 | +fi |
| 46 | + |
| 47 | +# Ensure explicit_log_dir is an absolute path before any cd commands |
| 48 | +if [[ "$explicit_log_dir" != /* ]]; then |
| 49 | + explicit_log_dir="${PWD}/${explicit_log_dir}" |
| 50 | +fi |
| 51 | +echo "Logging to ${explicit_log_dir}" |
| 52 | + |
| 53 | +if [[ -n "${TOKENIZER_PATH}" ]]; then |
| 54 | + echo "Getting tokenizer files" |
| 55 | + cp "${TOKENIZER_PATH}"/* . |
| 56 | + echo "" |
| 57 | +fi |
| 58 | + |
| 59 | +echo "Launching Torch distributed on the node rank $JOB_COMPLETION_INDEX out of $NNODES nodes" |
| 60 | + |
| 61 | +# Create the nsys directory. |
| 62 | +mkdir -p "${explicit_log_dir}/nsys" |
| 63 | + |
| 64 | +# Collect diagnostics |
| 65 | +linux_kv="$(uname --kernel-release)" |
| 66 | +cuda_driver_v="" |
| 67 | +driver_v="" |
| 68 | +vbios_v="" |
| 69 | +if command -v nvidia-smi &> /dev/null; then |
| 70 | + cuda_driver_v=$(nvidia-smi -q -x | grep -Po '(?<=<cuda_version>).*(?=</cuda_version>)' || true) |
| 71 | + driver_v=$(nvidia-smi -q -x | grep -Po '(?<=<driver_version>).*(?=</driver_version>)' || true) |
| 72 | + vbios_v=$(nvidia-smi -q -x | grep -Po '(?<=<vbios_version>).*(?=</vbios_version>)' | head -n1 || true) |
| 73 | +fi |
| 74 | +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") |
| 75 | +cuda_container_v=$(python3 -c "import torch; print(torch.version.cuda)" || echo "unknown") |
| 76 | + |
| 77 | +kv="{\"linux_kernel_version\": \"${linux_kv}\"" |
| 78 | +kv="${kv}, \"cuda_driver_version\": \"${cuda_driver_v}\"" |
| 79 | +kv="${kv}, \"cuda_container_version\": \"${cuda_container_v}\"" |
| 80 | +kv="${kv}, \"gpu_driver_version\": \"${driver_v}\"" |
| 81 | +kv="${kv}, \"vbios_version\": \"${vbios_v}\"" |
| 82 | +kv="${kv}, \"nccl_version\": \"${nccl_v}\"}" |
| 83 | + |
| 84 | +echo "VERSION_DIAGNOSTICS: ${kv}" |
| 85 | + |
| 86 | + |
| 87 | +export HF_TOKEN=${HF_TOKEN:-"<YOUR_HF_TOKEN>"} |
| 88 | +export CUDA_DEVICE_MAX_CONNECTIONS=32 |
| 89 | +export CUDNNFE_CLUSTER_OVERLAP_MARGIN=8 |
| 90 | +export HF_HUB_OFFLINE=0 |
| 91 | +export NCCL_CROSS_NIC=1 |
| 92 | +export NCCL_GRAPH_REGISTER=0 |
| 93 | +export NCCL_IB_SPLIT_DATA_ON_QPS=1 |
| 94 | +export NCCL_NVLS_ENABLE=0 |
| 95 | +export NCCL_P2P_NET_CHUNKSIZE=2097152 |
| 96 | +export NCCL_RAS_ENABLE=0 |
| 97 | +export NVLINK_DOMAIN_SIZE=72 |
| 98 | +export NVTE_BWD_LAYERNORM_SM_MARGIN=16 |
| 99 | +export NVTE_FWD_LAYERNORM_SM_MARGIN=16 |
| 100 | +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True |
| 101 | +export TOKENIZERS_PARALLELISM=False |
| 102 | +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 |
| 103 | +export TORCH_NCCL_ENABLE_MONITORING=0 |
| 104 | +export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 |
| 105 | +export TORCH_NCCL_HIGH_PRIORITY=1 |
| 106 | +export TRANSFORMERS_OFFLINE=1 |
| 107 | + |
| 108 | +cd /opt |
| 109 | +rm -rf Megatron-Bridge |
| 110 | +git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git |
| 111 | +cd Megatron-Bridge |
| 112 | +git checkout 5cb3444c43f7499cf3872b2d46870cf8bc2e00ce |
| 113 | +git submodule update --init --recursive |
| 114 | +sed -i -e '/pretrain(config=recipe/i \ recipe.dist.distributed_timeout_minutes = 10' scripts/performance/run_script.py |
| 115 | +ls |
| 116 | + |
| 117 | + |
| 118 | + |
| 119 | +worker_command=$(cat <<- EOM |
| 120 | + if [ "\$RANK" -eq "0" ]; then |
| 121 | + echo "Worker 0 is stalling for a few seconds.." ; |
| 122 | + sleep 3 ; |
| 123 | + echo "The detected environment within worker rank 0 is:" ; |
| 124 | + env | sed 's/^/ /' ; |
| 125 | + fi ; |
| 126 | +
|
| 127 | + cd /opt/Megatron-Bridge ; |
| 128 | +
|
| 129 | + export MEGATRON_CONFIG_LOCK_DIR="/tmp/\$LOCAL_RANK" |
| 130 | + mkdir -p "\$MEGATRON_CONFIG_LOCK_DIR" |
| 131 | +
|
| 132 | + numactl \ |
| 133 | + --cpunodebind=\$((LOCAL_RANK/2)) \ |
| 134 | + --membind=\$((LOCAL_RANK/2)) \ |
| 135 | + nice -10 \ |
| 136 | + python scripts/performance/run_script.py \ |
| 137 | + --model_family_name llama \ |
| 138 | + --model_recipe_name llama31_405b \ |
| 139 | + --config_variant v2 \ |
| 140 | + --gpu gb200 \ |
| 141 | + --num_gpus 256 \ |
| 142 | + --gpus_per_node 4 \ |
| 143 | + --compute_dtype nvfp4 \ |
| 144 | + --seq_length 8192 \ |
| 145 | + --global_batch_size 1536 \ |
| 146 | + --micro_batch_size 1 \ |
| 147 | + --tensor_model_parallel_size 4 \ |
| 148 | + --pipeline_model_parallel_size 16 \ |
| 149 | + --virtual_pipeline_model_parallel_size 8 \ |
| 150 | + --context_parallel_size 1 \ |
| 151 | + --expert_model_parallel_size 1 \ |
| 152 | + --expert_tensor_parallel_size 1 \ |
| 153 | + --cuda_graph_impl none \ |
| 154 | + --max_steps 30 \ |
| 155 | + logger.log_throughput=True \ |
| 156 | + dataset.num_workers=0 \ |
| 157 | + train.manual_gc_interval=100 |
| 158 | +
|
| 159 | +EOM |
| 160 | +) |
| 161 | + |
| 162 | +echo "$worker_command" > worker_command.sh |
| 163 | +chmod 777 worker_command.sh |
| 164 | + |
| 165 | +mkdir -p "/runtime-logs/${JOBSET_NAME}/logs" |
| 166 | + |
| 167 | +set -o pipefail |
| 168 | +torchrun \ |
| 169 | +--nproc-per-node="4" \ |
| 170 | +--nnodes="64" \ |
| 171 | +--node_rank="${JOB_COMPLETION_INDEX}" \ |
| 172 | +--rdzv_id="${JOB_IDENTIFIER}" \ |
| 173 | +--master_addr="${MASTER_ADDR}" \ |
| 174 | +--master_port="${MASTER_PORT}" \ |
| 175 | +--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 -a "/runtime-logs/${JOBSET_NAME}/logs/nemo_mb_RANK_${JOB_COMPLETION_INDEX}.log" |
| 176 | + |
| 177 | + |
| 178 | +if [[ "$JOB_COMPLETION_INDEX" == "0" && -n "${ARTIFACT_DIR}" ]]; then |
| 179 | + if mkdir -p "${ARTIFACT_DIR}" 2>/dev/null; then |
| 180 | + cp -r "${explicit_log_dir}"/* "${ARTIFACT_DIR}/" 2>/dev/null || true |
| 181 | + env > "${ARTIFACT_DIR}/environ.txt" 2>/dev/null || true |
| 182 | + ls "${ARTIFACT_DIR}" 2>/dev/null || true |
| 183 | + fi |
| 184 | +fi |
| 185 | + |
| 186 | +# All nodes copy their local nsys profiles to GCS |
| 187 | +if [[ -d /root/.nemo_run/experiments && -n "${ARTIFACT_DIR}" ]]; then |
| 188 | + if mkdir -p "${ARTIFACT_DIR}/nsys" 2>/dev/null; then |
| 189 | + find /root/.nemo_run/experiments -name "*.nsys-rep" -exec cp {} "${ARTIFACT_DIR}/nsys/" \; 2>/dev/null || true |
| 190 | + fi |
| 191 | +fi |
| 192 | +echo "Training completed" |
| 193 | +echo "Pod on $(hostname --fqdn) is exiting" |
0 commit comments