|
| 1 | +#!/usr/bin/env bash |
| 2 | +set -euo pipefail |
| 3 | +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY all_proxy |
| 4 | +export no_proxy="${no_proxy:-localhost,127.0.0.1}" |
| 5 | +export NO_PROXY="${no_proxy}" |
| 6 | + |
| 7 | +# Submit to an existing 8-node x 8-H100 Ray cluster. All paths must be shared. |
| 8 | +# Assume the BF16 GLM-5.3-Flash checkpoint and dapo-math-17k are already |
| 9 | +# downloaded under BASE_DIR, which must be mounted on every Ray node. |
| 10 | +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" |
| 11 | +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)" |
| 12 | +source "${SCRIPT_DIR}/models/glm5.3-flash-300B.sh" |
| 13 | +BASE_DIR="${BASE_DIR:-/data}" |
| 14 | +HF_CHECKPOINT="${HF_CHECKPOINT:-${BASE_DIR}/GLM-5.3-Flash}" |
| 15 | +PROMPT_DATA="${PROMPT_DATA:-${BASE_DIR}/dapo-math-17k/dapo-math-17k.jsonl}" |
| 16 | +RUN_DIR="${RUN_DIR:-${REPO_ROOT}/local/glm53/$(date -u +%Y%m%dT%H%M%SZ)}" |
| 17 | +mkdir -p "${RUN_DIR}" |
| 18 | +export REPO_ROOT HF_CHECKPOINT |
| 19 | +# Let NCCL/Gloo choose their interfaces unless the cluster needs an override. |
| 20 | +export SOCKET_IFNAME="${SOCKET_IFNAME:-}" |
| 21 | +export MEGATRON_PATH="${MEGATRON_PATH:-/root/Megatron-LM}" |
| 22 | + |
| 23 | +RUNTIME_ENV_JSON="$(python - <<'PY' |
| 24 | +import json, os |
| 25 | +env = { |
| 26 | + 'PYTHONPATH': os.environ['REPO_ROOT'] + ':' + os.environ['MEGATRON_PATH'], |
| 27 | + 'PYTHONUNBUFFERED': '1', 'OMP_NUM_THREADS': '4', |
| 28 | + 'CUDA_DEVICE_MAX_CONNECTIONS': '1', |
| 29 | + 'NCCL_IB_TIMEOUT': '22', 'NCCL_IB_QPS_PER_CONNECTION': '2', |
| 30 | + 'NCCL_NVLS_ENABLE': '0', 'NCCL_CUMEM_ENABLE': '0', |
| 31 | + # Ray workers inherit the head/worker daemon environment as well. |
| 32 | + 'http_proxy': '', 'https_proxy': '', 'HTTP_PROXY': '', 'HTTPS_PROXY': '', |
| 33 | + 'all_proxy': '', 'ALL_PROXY': '', |
| 34 | + 'no_proxy': os.environ['no_proxy'], |
| 35 | + 'NO_PROXY': os.environ['NO_PROXY'], |
| 36 | + 'TOKENIZERS_PARALLELISM': 'false', |
| 37 | +} |
| 38 | +if os.environ['SOCKET_IFNAME']: |
| 39 | + env['NCCL_SOCKET_IFNAME'] = env['GLOO_SOCKET_IFNAME'] = os.environ['SOCKET_IFNAME'] |
| 40 | +print(json.dumps({'env_vars': env})) |
| 41 | +PY |
| 42 | +)" |
| 43 | + |
| 44 | +cd "${REPO_ROOT}" |
| 45 | +# Populate remote-code caches once per node before concurrent worker startup. |
| 46 | +python - <<'PY' |
| 47 | +import os |
| 48 | +import ray |
| 49 | +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy |
| 50 | +
|
| 51 | +ray.init(address="auto", log_to_driver=False) |
| 52 | +
|
| 53 | +@ray.remote |
| 54 | +def prepare_processor(path): |
| 55 | + from transformers import AutoProcessor, AutoTokenizer |
| 56 | + AutoProcessor.from_pretrained(path, trust_remote_code=True) |
| 57 | + AutoTokenizer.from_pretrained(path, trust_remote_code=True) |
| 58 | +
|
| 59 | +ray.get([ |
| 60 | + prepare_processor.options(scheduling_strategy=NodeAffinitySchedulingStrategy(node["NodeID"], soft=False)).remote(os.environ["HF_CHECKPOINT"]) |
| 61 | + for node in ray.nodes() if node["Alive"] |
| 62 | +]) |
| 63 | +ray.shutdown() |
| 64 | +PY |
| 65 | + |
| 66 | +ray job submit --address="${RAY_DASHBOARD:-http://127.0.0.1:8265}" \ |
| 67 | + --runtime-env-json="${RUNTIME_ENV_JSON}" \ |
| 68 | + -- python train.py \ |
| 69 | + --actor-num-nodes 8 --actor-num-gpus-per-node 8 --colocate \ |
| 70 | + "${MODEL_ARGS[@]}" \ |
| 71 | + --hf-checkpoint "${HF_CHECKPOINT}" --load "${HF_CHECKPOINT}" \ |
| 72 | + --megatron-to-hf-mode raw \ |
| 73 | + --prompt-data "${PROMPT_DATA}" --input-key prompt --label-key label --apply-chat-template \ |
| 74 | + --rm-type "${RM_TYPE:-deepscaler}" \ |
| 75 | + --num-rollout "${NUM_ROLLOUT:-3}" --rollout-batch-size 16 --n-samples-per-prompt 8 \ |
| 76 | + --global-batch-size 128 --micro-batch-size 1 \ |
| 77 | + --rollout-max-response-len "${RESPONSE_LENGTH:-32768}" \ |
| 78 | + --rollout-temperature 1.0 --rollout-top-p 0.95 \ |
| 79 | + --use-rollout-routing-replay --use-score-centering --pg-loss-type reinforce \ |
| 80 | + --advantage-estimator grpo \ |
| 81 | + --use-tis --tis-clip-low 0.5 --tis-clip 2.0 \ |
| 82 | + --tensor-model-parallel-size 1 --pipeline-model-parallel-size 4 \ |
| 83 | + --decoder-first-pipeline-num-layers 12 --decoder-last-pipeline-num-layers 11 \ |
| 84 | + --context-parallel-size "${CP_SIZE:-2}" --expert-model-parallel-size 16 --expert-tensor-parallel-size 1 \ |
| 85 | + --allgather-cp \ |
| 86 | + --moe-token-dispatcher-type alltoall \ |
| 87 | + --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \ |
| 88 | + --seq-length 65536 --max-position-embeddings 65536 \ |
| 89 | + --use-dynamic-batch-size --max-tokens-per-gpu 2048 --data-pad-size-multiplier 128 \ |
| 90 | + --log-probs-chunk-size 256 \ |
| 91 | + --optimizer adam --lr 1e-6 --lr-decay-style constant \ |
| 92 | + --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \ |
| 93 | + --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer \ |
| 94 | + --no-load-optim --no-load-rng --no-save-optim \ |
| 95 | + --freeze-params-name-list 'model.visual.*' \ |
| 96 | + --rollout-num-gpus-per-engine 16 \ |
| 97 | + --sglang-dp-size 16 --sglang-ep-size 16 --sglang-enable-dp-attention \ |
| 98 | + --sglang-enable-dp-lm-head --sglang-moe-dense-tp-size 1 \ |
| 99 | + --sglang-mm-enable-dp-encoder --sglang-disable-shared-experts-fusion \ |
| 100 | + --sglang-dsa-prefill-backend fa3 --sglang-dsa-decode-backend fa3 \ |
| 101 | + --sglang-mem-fraction-static 0.75 --sglang-context-length 65536 \ |
| 102 | + --sglang-max-running-requests 32 --sglang-chunked-prefill-size 4096 \ |
| 103 | + --sglang-disable-overlap-schedule --sglang-cuda-graph-max-bs-decode 8 \ |
| 104 | + --sglang-mamba-ssm-dtype float32 --sglang-page-size 64 \ |
| 105 | + --sglang-kv-cache-dtype bfloat16 \ |
| 106 | + --sglang-watchdog-timeout 3600 \ |
| 107 | + --attention-dropout 0 --hidden-dropout 0 --bf16 \ |
| 108 | + --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 \ |
| 109 | + --distributed-timeout-minutes 60 \ |
| 110 | + --update-weight-buffer-size 1073741824 \ |
| 111 | + --save-debug-train-data "${RUN_DIR}/train-{rollout_id}.pt" \ |
| 112 | + "$@" |
0 commit comments