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
  • Single-Node
    • RL Placement
    • SFT and Torchrun
  • Multi-Node
• Parallelism Knobs
  • FSDP
  • Expert Parallelism
  • Context Parallelism
  • Activation Checkpointing and Offloading
  • Optimizer Offloading
  • LM Head Chunking
• Memory-Tight Recipe
• SLURM
  • Activation
  • [deployment] Block
  • Examples
  • Custom Templates
• Benchmarking

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 min(world_size // dp_replicate, 8): the FSDP island size, capped at 8. For non-MoE models it resolves to 1 (no-op). "auto" does not look at num_experts or cp; the resolved degree must still be a multiple of cp and divide num_experts, so set it explicitly when it does not. 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).

Context Parallelism

CP shards a single sequence across multiple GPUs along the token dimension — for long-context sequences. Prefer ulysses: it gets the most throughput and is the only style that works for hybrid linear-attention/Mamba models (Qwen3.5, NemotronH) and for VLMs. Each model class declares its supported styles in cp_support; an unsupported cp_style, or a model with no CP support at all, is rejected at setup.

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
cp_style = "ulysses"         # "ring"

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.

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
nodes_per_fsdp_group = 1         # optional — controls FSDP island size

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.

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

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. benchmarks/scripts/run_single_benchmark.py does exactly this and is what the CI benchmark matrix runs.