Skip to content

scaling.md

Latest commit

 

History

History
318 lines (215 loc) · 20.2 KB

scaling.md

File metadata and controls

318 lines (215 loc) · 20.2 KB
Raw

Scaling

This page covers how to scale prime-rl from a single GPU to a 1000-GPU cluster: single-node and multi-node deployments, FSDP / expert parallelism / context parallelism, and throughput benchmarking. See Training for detailed documentation of the trainer configuration and Inference for the inference configuration.

Table of Contents

Single-Node vs. Multi-Node Deployment

The rl, sft, and inference entrypoints all accept a [deployment] block (type = "single_node" or "multi_node") that picks how the trainer / orchestrator / inference processes are placed across hardware. Single-node runs locally; multi-node currently goes through SLURM — the launcher writes an sbatch script that places inference replicas, the orchestrator, and the trainer with the right rendezvous endpoints, IPs, ports, and shared-filesystem paths wired in.

Single-Node

RL Placement

rl defaults to 1 trainer GPU and 1 inference GPU. To give inference 6 GPUs with data parallelism and the trainer the remaining 2 on an 8-GPU node:

uv run rl @ rl.toml \
  --deployment.num-infer-gpus 6 \
  --deployment.num-train-gpus 2 \
  --inference.vllm.data-parallel-size 6

The launcher allocates GPUs in order from CUDA_VISIBLE_DEVICES (or all visible GPUs): inference first, trainer next. To target a specific physical subset, pin CUDA_VISIBLE_DEVICES before launching.

For quick A/B ablations on the same node, run two RL instances side-by-side in separate shells, each pinned to half the GPUs and a separate inference port, and track both in one dashboard:

# shell 1, GPUs 0–1, default port 8000
CUDA_VISIBLE_DEVICES=0,1 uv run rl @ rl.toml --run.name exp1

# shell 2, GPUs 2–3, port 8001
CUDA_VISIBLE_DEVICES=2,3 uv run rl @ rl.toml \
  --inference.server.port 8001 \
  --orchestrator.model.client.base-url http://localhost:8001/v1 \
  --run.name exp2

# watch both side by side
uv run dashboard outputs

SFT and Torchrun

uv run sft handles distributed launch internally. To scale from 1 to N GPUs, set the deployment GPU count (or just let it pick up WORLD_SIZE). For non-default layouts, the manual equivalent is:

uv run torchrun \
  --nproc-per-node 8 \
  --local-ranks-filter 0 \
  src/prime_rl/trainer/sft/train.py @ sft.toml

--local-ranks-filter 0 keeps console output to rank 0 only; per-rank stdout/stderr is still captured in <run_dir>/logs/latest/trainer/torchrun/.

Multi-Node

Multi-node deployments (RL or SFT) are launched via SLURM — set [deployment] type = "multi_node" plus the matching [slurm] block, and the launcher writes the sbatch script that places inference, orchestrator, and trainer across the requested nodes with the inter-process wiring set up correctly. See SLURM § Examples for full configs.

Parallelism Knobs

FSDP

FSDP2 is the default model sharding strategy. By default the trainer fully shards parameters, gradients, and optimizer state across the data-parallel mesh. Tweakable knobs:

Knob Effect
trainer.model.dp_replicate Number of dimensions to replicate instead of shard. Set to 2 to run 2-way DP replication × FSDP sharding within each replica — useful for very large clusters where pure FSDP communication dominates.
trainer.model.reshard_after_forward If true (default), parameters are resharded after the forward pass to free memory; the backward pass re-gathers. Set false to keep params resident — faster but more memory.
trainer.model.offload CPU offloading: optimizer (default) offloads optimizer state between steps; full offloads gradients, FP32 masters, and optimizer state and runs the optimizer (AdamW or SignSGD) on CPU during backward; none disables offloading.

Expert Parallelism

EP shards MoE expert weights across the EP mesh, dramatically reducing the FSDP communication volume per layer and improving the training throughput.

