From 6401ef572400bbb5eac510e9f1a34a075ef00033 Mon Sep 17 00:00:00 2001 From: Eric Tang Date: Fri, 19 Jun 2026 21:33:13 +0000 Subject: [PATCH 1/5] [megatron] Enable Nemotron-3-Ultra-550B GRPO RL + fix multi-rank weight sync Adds an end-to-end full-finetuning GRPO recipe for NVIDIA-Nemotron-3-Ultra-550B (NemotronH hybrid Mamba2+attention, latent MoE, reasoning) colocated with vLLM on 8x8 H200 (EFA), plus the weight-sync/reload correctness fixes it depends on. Validated: avg_raw_reward ~0.9, GSM8K eval ~0.94, grad_norm > 0. Core fix (general, affects any MoE synced at EP>16 / PP>2): the CUDA-IPC weight transport sent only rank-0's per-param slicing metadata, but each Megatron rank packs its own (per-rank-divergent) buffer. Each vLLM worker rebuilt its own GPU's buffer yet sliced it with rank-0's metadata -> correct bytes loaded under wrong names -> coherent- but-garbage generations and reward stuck at 0. Now sends per-GPU metadata and each worker slices its own buffer (cuda_ipc_strategy.py, new_inference_worker_wrap.py). Also included: - Preserve fp32 for the MoE router bias (gate.e_score_correction_bias) through sync; bf16 ULP at its ~25-57 magnitude collapses the per-expert offsets and corrupts routing. - Guard vLLM layerwise-reload get_numel_loaded (cf. vllm-project/vllm#44814): the composed-weight-loader double-count finalizes a layer early and drops Mamba mixer.D (uninitialized -> NaN). Mirrors the existing conv_weights reload workaround. - Forward HF_*/cache dirs and SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S to Ray worker actors (prepare_runtime_environment + GPU-CI conftest). - Reasoning-aware GSM8K reward: strip the trace, score strict `#### ` else last-number with comma/$ normalization. - Nemotron-Ultra logprob round-trip test params (EP16/PP2 baseline; EP32/PP4 regressions). - Example recipe (run_megatron_nemotron_ultra.sh), staging helper, and README. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../train/megatron/README_nemotron_ultra.md | 115 +++++++++++ .../megatron/run_megatron_nemotron_ultra.sh | 192 ++++++++++++++++++ .../train/megatron/stage_nemotron_ultra.py | 116 +++++++++++ skyrl-gym/skyrl_gym/envs/gsm8k/env.py | 24 ++- .../inference_servers/layerwise_reload.py | 32 +++ .../new_inference_worker_wrap.py | 13 ++ .../weight_sync/cuda_ipc_strategy.py | 50 +++-- .../workers/megatron/megatron_worker.py | 68 ++++--- skyrl/train/utils/utils.py | 37 ++++ .../skyrl_train/gpu/gpu_ci/conftest.py | 38 +++- .../gpu_ci/megatron/test_megatron_models.py | 140 ++++++++++++- tests/backends/skyrl_train/gpu/utils.py | 8 +- 12 files changed, 780 insertions(+), 53 deletions(-) create mode 100644 examples/train/megatron/README_nemotron_ultra.md create mode 100644 examples/train/megatron/run_megatron_nemotron_ultra.sh create mode 100644 examples/train/megatron/stage_nemotron_ultra.py diff --git a/examples/train/megatron/README_nemotron_ultra.md b/examples/train/megatron/README_nemotron_ultra.md new file mode 100644 index 0000000000..12ce1f2c90 --- /dev/null +++ b/examples/train/megatron/README_nemotron_ultra.md @@ -0,0 +1,115 @@ +# Nemotron-3-Ultra-550B GRPO RL on GSM8K (Megatron, multi-node) + +Full-finetuning GRPO RL of **NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16** (NemotronH hybrid +Mamba2 + attention, latent MoE with 512 experts, reasoning model) colocated with vLLM on +**8× nodes of 8×H200-141GB (64 GPUs, EFA)**. + +Recipe: [`run_megatron_nemotron_ultra.sh`](./run_megatron_nemotron_ultra.sh). +Staging helper: [`stage_nemotron_ultra.py`](./stage_nemotron_ultra.py). + +**Validated:** trains end-to-end with this config — `avg_raw_reward ≈ 0.9`, GSM8K +`eval ≈ 0.94`, `grad_norm > 0` (genuinely learning). Megatron mesh TP8 / PP4 / EP16 / ETP1 +(DP2); vLLM TP8 × PP4 (2 engines). + +--- + +## Replicating on a fresh cluster + +The cluster needs: 8 nodes × 8×H200-141GB, EFA, a Ray cluster, a large **node-local** disk at +`/mnt/local_storage` (~28 TB), and a small shared `/home` (which the 1.1 TB model must NOT touch). + +### 1. Make sure this PR's code is present +The recipe depends on several fixes in this PR (see [Why these fixes](#why-these-fixes-are-needed)). +On stock SkyRL/vLLM without them you get coherent-looking **garbage** generations and `reward=0`. + +### 2. Stage the model + data on every GPU node +Everything lives on node-local `/mnt/local_storage` (the model is too big for `/home`, and every +rank needs its data locally). One command does both, on all nodes, via Ray: + +```bash +export HF_TOKEN=$(cat ~/.HF_TOKEN) # fast authenticated download; unauthenticated is throttled +uv run --isolated --with ray --with huggingface_hub --with hf_transfer --with datasets \ + python examples/train/megatron/stage_nemotron_ultra.py +``` + +This downloads the HF snapshot to `/mnt/local_storage/hf_cache` **including `chat_template.jinja`** +and writes the GSM8K parquets to `/mnt/local_storage/data/gsm8k` on each node. Re-run it if the +autoscaler churns in a fresh (un-staged) node. + +> The `*.jinja` is essential. The tokenizer ships **no** chat template inline; the official ChatML + +> reasoning template lives in `chat_template.jinja`. Without it the instruct model is prompted +> off-distribution and never produces a parseable answer (reward stays 0). + +### 3. Caches go to `/mnt/local_storage` +Handled by the script: it exports `HF_HOME`, `XDG_CACHE_HOME`, `UV_CACHE_DIR`, `TRITON_CACHE_DIR`, +`TORCHINDUCTOR_CACHE_DIR`, `VLLM_CACHE_ROOT` → `/mnt/local_storage/...`, and SkyRL's +`prepare_runtime_environment` (this PR) forwards them to the Ray worker actors. Otherwise workers +default to `~/.cache` on the small `/home`, fill it, and take the node down (looks like an OOM / +preemption, but it's disk). + +### 4. Launch +```bash +export WANDB_API_KEY= +export HF_TOKEN=$(cat ~/.HF_TOKEN) # for churn resilience +bash examples/train/megatron/run_megatron_nemotron_ultra.sh +``` +EFA + multi-node specifics (all set by the script): `LD_LIBRARY_PATH=/opt/amazon/efa/lib`, +`SKYRL_LD_LIBRARY_PATH_EXPORT=1`, `VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1`, `NVTE_FLASH_ATTN=0`, and +raised placement-group / inference-server health timeouts (the 550B takes >600 s to come up). + +That's it — **stage model+data on every node, keep caches on `/mnt/local_storage`, and run.** + +--- + +## Why these fixes are needed + +The hard part was that vLLM generated **garbage** (multilingual token-salad / degenerate +repetition) after every weight sync → all rewards 0 → no learning. The root-cause chain and the +fixes (all in this PR): + +1. **CUDA-IPC weight sync used only rank-0's slicing metadata** (the core bug, general to MoE). + Each Megatron rank packs its *own* contiguous buffer (different params/order per rank — expert + chunks carry per-EP-rank names) and registers one IPC handle per physical GPU, but only rank 0's + `names`/`sizes`/`shapes` were sent. Each vLLM worker rebuilt *its own* GPU's buffer yet sliced it + with rank-0's metadata → correct bytes loaded under the wrong names → coherent-but-garbage, no + crash. Identical across PP ranks (so it worked at PP=2) but divergent at **PP>2 / EP>16**. + *Fix:* send per-GPU metadata; each worker slices its own buffer with its own + (`cuda_ipc_strategy.py`, `new_inference_worker_wrap.py`). Verified: EP16/PP4 post-sync logprob + diff `2.0 → 0.15`. + +2. **fp32 MoE router bias must not be down-cast.** `gate.e_score_correction_bias` is large + (~25–57) with tiny per-expert offsets (std ~7e-4) far below bf16 ULP at that scale; bf16 rounding + collapses the offsets and corrupts routing. *Fix:* transfer it in native fp32 through the sync + (`megatron_worker.py`). + +3. **vLLM layerwise-reload dropped Mamba `mixer.D`** (cf. vllm-project/vllm#44814). The reload + element-counter double-counts `composed_weight_loader` params (Mamba `A`), finalizing the layer + early and leaving `mixer.D` uninitialized → NaN. *Fix:* a guarded monkeypatch capping the count at + `param.numel()` (`layerwise_reload.py`), alongside SkyRL's existing `conv_weights` reload patch. + Remove once on a vLLM that includes #44814. + +4. **Chat template staging** (`*.jinja`) — see step 2 above. + +5. **Robust reasoning-aware GSM8K reward** — strip the `` trace, then score the answer with + strict `#### ` else last-number-with-normalization, so boxed/natural-language answers from a + reasoning model are scored correctly (`skyrl_gym/envs/gsm8k/env.py`). + +6. **Worker env forwarding** — `prepare_runtime_environment` (training) and the GPU-CI conftest + forward `HF_*` / cache dirs / `VLLM_USE_RAY_V2_EXECUTOR_BACKEND` / + `SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S` to the Ray worker actors. + +## Memory & parallelism notes +- Full-FT adds bf16 grads (~= weights) + the AdamW master. At EP16/**PP2** that's ~69+69 GiB/GPU → + OOMs the 141 GiB H200, so we use **PP=4** (halves per-GPU weights and grads → ~34+34). The + optimizer (fp32 master + Adam moments) is **CPU-offloaded** (host RAM, ~2 TB/node). +- With fix #1 in place, weight sync is correct at **any** EP/PP; EP is now just a memory/throughput + knob (e.g. EP=32 → 16 experts/rank for more headroom). EP must divide TP×DP. +- vLLM PP=4 keeps its weights ~34 GiB/GPU so both vLLM (woken) and the resident policy shard fit + during the colocated sync. + +## Known issues +- The model emits a `…` reasoning block; `max_generate_length=2048` gives room to + finish reasoning AND emit the answer (batched mode can't toggle `enable_thinking`). +- Node churn on large pools: a vLLM worker dying ("Executor failed") kills the run; raise + `trainer.ckpt_interval` resilience and re-stage churned-in nodes. `HF_HUB_OFFLINE=0` lets an + un-staged node self-download to `/mnt` instead of erroring. diff --git a/examples/train/megatron/run_megatron_nemotron_ultra.sh b/examples/train/megatron/run_megatron_nemotron_ultra.sh new file mode 100644 index 0000000000..41e173ed0b --- /dev/null +++ b/examples/train/megatron/run_megatron_nemotron_ultra.sh @@ -0,0 +1,192 @@ +set -x + +# Colocated GRPO training+generation for NVIDIA-Nemotron-3-Ultra-550B-A55B on GSM8K with Megatron. +# Runs on 8 nodes of 8xH200-141GB (64 GPUs), EFA interconnect. +# +# This is *full-finetuning* RL (no LoRA, no ref/KL model). It builds on the configs proven by +# the logprob round-trip test (tests/.../gpu_ci/megatron/test_megatron_models.py::[nemotron3-ultra]) +# but the test was forward-only (inference_only_init), so it had no optimizer/grads. Training adds +# bf16 grads (~same size as the weights) + the AdamW master state, so to fit the 141 GiB H200 we: +# (a) shard depth with PP=4 (halves per-GPU weights AND grads vs the test's PP=2 -> ~34+34 GiB), +# (b) CPU-offload the optimizer (fp32 master + Adam moments live on host RAM, not GPU), +# (c) recompute activations, (d) bin-pack microbatches by token count, and +# (e) drop the KL/ref model (no second 550B copy). +# VALIDATED working: this exact config trains end-to-end (reward ~0.9, gsm8k eval ~0.94, grad_norm>0). +# +# NOTE on correctness: getting coherent generations required two SkyRL fixes that are now in-tree +# (not knobs here): (1) the CUDA-IPC weight-sync sends per-GPU slicing metadata so each vLLM worker +# slices its own packed buffer correctly -- without it, weight sync corrupts vLLM at PP>2 / EP>16 +# (the policy stays fine, but vLLM generates token-salad and reward stays 0); (2) a vLLM +# layerwise-reload patch (cf. vllm-project/vllm#44814) so the NemotronH Mamba `mixer.D` isn't +# dropped during reload. If you run on a SkyRL/vLLM without these, expect garbage generations. +# +# Prereqs: +# uv run examples/train/gsm8k/gsm8k_dataset.py --output_dir $HOME/data/gsm8k +# Stage the model on every node's local disk (1.1TB; /home is too small): +# see the staging helper used for the test (HF_HOME=/mnt/local_storage/hf_cache). +# IMPORTANT: stage chat_template.jinja too (include *.jinja in allow_patterns). It is +# the model's official ChatML+reasoning template; without it the tokenizer/vLLM have NO +# chat template, the instruct model is prompted off-distribution, and reward stays 0. +# export WANDB_API_KEY= +# bash examples/train/megatron/run_megatron_nemotron_ultra.sh + +# --------------------------------------------------------------------------- +# Environment (must reach the Ray workers). These mirror the test's run env. +# --------------------------------------------------------------------------- +# Model is staged on each node's large local disk (1.1TB won't fit /home/ray's 255GB). +export HF_HOME=${HF_HOME:-/mnt/local_storage/hf_cache} +# Redirect ALL caches off the small home disk (255GB) to the big local disk (28TB). +# Workers write uv build envs, Triton/Inductor/vLLM/FlashInfer JIT caches, etc. to +# ~/.cache by default; on a small home disk that fills up and takes the node down. +# These are forwarded to Ray workers by prepare_runtime_environment. +export XDG_CACHE_HOME=${XDG_CACHE_HOME:-/mnt/local_storage/.cache} +export UV_CACHE_DIR=${UV_CACHE_DIR:-/mnt/local_storage/.cache/uv} +export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/mnt/local_storage/.cache/triton} +export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-/mnt/local_storage/.cache/inductor} +export VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT:-/mnt/local_storage/.cache/vllm} +# Use the local cache only (avoids re-downloading / HF rate limits). Unset if you +# want to allow downloads on first run. +export HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-1} +# EFA: NCCL must see the EFA libs, and SkyRL must forward LD_LIBRARY_PATH to Ray workers. +export LD_LIBRARY_PATH=/opt/amazon/efa/lib:${LD_LIBRARY_PATH:-} +export SKYRL_LD_LIBRARY_PATH_EXPORT=1 +# vLLM multi-node executor: the default Ray compiled-DAG (shm channel) crashes the raylet +# on the cross-node hop; the V2 (MultiprocExecutor/MessageQueue) backend avoids it. +export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 +# Megatron attention backend (TE flash attn off; see .claude/docs/backends/megatron.md). +export NVTE_FLASH_ATTN=0 +# 8-node uv cache warmup + 550B load can exceed the default placement-group timeout. +export SKYRL_RAY_PG_TIMEOUT_IN_S=${SKYRL_RAY_PG_TIMEOUT_IN_S:-1800} +# The 550B vLLM engines take a while to come up; raise the health-wait timeout. +export SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S=${SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S:-2400} +# Set HF_TOKEN (e.g. `export HF_TOKEN=$(cat ~/.HF_TOKEN)`) for fast authenticated staging. +# HF_HUB_OFFLINE=0 (instead of 1) makes workers re-download a missing shard to the big +# disk if a node churns in un-staged, instead of erroring; with a stable staged pool, 1 is fine. +# Surface vLLM/worker logs to stdout (helpful while bringing this up; comment out later). +export SKYRL_DUMP_INFRA_LOG_TO_STDOUT=${SKYRL_DUMP_INFRA_LOG_TO_STDOUT:-1} +# export NCCL_DEBUG=WARN + +# Data must be present on ALL nodes (node-local) for multi-node training. gsm8k is tiny; +# stage it to each node's local disk (e.g. copy $HOME/data/gsm8k -> here on every node). +DATA_DIR="/mnt/local_storage/data/gsm8k" +LOGGER="wandb" # change to "console" to print to stdout +MODEL_NAME="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16" + +INFERENCE_BACKEND="vllm" # currently only vllm is supported for megatron + +NUM_NODES=8 +NUM_GPUS=8 # per node + +### Megatron (policy) parallelism. world = TP*PP*DP = 64. +# TP within the NVLink domain; NemotronH Mamba requires TP | n_groups(=8), so TP in {1,2,4,8}. +MEGATRON_TP=8 +# PP=4 (vs the forward-only test's 2): training adds bf16 grads (~same size as weights), so at +# EP=16/PP=2 the ~69 GiB weights/GPU + ~69 GiB grads ~= 138 GiB doesn't fit the 141 GiB H200. +# PP=4 halves the layers (hence weights AND grads) per GPU to ~34+34 GiB, which fits. +MEGATRON_PP=4 +MEGATRON_CP=1 +# EP=16, ETP=1 -> EDP=1 (world = TP*PP*DP = 8*4*2 = 64; EP*ETP = 16 = TP*DP = 8*2). +# This is the validated config. Earlier runs at EP=32 produced garbage vLLM generations, but that +# was the CUDA-IPC weight-sync TRANSPORT bug (rank-0 slicing metadata reused for every GPU's +# divergent buffer), NOT the expert sharding itself -- the bridge's expert export is bit-correct at +# every EP. With the per-GPU-metadata fix now in-tree, any valid EP syncs correctly, so EP is purely +# a memory/throughput knob: e.g. EP=32 (16 experts/rank vs 32) further lowers per-GPU expert memory +# if you need more headroom. EP must divide TP*DP. +MEGATRON_EP=16 +MEGATRON_ETP=1 + +# Activation recompute (gated by trainer.gradient_checkpointing=true, which is the default). +RECOMPUTE_GRANULARITY="full" +RECOMPUTE_METHOD="uniform" +RECOMPUTE_NUM_LAYERS=1 + +# CPU-offload the optimizer (fp32 master + AdamW) so it doesn't sit on the GPU. +OPTIMIZER_OFFLOAD=true +OPTIMIZER_OFFLOAD_FRACTION=1.0 + +# Bin-pack microbatches by token count (with remove_microbatch_padding). When >0, +# micro_*_batch_size_per_gpu are ignored. Bounds activation memory; a single sequence +# longer than this still gets its own microbatch. longest seq here ~= 512+1024. +MAX_TOKENS_PER_MICROBATCH=4096 + +### Inference engine (vLLM), colocated over the same 64 GPUs. +# TP=8 (intra-node, NVLink) x PP=4 (cross-node, EFA) = 32 GPUs/engine, 2 engines -> 64 GPUs. +# vLLM TP must divide Mamba n_groups(=8); cross-node scale comes from PP. PP=4 (not 2) keeps +# vLLM weights ~34GB/GPU so during the colocated weight sync (vLLM woken alongside the resident +# policy shard) both fit on the 141 GiB H200 (PP=2 -> ~69+69 OOMs). +NUM_INFERENCE_ENGINES=2 +INFERENCE_ENGINE_TP=8 +INFERENCE_ENGINE_PP=4 +# Cap context: the model's native max is huge and vLLM sizes the KV pool for 1 max-len request. +INFERENCE_ENGINE_MAX_MODEL_LEN=4096 +# Nemotron-3-Ultra is a REASONING model: its official chat_template.jinja defaults to +# enable_thinking=true, so each rollout emits a ... block before the answer. +# In batched mode chat templating is done server-side by vLLM (chat_template_kwargs is not +# supported), so we cannot disable thinking from here -- instead we give generation enough +# budget to finish reasoning AND emit the final `#### ` the gsm8k reward parser wants. +# (Earlier runs got reward=0 because the chat template wasn't staged at all -> the instruct +# model was prompted off-distribution and never produced a parseable answer.) +GEN_MAX_LEN=2048 +# vLLM and the policy alternate on-GPU (sleep/wake); leave headroom for the policy shard. +GPU_MEMORY_UTILIZATION=0.6 + +uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \ + data.train_data="['$DATA_DIR/train.parquet']" \ + data.val_data="['$DATA_DIR/validation.parquet']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.policy.model.path=$MODEL_NAME \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TP \ + generator.inference_engine.pipeline_parallel_size=$INFERENCE_ENGINE_PP \ + generator.inference_engine.distributed_executor_backend=ray \ + generator.inference_engine.use_expandable_segments=true \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.transformer_config_kwargs.mtp_num_layers=0 \ + trainer.policy.megatron_config.transformer_config_kwargs.mtp_hybrid_override_pattern=null \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_granularity=$RECOMPUTE_GRANULARITY \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_method=$RECOMPUTE_METHOD \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_num_layers=$RECOMPUTE_NUM_LAYERS \ + trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \ + trainer.remove_microbatch_padding=true \ + trainer.max_tokens_per_microbatch=$MAX_TOKENS_PER_MICROBATCH \ + trainer.epochs=20 \ + trainer.eval_batch_size=1024 \ + trainer.eval_before_train=false \ + trainer.eval_interval=5 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=64 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.max_prompt_length=512 \ + generator.sampling_params.max_generate_length=$GEN_MAX_LEN \ + trainer.policy.optimizer_config.lr=1.0e-6 \ + trainer.algorithm.use_kl_loss=false \ + generator.inference_engine.backend=$INFERENCE_BACKEND \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.inference_engine.engine_init_kwargs.max_model_len=$INFERENCE_ENGINE_MAX_MODEL_LEN \ + generator.inference_engine.gpu_memory_utilization=$GPU_MEMORY_UTILIZATION \ + generator.batched=true \ + environment.env_class=gsm8k \ + generator.n_samples_per_prompt=5 \ + trainer.logger="$LOGGER" \ + trainer.project_name="gsm8k_nemotron_ultra" \ + trainer.run_name="gsm8k_nemotron_ultra_tp${MEGATRON_TP}_pp${MEGATRON_PP}_ep${MEGATRON_EP}" \ + trainer.resume_mode=latest \ + trainer.max_ckpts_to_keep=3 \ + trainer.ckpt_interval=20 \ + trainer.ckpt_path="$HOME/ckpts/gsm8k_nemotron_ultra_ckpt" \ + $@ diff --git a/examples/train/megatron/stage_nemotron_ultra.py b/examples/train/megatron/stage_nemotron_ultra.py new file mode 100644 index 0000000000..3c40604e9c --- /dev/null +++ b/examples/train/megatron/stage_nemotron_ultra.py @@ -0,0 +1,116 @@ +"""Pre-stage NVIDIA-Nemotron-3-Ultra-550B + the GSM8K dataset onto every GPU node's +local disk, for the multi-node recipe in ``run_megatron_nemotron_ultra.sh``. + +Why this exists: the model is ~1.1 TB and the cluster's shared ``/home`` is small +(~255 GB), so the checkpoint must live on each node's large node-local disk +(``/mnt/local_storage``, ~28 TB). Multi-node training also needs the data present on +every node (each rank reads its data locally). This launches one Ray task per distinct +GPU node that (1) downloads the HF snapshot and (2) writes the GSM8K parquets, both under +``/mnt/local_storage``. + +IMPORTANT: ``allow_patterns`` includes ``*.jinja`` so the model's official +``chat_template.jinja`` is staged. Without it the tokenizer/vLLM have NO chat template, +the instruct/reasoning model is prompted off-distribution, and reward stays 0. + +Usage (from the head node, on a running Ray cluster): + export HF_TOKEN=$(cat ~/.HF_TOKEN) # fast authenticated download; unauth is throttled + uv run --isolated --with ray --with huggingface_hub --with hf_transfer --with datasets \ + python examples/train/megatron/stage_nemotron_ultra.py +""" + +import os + +import ray +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy +from ray.util.state import list_nodes + +REPO = "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16" +HF_HOME = "/mnt/local_storage/hf_cache" +DATA_DIR = "/mnt/local_storage/data/gsm8k" + +ray.init(address="auto", ignore_reinit_error=True, log_to_driver=True) + +_ENV = {"HF_HOME": HF_HOME, "HF_HUB_ENABLE_HF_TRANSFER": "1"} +if os.environ.get("HF_TOKEN"): + _ENV["HF_TOKEN"] = os.environ["HF_TOKEN"] + + +@ray.remote( + num_cpus=8, + runtime_env={"pip": ["huggingface_hub", "hf_transfer", "datasets"], "env_vars": _ENV}, +) +def stage(node_ip): + import socket + import time + + from huggingface_hub import snapshot_download + + # 1) Model snapshot (weights + config + tokenizer + chat_template.jinja). + model_path = None + last_err = None + for _ in range(8): + try: + model_path = snapshot_download( + REPO, + cache_dir=f"{HF_HOME}/hub", + max_workers=16, + # NOTE: *.jinja is required (the chat template); *.py pulls any custom code. + allow_patterns=["*.safetensors", "*.json", "*.txt", "tokenizer*", "*.model", "*.jinja", "*.py"], + ) + break + except Exception as e: # noqa: BLE001 + last_err = f"{type(e).__name__}: {str(e)[:300]}" + time.sleep(15) + if model_path is None: + return (socket.gethostname(), node_ip, "MODEL_FAILED", last_err) + + # 2) GSM8K parquets (tiny; one copy per node so every rank reads locally). + try: + import re + + import datasets + + os.makedirs(DATA_DIR, exist_ok=True) + instruction = 'Let\'s think step by step and output the final answer after "####".' + + def to_row(example, idx, split): + q = example["question"] + sol = re.search(r"#### (\-?[0-9\.\,]+)", example["answer"]).group(0).split("#### ")[1].replace(",", "") + return { + "data_source": "openai/gsm8k", + "prompt": [{"role": "user", "content": q + " " + instruction}], + "env_class": "gsm8k", + "reward_spec": {"method": "rule", "ground_truth": sol}, + "extra_info": {"split": split, "index": idx}, + } + + ds = datasets.load_dataset("openai/gsm8k", "main") + ds["train"].map(lambda e, i: to_row(e, i, "train"), with_indices=True).to_parquet(f"{DATA_DIR}/train.parquet") + ds["test"].map(lambda e, i: to_row(e, i, "test"), with_indices=True).to_parquet( + f"{DATA_DIR}/validation.parquet" + ) + except Exception as e: # noqa: BLE001 + return (socket.gethostname(), node_ip, "DATA_FAILED", f"{type(e).__name__}: {str(e)[:300]}") + + return (socket.gethostname(), node_ip, "DONE", model_path) + + +nodes = [ + n + for n in list_nodes(detail=True) + if (n.get("resources_total") or {}).get("GPU", 0) > 0 and n.get("state") == "ALIVE" +] +print("staging on %d GPU nodes: %s" % (len(nodes), [n["node_ip"] for n in nodes]), flush=True) +results = ray.get( + [ + stage.options(scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=n["node_id"], soft=False)).remote( + n["node_ip"] + ) + for n in nodes + ] +) +print("\n===== RESULTS =====", flush=True) +ok = sum(1 for r in results if r[2] == "DONE") +for host, ip, status, detail in results: + print(f"[{ip} {host}] {status}: {detail}", flush=True) +print(f"\n{ok}/{len(results)} nodes staged OK", flush=True) diff --git a/skyrl-gym/skyrl_gym/envs/gsm8k/env.py b/skyrl-gym/skyrl_gym/envs/gsm8k/env.py index d4fe5b3d88..d8266312d2 100644 --- a/skyrl-gym/skyrl_gym/envs/gsm8k/env.py +++ b/skyrl-gym/skyrl_gym/envs/gsm8k/env.py @@ -1,6 +1,15 @@ from skyrl_gym.envs.base_text_env import BaseTextEnv, BaseTextEnvStepOutput from skyrl_gym.envs.gsm8k import utils from typing import Dict, Any +import re as _re + +# Matches an integer/decimal (optionally signed, with thousands separators), e.g. "1,800", "-5", "72.0". +_NUM_RE = _re.compile(r"-?[0-9][0-9,]*(?:\.[0-9]+)?") + + +def _norm_num(s: str) -> str: + """Normalize a parsed number to compare against the (comma-free integer) ground truth.""" + return s.strip().rstrip(".").replace(",", "").replace("$", "") class GSM8kEnv(BaseTextEnv): @@ -16,7 +25,20 @@ def __init__(self, env_config: Any = None, extras: Dict[str, Any] = {}): self.ground_truth = extras["reward_spec"]["ground_truth"] def _get_reward(self, action: str) -> float: - return utils.compute_score(action, self.ground_truth) + # Reasoning models (e.g. Nemotron-3-Ultra) emit `` and may write + # the final answer in their own style (boxed / natural language) rather than the exact + # `#### ` the strict parser wants. Drop the `` trace (so its intermediate + # numbers can't be mistaken for the answer), then reward the exact `#### ` format if + # present (strict), else fall back to the last number in the answer (robust to boxed / + # natural-language answers and comma/$ formatting). Non-reasoning models (no ``) keep + # the original strict scoring on the full output, so this is a no-op for them. + answer_segment = action.split("")[-1] if "" in action else action + reward = utils.compute_score(answer_segment, self.ground_truth, method="strict") + if reward == 0.0: + nums = _NUM_RE.findall(answer_segment) + if nums and _norm_num(nums[-1]) == _norm_num(str(self.ground_truth)): + reward = 1.0 + return reward def step(self, action: str) -> BaseTextEnvStepOutput: done = True # always done after one step diff --git a/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py b/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py index d3aa2918d5..c2e0ca484c 100644 --- a/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py +++ b/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py @@ -42,6 +42,38 @@ except ImportError: pass +# Workaround for a SECOND vLLM layerwise-reload corruption affecting NemotronH/Mamba. +# `get_numel_loaded` counts elements copied by a weight loader to decide when a layer is +# fully loaded (so it can be finalized). A `composed_weight_loader` copies into its param +# twice (initial load + in-place transform), so the counter reports 2x the param's numel. +# That inflated count can reach the layer's total before all params arrive, finalizing the +# layer early and leaving trailing params (e.g. Mamba `mixer.D`) UNINITIALIZED -> NaN logits. +# Whether it trips depends on the order params arrive (i.e. the weight-sync chunk/bucket +# layout), which is why it corrupts at some Megatron EP/PP configs (PP>2, EP>16) but not +# others. Fix mirrors vLLM PR #44814: cap each loader call's counted elements at the +# destination param's numel(). Patch both the `meta` source and the `layerwise` import-site +# reference (layerwise.py does `from .meta import get_numel_loaded` at module load). +# Remove once the vLLM version includes https://github.com/vllm-project/vllm/pull/44814 . +try: + from vllm.model_executor.model_loader.reload import ( + layerwise as _vllm_reload_layerwise, + ) + from vllm.model_executor.model_loader.reload import meta as _vllm_reload_meta + + _orig_get_numel_loaded = _vllm_reload_meta.get_numel_loaded + + def _capped_get_numel_loaded(weight_loader, args): + numel, return_value = _orig_get_numel_loaded(weight_loader, args) + param = getattr(args, "arguments", {}).get("param", None) + if isinstance(param, torch.Tensor): + numel = min(numel, param.numel()) + return numel, return_value + + _vllm_reload_meta.get_numel_loaded = _capped_get_numel_loaded + _vllm_reload_layerwise.get_numel_loaded = _capped_get_numel_loaded +except (ImportError, AttributeError): + pass + class LayerwiseReloadWorkerMixin: """Bracket a multi-chunk weight sync with one vLLM layerwise-reload init/finalize. diff --git a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py index 92f5906211..4d3831f536 100644 --- a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py +++ b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py @@ -91,6 +91,19 @@ def update_weights_ipc(self, update_info: dict) -> None: physical_gpu_id = str(torch.cuda.get_device_properties(device_index).uuid) if physical_gpu_id not in handles: raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. " f"Available: {list(handles)}") + + # Each GPU's packed buffer was built from a DIFFERENT set/order of params, so slice + # it with THIS GPU's own metadata when available (falling back to the flat rank-0 + # metadata for older senders). Using rank 0's metadata for every GPU silently + # mis-slices buffers whose layout differs from rank 0's (EP>16 / PP>2 expert chunks). + per_gpu_pickled = update_info.get("per_gpu_meta_pickled") + if per_gpu_pickled is not None: + per_gpu = pickle.loads(base64.b64decode(per_gpu_pickled)) + meta = per_gpu.get(physical_gpu_id) + if meta is not None: + names = meta["names"] + shapes = meta["shapes"] + sizes = meta["sizes"] func, args = handles[physical_gpu_id] # Remap device index to the LOCAL current-device. list_args = list(args) diff --git a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py index bdf2c7408e..fbf7e1337b 100644 --- a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py @@ -197,8 +197,6 @@ async def _send_chunks_vllm_native( world_size = torch.distributed.get_world_size() device = torch.cuda.current_device() gpu_uuid = str(torch.cuda.get_device_properties(device).uuid) - dtype = str_to_torch_dtype(self._init_info.model_dtype_str) - dtype_name = self._init_info.model_dtype_str.split(".")[-1] if rank == 0: await self._inference_client.start_weight_update(is_checkpoint_format=True) @@ -206,8 +204,13 @@ async def _send_chunks_vllm_native( for chunk in chunks: # --- pack all tensors in this chunk into one contiguous buffer --- - # Chunk tensors share a single dtype by construction (see - # weight_extractor_utils.py), so offsets in element units are safe. + # Chunk tensors share a single dtype by construction (the weight extractor + # emits one chunk per dtype), so offsets in element units are safe. Derive + # the packed dtype from the chunk itself rather than assuming model_dtype, + # so fp32 params (e.g. the MoE router bias) are transferred without bf16 + # down-casting; the receiver rebuilds at this dtype from the IPC handle. + chunk_dtype = chunk.tensors[0].dtype + dtype_name = str(chunk_dtype).split(".")[-1] names: List[str] = [] dtype_names: List[str] = [] shapes: List[List[int]] = [] @@ -217,7 +220,7 @@ async def _send_chunks_vllm_native( packed_tensor = torch.empty( total_numel, device=device, - dtype=dtype, + dtype=chunk_dtype, requires_grad=False, ) @@ -231,28 +234,49 @@ async def _send_chunks_vllm_native( shapes.append(list(shape) if not isinstance(shape, list) else shape) sizes.append(size) - # --- one IPC handle per rank for the packed buffer --- + # --- one IPC handle per rank for the packed buffer, plus THIS rank's own + # per-param metadata. Each rank packs a DIFFERENT set/order of params into its + # buffer (e.g. expert chunks carry per-EP-rank expert names), so the receiver + # must slice each GPU's buffer with THAT GPU's own names/sizes — not rank 0's. + # Using rank 0's metadata for every GPU mis-slices any buffer whose layout + # differs from rank 0's, which happens under EP>16 / PP>2 and silently loads + # the wrong tensor under each name (no crash; byte totals still match). See + # NEMOTRON_ULTRA_FINDINGS.md. ipc_handle: IpcHandle = reduce_tensor(packed_tensor) - local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} - gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size - torch.distributed.all_gather_object(gathered, local_handle_dict) + local_entry: Dict[str, Any] = { + gpu_uuid: { + "handle": ipc_handle, + "names": names, + "dtype_names": dtype_names, + "shapes": shapes, + "sizes": sizes, + } + } + gathered: List[Optional[Dict[str, Any]]] = [None] * world_size + torch.distributed.all_gather_object(gathered, local_entry) torch.distributed.barrier() torch.cuda.synchronize() if rank == 0: - merged_handles: Dict[str, IpcHandle] = {} + merged: Dict[str, Any] = {} for d in gathered: if d is not None: - merged_handles.update(d) + merged.update(d) - pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") + # {gpu_uuid: handle} (back-compat) and {gpu_uuid: per-GPU metadata} (the fix). + merged_handles: Dict[str, IpcHandle] = {uuid: e["handle"] for uuid, e in merged.items()} + per_gpu_meta: Dict[str, Any] = { + uuid: {k: e[k] for k in ("names", "dtype_names", "shapes", "sizes")} for uuid, e in merged.items() + } chunk_update_info: Dict[str, Any] = { + # rank 0's flat metadata kept for receivers not yet using per-GPU metadata "names": names, "dtype_names": dtype_names, "shapes": shapes, "sizes": sizes, - "ipc_handles_pickled": pickled, + "ipc_handles_pickled": base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8"), + "per_gpu_meta_pickled": base64.b64encode(pickle.dumps(per_gpu_meta)).decode("utf-8"), } await self._inference_client.update_weights_ipc(chunk_update_info) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 6e1f098653..61fe46f2a6 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -85,6 +85,19 @@ maybe_force_qwen35_text_bridge, ) +# Parameters that must be weight-synced in their native (fp32) precision rather than +# down-cast to the inference dtype (bf16). The MoE router score-correction bias +# (HF `*.gate.e_score_correction_bias`) has large magnitude (~25-57) with tiny +# per-expert offsets (std ~7e-4) far below bf16's ULP at that scale, so bf16 rounding +# collapses the offsets and corrupts expert routing after weight sync. vLLM keeps the +# router in fp32 (moe_router_dtype="fp32"), so we transfer this param in fp32. +_FP32_SYNC_SUFFIXES = ("e_score_correction_bias",) + + +def _sync_dtype_for(name: str, default_dtype: "torch.dtype", native_dtype: "torch.dtype") -> "torch.dtype": + """Dtype to weight-sync ``name`` in: native (fp32) for routing-critical params, else default.""" + return native_dtype if name.endswith(_FP32_SYNC_SUFFIXES) else default_dtype + class MegatronWeightExtractor(WeightExtractor): """Extracts weights from Megatron model-parallel models. @@ -216,7 +229,6 @@ def get_weight_metadata(self, dtype: torch.dtype) -> dict: names = [] dtype_names = [] shapes = [] - dtype_name = str(dtype).split(".")[-1] # Collect parameter metadata in the same order # as provided by `.extract_weights`. if not self.enable_bucketing: @@ -226,7 +238,7 @@ def get_weight_metadata(self, dtype: torch.dtype) -> dict: conversion_tasks=None, ): names.append(name) - dtype_names.append(dtype_name) + dtype_names.append(str(_sync_dtype_for(name, dtype, tensor.dtype)).split(".")[-1]) shapes.append(list(tensor.shape)) del tensor else: @@ -242,7 +254,7 @@ def get_weight_metadata(self, dtype: torch.dtype) -> dict: ): names.append(name) shapes.append(list(tensor.shape)) - dtype_names.append(dtype_name) + dtype_names.append(str(_sync_dtype_for(name, dtype, tensor.dtype)).split(".")[-1]) del tensor self._weight_metadata_cache = {"names": names, "dtype_names": dtype_names, "shapes": shapes} @@ -277,11 +289,12 @@ def extract_weights(self, dtype: torch.dtype): ) for name, tensor in hf_params_generator: - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) + out_dtype = _sync_dtype_for(name, dtype, tensor.dtype) + tensor = tensor.to(device=device, dtype=out_dtype, non_blocking=True) yield WeightChunk( names=[name], - dtypes=[str(dtype)], + dtypes=[str(out_dtype)], shapes=[list(tensor.shape)], tensors=[tensor], ) @@ -298,29 +311,30 @@ def extract_weights(self, dtype: torch.dtype): conversion_tasks=bucket_tasks, ) - # Collect all parameters in this bucket into one chunk - names = [] - dtypes_list = [] - shapes = [] - tensors = [] - + # Collect this bucket's params, grouped by output dtype. Each emitted + # chunk must be uniform-dtype because the CUDA-IPC sender packs a chunk + # into a single contiguous buffer of one dtype. Most params down-cast to + # `dtype` (bf16); routing-critical fp32 params (see _FP32_SYNC_SUFFIXES) + # keep their native dtype and ride a separate fp32 chunk. + groups: dict = {} # out_dtype -> (names, dtypes_list, shapes, tensors) for name, tensor in hf_params_generator: - # Move to device and convert dtype - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) - - names.append(name) - dtypes_list.append(str(dtype)) - shapes.append(list(tensor.shape)) - tensors.append(tensor) - - # Yield one chunk containing all parameters in this bucket - if tensors: - yield WeightChunk( - names=names, - dtypes=dtypes_list, - shapes=shapes, - tensors=tensors, - ) + out_dtype = _sync_dtype_for(name, dtype, tensor.dtype) + tensor = tensor.to(device=device, dtype=out_dtype, non_blocking=True) + g = groups.setdefault(out_dtype, ([], [], [], [])) + g[0].append(name) + g[1].append(str(out_dtype)) + g[2].append(list(tensor.shape)) + g[3].append(tensor) + + # Yield one chunk per dtype group in this bucket. + for names, dtypes_list, shapes, tensors in groups.values(): + if tensors: + yield WeightChunk( + names=names, + dtypes=dtypes_list, + shapes=shapes, + tensors=tensors, + ) class MegatronWorker: diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index f0b321e9ba..cf8c795e14 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -724,6 +724,43 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: logger.info(f"Exporting `SKYRL_RAY_PG_TIMEOUT_IN_S` to ray runtime env: {pg_timeout}") env_vars["SKYRL_RAY_PG_TIMEOUT_IN_S"] = pg_timeout + # The inference-server health-wait timeout is read at import time inside the + # VLLMServerActor process (a Ray worker actor), so it must be forwarded to take + # effect. Large models (e.g. 550B) loaded across multi-node PP engines from disk + # can take well over the 600s default to become healthy. + if health_timeout := os.environ.get("SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S"): + logger.info( + f"Exporting `SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S` to ray runtime env: {health_timeout}" + ) + env_vars["SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S"] = health_timeout + + # Forward HuggingFace + cache-location env vars to the Ray workers. Without this, + # workers fall back to the default home dir (`~/.cache`) for the HF model cache and + # uv/Triton/Inductor/vLLM build caches. For large models on nodes with a small home + # disk, that means re-downloading the checkpoint to (and filling) the home disk, + # which can take the node down. Point these at a large disk (e.g. /mnt/local_storage) + # via the launch env so workers read the pre-staged model and write caches there. + for var_name in ( + "HF_HOME", + "HF_HUB_OFFLINE", + "HF_TOKEN", + "HF_HUB_ENABLE_HF_TRANSFER", + "HUGGING_FACE_HUB_TOKEN", + "XDG_CACHE_HOME", + "UV_CACHE_DIR", + "TRITON_CACHE_DIR", + "TORCHINDUCTOR_CACHE_DIR", + "VLLM_CACHE_ROOT", + "TMPDIR", + # vLLM executor selection: cross-node serving needs the V2 (MultiprocExecutor) + # backend; the default compiled-DAG executor both crashes the raylet on the + # cross-node hop and takes a KV-cache-init path that KeyErrors on NemotronH. + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND", + "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", + ): + if value := os.environ.get(var_name): + env_vars.setdefault(var_name, value) + return env_vars diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/conftest.py b/tests/backends/skyrl_train/gpu/gpu_ci/conftest.py index b3d154c134..a6e3b5d98b 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/conftest.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/conftest.py @@ -6,7 +6,11 @@ import ray from loguru import logger -from skyrl.env_vars import _SKYRL_USE_NEW_INFERENCE, SKYRL_PYTHONPATH_EXPORT +from skyrl.env_vars import ( + _SKYRL_USE_NEW_INFERENCE, + SKYRL_LD_LIBRARY_PATH_EXPORT, + SKYRL_PYTHONPATH_EXPORT, +) from skyrl.train.utils.utils import peer_access_supported @@ -47,6 +51,38 @@ def _build_ray_env_vars(): raise RuntimeError("SKYRL_PYTHONPATH_EXPORT is set but PYTHONPATH is not defined in environment") env_vars["PYTHONPATH"] = pythonpath + # Mirror prepare_runtime_environment: for multi-node tests over EFA, the + # driver's LD_LIBRARY_PATH (e.g. /opt/amazon/efa/lib) must reach the Ray + # workers so NCCL picks up the EFA provider. Set SKYRL_LD_LIBRARY_PATH_EXPORT=1. + if SKYRL_LD_LIBRARY_PATH_EXPORT: + ld_library_path = os.environ.get("LD_LIBRARY_PATH") + if ld_library_path is None: + raise RuntimeError("SKYRL_LD_LIBRARY_PATH_EXPORT is set but LD_LIBRARY_PATH is not defined in environment") + env_vars["LD_LIBRARY_PATH"] = ld_library_path + + # Forward debugging/observability env vars to the Ray workers when set on the + # driver. Useful for multi-node bring-up: NCCL_DEBUG / FI_* surface + # NCCL+EFA init failures, and SKYRL_DUMP_INFRA_LOG_TO_STDOUT=1 stops the + # inference-server actors from redirecting their stdout to per-node log + # files (so crash tracebacks reach the driver). vLLM additionally copies + # NCCL_*/FI_* prefixed vars from the engine to its TP/PP workers. + # HF_* / HUGGING_FACE_* let the workers find a pre-staged model cache + # (e.g. HF_HOME on a large node-local disk for very big models) and use an + # HF token; vLLM also copies HF_-prefixed vars to its TP/PP workers. + # SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S is read at import time inside + # the VLLMServerActor worker process, so it must be forwarded to take effect. Very + # large models (e.g. 550B) loaded across multi-node PP engines can take well over + # the 600s default to become healthy. + _DEBUG_PASSTHROUGH = ( + "SKYRL_DUMP_INFRA_LOG_TO_STDOUT", + "PYTORCH_CUDA_ALLOC_CONF", + "SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S", + ) + _DEBUG_PREFIXES = ("NCCL_", "FI_", "HF_", "HUGGING_FACE_", "VLLM_") + for name, value in os.environ.items(): + if name in _DEBUG_PASSTHROUGH or name.startswith(_DEBUG_PREFIXES): + env_vars.setdefault(name, value) + return env_vars diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py index b92ef1aa73..5982ef3942 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py @@ -25,6 +25,7 @@ from tests.backends.skyrl_train.gpu.utils import ( InferenceEngineState, Timer, + _ensure_chat_template, get_test_generator_input, init_worker_with_type, ) @@ -83,11 +84,25 @@ def get_test_actor_config(model_name) -> SkyRLTrainConfig: # the fp32 master + AdamW state on GPU at init (~6x model size), which # OOMs on 4xH100 before forward ever runs. These tests only forward + # weight-sync, so skip optimizer construction entirely. - is_large_moe = ("qwen3.5-35b" in model_name.lower() and "tiny" not in model_name.lower()) or ( - "nemotron-3-nano" in model_name.lower() + is_large_moe = ( + ("qwen3.5-35b" in model_name.lower() and "tiny" not in model_name.lower()) + or ("nemotron-3-nano" in model_name.lower()) + or ("nemotron-3-ultra" in model_name.lower()) ) if is_large_moe: cfg.trainer.policy.inference_only_init = True + if "nemotron-3-ultra" in model_name.lower(): + # Nemotron-Ultra ships MTP layers (num_nextn_predict_layers=1). Megatron-Bridge's + # Mamba provider builds an `mtp_hybrid_override_pattern` from that, and its + # finalize() does `[pattern] * mtp_num_layers`. SkyRL disables MTP for + # training by nulling mtp_num_layers, but for the Mamba provider that leaves + # the pattern set -> `[pattern] * None` -> TypeError. Clear both up front + # (transformer_config_kwargs are applied right before provider.finalize()). + if cfg.trainer.policy.megatron_config.transformer_config_kwargs is None: + cfg.trainer.policy.megatron_config.transformer_config_kwargs = {} + cfg.trainer.policy.megatron_config.transformer_config_kwargs.update( + {"mtp_num_layers": 0, "mtp_hybrid_override_pattern": None} + ) validate_cfg(cfg) return cfg @@ -108,6 +123,14 @@ def _engine_overrides_for_model(model_name: str) -> dict: # Megatron policy init also needs room alongside vLLM on the same # GPU, so lower vLLM's pool footprint. overrides["gpu_memory_utilization"] = 0.5 + if "Nemotron-3-Ultra" in model_name: + # 550B sharded 16-way (vLLM TP8 x PP2) is ~69 GB of weights per GPU, so + # the KV pool needs gmu well above 0.5 just to leave cache room. vLLM + # and the Megatron policy alternate on-GPU via sleep/wake (sleep_level=2), + # so vLLM can claim most of the H200 while loading. Cap context to 4096. + # Starting point -- tune alongside the parallelism. + overrides["engine_init_kwargs"]["max_model_len"] = 4096 + overrides["gpu_memory_utilization"] = 0.85 # Large MoE: Megatron policy init also needs room alongside vLLM on the # same GPU, so lower vLLM's pool footprint. if "qwen3.5-35b" in model_name.lower() and "tiny" not in model_name.lower(): @@ -195,10 +218,10 @@ async def construct_training_input_from_generator_output(generator_output, token @pytest.mark.asyncio @pytest.mark.megatron_models @pytest.mark.parametrize( - "tp,pp,cp,ep,etp,inference_tp,num_gpus,model_name,vllm_threshold,megatron_threshold", + "tp,pp,cp,ep,etp,inference_tp,inference_pp,num_gpus,num_nodes,model_name,vllm_threshold,megatron_threshold", [ - pytest.param(2, 1, 1, 2, 1, 2, 4, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, id="qwen3-moe_tp2_ep2"), - pytest.param(1, 2, 2, 1, None, 2, 4, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, id="qwen3-moe_pp2_cp2"), + pytest.param(2, 1, 1, 2, 1, 2, 1, 4, 1, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, id="qwen3-moe_tp2_ep2"), + pytest.param(1, 2, 2, 1, None, 2, 1, 4, 1, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, id="qwen3-moe_pp2_cp2"), pytest.param( 2, 1, @@ -206,7 +229,9 @@ async def construct_training_input_from_generator_output(generator_output, token 2, 1, 2, + 1, 4, + 1, "eatang/glm-4.7-flash-tiny-random", 1e-1, 2e-2, @@ -220,7 +245,9 @@ async def construct_training_input_from_generator_output(generator_output, token 2, 1, 4, + 1, 4, + 1, "eatang/qwen3.5-moe-tiny-random", 1e-1, 2e-1, @@ -237,7 +264,9 @@ async def construct_training_input_from_generator_output(generator_output, token 1, None, 2, + 1, 2, + 1, "Qwen/Qwen3.5-0.8B", 1e-1, 5e-2, @@ -257,7 +286,9 @@ async def construct_training_input_from_generator_output(generator_output, token 4, 1, 4, + 1, 4, + 1, "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", 5e-1, 5e-2, @@ -275,26 +306,110 @@ async def construct_training_input_from_generator_output(generator_output, token 4, 1, 4, + 1, 4, + 1, "Qwen/Qwen3.5-35B-A3B", 3e-1, 5e-2, id="qwen3.5-35b-a3b_h100_tp4_ep4", marks=pytest.mark.h100, ), + # Nemotron-3-Ultra (550B MoE, ~55B activated, bf16) on 8 nodes x 8xH200 + # = 64 GPUs. Megatron mesh: TP=8 PP=2 CP=1 EP=16 ETP=1 -> DP=4 + # (8*2 = 16 GPUs per model replica, 64/16 = 4 DP, ~69 GiB weights/GPU). + # vLLM: TP=8 (intra-node, NVLink) x PP=4 (across 4 nodes, EFA) = 32 + # GPUs/engine, num_engines = 64/32 = 2 (colocated over the same 64 GPUs). + # vLLM TP must divide NemotronH's Mamba n_groups (=8), so TP=16 is invalid; + # the cross-node parallelism comes from PP instead. vLLM PP=4 (not 2) shards + # vLLM weights ~34 GiB/GPU so that during the colocated Megatron->vLLM weight + # broadcast the policy shard (~69 GiB) + woken vLLM weights (~34 GiB) fit + # alongside the broadcast buffers on the 141 GiB H200 (PP=2 -> ~69+69 OOMs). + # Needs VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 (compiled-DAG shm channel crashes + # the raylet cross-node) and the ray distributed executor backend. + pytest.param( + 8, + 2, + 1, + 16, + 1, + 8, + 4, + 64, + 8, + "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16", + 3e-1, + # 550B hybrid-MoE diverges more between vLLM and Megatron than the + # smaller MoE cases (observed diff mean ~0.059, driven by a minority + # of high-divergence tokens: Megatron logprob std ~0.445 vs vLLM ~0.178). + 8e-2, + id="nemotron3-ultra_tp8pp4_mega_tp8pp2ep16_8node", + marks=pytest.mark.h100, + ), + # Same as above but EP=32 (full expert sharding, ETP=1 -> EDP=1). This is the + # parallelism the full-FT training recipe uses (EP=16 OOMs there once the + # optimizer is added). Diagnostic: training at EP=32 produced GARBAGE vLLM + # generations from step 1, while EP=16 (above) passes -> this isolates whether + # the Megatron->vLLM expert weight gather is correct at EP=32. A blown-up + # post-sync diff (vs pre-sync OK) would pin the bug to the EP=32 sync gather. + pytest.param( + 8, + 2, + 1, + 32, + 1, + 8, + 4, + 64, + 8, + "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16", + 3e-1, + 8e-2, + id="nemotron3-ultra_tp8pp4_mega_tp8pp2ep32_8node", + marks=pytest.mark.h100, + ), + # Same as the validated EP=16 case but Megatron PP=4 (vs 2). DISCRIMINATOR: full-FT + # training at EP=16/PP=4 produced GARBAGE vLLM generations, while this inference-only + # path at EP=16/PP=2 passes (0.298). PP=4 is the only parallelism change from the + # validated case, so if this PASSES, PP=4 sync is fine and the training garbage comes + # from the full-FT optimizer init; if it FAILS, PP=4 itself breaks the weight-sync gather. + pytest.param( + 8, + 4, + 1, + 16, + 1, + 8, + 4, + 64, + 8, + "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16", + 3e-1, + 8e-2, + id="nemotron3-ultra_tp8pp4_mega_tp8pp4ep16_8node", + marks=pytest.mark.h100, + ), ], ) async def test_logprobs_matching_roundtrip( - tp, pp, cp, ep, etp, inference_tp, num_gpus, model_name, vllm_threshold, megatron_threshold + tp, pp, cp, ep, etp, inference_tp, inference_pp, num_gpus, num_nodes, model_name, vllm_threshold, megatron_threshold ): """ Check that logprob diff matches acrosss vllm and megatron. """ + assert num_gpus % num_nodes == 0, f"num_gpus ({num_gpus}) must be divisible by num_nodes ({num_nodes})" + num_gpus_per_node = num_gpus // num_nodes with ray_init(extra_env_vars=_extra_env_vars_for_model(model_name)): cfg = get_test_actor_config(model_name=model_name) cfg.trainer.strategy = "megatron" cfg.generator.inference_engine.tensor_parallel_size = inference_tp - cfg.generator.inference_engine.num_engines = num_gpus // inference_tp + cfg.generator.inference_engine.pipeline_parallel_size = inference_pp + cfg.generator.inference_engine.num_engines = num_gpus // (inference_tp * inference_pp) + # Colocated weight sync keeps the Megatron policy shard and the woken vLLM + # weights on the same GPUs simultaneously. For large models this is tight; + # the expandable_segments allocator reclaims fragmentation so the weight + # broadcast's working buffers fit. + cfg.generator.inference_engine.use_expandable_segments = True cfg.generator.sampling_params = SamplingParams( max_generate_length=MAX_GENERATE_LENGTH, logprobs=1, @@ -305,6 +420,11 @@ async def test_logprobs_matching_roundtrip( tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token + # Some checkpoints (e.g. NemotronH-Ultra) ship no chat template; install a + # minimal one so the generator can format prompts. No-op if one exists. + # vLLM and Megatron both see the same formatted input, so the logprob + # comparison stays valid. + _ensure_chat_template(tokenizer) engine_overrides = _engine_overrides_for_model(model_name) async with InferenceEngineState.create( @@ -331,7 +451,8 @@ async def test_logprobs_matching_roundtrip( generator, client, model_name, tokenizer, return_training_input=True ) await client.sleep() - cfg.trainer.placement.policy_num_gpus_per_node = num_gpus + cfg.trainer.placement.policy_num_nodes = num_nodes + cfg.trainer.placement.policy_num_gpus_per_node = num_gpus_per_node cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = pp cfg.trainer.policy.megatron_config.context_parallel_size = cp @@ -344,7 +465,8 @@ async def test_logprobs_matching_roundtrip( "policy", shared_pg=pg, colocate_all=True, - num_gpus_per_node=num_gpus, + num_gpus_per_node=num_gpus_per_node, + num_nodes=num_nodes, cfg=cfg, ) ray.get( diff --git a/tests/backends/skyrl_train/gpu/utils.py b/tests/backends/skyrl_train/gpu/utils.py index b4410e4890..19a557350b 100644 --- a/tests/backends/skyrl_train/gpu/utils.py +++ b/tests/backends/skyrl_train/gpu/utils.py @@ -44,7 +44,11 @@ TrainingInputBatch, ) from skyrl.backends.skyrl_train.workers.worker import PPORayActorGroup -from skyrl.env_vars import _SKYRL_USE_NEW_INFERENCE, SKYRL_PYTHONPATH_EXPORT +from skyrl.env_vars import ( + _SKYRL_USE_NEW_INFERENCE, + SKYRL_PYTHONPATH_EXPORT, + SKYRL_RAY_PG_TIMEOUT_IN_S, +) from skyrl.train.config import SkyRLTrainConfig from skyrl.train.dataset import PromptDataset from skyrl.train.dataset.replay_buffer import Experience @@ -584,7 +588,7 @@ def create( [{"GPU": 1, "CPU": 1}] * total_gpu_slots, strategy="PACK", ) - get_ray_pg_ready_with_timeout(raw_pg, timeout=60) + get_ray_pg_ready_with_timeout(raw_pg, timeout=SKYRL_RAY_PG_TIMEOUT_IN_S) shared_pg = ResolvedPlacementGroup(raw_pg) sleep = True else: From 8c5dd76226c3f6e8b0c49317542cb6d903b137b2 Mon Sep 17 00:00:00 2001 From: Eric Tang Date: Sat, 20 Jun 2026 12:43:28 +0000 Subject: [PATCH 2/5] [megatron] Nemotron-Ultra-550B throughput/memory sweep: harness + findings Adds a sweep harness (real Megatron fwd+bwd on fabricated rollouts, no vLLM generation) to map, on 64xH200, the max max_tokens_per_microbatch and the parallelism (TP/PP/CP/EP/DP) that maximizes full-FT GRPO training throughput for NVIDIA-Nemotron-3-Ultra-550B, plus a long-context (variable-length) study. Findings (examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md): - Max MTPM ~= 64k tokens/microbatch at the validated TP8/PP4/EP16/DP2 config. - Highest throughput: TP8/PP2/EP32/DP4 (~8.5k tok/s, +11% over PP4/DP2) for short/medium seqs; config space is pinned (PP8 invalid for 108 layers, EP8 OOMs, TP4 doubles activations via sequence parallelism). - Long context: single-sequence ceiling ~40-48k tokens (CP1/PP4/DP2); CP gives little net benefit (CP=2 forces PP2 whose 2x weights cancel the savings). Long seqs are more throughput-efficient per token (~12k tok/s at ~39k mean). New files: - examples/train_scripts/full_context/{trainer_ultra_sweep,main_ultra_sweep,analyze_sweep}.py - examples/train/megatron/run_ultra_sweep.sh - examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md worker.py: get_cuda_memory() now also returns max_allocated/max_reserved high-water marks (capture in-step peak even when queried after offload). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../megatron/NEMOTRON_ULTRA_THROUGHPUT.md | 155 +++++++++++++ examples/train/megatron/run_ultra_sweep.sh | 164 ++++++++++++++ .../full_context/analyze_sweep.py | 62 ++++++ .../full_context/main_ultra_sweep.py | 77 +++++++ .../full_context/trainer_ultra_sweep.py | 210 ++++++++++++++++++ skyrl/backends/skyrl_train/workers/worker.py | 4 + 6 files changed, 672 insertions(+) create mode 100644 examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md create mode 100644 examples/train/megatron/run_ultra_sweep.sh create mode 100644 examples/train_scripts/full_context/analyze_sweep.py create mode 100644 examples/train_scripts/full_context/main_ultra_sweep.py create mode 100644 examples/train_scripts/full_context/trainer_ultra_sweep.py diff --git a/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md b/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md new file mode 100644 index 0000000000..b1da2967ba --- /dev/null +++ b/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md @@ -0,0 +1,155 @@ +# Nemotron-3-Ultra-550B — Megatron training throughput & memory sweep + +Goal: find, on **8× nodes of 8×H200-141GB (64 GPUs, EFA)**, (1) the maximum +`trainer.max_tokens_per_microbatch` (MTPM) that fits for training fwd+bwd, and +(2) the Megatron parallelism (TP/PP/CP/EP/DP) that maximizes training throughput +for full-finetuning GRPO of Nemotron-3-Ultra-550B. + +## Method + +Runs use a dedicated harness that executes the **real** Megatron fwd+bwd training +path on fabricated rollouts (no vLLM generation), so numbers reflect genuine +training cost while iterating fast: + +- Trainer: `examples/train_scripts/full_context/trainer_ultra_sweep.py` + (extends the dummy `FullCtxTrainer`; logs per-step wall time + peak CUDA + reserved/allocated memory across policy workers to a JSONL file). +- Launcher: `examples/train/megatron/run_ultra_sweep.sh` (all knobs are env vars). +- Analysis: `examples/train_scripts/full_context/analyze_sweep.py`. + +Colocated config matches the validated recipe (`run_megatron_nemotron_ultra.sh`): +optimizer CPU-offloaded, `recompute_granularity=full`, `remove_microbatch_padding=true`, +vLLM colocated but **asleep during training** (engines sleep at init in colocate +mode, and the harness never wakes them — exactly the memory state of the real +training step). Peak reserved ≈ caching-allocator high-water mark during the step. + +`max_tokens_per_microbatch` bin-packs each microbatch up to that many tokens **per +DP rank**; a single sequence longer than the budget gets its own microbatch. So the +MTPM memory ceiling is per-DP-rank and independent of DP size. + +## TL;DR + +- **Max `max_tokens_per_microbatch`** at the validated config (TP8/PP4/EP16/DP2): **~64k** tokens/microbatch + (per DP rank). 64k fits; 72k+ OOMs. The model + bf16 grads take ~38 GiB/GPU (optimizer CPU-offloaded) + and a sleeping vLLM holds ~5–8 GiB, leaving ~95 GiB for activations even with `recompute_granularity=full`. +- **Highest training throughput**: **TP8 / PP2 / EP32 / DP4** at MTPM≈32k → **~8540 tok/s**, vs **7720 tok/s** + for the PP4/DP2 baseline at MTPM=48k — a ~11% gain from doubling data parallelism (DP2→DP4). PP2 is what + frees the GPUs for DP4; it costs MTPM headroom (caps ~32–40k) but nets faster. +- **The config space is tightly pinned** by the model (108 layers, MoE with 512 experts) and 141 GiB H200s: + PP must divide 108 (**PP=8 is invalid**), EP8 doubles expert memory and **OOMs at load**, and TP4 doubles + activation memory (sequence-parallel shards activations by TP) and **OOMs**. Viable full-FT configs are + essentially {TP8/PP4/EP16/DP2, TP8/PP2/EP32/DP4}. +- **Long context is activation-bound by the single longest sequence**: with `remove_microbatch_padding`, any + sequence longer than MTPM gets its own microbatch that must fit alone. The *single-sequence* ceiling is + **~40–48k tokens** (CP1/PP4/DP2) — well below the ~64k *packed* ceiling, because a long contiguous sequence + has a much larger per-microbatch footprint than an equal-token pack of short ones — rising to ~96k only at + CP2/PP2/DP2. **CP does not help much**: CP=2 forces PP2 (to keep EP16/DP2), and PP2's 2× weights cancel the + activation memory CP frees. The full 60k±30k distribution (tail → 131k) is **not trainable on 64 GPUs** + without truncating sequences to ~40k. +- **Long sequences are *more* throughput-efficient per token** (~12k tok/s on a ~39k-mean distribution at + PP4/DP2, vs ~7.7k for uniform 10 240-token seqs): bigger microbatches use the GEMMs better and incur less + per-microbatch/pipeline overhead. +- **The throughput-optimal config is sequence-length-dependent**: PP2/EP32/DP4 for short/medium sequences + (≤~48k); PP4/EP16/DP2 for long sequences (higher single-sequence ceiling). + +## Cluster / model facts (measured) +- 64× H200-141 GiB (8 nodes), EFA. GPU usable ≈ 139.8 GiB; sleeping colocated vLLM holds ~5–8 GiB during training. +- Model: 108 layers (hybrid Mamba2 + attention + latent MoE, 512 experts). At TP8/PP4/EP16: **9.44B params/GPU** + → ~18.9 GiB bf16 weights + ~18.9 GiB bf16 grads (~38 GiB fixed; AdamW master/moments CPU-offloaded). +- MoE expert memory/GPU ∝ (108/PP)·(512/EP); **PP·EP is the invariant**. Baseline PP4·EP16 = 64. + PP2 needs EP32 to match (PP2·EP32=64); PP4·EP8=32 ⇒ 2× expert memory ⇒ OOM. + +## Stage 1 — max tokens per microbatch (TP8/PP4/EP16/DP2, uniform 10240-token seqs) + +| MTPM (setting) | largest microbatch | result | +|---:|---:|:--| +| 65536 (64k) | 61 440 tok | **FITS** — steady 77 s/step @ 327 680 tok | +| 73728 (72k) | 71 680 tok | FAIL (DistBackend; one rank OOM aborts NCCL — boundary) | +| 81920 (80k) | 81 920 tok | OOM (forward_backward) | +| 98304 (96k) | 92 160 tok | OOM | +| 131072 (128k) | 131 072 tok | OOM (needed +14.1 GiB; only 8.2 GiB free) | + +**Max safe MTPM ≈ 64k** for packed short/medium sequences. + +## Stage 2 — parallelism sweep for throughput (fixed 655 360-token workload, uniform 10240-token seqs) + +| Config (TP/PP/CP/EP, DP) | MTPM | step time | **throughput** | in-step peak | result | +|:--|---:|---:|---:|---:|:--| +| **TP8/PP2/EP32, DP4** | 32k | 76.8 s | **8 539 tok/s** | 109.9 GB | **OK — fastest** | +| TP8/PP4/EP16, DP2 (baseline) | 48k | 84.9 s | 7 719 tok/s | 108.3 GB | OK | +| TP8/PP2/EP32, DP4 | 48k | — | — | — | OOM (PP2 caps MTPM ~32–40k) | +| TP4/PP4/EP16, DP4 | 32k | — | — | — | OOM (TP4 ⇒ 2× activation via SP) | +| TP8/PP4/EP8, DP2 | 48k | — | — | — | OOM at model load (2× expert mem) | +| TP8/PP8/EP8, DP1 | — | — | — | — | INVALID (108 not divisible by 8) | +| TP4/PP8/EP8, DP2 | — | — | — | — | INVALID (108 not divisible by 8) | + +Doubling data parallelism (DP2→DP4) is the throughput lever; PP2 is the only way to free GPUs for DP4 +given the layer-count and EP constraints. TP must stay 8 (sequence parallelism shards activations by TP). + +## Stage 3 — long context (variable length) & context parallelism + +Single-sequence ceiling (one sequence alone in its microbatch): + +| Config | single-seq ceiling | note | +|:--|---:|:--| +| CP1 / PP4 / DP2 | **~40–48k tok** | single 40 960 fits (peak 110 GB); single 49 152 OOMs. A contiguous long seq has a much bigger per-microbatch footprint than the packed 6×10 240 that fit in stage 1, so the *single-sequence* ceiling (~40–48k) is well below the *packed* 64k ceiling. | +| CP2 / PP2 / EP16 / DP2 | ~96k tok | 128k OOMs: needs 60 GiB activation, only 54 free (PP2 weights eat the headroom CP frees) | + +**CP gives little net benefit here**: CP=2 forces PP2 (to keep EP16/DP2), whose 2× weights cancel CP's +activation savings; CP=4 is infeasible (forces EP8-OOM or PP1-OOM). So with `remove_microbatch_padding`, +**individual training sequences are capped at ~40–48k tokens** (CP1/PP4/DP2), or ~96k at CP2/PP2/DP2, +on 64 GPUs — and the requested 60k±30k distribution (tail → 131k) must be truncated. + +### Throughput on a variable-length long-context distribution + +The requested distribution was 256 samples ~ N(60k, 30k) tokens. Because the single-sequence +ceiling is ~40–48k (above), the distribution must be **truncated** to a value that fits as a single +microbatch; the 60k mean is not trainable per-sequence on 64 GPUs. Measured on PP4/EP16/DP2 with +sequences ~ N(60k, 30k) **clamped to [1k, 40 960]** (realized mean ~39k, max 40 960), MTPM=40 960: + +| Config | distribution (realized) | MTPM | step time | **throughput** | peak | +|:--|:--|---:|---:|---:|---:| +| TP8/PP4/EP16/DP2 | 64 varlen seqs, mean ~39k, max 40 960 | 40 960 | 207 s | **~12 070 tok/s** | 112.9 GB | + +(Throughput is per-token and count-independent; 64 samples used for tractable wall-clock. Earlier +256-sample runs confirmed the packing/fit behaviour and the OOM on sequences >~48k.) + +**Long sequences are *more* throughput-efficient per token** than short ones (~12.1k tok/s here vs +~7.7k for uniform 10 240-token seqs at the same PP4/DP2): a microbatch of one ~40k sequence (or a few +medium ones) uses the GEMMs far better and incurs less per-microbatch / pipeline overhead than packing +many tiny 10 240-token sequences. So for long-context RL the throughput ceiling is set by **fitting the +longest single sequence** (clamp responses to ~40k), not by aggregate tokens. + +## Recommendations + +- **Short/medium sequences (≤~40k), max throughput:** `TP8 / PP2 / EP32 / ETP1 / DP4`, MTPM≈32k + (~8.5k tok/s on uniform short seqs; ~11% over the PP4/DP2 baseline). +- **Long context / largest single sequences:** `TP8 / PP4 / EP16 / ETP1 / DP2`, MTPM≈40–64k + (single-sequence ceiling ~40–48k; ~12k tok/s on a ~39k-mean distribution). +- **Cap RL response length to ~40k tokens.** The full 60k±30k distribution is not trainable + per-sequence on 64×H200; CP cannot rescue it (CP=2 forces PP2, whose 2× weights cancel CP's savings). +- Keep `TP=8` (sequence parallelism shards activations by TP — TP4 doubles activation memory and OOMs), + optimizer CPU-offload on, and `recompute_granularity=full`. + + +## Reproducing + +```bash +# Stage 1 — MTPM ceiling at the baseline config (uniform 10k-token seqs). +TP=8 PP=4 CP=1 EP=16 ETP=1 MTPM=65536 MODE=uniform SEQ_LEN=10240 \ + TAG=s1_tp8pp4ep16_mtpm65536 SWEEP_RESULTS_FILE=/home/ray/ultra_sweep/results.jsonl \ + bash examples/train/megatron/run_ultra_sweep.sh + +# Stage 2 — highest-throughput config (DP4). +TP=8 PP=2 CP=1 EP=32 ETP=1 MTPM=32768 MODE=uniform SEQ_LEN=10240 NUM_SEQ=64 \ + TAG=s2_tp8pp2ep32_dp4 SWEEP_RESULTS_FILE=/home/ray/ultra_sweep/results.jsonl \ + bash examples/train/megatron/run_ultra_sweep.sh + +# Stage 3 — long-context varlen distribution (clamped to the ~40k single-seq ceiling). +TP=8 PP=4 CP=1 EP=16 ETP=1 MODE=varlen AVG_LEN=60000 STD_LEN=30000 MAX_LEN=40960 \ + MTPM=40960 NUM_SEQ=64 TAG=s3_varlen_clamp40k \ + SWEEP_RESULTS_FILE=/home/ray/ultra_sweep/results.jsonl \ + bash examples/train/megatron/run_ultra_sweep.sh + +python examples/train_scripts/full_context/analyze_sweep.py /home/ray/ultra_sweep/results.jsonl +``` diff --git a/examples/train/megatron/run_ultra_sweep.sh b/examples/train/megatron/run_ultra_sweep.sh new file mode 100644 index 0000000000..e5de472379 --- /dev/null +++ b/examples/train/megatron/run_ultra_sweep.sh @@ -0,0 +1,164 @@ +#!/usr/bin/env bash +# Throughput / memory sweep harness for Nemotron-3-Ultra-550B (Megatron). +# +# Drives examples.train_scripts.full_context.main_ultra_sweep, which runs the real +# fwd+bwd training path on fabricated rollouts (no vLLM generation) and logs per-step +# peak CUDA memory + step time to $SWEEP_RESULTS_FILE. +# +# All knobs are env vars (with defaults). Example: +# TP=8 PP=4 CP=1 EP=16 ETP=1 MTPM=131072 MODE=uniform SEQ_LEN=10240 NUM_SEQ=64 \ +# TAG=tp8pp4ep16_mtpm128k SWEEP_RESULTS_FILE=/path/results.jsonl \ +# bash examples/train/megatron/run_ultra_sweep.sh +set -x + +# ---- Megatron parallelism (world = TP*PP*CP*DP = 64) ---- +TP=${TP:-8} +PP=${PP:-4} +CP=${CP:-1} +EP=${EP:-16} +ETP=${ETP:-1} + +# ---- Microbatch token budget (per DP rank) ---- +MTPM=${MTPM:-131072} + +# ---- Workload ---- +MODE=${MODE:-uniform} # uniform | varlen +SEQ_LEN=${SEQ_LEN:-10240} # uniform: total tokens/seq +PROMPT_LEN=${PROMPT_LEN:-512} +AVG_LEN=${AVG_LEN:-60000} # varlen +STD_LEN=${STD_LEN:-30000} +MIN_LEN=${MIN_LEN:-1024} +MAX_LEN=${MAX_LEN:-131072} +NUM_STEPS=${NUM_STEPS:-3} + +NUM_NODES=8 +NUM_GPUS=8 +WORLD=$((NUM_NODES * NUM_GPUS)) +DP=$((WORLD / (TP * PP * CP))) + +# ---- sequence count & batch sizes ---- +# Default: enough sequences so each DP rank forms ~2 full microbatches (uniform), +# or exactly 256 for varlen (the stage-3 target distribution). +if [ -z "${NUM_SEQ:-}" ]; then + if [ "$MODE" = "varlen" ]; then + NUM_SEQ=256 + else + # ceil(2 * MTPM * DP / SEQ_LEN), rounded up to a multiple of 8 + NUM_SEQ=$(( (2 * MTPM * DP + SEQ_LEN - 1) / SEQ_LEN )) + NUM_SEQ=$(( ((NUM_SEQ + 7) / 8) * 8 )) + if [ "$NUM_SEQ" -lt 8 ]; then NUM_SEQ=8; fi + fi +fi +# Round NUM_SEQ up to a multiple of (2*DP) so n_samples=2 and DP divisibility both hold. +LCM=$((2 * DP)) +NUM_SEQ=$(( ((NUM_SEQ + LCM - 1) / LCM) * LCM )) +N_SAMPLES=2 +TBS=$(( NUM_SEQ / N_SAMPLES )) +MINI=$TBS + +TAG=${TAG:-tp${TP}pp${PP}cp${CP}ep${EP}_mtpm${MTPM}_${MODE}} +SWEEP_RESULTS_FILE=${SWEEP_RESULTS_FILE:-/home/ray/ultra_sweep/results.jsonl} +mkdir -p "$(dirname "$SWEEP_RESULTS_FILE")" + +# ---- Environment (mirror the validated nemotron recipe) ---- +export HF_HOME=${HF_HOME:-/mnt/local_storage/hf_cache} +export XDG_CACHE_HOME=${XDG_CACHE_HOME:-/mnt/local_storage/.cache} +export UV_CACHE_DIR=${UV_CACHE_DIR:-/mnt/local_storage/.cache/uv} +export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/mnt/local_storage/.cache/triton} +export TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR:-/mnt/local_storage/.cache/inductor} +export VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT:-/mnt/local_storage/.cache/vllm} +export HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-1} +export LD_LIBRARY_PATH=/opt/amazon/efa/lib:${LD_LIBRARY_PATH:-} +export SKYRL_LD_LIBRARY_PATH_EXPORT=1 +export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 +export NVTE_FLASH_ATTN=0 +export SKYRL_RAY_PG_TIMEOUT_IN_S=${SKYRL_RAY_PG_TIMEOUT_IN_S:-1800} +export SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S=${SKYRL_WAIT_UNTIL_INFERENCE_SERVER_HEALTHY_TIMEOUT_S:-2400} +export SKYRL_DUMP_INFRA_LOG_TO_STDOUT=${SKYRL_DUMP_INFRA_LOG_TO_STDOUT:-1} + +# ---- sweep trainer params (read by trainer_ultra_sweep.py) ---- +export SWEEP_RESULTS_FILE SWEEP_TAG="$TAG" SWEEP_MODE="$MODE" +export SWEEP_NUM_STEPS="$NUM_STEPS" SWEEP_NUM_SEQ="$NUM_SEQ" SWEEP_PROMPT_LEN="$PROMPT_LEN" +export SWEEP_SEQ_LEN="$SEQ_LEN" SWEEP_AVG_LEN="$AVG_LEN" SWEEP_STD_LEN="$STD_LEN" +export SWEEP_MIN_LEN="$MIN_LEN" SWEEP_MAX_LEN="$MAX_LEN" + +DATA_DIR="/mnt/local_storage/data/gsm8k" +MODEL_NAME="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16" + +# vLLM colocated config (matches the recipe; vLLM sleeps during the training step). +NUM_INFERENCE_ENGINES=2 +INFERENCE_ENGINE_TP=8 +INFERENCE_ENGINE_PP=4 + +echo "[run_ultra_sweep] TAG=$TAG TP=$TP PP=$PP CP=$CP EP=$EP ETP=$ETP DP=$DP MTPM=$MTPM MODE=$MODE NUM_SEQ=$NUM_SEQ TBS=$TBS MINI=$MINI" + +uv run --isolated --extra megatron -m examples.train_scripts.full_context.main_ultra_sweep \ + data.train_data="['$DATA_DIR/train.parquet']" \ + data.val_data="['$DATA_DIR/validation.parquet']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.policy.model.path=$MODEL_NAME \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TP \ + generator.inference_engine.pipeline_parallel_size=$INFERENCE_ENGINE_PP \ + generator.inference_engine.distributed_executor_backend=ray \ + generator.inference_engine.use_expandable_segments=true \ + trainer.policy.megatron_config.tensor_model_parallel_size=$TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$PP \ + trainer.policy.megatron_config.context_parallel_size=$CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$ETP \ + trainer.policy.megatron_config.transformer_config_kwargs.mtp_num_layers=0 \ + trainer.policy.megatron_config.transformer_config_kwargs.mtp_hybrid_override_pattern=null \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_granularity=full \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_method=uniform \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_num_layers=1 \ + trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=true \ + trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=true \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=true \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=1.0 \ + trainer.remove_microbatch_padding=true \ + trainer.max_tokens_per_microbatch=$MTPM \ + trainer.epochs=1 \ + trainer.eval_before_train=false \ + trainer.eval_interval=100000 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TBS \ + trainer.policy_mini_batch_size=$MINI \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.max_prompt_length=$PROMPT_LEN \ + generator.max_input_length=$PROMPT_LEN \ + generator.sampling_params.max_generate_length=$((SEQ_LEN - PROMPT_LEN)) \ + trainer.policy.optimizer_config.lr=1.0e-6 \ + trainer.algorithm.use_kl_loss=false \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.inference_engine.engine_init_kwargs.max_model_len=4096 \ + generator.inference_engine.gpu_memory_utilization=0.6 \ + generator.batched=true \ + environment.env_class=gsm8k \ + generator.n_samples_per_prompt=$N_SAMPLES \ + trainer.num_dummy_steps=$NUM_STEPS \ + trainer.sweep_results_file="$SWEEP_RESULTS_FILE" \ + trainer.sweep_tag="$TAG" \ + trainer.sweep_mode="$MODE" \ + trainer.sweep_num_seq=$NUM_SEQ \ + trainer.sweep_prompt_len=$PROMPT_LEN \ + trainer.sweep_seq_len=$SEQ_LEN \ + trainer.sweep_avg_len=$AVG_LEN \ + trainer.sweep_std_len=$STD_LEN \ + trainer.sweep_min_len=$MIN_LEN \ + trainer.sweep_max_len=$MAX_LEN \ + trainer.logger="console" \ + trainer.project_name="ultra_sweep" \ + trainer.run_name="$TAG" \ + trainer.resume_mode=none \ + trainer.ckpt_interval=100000 \ + trainer.ckpt_path="/mnt/local_storage/ultra_sweep_ckpt" \ + "$@" diff --git a/examples/train_scripts/full_context/analyze_sweep.py b/examples/train_scripts/full_context/analyze_sweep.py new file mode 100644 index 0000000000..6471ce7221 --- /dev/null +++ b/examples/train_scripts/full_context/analyze_sweep.py @@ -0,0 +1,62 @@ +"""Summarize ultra-sweep results.jsonl into a throughput/memory table. + +Usage: python analyze_sweep.py /home/ray/ultra_sweep/results.jsonl +For each tag, reports the steady-state (non-warmup, non-error) median step time, +tokens/s, max peak reserved GB, min free GB, and OOM/error status. +""" + +import json +import sys +from collections import defaultdict +from statistics import median + + +def main(path): + by_tag = defaultdict(list) + order = [] + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + r = json.loads(line) + if r["tag"] not in by_tag: + order.append(r["tag"]) + by_tag[r["tag"]].append(r) + + hdr = ( + f"{'tag':<34} {'tp':>2} {'pp':>2} {'cp':>2} {'ep':>3} {'dp':>2} " + f"{'mtpm':>8} {'mode':>7} {'maxseq':>7} {'tok/step':>9} " + f"{'step_s':>7} {'tok/s':>9} {'pkRsv':>6} {'minFree':>7} {'status':>8}" + ) + print(hdr) + print("-" * len(hdr)) + for tag in order: + recs = by_tag[tag] + steady = [r for r in recs if not r.get("warmup") and not r.get("error")] + oom = any(r.get("oom") for r in recs) + err = next((r["error"] for r in recs if r.get("error")), None) + any0 = recs[0] + if steady: + st = median([r["step_time_s"] for r in steady]) + tps = median([r["tokens_per_s"] for r in steady if r["tokens_per_s"]]) + pk = max([r["peak_reserved_gb"] for r in steady if r["peak_reserved_gb"]] or [0]) + mf = min([r["min_free_gb"] for r in steady if r["min_free_gb"]] or [0]) + status = "OK" + else: + st = tps = pk = mf = 0 + # use warmup peak if present + pk = max([r.get("peak_reserved_gb") or 0 for r in recs] or [0]) + mf = min([r.get("min_free_gb") or 99 for r in recs] or [0]) + status = "OOM" if oom else ("ERR" if err else "NONE") + print( + f"{tag:<34} {any0['tp']:>2} {any0['pp']:>2} {any0['cp']:>2} {any0['ep']:>3} {any0['dp']:>2} " + f"{any0['mtpm']:>8} {any0.get('mode',''):>7} {any0.get('max_seqlen',0):>7} " + f"{any0.get('total_tokens',0):>9} {st:>7.1f} {tps:>9.0f} {pk:>6.1f} {mf:>7.1f} {status:>8}" + ) + if err and status != "OK": + print(f" └─ {err[:160]}") + + +if __name__ == "__main__": + main(sys.argv[1] if len(sys.argv) > 1 else "/home/ray/ultra_sweep/results.jsonl") diff --git a/examples/train_scripts/full_context/main_ultra_sweep.py b/examples/train_scripts/full_context/main_ultra_sweep.py new file mode 100644 index 0000000000..477f70462b --- /dev/null +++ b/examples/train_scripts/full_context/main_ultra_sweep.py @@ -0,0 +1,77 @@ +"""Entrypoint for the Nemotron-Ultra throughput/memory sweep. + +uv run --isolated --extra megatron -m examples.train_scripts.full_context.main_ultra_sweep ... +""" + +import sys +from dataclasses import dataclass + +import ray + +from skyrl.train.config import TrainerConfig, make_config +from skyrl.train.entrypoints.main_base import BasePPOExp +from skyrl.train.utils import initialize_ray, validate_cfg + +from .trainer_ultra_sweep import UltraSweepTrainer + + +@dataclass +class UltraSweepTrainerConfig(TrainerConfig): + num_dummy_steps: int = 3 + # Sweep params passed via CLI (trainer.=...) so they reach the Ray + # worker through cfg (shell env vars do NOT propagate to the entrypoint actor). + sweep_results_file: str = "/home/ray/ultra_sweep/results.jsonl" + sweep_tag: str = "run" + sweep_mode: str = "uniform" # uniform | varlen + sweep_num_seq: int = -1 # -1 -> train_batch_size * n_samples_per_prompt + sweep_prompt_len: int = 512 + sweep_seq_len: int = 10240 # uniform: total tokens/seq + sweep_avg_len: int = 60000 # varlen + sweep_std_len: int = 30000 + sweep_min_len: int = 1024 + sweep_max_len: int = 131072 + sweep_seed: int = 1234 + + +UltraSweepConfig = make_config(trainer_cls=UltraSweepTrainerConfig) + + +class UltraSweepPPOExp(BasePPOExp): + def get_trainer( + self, + cfg, + tracker, + tokenizer, + train_dataset, + eval_dataset, + inference_engine_client, + generator, + colocate_pg, + ): + return UltraSweepTrainer( + cfg=cfg, + tracker=tracker, + tokenizer=tokenizer, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + inference_engine_client=inference_engine_client, + generator=generator, + colocate_pg=colocate_pg, + ) + + +@ray.remote(num_cpus=1) +def skyrl_entrypoint(cfg): + exp = UltraSweepPPOExp(cfg) + exp.run() + + +def main() -> None: + cfg = UltraSweepConfig.from_cli_overrides(sys.argv[1:]) + validate_cfg(cfg) + initialize_ray(cfg) + ray.get(skyrl_entrypoint.remote(cfg)) + + +if __name__ == "__main__": + main() diff --git a/examples/train_scripts/full_context/trainer_ultra_sweep.py b/examples/train_scripts/full_context/trainer_ultra_sweep.py new file mode 100644 index 0000000000..eec3ed6d0d --- /dev/null +++ b/examples/train_scripts/full_context/trainer_ultra_sweep.py @@ -0,0 +1,210 @@ +"""Throughput / memory sweep trainer for Nemotron-Ultra-550B (Megatron). + +Extends the dummy ``FullCtxTrainer`` to (a) build either uniform-length or +variable-length synthetic batches, (b) record per-step wall time and the peak +CUDA reserved memory across policy workers, and (c) append one JSON record per +step to ``$SWEEP_RESULTS_FILE``. It runs the *real* fwd+bwd training path +(``fwd_logprobs_values_reward`` + ``train_critic_and_policy``) so the numbers +reflect genuine training cost, but fabricates the rollout so no vLLM generation +is needed. + +Driven entirely by env vars (so the same module serves every sweep config): + + SWEEP_RESULTS_FILE path to append JSONL results to (required) + SWEEP_TAG label for this run (e.g. "tp8pp4ep16_mtpm128k") + SWEEP_MODE "uniform" (default) or "varlen" + SWEEP_NUM_STEPS number of measured steps (default 3; first is warmup) + SWEEP_NUM_SEQ total sequences per step (default = train_batch_size*n_samples) + SWEEP_PROMPT_LEN prompt length per sequence (default 512) + # uniform mode: + SWEEP_SEQ_LEN total tokens per sequence (prompt+response), default 10240 + # varlen mode: + SWEEP_AVG_LEN mean total tokens/sequence (default 60000) + SWEEP_STD_LEN stddev of total tokens/sequence (default 30000) + SWEEP_MIN_LEN clamp floor for total length (default 1024) + SWEEP_MAX_LEN clamp ceiling for total length (default 131072) + SWEEP_SEED RNG seed for reproducible varlen draws (default 1234) +""" + +import json +import random +import time + +from loguru import logger + +from skyrl.train.utils.utils import Timer + +from .trainer_full_ctx import FullCtxTrainer + + +class UltraSweepTrainer(FullCtxTrainer): + def _build_lengths(self, num_seq): + """Return a list of (prompt_len, response_len) per sequence.""" + t = self.cfg.trainer + mode = t.sweep_mode + prompt_len = t.sweep_prompt_len + if mode == "uniform": + seq_len = t.sweep_seq_len + resp = max(1, seq_len - prompt_len) + return [(prompt_len, resp)] * num_seq, mode, seq_len, 0 + # varlen + rng = random.Random(t.sweep_seed) + out = [] + for _ in range(num_seq): + total = int(round(rng.gauss(t.sweep_avg_len, t.sweep_std_len))) + total = max(t.sweep_min_len, min(t.sweep_max_len, total)) + total = max(total, prompt_len + 1) + out.append((prompt_len, total - prompt_len)) + return out, mode, t.sweep_avg_len, t.sweep_std_len + + def _peak_reserved_gb(self): + """Peak CUDA high-water (reserved, allocated) and min free, GB, across policy workers. + + Uses max_reserved/max_allocated (high-water marks that survive empty_cache/offload) + so we capture the in-step fwd/bwd peak even though this is queried after the policy + has been offloaded back to CPU. Falls back to current reserved/allocated on older + workers that don't return the max_* keys. + """ + import ray + + try: + mems = ray.get( + self.policy_model.async_run_ray_method("pass_through", "get_cuda_memory"), + timeout=60, + ) + except Exception as e: # noqa: BLE001 + logger.warning(f"get_cuda_memory failed: {e}") + return None, None, None + reserved = [m.get("max_reserved", m["reserved"]) / 1e9 for m in mems] + allocated = [m.get("max_allocated", m["allocated"]) / 1e9 for m in mems] + free = [m["free"] / 1e9 for m in mems] + return max(reserved), max(allocated), min(free) + + def _record(self, rec): + path = self.cfg.trainer.sweep_results_file + if path: + try: + import os + + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "a") as f: + f.write(json.dumps(rec) + "\n") + except Exception as e: # noqa: BLE001 + logger.warning(f"could not write results file {path}: {e}") + logger.info("SWEEP_RESULT " + json.dumps(rec)) + + async def train(self): + tag = self.cfg.trainer.sweep_tag + num_steps = self.cfg.trainer.num_dummy_steps + n_samples = self.cfg.generator.n_samples_per_prompt + default_num_seq = self.cfg.trainer.train_batch_size * n_samples + cfg_num_seq = self.cfg.trainer.sweep_num_seq + num_seq = cfg_num_seq if cfg_num_seq and cfg_num_seq > 0 else default_num_seq + + mc = self.cfg.trainer.policy.megatron_config + world = self.cfg.trainer.placement.policy_num_nodes * self.cfg.trainer.placement.policy_num_gpus_per_node + dp = world // (mc.pipeline_model_parallel_size * mc.context_parallel_size * mc.tensor_model_parallel_size) + cfg_hdr = { + "tag": tag, + "tp": mc.tensor_model_parallel_size, + "pp": mc.pipeline_model_parallel_size, + "cp": mc.context_parallel_size, + "ep": mc.expert_model_parallel_size, + "etp": mc.expert_tensor_parallel_size, + "dp": dp, + "mtpm": self.cfg.trainer.max_tokens_per_microbatch, + "num_seq": num_seq, + } + logger.info(f"[ultra-sweep] starting: {json.dumps(cfg_hdr)}") + + self.global_step = 0 + with Timer("init_weight_sync_state", self.all_timings): + self.init_weight_sync_state() + + lengths, mode, p0, p1 = self._build_lengths(num_seq) + total_tokens = sum(p + r for p, r in lengths) + max_seqlen = max(p + r for p, r in lengths) + logger.info( + f"[ultra-sweep] mode={mode} num_seq={num_seq} total_tokens={total_tokens} " + f"max_seqlen={max_seqlen} dp={dp} per_dp_tokens={total_tokens // dp}" + ) + + # uids group sequences into n_samples_per_prompt groups (for grpo advantage std). + uids = [] + for i in range(num_seq): + uids.append(str(i // n_samples)) + + self.global_step += 1 + for step in range(num_steps): + oom = False + err = None + t0 = time.time() + try: + prompt_token_ids = [[random.randint(0, self.tokenizer.vocab_size - 1)] * p for (p, r) in lengths] + response_ids = [[random.randint(0, self.tokenizer.vocab_size - 1)] * r for (p, r) in lengths] + rewards = [[0.0] * (r - 1) + [float(random.randint(0, 1))] for (p, r) in lengths] + loss_masks = [[1] * r for (p, r) in lengths] + dummy = { + "prompt_token_ids": prompt_token_ids, + "response_ids": response_ids, + "rewards": rewards, + "loss_masks": loss_masks, + } + training_input = self.convert_to_training_input(dummy, uids) + with Timer("step", self.all_timings): + with Timer("fwd_logprobs_values_reward", self.all_timings): + training_input = self.fwd_logprobs_values_reward(training_input) + with Timer("compute_advantages_and_returns", self.all_timings): + training_input = self.compute_advantages_and_returns(training_input) + for key in ["rewards"]: + training_input.pop(key) + training_input.metadata.pop("uids") + with Timer("train_critic_and_policy", self.all_timings): + self.train_critic_and_policy(training_input) + except Exception as e: # noqa: BLE001 + err = f"{type(e).__name__}: {str(e)[:500]}" + if "out of memory" in str(e).lower() or "OutOfMemory" in str(e): + oom = True + logger.error(f"[ultra-sweep] step {step} FAILED: {err}") + + step_time = time.time() - t0 + if err: + # Workers can be wedged after a CUDA OOM; querying memory may hang. + peak_res, peak_alloc, min_free = None, None, None + else: + peak_res, peak_alloc, min_free = self._peak_reserved_gb() + rec = dict(cfg_hdr) + rec.update( + { + "step": step, + "warmup": step == 0, + "mode": mode, + "total_tokens": total_tokens, + "max_seqlen": max_seqlen, + "per_dp_tokens": total_tokens // dp, + "step_time_s": round(step_time, 2), + "tokens_per_s": round(total_tokens / step_time, 1) if step_time > 0 and not err else None, + "peak_reserved_gb": round(peak_res, 2) if peak_res else None, + "peak_alloc_gb": round(peak_alloc, 2) if peak_alloc else None, + "min_free_gb": round(min_free, 2) if min_free else None, + "oom": oom, + "error": err, + } + ) + self._record(rec) + try: + self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step) + except Exception: # noqa: BLE001 + pass + self.all_timings = {} + self.all_metrics = {} + self.global_step += 1 + if err: + # No point continuing this config once it fails. + break + + try: + self.tracker.finish() + except Exception: # noqa: BLE001 + pass + logger.info(f"[ultra-sweep] done: {tag}") diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index 6bc218313b..2eabffd02c 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -337,6 +337,10 @@ def get_cuda_memory(self) -> Dict[str, Any]: return { "allocated": torch.cuda.memory_allocated(), "reserved": torch.cuda.memory_reserved(), + # High-water marks (persist across empty_cache / offload, until reset_peak_memory_stats), + # so they capture the in-step fwd/bwd peak even when queried after the policy is offloaded. + "max_allocated": torch.cuda.max_memory_allocated(), + "max_reserved": torch.cuda.max_memory_reserved(), "free": free, "total": total, } From 1d64a64c705c7821cdf9d79ab5af64b8a4fb3dc5 Mon Sep 17 00:00:00 2001 From: Eric Tang Date: Sat, 20 Jun 2026 17:22:18 +0000 Subject: [PATCH 3/5] [megatron] Correct long-context/CP findings: CP composes with EP Follow-up to the throughput sweep. CP composes with EP in Megatron-Core (EP divides TP*CP*DP, with ETP=1), so adding CP does NOT force EP down. Measured: - TP8/PP4/CP2/EP16/DP1 fits a single 96k sequence (128k OOMs) -- CP2 roughly doubles the single-sequence ceiling (~40-48k -> ~96k) while keeping PP4's low weights and baseline expert memory. Best long-context config. - TP8/PP2/CP4/EP32/DP1 is valid and loads but still OOMs at 128k: dropping to PP2 to free GPUs for CP4 doubles the weights and eats the budget CP frees. So the 60k+-30k distribution is mostly trainable with PP4/CP2 (clamp ~96k, ~10% truncated) at the cost of DP->1; the full 131k tail still OOMs. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../megatron/NEMOTRON_ULTRA_THROUGHPUT.md | 48 ++++++++++++------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md b/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md index b1da2967ba..e7268602f1 100644 --- a/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md +++ b/examples/train/megatron/NEMOTRON_ULTRA_THROUGHPUT.md @@ -41,11 +41,15 @@ MTPM memory ceiling is per-DP-rank and independent of DP size. essentially {TP8/PP4/EP16/DP2, TP8/PP2/EP32/DP4}. - **Long context is activation-bound by the single longest sequence**: with `remove_microbatch_padding`, any sequence longer than MTPM gets its own microbatch that must fit alone. The *single-sequence* ceiling is - **~40–48k tokens** (CP1/PP4/DP2) — well below the ~64k *packed* ceiling, because a long contiguous sequence - has a much larger per-microbatch footprint than an equal-token pack of short ones — rising to ~96k only at - CP2/PP2/DP2. **CP does not help much**: CP=2 forces PP2 (to keep EP16/DP2), and PP2's 2× weights cancel the - activation memory CP frees. The full 60k±30k distribution (tail → 131k) is **not trainable on 64 GPUs** - without truncating sequences to ~40k. + **~40–48k tokens** at CP1/PP4/DP2 (well below the ~64k *packed* ceiling — a long contiguous sequence has a + much larger per-microbatch footprint than an equal-token pack of short ones). +- **Context parallelism roughly doubles that ceiling to ~96k.** CP *composes* with EP (in Megatron-Core + `EP` divides `TP·CP·DP`, so CP does not steal from EP's budget): **`TP8/PP4/CP2/EP16/DP1` fits a single 96k + sequence** (128k OOMs) while keeping baseline expert memory. CP=4 via `TP8/PP2/CP4/EP32/DP1` is *valid* but + *worse* — it still OOMs at 128k because dropping to PP2 (needed to free GPUs for CP4) doubles the weights and + eats the budget CP frees. So the practical long-context recipe is **PP4 + CP2** (≤~96k/seq, at the cost of + DP→1), and the 60k±30k distribution becomes mostly trainable (clamp ~96k truncates only the ~10% tail vs + ~half at CP1's ~40k clamp). - **Long sequences are *more* throughput-efficient per token** (~12k tok/s on a ~39k-mean distribution at PP4/DP2, vs ~7.7k for uniform 10 240-token seqs): bigger microbatches use the GEMMs better and incur less per-microbatch/pipeline overhead. @@ -88,17 +92,21 @@ given the layer-count and EP constraints. TP must stay 8 (sequence parallelism s ## Stage 3 — long context (variable length) & context parallelism -Single-sequence ceiling (one sequence alone in its microbatch): +**CP composes with EP** (not with DP-budget-for-EP as first assumed): in Megatron-Core the expert group is +formed over `TP·CP·DP`, so with `ETP=1`, **`EP` must divide `TP·CP·DP`** — adding CP does *not* force EP down. +This makes CP genuinely useful here. Single-sequence ceiling (one sequence alone in its microbatch): -| Config | single-seq ceiling | note | -|:--|---:|:--| -| CP1 / PP4 / DP2 | **~40–48k tok** | single 40 960 fits (peak 110 GB); single 49 152 OOMs. A contiguous long seq has a much bigger per-microbatch footprint than the packed 6×10 240 that fit in stage 1, so the *single-sequence* ceiling (~40–48k) is well below the *packed* 64k ceiling. | -| CP2 / PP2 / EP16 / DP2 | ~96k tok | 128k OOMs: needs 60 GiB activation, only 54 free (PP2 weights eat the headroom CP frees) | +| Config | world | single-seq ceiling | note | +|:--|:--|---:|:--| +| CP1 / PP4 / EP16 / DP2 | 8·4·1·2 | **~40–48k tok** | single 40 960 fits (peak 110 GB); 49 152 OOMs. Long contiguous seq footprint ≫ equal-token pack of short seqs, so single-seq ceiling < packed 64k. | +| **CP2 / PP4 / EP16 / DP1** | 8·4·2·1 | **~96k tok** | EP16 still valid (`EP \| TP·CP·DP = 16`); **single 98 304 FITS**, 131 072 OOMs. Keeps PP4's low weights (~38 GiB) and shards the seq 2×. Best long-context config. | +| CP4 / PP2 / EP32 / DP1 | 8·2·4·1 | <128k | *valid* (`EP32 \| 8·4·1=32`) and loads, but **131 072 OOMs**: dropping to PP2 to free GPUs for CP4 doubles weights (~76 GiB) and eats the budget CP frees — worse than PP4/CP2. | -**CP gives little net benefit here**: CP=2 forces PP2 (to keep EP16/DP2), whose 2× weights cancel CP's -activation savings; CP=4 is infeasible (forces EP8-OOM or PP1-OOM). So with `remove_microbatch_padding`, -**individual training sequences are capped at ~40–48k tokens** (CP1/PP4/DP2), or ~96k at CP2/PP2/DP2, -on 64 GPUs — and the requested 60k±30k distribution (tail → 131k) must be truncated. +So **CP roughly doubles the usable single-sequence length** (~40–48k → ~96k) via **PP4 + CP2**, at the cost of +collapsing DP to 1 (≈ half the data-parallel throughput). The 60k±30k distribution is then **mostly trainable**: +clamping at ~96k truncates only the ~10% upper tail (vs truncating ~half at CP1's ~40k clamp). The extreme +131k tail still OOMs (the LM-head logits / non-CP-sharded buffers don't shrink enough); CP4 doesn't fix it +because of the PP2 weight penalty. ### Throughput on a variable-length long-context distribution @@ -124,10 +132,14 @@ longest single sequence** (clamp responses to ~40k), not by aggregate tokens. - **Short/medium sequences (≤~40k), max throughput:** `TP8 / PP2 / EP32 / ETP1 / DP4`, MTPM≈32k (~8.5k tok/s on uniform short seqs; ~11% over the PP4/DP2 baseline). -- **Long context / largest single sequences:** `TP8 / PP4 / EP16 / ETP1 / DP2`, MTPM≈40–64k - (single-sequence ceiling ~40–48k; ~12k tok/s on a ~39k-mean distribution). -- **Cap RL response length to ~40k tokens.** The full 60k±30k distribution is not trainable - per-sequence on 64×H200; CP cannot rescue it (CP=2 forces PP2, whose 2× weights cancel CP's savings). +- **Long context, no CP:** `TP8 / PP4 / EP16 / ETP1 / DP2`, MTPM≈40–64k (single-sequence ceiling ~40–48k; + ~12k tok/s on a ~39k-mean distribution — long seqs are more throughput-efficient per token). +- **Longest single sequences (up to ~96k):** `TP8 / PP4 / CP2 / EP16 / ETP1 / DP1` — CP2 ~doubles the + single-seq ceiling to ~96k (EP16 stays valid since `EP \| TP·CP·DP`). Costs DP→1 (≈ half the DP throughput), + so use it only when sequences actually exceed ~48k. Prefer **PP4+CP2** over PP2+CP4 (PP2's weights negate CP). +- **For the 60k±30k distribution:** clamp responses to ~96k with PP4/CP2 (≈10% of samples truncated), or to + ~40k with PP4/DP2 (≈half truncated, but full DP throughput). The full untruncated 131k tail is not trainable + on 64×H200. - Keep `TP=8` (sequence parallelism shards activations by TP — TP4 doubles activation memory and OOMs), optimizer CPU-offload on, and `recompute_granularity=full`. From cc686bbff6e5a54a193672b4796fb2ce59a58246 Mon Sep 17 00:00:00 2001 From: Eric Tang Date: Wed, 1 Jul 2026 23:42:32 +0000 Subject: [PATCH 4/5] x --- .../train/megatron/README_nemotron_ultra.md | 53 +++++++++++++++++-- pyproject.toml | 11 +++- 2 files changed, 58 insertions(+), 6 deletions(-) diff --git a/examples/train/megatron/README_nemotron_ultra.md b/examples/train/megatron/README_nemotron_ultra.md index 12ce1f2c90..f68ae63ffe 100644 --- a/examples/train/megatron/README_nemotron_ultra.md +++ b/examples/train/megatron/README_nemotron_ultra.md @@ -15,19 +15,53 @@ Staging helper: [`stage_nemotron_ultra.py`](./stage_nemotron_ultra.py). ## Replicating on a fresh cluster -The cluster needs: 8 nodes × 8×H200-141GB, EFA, a Ray cluster, a large **node-local** disk at -`/mnt/local_storage` (~28 TB), and a small shared `/home` (which the 1.1 TB model must NOT touch). +The cluster needs: 8 nodes × 8×H200-141GB, EFA, a Ray cluster, and a large **node-local** disk at +`/mnt/local_storage` (~28 TB) for the model + caches. On this cluster `/mnt/local_storage` is also the +overlay filesystem on the worker nodes (a worker's `/` and `/home` sit on the same big 28 TB disk), so +writing to `/home` on the *workers* is no longer a hazard. The recipe still keeps the model + all caches +on `/mnt/local_storage` (node-local) so every rank reads locally, and the head node's `/home` is still +small — keep the 1.1 TB model off it. ### 1. Make sure this PR's code is present The recipe depends on several fixes in this PR (see [Why these fixes](#why-these-fixes-are-needed)). On stock SkyRL/vLLM without them you get coherent-looking **garbage** generations and `reward=0`. +### 1b. Prep the megatron uv env on every node (transformer-engine wheel) +`uv run --isolated --extra megatron` builds `transformer-engine-torch==2.11.0` from its PyPI **sdist** +(NVIDIA publishes no wheel for it), and that sdist cannot build standalone — it fails with +`ModuleNotFoundError: No module named 'build_tools'`. The env only comes up if uv **reuses a prebuilt TE +wheel** instead of rebuilding, which requires two things to line up: + +- **Use uv 0.9.x on every node, the head included** (matches `docker/Dockerfile.megatron`). uv 0.11.x + computes a different build cache-key and will *not* reuse the cached wheel → rebuild → the `build_tools` + failure. Pin the head with: + `curl -LsSf https://astral.sh/uv/0.9.6/install.sh | env UV_UNMANAGED_INSTALL=$HOME/.local/bin sh` +- **This PR's `pyproject.toml`** adds `nvidia-cudnn-cu12` to transformer-engine's + `[tool.uv.extra-build-dependencies]` and an `[tool.uv.extra-build-variables]` cudnn + `CPATH`/`LIBRARY_PATH`/`CUDNN_PATH` block, so the build-key matches the prebuilt wheel. + +Cache the prebuilt wheel at `/mnt/local_storage/wheels/` on every node (build it once per +`docker/Dockerfile.megatron`, or copy it around) and add to `~/.bashrc` on **every** node: + +```bash +export XDG_CACHE_HOME=/mnt/local_storage/.cache # uv cache -> /mnt/local_storage/.cache/uv +export UV_FIND_LINKS=/mnt/local_storage/wheels +``` + +Then pre-warm the uv cache on every node so the first launch doesn't stall (or fail on an un-warmed node): +```bash +uv run --isolated --extra megatron python -c "import transformer_engine.pytorch, megatron.core, vllm; print('ok')" +``` +Ray ships the working dir + this `uv run` invocation to the workers (`RAY_RUNTIME_ENV_HOOK`), so the +`pyproject.toml` fix propagates cluster-wide automatically at launch. + ### 2. Stage the model + data on every GPU node Everything lives on node-local `/mnt/local_storage` (the model is too big for `/home`, and every rank needs its data locally). One command does both, on all nodes, via Ray: ```bash -export HF_TOKEN=$(cat ~/.HF_TOKEN) # fast authenticated download; unauthenticated is throttled +export HF_TOKEN= # fast authenticated download; unauthenticated is throttled +export HF_XET_HIGH_PERFORMANCE=1 # much faster Xet download of the ~1.1 TB/node snapshot uv run --isolated --with ray --with huggingface_hub --with hf_transfer --with datasets \ python examples/train/megatron/stage_nemotron_ultra.py ``` @@ -36,6 +70,14 @@ This downloads the HF snapshot to `/mnt/local_storage/hf_cache` **including `cha and writes the GSM8K parquets to `/mnt/local_storage/data/gsm8k` on each node. Re-run it if the autoscaler churns in a fresh (un-staged) node. +> The staging script targets **GPU (worker) nodes** only. The **driver** (head) also loads the +> tokenizer/config and the dataset, so it needs — on its own `/mnt/local_storage` — the model *metadata* +> (`config.json` + `tokenizer*` + `chat_template.jinja`, no weights) at `/mnt/local_storage/hf_cache` +> and the GSM8K parquets at `/mnt/local_storage/data/gsm8k`. Either run a metadata-only +> `snapshot_download` (`allow_patterns=["*.json","tokenizer*","*.model","*.jinja","*.txt","*.py"]`) on the +> head and copy the parquets over, or set `HF_HUB_OFFLINE=0` for the run so the driver fetches the small +> metadata on demand. + > The `*.jinja` is essential. The tokenizer ships **no** chat template inline; the official ChatML + > reasoning template lives in `chat_template.jinja`. Without it the instruct model is prompted > off-distribution and never produces a parseable answer (reward stays 0). @@ -50,9 +92,12 @@ preemption, but it's disk). ### 4. Launch ```bash export WANDB_API_KEY= -export HF_TOKEN=$(cat ~/.HF_TOKEN) # for churn resilience +export HF_TOKEN= # for churn resilience bash examples/train/megatron/run_megatron_nemotron_ultra.sh ``` +Prereq: uv 0.9.x + the staged TE wheel + `UV_FIND_LINKS`/`XDG_CACHE_HOME` on every node (see step 1b), +otherwise the megatron env fails to build with `No module named 'build_tools'`. + EFA + multi-node specifics (all set by the script): `LD_LIBRARY_PATH=/opt/amazon/efa/lib`, `SKYRL_LD_LIBRARY_PATH_EXPORT=1`, `VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1`, `NVTE_FLASH_ATTN=0`, and raised placement-group / inference-server health timeouts (the 550B takes >600 s to come up). diff --git a/pyproject.toml b/pyproject.toml index 48effda17b..06561af0a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -244,11 +244,18 @@ override-dependencies = [ [tool.uv.extra-build-dependencies] flash-attn = [{requirement = "torch", match-runtime = true}] -transformer-engine = [{requirement = "torch", match-runtime = true}, "build_tools", "ninja"] -transformer-engine-torch = [{requirement = "torch", match-runtime = true}, "build_tools", "ninja"] +transformer-engine = [{requirement = "torch", match-runtime = true}, "build_tools", "ninja", "nvidia-cudnn-cu12"] +transformer-engine-torch = [{requirement = "torch", match-runtime = true}, "build_tools", "ninja", "nvidia-cudnn-cu12"] [tool.uv.extra-build-variables] flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE"} +# Point the C++ build env to anaconda's nvidia-cudnn-cu12 includes/libs. +# AWS H200 nodes don't ship cudnn under /usr/local/cuda; without these +# CPATH/LIBRARY_PATH/CUDNN_PATH entries, transformer-engine-torch's setup +# can't find cudnn.h and the build fails. Anaconda is at the same path on +# the head and all workers (Anyscale image). +transformer-engine-torch = { CPATH = "/home/ray/anaconda3/lib/python3.12/site-packages/nvidia/cudnn/include", LIBRARY_PATH = "/home/ray/anaconda3/lib/python3.12/site-packages/nvidia/cudnn/lib", CUDNN_PATH = "/home/ray/anaconda3/lib/python3.12/site-packages/nvidia/cudnn" } +transformer-engine = { CPATH = "/home/ray/anaconda3/lib/python3.12/site-packages/nvidia/cudnn/include", LIBRARY_PATH = "/home/ray/anaconda3/lib/python3.12/site-packages/nvidia/cudnn/lib", CUDNN_PATH = "/home/ray/anaconda3/lib/python3.12/site-packages/nvidia/cudnn" } [[tool.uv.index]] name = "pytorch-cu128" From 290466b87be475419bb87582fe11612f454c5e53 Mon Sep 17 00:00:00 2001 From: SumanthRH Date: Thu, 2 Jul 2026 06:54:43 +0000 Subject: [PATCH 5/5] [docs] GLM-4.7 355B 128K max_tokens / throughput re-tune (FP32 grads) + varlen full-ctx harness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-tune of glm47_355b_128k_megatron_ablation after peak-mem improvements + the switch to token-based dynamic micro-batching (trainer.max_tokens_per_microbatch). - ai_docs/glm47_355b_128k_max_tokens_throughput_0702.md: full writeup (Stages 1-3). - examples/train/megatron/stage_glm47.py: stage GLM-4.7 (~667GB) to node-local NVMe on all nodes (subprocess watchdog + resume — a single snapshot_download hung 6/8 nodes). - examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh: parameterized 355B full-ctx benchmark harness (all TP/PP/CP/EP/grad/recompute/batch/max_tokens knobs env-overridable). - full_context/{main_full_ctx,trainer_full_ctx}.py: mature peak-mem + step-timing logging, stubbed inference, and Stage-3 variable-length dummy batch (trainer.dummy_variable_length). Result (FP32 grads): winner TP8/PP4/CP2/EP16/ETP1/DP1. Realistic 256-sample varlen batch (avg ~70K) ~8.9K tok/s/cluster (139 tok/s/GPU) at max_tokens=256K, peak 66GB. all-128K train_batch=32 ~7.0K tok/s. max_tokens sweet spot 256K (128K/GPU); 384K OOMs at full pipeline. DP2/DP4 do not help under FP32 grads (grad-reduce 2x/4x bytes; PP2 OOMs) — need bf16 grads. Co-Authored-By: Claude Opus 4.8 (1M context) --- ...47_355b_128k_max_tokens_throughput_0702.md | 144 ++++++++++++++ examples/train/megatron/stage_glm47.py | 157 +++++++++++++++ .../full_context/main_full_ctx.py | 35 ++++ .../full_context/run_full_ctx_glm47_355b.sh | 179 ++++++++++++++++++ .../full_context/trainer_full_ctx.py | 150 ++++++++++++--- 5 files changed, 634 insertions(+), 31 deletions(-) create mode 100644 ai_docs/glm47_355b_128k_max_tokens_throughput_0702.md create mode 100644 examples/train/megatron/stage_glm47.py create mode 100644 examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh diff --git a/ai_docs/glm47_355b_128k_max_tokens_throughput_0702.md b/ai_docs/glm47_355b_128k_max_tokens_throughput_0702.md new file mode 100644 index 0000000000..2703ffacc3 --- /dev/null +++ b/ai_docs/glm47_355b_128k_max_tokens_throughput_0702.md @@ -0,0 +1,144 @@ +# GLM-4.7 (355B) — 128K full-context Megatron throughput & `max_tokens_per_microbatch` re-tune (8×8×H200) + +**Status:** complete (2026-07-02). Re-tune of `ai_docs/glm47_355b_128k_megatron_ablation.md` after +(a) peak-memory improvements in the Megatron backend, and (b) the switch from +`micro_train_batch_size_per_gpu` to **token-based dynamic micro-batching** (`trainer.max_tokens_per_microbatch`). + +**TL;DR (FP32 grads — the configured default this session):** +- Best config: **TP=8 PP=4 CP=2 EP=16 ETP=1 (DP=1)** + full recompute + full CPU optimizer offload + FP32 grads. + On a **realistic 256-sample variable-length batch** (avg ~70K, ~17.4 M tokens) it runs at **~8.9 K + tok/s/cluster (139 tok/s/GPU)**, peak **66 GB/GPU**. On the all-128K stress test (train_batch=32) it is + **~7.0 K tok/s** (599 s/step) — still the winner and ~7 % faster than the old ablation's D1 (647 s). +- **`max_tokens_per_microbatch`: use ~256K/rank (128K/GPU at CP=2).** That fits at 66 GB even at a full + pipeline and is one step below the **384K OOM cliff** (a per-GPU-token *transient* in the MoE all-to-all + dispatcher, amplified when the PP-4 pipeline holds several microbatches concurrently). The Stage-1 "384K + fits" was an artifact of an underfilled pipeline (3 microbatches); at a realistic batch the ceiling is 256K. + Note: max_tokens is a *packing/fit* knob, not a raw-speed knob — at fixed batch, more/smaller microbatches + fill the pipeline better. +- **DP2 / DP4 do not help under FP32 grads** (unlike the Nemotron-Ultra recipe): the FP32 gradient all-reduce + is 2×/4× the bytes and eats the DP benefit even with `overlap_grad_reduce=true`; PP=2 (required for DP4) + OOMs under FP32. The Nemotron DP2/DP4 wins depend on **bf16 grads**, which were disabled here by request — + revisit `PP4/CP1/DP2` if bf16 grads are later allowed. + +## Setup + +- Cluster: 8 nodes × 8× H200-141G = 64 GPUs (Anyscale, EFA ~363 GBps busbw). +- Weights: **real** `zai-org/GLM-4.7` (~667 GB BF16, 94 safetensors) staged to `/mnt/local_storage/hf_cache` + on all 8 nodes via `examples/train/megatron/stage_glm47.py` (subprocess-watchdog + resume; see gotchas). +- Driver: `examples/train_scripts/full_context/main_full_ctx.py` (dummy trainer, inference stubbed). + Harness: `examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh` (all knobs env-overridable). +- GLM-4.7 arch: 92 layers, 5120 hidden, 96 attn / 8 KV heads (GQA), 160 routed + 1 shared expert, top-8, + moe_intermediate 1536, first_k_dense=3, 1 MTP layer (disabled via `mtp_num_layers=null`). + ⇒ **PP ∈ {1,2,4}** (92 not divisible by 8), **TP ≤ 8** (8 KV heads), **EP | 160**, **EP×ETP = TP·CP·DP**. +- All runs: full recompute (`granularity=full, method=uniform, num_layers=1`), CPU-offloaded optimizer + (`offload_fraction=1.0`, precision-aware, d2h/h2d overlap), `remove_microbatch_padding=true` (THD), flash-attn, + **FP32 grads** (`ddp_config.grad_reduce_in_fp32=true`). Reported number = **step 2** (step 1 warmup-tainted). +- `max_tokens_per_microbatch` semantics: **per-DP-rank** bin-packing budget (each DP worker packs its own shard; + padding microbatches sync counts across DP). Per-GPU tokens ≈ `max_tokens / CP`. + +## Stage 1 — `max_tokens_per_microbatch` ceiling at the starter config + +Starter = TP8/PP4/CP2/EP16/ETP1 (DP1), FP32 grads, full recompute. train_batch=8 all-128K = 1.05 M tok. + +| max_tokens | tokens/GPU | #microbatch | step2 (s) | fwd (s) | train (s) | peak max_used (GB) | result | +|---|---|---|---|---|---|---|---| +| 128 000 | 64K | 8 | 205.0 | 37.8 | 167.2 | 60.2 | ✅ baseline | +| 256 000 | 128K | 4 | 257.0 | 51.5 | 205.3 | 64.9 | ✅ (alloc 42.9) | +| 384 000 | 192K | 3 | 286.3 | 57.6 | 228.7 | 65.7 | ✅ (alloc 43.5) | +| 512 000 | 256K | 2 | — | — | — | OOM | ❌ transient MoE buffer @ expert-bias all-reduce | + +**Findings.** Steady allocation is nearly flat (peak_alloc 38→44 GB from 64K→192K tokens/GPU) — full recompute +keeps activations tiny. The 512K OOM is a **transient** (MoE all-to-all dispatch / grouped-GEMM workspace) that +scales with tokens/GPU and spikes between 192K and 256K/GPU. So the packing ceiling ≈ **384K/rank** at CP=2 with +FP32 grads. Throughput note: **lower max_tokens is faster** at fixed batch (205 s @ 8 microbatches < 257 s @ 4 < +286 s @ 3) because more/smaller microbatches fill the PP-4 pipeline better. So max_tokens is a *packing/fit* knob, +not a speed knob — you want it just large enough to pack your longest sequences, not maxed out. + +## Stage 2 — parallelism sweep for throughput (FP32 grads) + +All: FP32 grads, full recompute, offload=1.0, all-128K sequences, `max_tokens=128000`/rank (1 seq/microbatch = +most microbatches = best pipeline fill), **train_batch=32 = 4.19 M tok**. tok/s = 4,194,304 / step2. + +| tag | TP | PP | CP | EP | ETP | DP | overlap | step2 (s) | tok/s | peak (GB) | result | +|---|---|---|---|---|---|---|---|---|---|---|---| +| **c0** | 8 | 4 | 2 | 16 | 1 | 1 | — | **599.4** | **6997** | 60.2 | ✅ **winner** (beats old D1's 647 s) | +| c1 | 8 | 4 | 1 | 16 | 1 | 2 | no | 617.7 | 6789 | 64.4 | ✅ slower — FP32 DP-reduce not overlapped | +| c1b | 8 | 4 | 1 | 16 | 1 | 2 | yes | 609.9 | 6877 | 64.3 | ✅ overlap helps +1.3%, still < c0 | +| c3 | 8 | 2 | 1 | 32 | 1 | 4 | yes | — | — | OOM | ❌ PP2 2× FP32 grad buffer + 128K/GPU | + +**Findings.** +- **c0 (PP4/CP2/DP1) wins.** DP=1 pays *no* gradient all-reduce; CP=2 splits each 128K sequence to 64K/GPU and + balances load well. +- **DP2 (c1/c1b) is ~2–3% slower.** With FP32 grads the DP all-reduce of the full 355B gradient is 2× the bytes; + `overlap_grad_reduce=true` recovers ~1.3% but not enough. CP=1 also puts the whole 128K sequence on one GPU. +- **DP4 (c3) OOMs.** DP4 requires PP=2 (TP8·PP2·CP1·DP4=64), and PP=2's ~2× dense FP32 grad buffer + 128K/GPU + activation exceeds 141 GB — matching the old ablation's "PP<4 OOMs with FP32 grads". +- **Why this differs from Nemotron-Ultra** (`NEMOTRON_ULTRA_THROUGHPUT.md`, which found PP4/CP1/DP2 and PP2/DP4 + fastest): those wins ride on **bf16 grads** (halved grad buffer ⇒ PP2 fits, and half the DP-reduce bytes). + Under the FP32-grad constraint set for this session, DP scaling is net-negative and PP=2 is infeasible, so the + DP1 shape wins. If bf16 grads are later allowed, PP4/CP1/DP2 is the config to revisit. + +## Stage 3 — realistic variable-length throughput (256 samples, avg 70K, std 30K) + +Winner config c0 (TP8/PP4/CP2/EP16/DP1), FP32 grads. Dummy batch = 256 samples, per-sample length ~ +clamped Normal(70K, 30K) in [2048, 128000]; realized mean **71,406**, std **30,002**, **18.28 M tokens/step**. +Generated by the new `trainer.dummy_variable_length` path in `trainer_full_ctx.py`. + +| max_tokens | tokens/GPU | ~#microbatch | step2 (s) | tok/s/cluster | tok/s/GPU | peak (GB) | result | +|---|---|---|---|---|---|---|---| +| 256 000 | 128K | ~72 | **1954.2** | **8 927** | **139** | 65.9 | ✅ **best** (28% > all-128K baseline) | +| 384 000 | 192K | ~48 | — | — | — | OOM | ❌ MoE dispatcher transient (13.8 GB) in backward @ full pipeline | + +**Findings.** The realistic variable-length batch runs at **~8.9 K tok/s/cluster (139 tok/s/GPU)** — **28 % +faster** than the all-128K stress test (6,997 tok/s), from (a) ~72 microbatches ⇒ near-ideal PP-4 fill (~8 % +bubble vs 19 % at batch=32), (b) THD packing of multiple short sequences per microbatch, and (c) lower +quadratic attention cost at the 68K average length. Step = fwd_logprobs 447 s + train 1,507 s = 1,954 s for +17.44 M tokens; peak **65.9 GB**. + +**`max_tokens` ceiling at a realistic (full-pipeline) batch is 256K/rank (128K/GPU), NOT the 384K Stage-1 +suggested.** 384K OOMs here: at ~48 microbatches the PP-4 1F1B steady state holds several 192K/GPU +microbatches' MoE-dispatcher transients concurrently (the OOM is a 13.8 GB alloc inside +`token_dispatcher.combine_preprocess` during backward). Stage 1 only "fit" 384K because train_batch=8 gave +3 microbatches (underfilled pipeline). **So size `max_tokens` for the full-pipeline peak: 256K is the sweet +spot** — packs efficiently, fits at 66 GB, one step below the OOM cliff. (128K/rank also works, more +microbatches; not benchmarked here per request, expected similar-or-slightly-lower throughput.) + +## Recommended config (8-node, 128K, FP32 grads) + +``` +trainer.policy.megatron_config.tensor_model_parallel_size=8 +trainer.policy.megatron_config.pipeline_model_parallel_size=4 +trainer.policy.megatron_config.context_parallel_size=2 +trainer.policy.megatron_config.expert_model_parallel_size=16 +trainer.policy.megatron_config.expert_tensor_parallel_size=1 +trainer.policy.megatron_config.ddp_config.grad_reduce_in_fp32=true +trainer.policy.megatron_config.ddp_config.overlap_grad_reduce=false # DP=1: nothing to overlap +trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=true +trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=1.0 +trainer.policy.megatron_config.transformer_config_kwargs.recompute_granularity=full # +method=uniform, num_layers=1 +trainer.policy.megatron_config.transformer_config_kwargs.num_moe_experts=160 +trainer.policy.megatron_config.transformer_config_kwargs.mtp_num_layers=null +trainer.remove_microbatch_padding=true +trainer.max_tokens_per_microbatch=256000 # 128K/GPU at CP2; sweet spot — fits at full pipeline, 1 step below the 384K OOM cliff +``` + +**Headline numbers (8 nodes, 64×H200-141G, real GLM-4.7 weights, FP32 grads, full recompute):** + +| workload | config | step (s) | tokens/step | tok/s/cluster | tok/s/GPU | peak (GB) | +|---|---|---|---|---|---|---| +| all-128K, train_batch=32, max_tokens=128K | c0 (PP4/CP2/DP1) | 599 | 4.19 M | 6 997 | 109 | 60 | +| realistic varlen (256 samples, avg 70K), max_tokens=256K | c0 (PP4/CP2/DP1) | 1 954 | 17.44 M | **8 927** | **139** | 66 | + +## Setup gotchas hit this session + +- **GLM-4.7 is ~667 GB** (BF16), not 358 GB (HF API param-count looked like a byte count). 94 safetensors. +- **Staging hang:** a single `snapshot_download(max_workers=16)` on 8 nodes (128 concurrent HF connections) + hung on 6/8 nodes with dead TCP sockets (no timeout) after ~40 min — a *hang*, not an exception, so the + retry loop never fired. Fix in `stage_glm47.py`: run each download attempt in a **subprocess** killed after a + 20-min hard timeout, then retry — `hf_transfer` resumes from `.incomplete` blobs, so every attempt makes + forward progress and a hung socket can't wedge the stage. Also dropped `max_workers` 16→6. +- `HF_HOME=/mnt/local_storage/hf_cache` + `HF_HUB_OFFLINE=1` must be exported by the launcher; SkyRL propagates + them to workers (`skyrl/train/utils/utils.py`), which is how the Megatron bridge finds the staged snapshot. +- `trainer.use_sample_packing` and `trainer.remove_microbatch_padding` are mutually exclusive (the former is a + deprecated alias) — passing both raises. Use `remove_microbatch_padding` only. +- For Megatron + flash-attn, SkyRL now forces `NVTE_FUSED_ATTN=0`, so the old FA3 (~1%) trick no longer applies. diff --git a/examples/train/megatron/stage_glm47.py b/examples/train/megatron/stage_glm47.py new file mode 100644 index 0000000000..dec0171afb --- /dev/null +++ b/examples/train/megatron/stage_glm47.py @@ -0,0 +1,157 @@ +"""Pre-stage zai-org/GLM-4.7 (355B, ~667 GB BF16 checkpoint) + the GSM8K dataset onto every +GPU node's local disk, for the full-context Megatron throughput ablation. + +Adapted from ``stage_nemotron_ultra.py`` (nemotron-3-ultra-550b-rl branch). Why this exists: +the model is ~667 GB and the cluster's shared ``/home`` is small, so the checkpoint must live on +each node's large node-local disk (``/mnt/local_storage``, ~24 TB). Multi-node training also needs +the data present on every node (each rank reads its data locally). This launches one Ray task per +distinct GPU node that (1) downloads the HF snapshot and (2) writes the GSM8K parquets, both under +``/mnt/local_storage``. + +ROBUSTNESS: an earlier single-shot ``snapshot_download(max_workers=16)`` hung on 6/8 nodes with +dead TCP connections (no timeout) after ~40 min — likely Hub throttling from 128 concurrent +connections cluster-wide. Fix: run each download attempt in a *child process* killed after a hard +timeout, then retry — ``hf_transfer`` resumes from the ``.incomplete`` blobs, so every attempt makes +forward progress and a hung socket can never wedge the whole stage. Also fewer workers per node. + +IMPORTANT: ``allow_patterns`` includes ``*.jinja`` (chat template) and ``*.py`` (custom modeling +code, ``Glm4MoeForCausalLM``). + +Usage (from the head node, on a running Ray cluster): + set -a; source /home/ray/default/SkyRL-private/.env.apex; set +a # exports HF_TOKEN + uv run --isolated --with ray --with huggingface_hub --with hf_transfer --with datasets \ + python examples/train/megatron/stage_glm47.py +""" + +import os + +import ray +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy +from ray.util.state import list_nodes + +REPO = "zai-org/GLM-4.7" +HF_HOME = "/mnt/local_storage/hf_cache" +DATA_DIR = "/mnt/local_storage/data/gsm8k" +CACHE_DIR = f"{HF_HOME}/hub" +ALLOW = ["*.safetensors", "*.json", "*.txt", "tokenizer*", "*.model", "*.jinja", "*.py"] +MAX_WORKERS = 6 +ATTEMPT_TIMEOUT_S = 1200 # 20 min/attempt; a hang is killed and resumed +MAX_ATTEMPTS = 40 + +ray.init(address="auto", ignore_reinit_error=True, log_to_driver=True) + +_ENV = {"HF_HOME": HF_HOME, "HF_HUB_ENABLE_HF_TRANSFER": "1"} +if os.environ.get("HF_TOKEN"): + _ENV["HF_TOKEN"] = os.environ["HF_TOKEN"] + + +_CHILD_CODE = """ +import sys +from huggingface_hub import snapshot_download +p = snapshot_download( + {repo!r}, cache_dir={cache_dir!r}, max_workers={max_workers}, + allow_patterns={allow!r}, +) +print("SNAPSHOT_OK:" + p) +""".format(repo=REPO, cache_dir=CACHE_DIR, max_workers=MAX_WORKERS, allow=ALLOW) + + +@ray.remote( + num_cpus=8, + runtime_env={"pip": ["huggingface_hub", "hf_transfer", "datasets"], "env_vars": _ENV}, +) +def stage(node_ip): + import shutil + import socket + import subprocess + import sys + import time + + free_before_gb = shutil.disk_usage("/mnt/local_storage").free / 1024**3 + + # 1) Model snapshot with a per-attempt subprocess watchdog. subprocess.run(timeout=...) + # SIGKILLs a hung child (and its hf_transfer threads); the next attempt resumes the + # .incomplete blobs, so a dead socket can never wedge the whole stage. + model_path = None + last_err = None + for attempt in range(MAX_ATTEMPTS): + try: + r = subprocess.run( + [sys.executable, "-c", _CHILD_CODE], + capture_output=True, + text=True, + timeout=ATTEMPT_TIMEOUT_S, + ) + except subprocess.TimeoutExpired: + last_err = f"attempt {attempt}: {ATTEMPT_TIMEOUT_S}s timeout -> killed, resuming" + print(f"[{node_ip}] {last_err}", flush=True) + continue + ok_line = [ln for ln in (r.stdout or "").splitlines() if ln.startswith("SNAPSHOT_OK:")] + if r.returncode == 0 and ok_line: + model_path = ok_line[0].split("SNAPSHOT_OK:", 1)[1] + break + last_err = f"attempt {attempt}: rc={r.returncode} :: {(r.stderr or '')[-300:]}" + print(f"[{node_ip}] {last_err}", flush=True) + time.sleep(15) + if model_path is None: + return (socket.gethostname(), node_ip, "MODEL_FAILED", last_err, round(free_before_gb, 1), None) + + free_after_gb = shutil.disk_usage("/mnt/local_storage").free / 1024**3 + + # 2) GSM8K parquets (tiny; one copy per node so every rank reads locally). + try: + import re + + import datasets + + os.makedirs(DATA_DIR, exist_ok=True) + instruction = 'Let\'s think step by step and output the final answer after "####".' + + def to_row(example, idx, split): + q = example["question"] + sol = re.search(r"#### (\-?[0-9\.\,]+)", example["answer"]).group(0).split("#### ")[1].replace(",", "") + return { + "data_source": "openai/gsm8k", + "prompt": [{"role": "user", "content": q + " " + instruction}], + "env_class": "gsm8k", + "reward_spec": {"method": "rule", "ground_truth": sol}, + "extra_info": {"split": split, "index": idx}, + } + + ds = datasets.load_dataset("openai/gsm8k", "main") + ds["train"].map(lambda e, i: to_row(e, i, "train"), with_indices=True).to_parquet(f"{DATA_DIR}/train.parquet") + ds["test"].map(lambda e, i: to_row(e, i, "test"), with_indices=True).to_parquet( + f"{DATA_DIR}/validation.parquet" + ) + except Exception as e: # noqa: BLE001 + return ( + socket.gethostname(), + node_ip, + "DATA_FAILED", + f"{type(e).__name__}: {str(e)[:300]}", + round(free_before_gb, 1), + round(free_after_gb, 1), + ) + + return (socket.gethostname(), node_ip, "DONE", model_path, round(free_before_gb, 1), round(free_after_gb, 1)) + + +nodes = [ + n + for n in list_nodes(detail=True) + if (n.get("resources_total") or {}).get("GPU", 0) > 0 and n.get("state") == "ALIVE" +] +print("staging on %d GPU nodes: %s" % (len(nodes), [n["node_ip"] for n in nodes]), flush=True) +results = ray.get( + [ + stage.options(scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=n["node_id"], soft=False)).remote( + n["node_ip"] + ) + for n in nodes + ] +) +print("\n===== RESULTS =====", flush=True) +ok = sum(1 for r in results if r[2] == "DONE") +for host, ip, status, detail, free_before, free_after in results: + print(f"[{ip} {host}] {status}: free_before={free_before}GB free_after={free_after}GB :: {detail}", flush=True) +print(f"\n{ok}/{len(results)} nodes staged OK", flush=True) diff --git a/examples/train_scripts/full_context/main_full_ctx.py b/examples/train_scripts/full_context/main_full_ctx.py index 169ebee6d4..b939e70850 100644 --- a/examples/train_scripts/full_context/main_full_ctx.py +++ b/examples/train_scripts/full_context/main_full_ctx.py @@ -17,12 +17,47 @@ @dataclass class FullCtxTrainerConfig(TrainerConfig): num_dummy_steps: int = 5 + # --- Stage 3: variable-length dummy sequences (realistic throughput test) --- + # When enabled, each of the ``train_batch_size * n_samples_per_prompt`` dummy samples gets a + # total length drawn from a clamped Normal(dummy_mean_len, dummy_std_len), instead of all + # sequences being fully padded to the max context length. This exercises the token-based + # micro-batch packing (``max_tokens_per_microbatch``) the way a real RL batch would. + dummy_variable_length: bool = False + dummy_mean_len: int = 70000 + dummy_std_len: int = 30000 + dummy_min_len: int = 2048 + dummy_prompt_len: int = 512 + dummy_seed: int = 1234 FullCtxConfig = make_config(trainer_cls=FullCtxTrainerConfig) +class _StubInferenceEngineClient: + """Placeholder inference client for full-context perf testing. + + The dummy trainer never generates tokens or syncs weights, so this object + only needs to be a valid attribute holder. Any unexpected call raises so we + notice if some code path drifts and starts requiring real generation. + """ + + def __getattr__(self, name): # pragma: no cover - shouldn't be hit + raise RuntimeError( + f"_StubInferenceEngineClient.{name}() called from full-context perf test. " + "Generation is not supposed to run; the trainer should only exercise the " + "policy worker forward/backward path." + ) + + class FullCtxPPOExp(BasePPOExp): + def get_inference_client(self): + # Skip vLLM init: full-context dummy trainer never generates and we don't + # have weights for vLLM to load (random-init perf test). + return _StubInferenceEngineClient() + + def get_generator(self, cfg, tokenizer, inference_engine_client): + return None # Unused in dummy trainer. + def get_trainer( self, cfg, diff --git a/examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh b/examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh new file mode 100644 index 0000000000..6feac40315 --- /dev/null +++ b/examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh @@ -0,0 +1,179 @@ +#!/bin/bash +set -x +# ============================================================================= +# GLM-4.7 (355B) full-context Megatron throughput / MAX_TOKENS_PER_MICROBATCH +# ablation on 8x8xH200 (64 GPUs). +# +# Uses the dummy `main_full_ctx` trainer (stubbed inference, no generation) to run +# `num_dummy_steps` of fully-padded 128K-context batches and reports peak GPU memory +# + step timings, so we can find: +# (stage 1) the largest MAX_TOKENS_PER_MICROBATCH that fits, and +# (stage 2) the parallelism shape (TP/PP/CP/EP/DP) with the best throughput. +# +# Weights: real staged GLM-4.7 from /mnt/local_storage/hf_cache (see stage_glm47.py). +# Inference is stubbed; we keep it colocated (tp=8, 8 engines) only to satisfy the +# colocate_all placement-group GPU-count check. +# +# All knobs are env-overridable so the driver harness can sweep them: +# MAX_TOKENS_PER_MICROBATCH, MEGATRON_TP/PP/CP/EP/ETP, RECOMPUTE_*, GRAD_FP32, +# TRAIN_BATCH, MINI_BATCH, N_SAMPLES, CONTEXT_LEN, OPTIMIZER_*, NUM_DUMMY_STEPS. +# +# Usage: +# RUN_TAG=s1_mtpm256k MAX_TOKENS_PER_MICROBATCH=256000 \ +# bash examples/train_scripts/full_context/run_full_ctx_glm47_355b.sh +# ============================================================================= + +# --- AWS EFA networking (multi-node NCCL) --- +export LD_LIBRARY_PATH=/opt/amazon/efa/lib:${LD_LIBRARY_PATH:-} +export SKYRL_LD_LIBRARY_PATH_EXPORT=1 +export NCCL_SOCKET_IFNAME=eth0 +export FI_EFA_USE_DEVICE_RDMA=1 +export FI_PROVIDER=efa +export NCCL_DEBUG=WARN + +# --- HF cache: point at the staged snapshot on node-local NVMe --- +export HF_HOME=/mnt/local_storage/hf_cache +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 + +# --- Memory allocator + bigger PG/NCCL timeouts (355B weight load from local NVMe) --- +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export SKYRL_RAY_PG_TIMEOUT_IN_S=${SKYRL_RAY_PG_TIMEOUT_IN_S:-1800} +export SKYRL_WORKER_NCCL_TIMEOUT_IN_S=${SKYRL_WORKER_NCCL_TIMEOUT_IN_S:-3600} + +# --- WANDB (optional; we log to console by default) --- +if [ -f /home/ray/default/SkyRL-private/.env.apex ]; then + export WANDB_API_KEY=$(grep -E '^export WANDB_API_KEY=' /home/ray/default/SkyRL-private/.env.apex | head -1 | sed 's/^export WANDB_API_KEY=//') +fi + +DATA_DIR="/mnt/local_storage/data/gsm8k" +TRAIN_FILE="$DATA_DIR/train.parquet" +TEST_FILE="$DATA_DIR/validation.parquet" +MODEL_NAME="zai-org/GLM-4.7" + +# --- Cluster shape --- +NUM_NODES=${NUM_NODES:-8} +NUM_GPUS_PER_NODE=8 + +# --- Single-sequence context length (all samples fully padded to this in stage 1/2) --- +CONTEXT_LEN=${CONTEXT_LEN:-128000} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-8000} +MAX_RESPONSE_LENGTH=$(( CONTEXT_LEN - MAX_PROMPT_LENGTH )) + +# --- The knob to sweep (token-based dynamic micro-batching) --- +MAX_TOKENS_PER_MICROBATCH=${MAX_TOKENS_PER_MICROBATCH:-128000} + +# --- Batch sizing (global). DP-normalized constraints validated in utils.validate_batch_sizes --- +TRAIN_BATCH=${TRAIN_BATCH:-8} +MINI_BATCH=${MINI_BATCH:-8} +N_SAMPLES=${N_SAMPLES:-1} +NUM_DUMMY_STEPS=${NUM_DUMMY_STEPS:-2} + +# --- Megatron parallelism (starter = winner A1 from the 355B 128K ablation) --- +MEGATRON_TP=${MEGATRON_TP:-8} +MEGATRON_PP=${MEGATRON_PP:-4} +MEGATRON_CP=${MEGATRON_CP:-2} +MEGATRON_EP=${MEGATRON_EP:-16} +MEGATRON_ETP=${MEGATRON_ETP:-1} + +# --- MoE routing flags (GLM-4.7: sigmoid scoring + expert bias, DeepSeek-V3 family) --- +MOE_TOKEN_DISPATCHER="alltoall" +MOE_ROUTER_LB="none" +MOE_GROUPED_GEMM=true +MOE_ROUTER_SCORE_FN="sigmoid" +MOE_ROUTER_EXPERT_BIAS=true +NUM_MOE_EXPERTS=${NUM_MOE_EXPERTS:-160} + +# --- CPU optimizer offload (full offload = winner) --- +OPTIMIZER_CPU_OFFLOAD=${OPTIMIZER_CPU_OFFLOAD:-true} +OPTIMIZER_OFFLOAD_FRACTION=${OPTIMIZER_OFFLOAD_FRACTION:-1.0} + +# --- Activation recompute (full+uniform+1 required at 128K in the old ablation) --- +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} +RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-uniform} +RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-1} +RECOMPUTE_MODULES=${RECOMPUTE_MODULES:-'[core_attn]'} + +# --- Gradient accumulation dtype + DDP overlap --- +GRAD_FP32=${GRAD_FP32:-true} +DDP_OVERLAP_GRAD=${DDP_OVERLAP_GRAD:-false} +DDP_OVERLAP_PARAM=${DDP_OVERLAP_PARAM:-false} + +FLASH_ATTN=true + +# --- Inference (stubbed; colocated tp=8 just to satisfy placement-group GPU count) --- +INFERENCE_ENGINE_TP=8 +NUM_INFERENCE_ENGINES=$(( NUM_NODES * NUM_GPUS_PER_NODE / INFERENCE_ENGINE_TP )) + +RUN_TAG=${RUN_TAG:-glm47_355b} +RUN_NAME="fullctx_glm47_355b_${RUN_TAG}_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_ep${MEGATRON_EP}_etp${MEGATRON_ETP}_mtpm${MAX_TOKENS_PER_MICROBATCH}" +echo "=== RUN $RUN_NAME : nodes=$NUM_NODES TP=$MEGATRON_TP PP=$MEGATRON_PP CP=$MEGATRON_CP EP=$MEGATRON_EP ETP=$MEGATRON_ETP max_tokens=$MAX_TOKENS_PER_MICROBATCH ctx=$CONTEXT_LEN train_batch=$TRAIN_BATCH n_samples=$N_SAMPLES recompute=$RECOMPUTE_GRANULARITY grad_fp32=$GRAD_FP32 offload=$OPTIMIZER_CPU_OFFLOAD ===" + +uv run --isolated --extra megatron -m examples.train_scripts.full_context.main_full_ctx \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.policy.model.path=$MODEL_NAME \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TP \ + generator.inference_engine.enforce_eager=true \ + generator.inference_engine.engine_init_kwargs.max_model_len=$CONTEXT_LEN \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.moe_token_dispatcher_type=$MOE_TOKEN_DISPATCHER \ + trainer.policy.megatron_config.moe_router_load_balancing_type=$MOE_ROUTER_LB \ + trainer.policy.megatron_config.moe_grouped_gemm=$MOE_GROUPED_GEMM \ + trainer.policy.megatron_config.moe_router_score_function=$MOE_ROUTER_SCORE_FN \ + trainer.policy.megatron_config.moe_router_enable_expert_bias=$MOE_ROUTER_EXPERT_BIAS \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_CPU_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \ + trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=true \ + trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=true \ + trainer.policy.megatron_config.empty_cuda_cache=true \ + trainer.policy.megatron_config.transformer_config_kwargs.num_moe_experts=$NUM_MOE_EXPERTS \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_granularity=$RECOMPUTE_GRANULARITY \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_method=$RECOMPUTE_METHOD \ + trainer.policy.megatron_config.transformer_config_kwargs.recompute_num_layers=$RECOMPUTE_NUM_LAYERS \ + "trainer.policy.megatron_config.transformer_config_kwargs.recompute_modules=$RECOMPUTE_MODULES" \ + trainer.policy.megatron_config.transformer_config_kwargs.mtp_num_layers=null \ + trainer.policy.megatron_config.ddp_config.grad_reduce_in_fp32=$GRAD_FP32 \ + trainer.policy.megatron_config.ddp_config.overlap_grad_reduce=$DDP_OVERLAP_GRAD \ + trainer.policy.megatron_config.ddp_config.overlap_param_gather=$DDP_OVERLAP_PARAM \ + trainer.remove_microbatch_padding=true \ + trainer.flash_attn=$FLASH_ATTN \ + trainer.max_tokens_per_microbatch=$MAX_TOKENS_PER_MICROBATCH \ + trainer.epochs=1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TRAIN_BATCH \ + trainer.policy_mini_batch_size=$MINI_BATCH \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=999999 \ + trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ + generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.policy.optimizer_config.lr=1.0e-6 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.algorithm.use_kl_loss=false \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=false \ + generator.batched=true \ + environment.env_class=gsm8k \ + generator.n_samples_per_prompt=$N_SAMPLES \ + trainer.logger="console" \ + trainer.project_name="glm47_355b_full_ctx" \ + trainer.run_name="$RUN_NAME" \ + trainer.eval_before_train=false \ + trainer.eval_interval=999999 \ + trainer.num_dummy_steps=$NUM_DUMMY_STEPS \ + "$@" diff --git a/examples/train_scripts/full_context/trainer_full_ctx.py b/examples/train_scripts/full_context/trainer_full_ctx.py index 6022639d9d..9aeca647f1 100644 --- a/examples/train_scripts/full_context/trainer_full_ctx.py +++ b/examples/train_scripts/full_context/trainer_full_ctx.py @@ -1,6 +1,7 @@ from skyrl.train.trainer import RayPPOTrainer from loguru import logger import random +import numpy as np from skyrl.train.utils.utils import Timer @@ -20,9 +21,10 @@ async def train(self): self.global_step = 0 - # Initialize weight sync state - with Timer("init_weight_sync_state", self.all_timings): - self.init_weight_sync_state() + # NOTE(charlie/full_ctx): skip init_weight_sync_state — no real inference engine + # in this perf test. The trainer only exercises policy fwd/bwd paths. + # with Timer("init_weight_sync_state", self.all_timings): + # self.init_weight_sync_state() # Run a few training steps self.global_step += 1 # start from 1 @@ -31,34 +33,10 @@ async def train(self): # Run a single training step with Timer("step", self.all_timings): - # Create training input directly with max length sequences - num_samples = self.cfg.trainer.train_batch_size * self.cfg.generator.n_samples_per_prompt - uids = [str(i) for i in range(self.cfg.trainer.train_batch_size)] - prompt_token_ids = [ - [random.randint(0, self.tokenizer.vocab_size - 1)] * self.cfg.generator.max_input_length - ] * self.cfg.trainer.train_batch_size - prompt_token_ids = sum( - [ - [prompt_token_id] * self.cfg.generator.n_samples_per_prompt - for prompt_token_id in prompt_token_ids - ], - [], - ) - response_ids = [ - [random.randint(0, self.tokenizer.vocab_size - 1)] - * self.cfg.generator.sampling_params.max_generate_length - ] * num_samples - uids = sum([[uid] * self.cfg.generator.n_samples_per_prompt for uid in uids], []) - - dummy_generator_output = { - "prompt_token_ids": prompt_token_ids, - "response_ids": response_ids, - "rewards": [ - [0] * (self.cfg.generator.sampling_params.max_generate_length - 1) + [random.randint(0, 1)] - ] - * num_samples, - "loss_masks": [[1] * self.cfg.generator.sampling_params.max_generate_length] * num_samples, - } + if getattr(self.cfg.trainer, "dummy_variable_length", False): + dummy_generator_output, uids = self._build_variable_length_batch(step) + else: + dummy_generator_output, uids = self._build_max_length_batch() training_input = self.convert_to_training_input(dummy_generator_output, uids) with Timer("fwd_logprobs_values_reward", self.all_timings): @@ -84,6 +62,13 @@ async def train(self): self.tracker.log(self.all_metrics, step=self.global_step) self.all_metrics = {} self.tracker.log({"timing/" + k: v for k, v in self.all_timings.items()}, step=self.global_step) + # NOTE(charlie/full_ctx): print step timings + peak GPU memory across all + # workers so the ablation harness can grep this line. + try: + mem_info = self._collect_peak_memory() + except Exception as e: + mem_info = f"" + logger.info(f"FULLCTX_STEP step={step + 1} timings={dict(self.all_timings)} peak_mem={mem_info}") self.all_timings = {} self.global_step += 1 @@ -91,3 +76,106 @@ async def train(self): self.tracker.finish() logger.info("Dummy training completed successfully!") + + def _build_max_length_batch(self): + """Original behavior: every sample fully padded to the max context length.""" + num_samples = self.cfg.trainer.train_batch_size * self.cfg.generator.n_samples_per_prompt + uids = [str(i) for i in range(self.cfg.trainer.train_batch_size)] + prompt_token_ids = [ + [random.randint(0, self.tokenizer.vocab_size - 1)] * self.cfg.generator.max_input_length + ] * self.cfg.trainer.train_batch_size + prompt_token_ids = sum( + [[prompt_token_id] * self.cfg.generator.n_samples_per_prompt for prompt_token_id in prompt_token_ids], + [], + ) + max_gen = self.cfg.generator.sampling_params.max_generate_length + response_ids = [[random.randint(0, self.tokenizer.vocab_size - 1)] * max_gen] * num_samples + uids = sum([[uid] * self.cfg.generator.n_samples_per_prompt for uid in uids], []) + dummy_generator_output = { + "prompt_token_ids": prompt_token_ids, + "response_ids": response_ids, + "rewards": [[0] * (max_gen - 1) + [random.randint(0, 1)]] * num_samples, + "loss_masks": [[1] * max_gen] * num_samples, + } + return dummy_generator_output, uids + + def _build_variable_length_batch(self, step: int): + """Stage-3 realistic batch: per-sample total length ~ clamped Normal(mean, std). + + Produces ``train_batch_size * n_samples_per_prompt`` samples. A fixed small prompt + (``dummy_prompt_len``) is prepended to each sample; the remaining tokens are the + response, so the packed micro-batch iterator (``max_tokens_per_microbatch``) sees a + realistic length distribution rather than uniform max-length sequences. + """ + num_prompts = self.cfg.trainer.train_batch_size + n_samples = self.cfg.generator.n_samples_per_prompt + num_samples = num_prompts * n_samples + + prompt_len = int(self.cfg.trainer.dummy_prompt_len) + mean_len = int(self.cfg.trainer.dummy_mean_len) + std_len = int(self.cfg.trainer.dummy_std_len) + min_len = int(self.cfg.trainer.dummy_min_len) + max_len = self.cfg.generator.max_input_length + self.cfg.generator.sampling_params.max_generate_length + + # Seed per step so step 1 (warmup) and step 2 (measured) use the same distribution family + # but different draws; reproducible across runs for a given config. + rng = np.random.default_rng(self.cfg.trainer.dummy_seed + step) + total_lens = rng.normal(loc=mean_len, scale=std_len, size=num_samples) + total_lens = np.clip(np.rint(total_lens), min_len, max_len).astype(int) + # response length = total - prompt (>= 1) + resp_lens = np.maximum(total_lens - prompt_len, 1).tolist() + + vocab = self.tokenizer.vocab_size - 1 + # One shared random prompt per prompt-group (content is irrelevant for a perf test). + prompt_token_ids = [] + for _ in range(num_prompts): + p = [random.randint(0, vocab)] * prompt_len + prompt_token_ids.extend([p] * n_samples) + + response_ids, rewards, loss_masks = [], [], [] + for rl in resp_lens: + response_ids.append([random.randint(0, vocab)] * rl) + rewards.append([0] * (rl - 1) + [random.randint(0, 1)]) + loss_masks.append([1] * rl) + + uids = sum([[str(i)] * n_samples for i in range(num_prompts)], []) + total_tokens = int(sum(total_lens)) + logger.info( + f"FULLCTX_VARLEN step={step + 1} num_samples={num_samples} " + f"total_tokens={total_tokens} mean_len={float(total_lens.mean()):.0f} " + f"std_len={float(total_lens.std()):.0f} min={int(total_lens.min())} max={int(total_lens.max())}" + ) + dummy_generator_output = { + "prompt_token_ids": prompt_token_ids, + "response_ids": response_ids, + "rewards": rewards, + "loss_masks": loss_masks, + } + return dummy_generator_output, uids + + def _collect_peak_memory(self): + """Aggregate per-rank peak/used GPU memory via worker.get_cuda_memory. + + Returns a compact summary {peak_alloc_GB, peak_reserved_GB, max_used_GB} + across all training workers — enough for ablation comparisons without + flooding logs with one entry per rank. + """ + import ray as _ray + + actor_groups = self.dispatch._actor_groups + policy_group = actor_groups.get("policy") + if policy_group is None: + return None + + results = _ray.get(policy_group.async_run_ray_method("pass_through", "get_cuda_memory")) + if not results: + return None + gb = 1024**3 + max_alloc = max(r["allocated"] for r in results) / gb + max_reserved = max(r["reserved"] for r in results) / gb + max_used = max((r["total"] - r["free"]) for r in results) / gb + return { + "peak_alloc_GB": round(max_alloc, 2), + "peak_reserved_GB": round(max_reserved, 2), + "max_used_GB": round(max_used, 2), + }