Table of Contents
Custom Modeling
The trainer only runsprime-rl’s own model implementations, selected from the HF config type. Besides dense Llama, Qwen3 and Qwen3.5, these cover the families below. Other architectures fail at trainer setup.
GPT-OSS uses FlashAttention 4 with learned attention sinks. Training requires SM90 or SM100/SM110 GPUs
and a BF16 checkpoint such as
unsloth/gpt-oss-20b-BF16; the original MXFP4 checkpoints are not supported.
Low-precision training
Dense linear precision and routed-expert precision are configured independently.[trainer.model.quantization] applies only to dense Linear modules:
type = "fp8"— DeepGEMM FP8 blockwise linears (requires SM90+).type = "mxfp8"— torchao MXFP8 linears (requires SM100).recipeismxfp8_rceilormxfp8_rceil_wgrad_with_hp.
[trainer.model.moe.compute] selects routed-expert compute independently:
type = "bf16"(default), withbackend = "torch"(default) or"sonicmoe".type = "deepgemm_fp8"(requires DeepGEMM and SM90+)type = "mxfp8"(requiresprime-kernels, torchao, and SM100)
sonic-moe package (uv sync --extra sonic-moe) for fused BF16 expert computation. The supported model is Qwen3 MoE with the gate_up model fusion enabled. Backend selection requires fused gate/up weights, standard SwiGLU, and bias-free experts; incompatible expert structures raise an error during setup. It uses the same router and local, torch EP, or DeepEP dispatch as other compute backends:
apply_to:
"all"(default) applies the backend to all expert groups."85%"applies it to the first 85% of model layers, rounded down. For a 48-layer model, this selects layers 0–39.[0, 1, 2, 3]selects explicit zero-based model layer indices;[]selects none.
apply_to. Check the selected modules on both sides before comparing trainer and rollout logprobs.
GLM-5.2 adds IndexShare: the DSA sparse-attention indexer runs only on a subset of layers and the remaining layers reuse the cached top-k indices. The trainer reads this schedule from the model’s indexer_types config field and enables the index cache automatically, so no extra config is needed.
Expert Parallelism Backends
[trainer.model.moe.dispatch] selects how routed tokens are dispatched and combined:
torch(default): torch all-to-all withtransport = "bf16"or, when MXFP8 expert compute is selected,transport = "mxfp8"on SM100.deepep: DeepEP custom dispatch/combine kernels. Setnum_smsand optionaltoken_chunk_sizein the same table. Pre-built H100/H200 binaries use CUDA 13.0 and are installed byuv sync --all-extras.
optim.max_norm is set to None automatically.)
Runtime Fusions
model.fusions packs parameters that are always computed together into one tensor, turning several GEMMs into one. Both fusions are on by default:
gate_up— each gated MoE expert’sgate_projandup_projbecome one[num_experts, 2 * intermediate_size, hidden_size]weight, halving the routed-expert grouped GEMMs.qkv— an attention layer’sq_proj,k_projandv_proj(and their biases) become one linear layer.
shard_fused_on_dim1 = true shards fused 2-D weights along dim 1 under FSDP, which makes weight loading and checkpointing zero-copy: fused weights and their optimizer state are read and written in place rather than assembled into a full copy on each rank first. It requires hidden_size to be divisible by the FSDP shard mesh size.
Multimodal Training
Supported Families
The built-in VLM registry covers:Enabling VLM Mode
Add[model.vlm] and bfloat16 dtypes:
{language_model_attr}.layers. automatically.
VLM training requires a registered custom PrimeRL implementation.
Limitations
- Vision encoder frozen by default. The default LoRA targets do not match Qwen3.5 vision modules. Set
freeze_vision_encoder = falseto fine-tune the encoder; this is incompatible with LoRA because LoRA freezes all non-adapter parameters. - bfloat16 mandatory. The trainer config validator refuses any other
optimization_dtype/reduce_dtypefor VLMs — vLLM serves VLMs in bfloat16 and a mismatch breaks the importance ratio. - Higher KL mismatch with multi-image inputs. Expect noisier
mismatch_klthan text-only; this is from minor numerical differences between the trainer’s and vLLM’s image processing. - Images aren’t logged to monitors. Sample logging captures the prompt text but not the actual images.
LoRA Training
LoRA is enabled by adding[model.lora]:
target_modules defaults to a reasonable cross-family set (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj, experts, plus a few latent-projection names for Nemotron). Unknown names are silently ignored, so the defaults work across architectures. Add architecture-specific names to extend coverage (e.g. in_proj / out_proj for Mamba).
LoRA is supported across SFT and RL. NCCL weight broadcast is not supported with LoRA — the default NCCL transport automatically falls back to filesystem when LoRA is enabled. Broadcast dirs of LoRA runs contain the raw adapter (adapter_model.safetensors + adapter_config.json). With LoRA on, the inference server registers the adapter under the model name and serves the base model as <model>-base.
Disaggregated Prefill/Decode Inference
For large MoE serving, splitting prefill and decode onto separate vLLM groups can substantially improve throughput. Pick the prefill:decode ratio based on workload shape:
Example config:
examples/advanced/glm-5.3/swe.toml — full RL run on GLM-5 with P/D disaggregation behind a vllm-router, FP8 inference, and NCCL weight broadcast, paired with an inference config from examples/advanced/glm-5.3/infer/.
Monitor live queue depths to detect imbalance:
signal 11: invalid permissions for mapped object in libucs.so) — reproduced on vLLM 0.22 and 0.23, with/without mooncake, with/without llm-d. Building NIXL against UCX 1.19.x from source is therefore required (not optional) for disaggregated P/D.
third_party/ucx/; the bundled sbatch templates prepend it to LD_LIBRARY_PATH so it overrides the system version. Re-run both commands after every uv sync, since the lock pins the wheel.