ep defaults to "auto". At startup it loads the model config to check whether the model is MoE. For MoE models it resolves to the largest degree up to 8 that divides both the FSDP island size (world_size // dp_replicate) and num_experts and is a multiple of cp (1 if none is). For example, 6 GPUs with 128 experts resolve to ep=2. For non-MoE models it resolves to 1 (no-op). Set ep to an explicit integer to override:

[trainer.model]
ep = 8  # explicit EP degree; must divide num_experts

[trainer.model.moe.dispatch]
type = "torch"
transport = "bf16"

For DeepEP, set type = "deepep" and tune num_sms plus optional token_chunk_size in the same dispatch table. Routed-expert precision is selected separately with [trainer.model.moe.compute] (bf16, deepgemm_fp8, or mxfp8).

Mega MoE (type = "mega_moe") fuses dispatch, the SwiGLU expert MLP, and combine into one persistent kernel per direction, overlapping the expert-parallel communication with the GEMMs. precision selects the expert GEMMs: bf16 (default, SM90+), mxfp8 (SM100), fp8 (blockwise FP8 forward and backward, SM90), or fp8_bwd (BF16 forward, blockwise FP8 backward, SM90); [trainer.model.moe.compute] does not apply to those layers. The kernels come from prime-mega-moe, Prime Intellect's closed-source DeepGEMM fork, built as a deep_gemm wheel that replaces the public one (uv pip install --no-deps <wheel> after syncing). Without it, selecting mega_moe fails at startup with a message naming the missing kernels. Mega MoE needs EP > 1, nvcc 13 on CUDA_HOME for the runtime JIT, gated experts without biases, hidden_size divisible by 256 and moe_intermediate_size by 128.

The kernels are persistent grids that synchronize across ranks, so an NCCL kernel on another stream (FSDP all-gathers, gradient reductions) has to fit in the SMs they leave free, or the two deadlock. num_reserved_sms is how many SMs they leave free. NCCL_MAX_CTAS caps the CTAs (one per SM) a NCCL kernel launches. In the worst case each NCCL kernel uses all NCCL_MAX_CTAS, so at most num_reserved_sms // NCCL_MAX_CTAS NCCL kernels can run concurrently. FSDP overlaps an all-gather and a reduce-scatter, and HSDP (dp_replicate > 1) adds an all-reduce, so the trainer sets NCCL_MAX_CTAS to num_reserved_sms // 2 (// 3 under HSDP) when it is unset, and refuses to start if it is set higher than that. The default 16 reserved SMs (8 CTAs, 5 under HSDP) are enough for FSDP's collectives; every reserved SM is taken from the GEMMs, so raise num_reserved_sms only if NCCL needs more. The symmetric buffer is sized once at startup from max_tokens_per_rank, which defaults to a rank's share of the micro batch under context parallelism (model.seq_len / model.cp for RL, data.micro_batch_size * data.seq_len / model.cp for SFT); a smaller value is rejected at startup.

[trainer.model.moe.dispatch]
type = "mega_moe"
precision = "bf16"     # or "mxfp8" (SM100), "fp8" / "fp8_bwd" (SM90)
num_reserved_sms = 16  # SMs left free for NCCL; NCCL_MAX_CTAS defaults to this // 2 (// 3 under HSDP)

[trainer.env_vars]
CUDA_HOME = "/usr/local/cuda-13.0"
PATH = "/usr/local/cuda-13.0/bin:${PATH}"

Context Parallelism

CP shards a single sequence across multiple GPUs along the token dimension — for long-context sequences. It uses Ulysses: an all-to-all turns sequence-sharded Q/K/V into head-sharded Q/K/V over the full sequence, attention runs locally, and a second all-to-all shards the output back. Hybrid linear-attention/Mamba models (Qwen3.5, NemotronH) redistribute their recurrent layers the same way. The sparse-attention models (DeepSeek-V4, GLM-5 DSA) keep queries sequence-sharded and all-gather their keys instead, which needs no head sharding.

Ulysses head-shards Q/K/V, so the CP degree must divide num_attention_heads. GQA models with fewer KV heads than the CP degree (e.g. NemotronH: 32 query heads, 2 KV heads) are supported via KV-head replication; the CP degree must then be a multiple of num_key_value_heads. Nemotron-H redistributes its Mamba activations between sequence- and head-parallel layouts in the model-owned Mamba path; the CP degree must divide mamba_num_heads and n_groups.

[trainer.model]
attn = "auto"                # auto = FA3 on Hopper, FA4 on datacenter Blackwell, FA2 otherwise; or flash_attention_2/3/4
cp = 2                       # CP degree

Activation Checkpointing and Offloading

Knob Memory ↓ Throughput ↓
trainer.model.ac large ~25%
trainer.model.ac.mode = "selective" medium small
trainer.model.ac_offloading extra a bit more

AC and AC offloading are enabled by default (full mode). Both AC modes retain stateful MoE routing updates and DeepEP communication that cannot safely replay. Selective AC uses the same transformer-block boundaries and additionally retains distributed communication, expensive matrix multiplications, grouped GEMMs, and supported attention kernels:

[trainer.model.ac]
mode = "selective"

Set targets to operator names (for example, "aten::mm") or namespaces (for example, "prime_rl_collectives") to replace the default selective targets. Correctness-required operations remain retained.

Activation offloading still applies to tensors saved by autograd, but tensors retained by the checkpoint policy remain on the accelerator.

ac_offloading is also on by default with max_inflight_activations = 5. We've observed this feature to be very effective, lowering the peak memory usage by 30-40% in some cases, while only losing ~3-5% of throughput. To disable either, set model.ac = "None" or model.ac_offloading = "None".

Optimizer Offloading

State-only optimizer offload is enabled by default (model.offload = "optimizer"). For full offload, set model.offload = "full"; this keeps BF16 compute weights on GPU and runs CPU optimizer chunks as gradients become ready during backward. Full offload only supports AdamW and SignSGD (optim.type = "sign_sgd") and disables gradient clipping. SignSGD is stateless, so it halves the host RAM footprint versus AdamW (8 instead of 16 bytes per parameter: FP32 master + FP32 accumulated gradient, no moments).

LM Head Chunking

The vanilla LM head materializes a [batch * seq, vocab] logits tensor on every step — a major memory tax when the vocabulary is large (often >100K). fused_lm_head_token_chunk_size swaps in a custom fused linear + logprob/entropy kernel that streams through chunk_size tokens at a time, avoiding the materialization. It defaults to 8192 for RL training:

[trainer.model]
fused_lm_head_token_chunk_size = 8192       # default
# fused_lm_head_token_chunk_size = "disabled"  # vanilla LM head

Drop the chunk size further when peak memory is still tight (e.g. with very long sequences); raise it to amortize kernel-launch overhead. SFT training silently disables this (not supported yet).

Memory-Tight Recipe

The kitchen-sink config for fitting large MoE on limited GPUs at acceptable throughput. AC, AC offloading, compile, fused LM head chunking, optimizer offload, and EP auto-resolution are on by default — only CP needs to be set explicitly (and EP overridden if auto-resolution is not desired):

[trainer.model]
ep = 8
cp = 2

[trainer.model.ac]
freq = 1

[trainer.model.ac_offloading]
max_inflight_activations = 1

The defaults already cover: fused LM head chunking (8192), torch.compile (fullgraph=False), AC (full mode), AC offloading (max_inflight_activations=5), and optimizer CPU offload. Walks through every memory lever in order: FSDP+EP shard the weights, CP shards the activations along the token dim, AC + AC offloading shrink the activation footprint, fused LM head chunks the loss, torch.compile reduces fragmentation, optim offload moves Adam state off GPU. Apply selectively — each knob has a throughput cost.

Max Trainer Throughput Recipe

The trainer defaults are tuned for memory, not speed: optimizer CPU offload, AC offloading, and torch.compile without fullgraph let large models fit on few GPUs, at a throughput cost. When you have enough GPUs to keep the model state and activations on GPU, turn them off:

[model]
offload = "none"            # optimizer states stay on GPU
ac_offloading = "None"      # activations stay on GPU

[model.compile]
fullgraph = true            # compile each block as a single graph

Then add nodes instead of CP (cp = 1) and run one full-length sequence per GPU (micro_batch_size = 1). Keep full AC unless selective AC fits. Router gates share the block's FSDP unit and use BF16 operands with FP32 logits.

Examples:

SLURM

The rl, sft, and inference entrypoints all submit to SLURM when a [slurm] table is present — there's no separate entrypoint.

The prime-rl checkout and its uv venv must live on a shared filesystem visible to every node. The generated sbatch script runs a single uv sync --all-extras --all-packages on the batch node (not once per node), so all ranks share that one environment — a node-local venv would leave the other nodes stale.

Compile and JIT caches are node-local by default: the launcher sets TRITON_CACHE_DIR=/tmp/triton-<uid> for every component, and VLLM_CACHE_ROOT=/tmp/.vllm-cache-<uid> and FLASHINFER_WORKSPACE_BASE=/tmp/.flashinfer-cache-<uid> for inference. Ranks writing one cache on a shared filesystem can crash or hang. Override them in [env_vars] or a component's env_vars.

Activation

A SLURM config is usually a thin overlay that adds [slurm] (and [deployment] for multi-node) on top of a base config. Configs are composed left-to-right via the @ CLI syntax — see Configuration § TOML Composition:

# my_slurm.toml
output_dir = "/shared/outputs/my-rl"

[slurm]
job_name = "my-rl-run"

Launch:

uv run rl @ base_rl.toml @ my_slurm.toml             # submits via sbatch
uv run rl @ base_rl.toml @ my_slurm.toml --dry-run   # writes launcher/rl.sbatch + resolved config, exits

Every SLURM entrypoint stores its generated scripts and coordination files under <run_dir>/launcher/. SLURM stdout and stderr go to launcher/logs/. Local launches do not create launcher/.

[deployment] Block

[deployment] is a discriminated union picked by type — single_node or multi_node for RL/SFT, with an extra disaggregated variant for inference. RL multi-node:

[deployment]
type = "multi_node"
num_train_nodes = 2
num_infer_nodes = 1              # optional when inference.deployment defines the node topology
gpus_per_node = 8                # default

SFT multi-node:

[deployment]
type = "multi_node"
num_train_nodes = 2
num_infer_nodes = 1  # required only for online evals
gpus_per_node = 8

Examples

Full multi-node configs ship under examples/advanced/:

  • nemotron-3-super/swe.toml — 4 trainer + 1 inference node RL on a 120B hybrid-Mamba MoE with NCCL weight broadcast.
  • glm-5.3/ — large-scale and P/D-disaggregated inference across the GLM-5 family.
  • glm-4.5-air/swe-2-node.toml — train a 100B+ MoE on only two nodes: a single-node sign-SGD trainer (full optimizer offload) paired with one inference replica at batch size 64. An overlay on the shared swe-budget.toml base — launch with uv run rl @ swe-budget.toml @ swe-2-node.toml.

For inference-only multi-node, set [deployment] type = "multi_node" on an inference TOML — each node runs an independent vLLM replica (TP and DP must fit within one node), fronted by a single global router on node 0. Point clients at the router URL the launcher prints.

ModelExpress weight broadcast

See ModelExpress for the modelexpress weight transport.

NIXL weight broadcast

Set [weight_broadcast] type = "nixl" to use receiver-driven NIXL weight transfer. Before the first SLURM run, install the NIXL/UCX build and the ModelExpress service binaries on the shared filesystem:

bash scripts/install_nixl_from_source.sh
uv pip install --reinstall --no-deps deps/nixl_cu13-*.whl
bash scripts/install_modelexpress.sh

The generated job starts a job-scoped ModelExpress server and Redis backend on the trainer head node and passes that address to every component. To use an existing service, set slurm.launch_modelexpress = false and configure weight_broadcast.host and weight_broadcast.port.

The launcher requires the CUDA and InfiniBand transports from third_party/ucx. Each NIXL process selects the active InfiniBand port nearest its GPU; an explicitly configured UCX_NET_DEVICES takes precedence. Inference ranks start their pulls at different trainer ranks so concurrent workers distribute traffic across all available source rails.

ModelExpress exchanges peer metadata during startup. Weight updates reuse prepared NIXL requests, post every trainer-rank read in a transfer group concurrently, and use versioned NIXL notifications for source readiness and buffer credits.

By default, the trainer and inference worker each allocate one transfer arena. Set weight_broadcast.overlap_transfer_and_replay = true to allocate two arenas on both sides and replay one weight group while receiving the next. The additional arena is the size of the largest transfer group per GPU; allocation errors are reported instead of silently disabling overlap.

Custom Templates

For unusual partitions, module loads, or environment setup, supply your own Jinja2 template:

uv run rl @ my_config.toml --slurm.template-path path/to/my_template.sbatch.j2

The default templates live under src/prime_rl/templates/ — copy one as a starting point.

Benchmarking

To benchmark a parallelism config before committing a multi-day run, run a short training with fake data and a step cap, and read throughput / MFU / step time / peak memory from the logs:

# SFT trainer alone
uv run sft @ sft.toml --data.type fake --max-steps 4

# RL trainer alone (no inference involved) -- launch it under torchrun, like the `rl` launcher does
uv run torchrun --nproc-per-node=8 src/prime_rl/trainer/rl/train.py @ train.toml --data.fake --max-steps 4

With --data.fake.generate-samples, fake RL data packs random-length samples into each row; add --data.fake.doc-len N to make every sample N tokens instead (the last one cut to fill the row), e.g. two 262k-token rollouts in a 524k-token row.

Every step logs Throughput, MFU, and Peak Mem. to the console. For machine-readable numbers, the file monitor writes monitors/file/metrics.jsonl under the run's output directory by default (monitors.file); aggregate perf/throughput, perf/mfu, time/step, and perf/peak_memory from the run's metrics.jsonl — skip the first step, it is warmup.