Repository navigation
Conversation
added 21 commits
September 21, 2026 15:06
This was referenced Oct 4, 2026
faresobeid
added a commit
that referenced
this pull request
Oct 10, 2026
Removes `model.index_cache` (`IndexCacheConfig.topk_freq` and `.topk_pattern`). The trainer still turns on IndexShare from the model's own `indexer_types` schedule (`src/prime_rl/trainer/model.py`), as before. Why: the override changes the trainer only. No inference-side config matches it, so setting it makes the trainer pick different top-k indices than vLLM unless the user also passes the same vLLM `hf_overrides` by hand. The setting appears in no config, example or test. The HF-level `use_index_cache` / `index_topk_freq` / `index_topk_pattern` attributes on `GlmMoeDsaConfig` stay, so a checkpoint's own `config.json` still drives them. Breaking: `[trainer.model.index_cache]` / `[model.index_cache]` now fails validation. Overlap with open PRs: - #3833 rewrites the same `index_cache` lines in `model.py` (`set_model_options(model_config, "index_cache", ...)`). Whichever lands second drops or rebases that block; the conflict is small. - #3879 is based on that code too. - #3616 already leaves `trainer.model.index_cache` unset for the reason above: a trainer frequency override replaces the checkpoint schedule and requests weights the checkpoint does not have. Test plan - `pytest tests/unit/orchestrator tests/unit/test_configs.py -m "not gpu"`: 334 passed. - `pytest tests/unit/utils/test_weights.py` (includes `GlmMoeDsaForCausalLM`): 7 passed. - `ruff check` and `ruff format --check` on the changed Python files: clean. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Breaking config validation for anyone using `index_cache`, and it changes sparse-attention index scheduling to follow the checkpoint only—important for trainer/inference logprob alignment on DSA models. > > **Overview** > **Removes the trainer-only `[trainer.model.index_cache]` override** (`topk_freq` / `topk_pattern`). Configs that still set it will fail validation. > > IndexShare for GLM-5.2-style DSA models is unchanged in behavior: when the checkpoint’s `indexer_types` includes `"shared"` layers, the trainer still sets `use_index_cache` on the HF config and logs the full/shared layer counts. It no longer applies manual `index_topk_freq` / `index_topk_pattern` or clears `indexer_types`, which could diverge from vLLM rollouts. > > Docs no longer mention manually overriding the schedule via `[trainer.model.index_cache]`. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit bf75c92. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY -->
samsja
added a commit
that referenced
this pull request
Oct 10, 2026
Replaces the three switches `model.fsdp_cpu_offload`, `model.optim_cpu_offload` and `model.full_offload` with one field: ```toml [model] offload = "optimizer" # default; also "full", "none" ``` This applies to both the RL trainer and SFT. The `fsdp` switch is not migrated: the mode is removed outright (rationale below). The `[model.full_offload] numa_bind` knob is dropped: no recipe, deployment, template or doc sets it. Full offload now always pins each rank's CPUs to its GPU's NUMA node. That was already the default. Ranks whose GPU reports no NUMA node still skip binding with a warning. On hosts without GPU sysfs topology (`/sys/bus/pci/devices/<bus>/numa_node`), binding is no longer optional: full offload fails there instead of being switched off. Why: the three switches could only be used one at a time, enforced by `cpu_offload_mutual_exclusion`. Because `optim_cpu_offload` defaulted to `true`, enabling full offload also meant writing `optim_cpu_offload = false`, otherwise validation failed (see the old `configs/advanced/deepseek-v4-flash/sft.toml` comment). With one value the conflicting combinations cannot be expressed, so that validator is deleted. The default behaviour does not change: optimizer-state offload stays on. Why drop `fsdp`: the mode was the stock torch FSDP2 `CPUOffloadPolicy` with no in-tree config, example, or test user; Muon rejected it and it broke NCCL weight broadcast. It was kept only for long-context runs whose fp32 masters/moments no longer fit on GPU (the reason it was restored in #3052); the `full` mode added in #3234 keeps FP32 masters, moments, and gradients in host RAM, covering that memory pressure instead. Gone with it: the `CPUOffloadPolicy` selection in `setup_fsdp`, CPU-side model init in `setup_model`/`load_dcp_from_hf`, `_move_buffers_to_cuda`, and the fsdp leg of the Gloo special case. Changes: - `ModelConfig.offload: Literal["none", "optimizer", "full"] = "optimizer"` replaces the three fields. `OptimizerInBackwardOffloadConfig` (and with it `numa_bind`) and its bool-normalizing `BeforeValidator` are gone. - Validators in `TrainerConfig` and `SFTConfig`: - full offload accepts only AdamW or SignSGD and disables clipping, now reading `offload == "full"`; - the `Muon optimizer does not support FSDP CPU offload` validators are deleted with the mode. - Runtime: - `setup_optimizer(..., offload=...)` replaces `cpu_offload` + `full_offload_config`, and its "both enabled" check is gone. - `FullCPUOffloadOptimizer` drops its unused `offload_config` argument. - `setup_full_cpu_optimizer_offload()` takes no arguments and always binds. - `model.py`: FSDP always uses the default `OffloadPolicy()`; CPU-side model init in `setup_model`/`load_dcp_from_hf` and `_move_buffers_to_cuda` are gone. - Tomls: `glm-4.5-air/swe-budget.toml` and `ci/integration/reverse-text-sft/full-offload-resume.toml` → `offload = "full"`; `glm-5.3/sft/h200/16k-12-node.toml` → `offload = "none"`; `glm-5.3/sft/h200/base.toml` → `offload = "optimizer"`. Also updated `docs/scaling.md`, the `deepseek-v4-flash/sft.toml` comment, and `skills/training/start-run/SKILL.md`, which documented the removed fields. Breaking: the old keys now fail with `extra_forbidden`, and `offload = "fsdp"` fails the `Literal`. There is no shim. Interaction with #3849 (one trainer-common config base): #3849 moves the full-offload and Muon/FSDP validators from `TrainerConfig`/`SFTConfig` onto `BaseTrainerConfig`. This PR edits the full-offload validator's condition and deletes the Muon/FSDP one, so the conflict is textual only: keep #3849's single copy of the full-offload validators with `self.model.offload == "full"` and drop its Muon/FSDP copy. Whichever lands second takes a small rebase. Other open PRs that set or read the old keys need a rebase: - tomls: #3883, #3879, #3766, #3765, #3729, #3629, #3627, #3616, #3544 - code: #3885 (`optim_cpu_offload`), #3739 (`full_offload` validators), #3724, #3719, #3699, #3610, #3487 (`fsdp_cpu_offload` / `enable_gloo`) Test plan - `pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train/test_cpu_adam.py tests/unit/train/test_state_offload.py -m "not gpu"`: 336 passed, 1 skipped (Qwen3-VL snapshot not cached on this host), 4 deselected (GPU; this host has no GPU). The existing offload config tests are updated to the new field. `test_load_configs` covers the 4 edited tomls. - `cli(TrainerConfig|SFTConfig, ...)`: `--model.offload full|none` and the default (`optimizer`) resolve as expected; `optim_cpu_offload` and `offload = "fsdp"` are rejected. - `import prime_rl.trainer.rl.train, prime_rl.trainer.sft.train` works. - `ruff check` and `ruff format --check` on the changed Python files: clean. - Not run: GPU or multi-node training with `offload = "full"`. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Breaking TOML/API surface touches distributed training memory paths (FSDP offload, optimizer-in-backward, Gloo); behavior should match prior modes when mapped correctly, but mis-migrated configs will fail at parse time. > > **Overview** > **Breaking config change:** `model.fsdp_cpu_offload`, `model.optim_cpu_offload`, and `model.full_offload` are replaced by a single **`model.offload`** setting (`"optimizer"` default, or `"full"`, `"fsdp"`, `"none"`). RL and SFT share the same field; example TOMLs and scaling docs are updated accordingly. > > The mutual-exclusion validator between the old flags is removed because only one mode can be expressed. **`[model.full_offload]` / `numa_bind` is dropped** — `"full"` always runs NUMA binding and CPU optimizer thread setup via parameterless `setup_full_cpu_optimizer_offload()`. Runtime paths (`setup_optimizer`, FSDP CPU offload in `model.py`, Gloo when offload is `"full"` or `"fsdp"`) key off the new enum; `FullCPUOffloadOptimizer` no longer takes an offload config object. > > Default training behavior is unchanged: optimizer-state CPU offload remains the default (`"optimizer"`). Old TOML keys are not aliased and should fail validation. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 35da73a. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: sami jaghouar <sami.jaghouar@gmail.com>
samsja
pushed a commit
that referenced
this pull request
Oct 19, 2026
Disaggregated P/D with per-role KV offload, so P/D + router replay + sampling replay + KV offload can be configured from the stack alone. **Stacked on #3831.** Supersedes the parts of #3616 needed here. **Per-role KV offload** - `inference.deployment.prefill_kv_cache_offload` / `decode_kv_cache_offload`: default `"inherit"` (use `inference.kv_cache_offload`), `None` (`"None"` in TOML) disables offload for that role, or a full offload config. - Mooncake `cpu.num_bytes = 0` connects a role to the shared store without starting a `mooncake_client` on its nodes (no storage contributed; the vLLM process still uses its staging buffer). Native offload still needs a positive size. - Under P/D the Mooncake connector gets `save_decode_cache = true`: decode also writes the KV it generates, so the next turn's prefill can load the previous completion. Existing keys are skipped by the store, so decode re-saving the prompt costs one existence check. Writes are not clamped; prefill loads are clamped at `routed_experts_prompt_start` by #3831, which keeps multi-turn completion blocks loadable. - `_mooncake_store.sh.j2` runs per role: one master on the head node; storage clients only on roles with capacity. Roles must agree on a Mooncake disk path (one master, one `-root_fs_dir`). - Mooncake keys carry model + TP/PCP/DCP/PP rank + content hash only, so both roles must keep the same KV layout (TP size, `kv_cache_dtype`, block size) to share the store. Documented. **Launcher.** One engine config per role (`inference-{prefill,decode}.json`, `InferenceConfig.for_pd_role` / `write_pd_configs`); each node picks its file by `$ROLE`. Inference nodes always run `ulimit -l unlimited` (NIXL/RDMA registers pinned host pools). Example (all storage on prefill; decode connects without capacity): ```toml [inference.deployment.prefill_kv_cache_offload] type = "mooncake" [inference.deployment.prefill_kv_cache_offload.cpu] num_bytes = 1_000_000_000_000 [inference.deployment.decode_kv_cache_offload] type = "mooncake" [inference.deployment.decode_kv_cache_offload.cpu] num_bytes = 0 ``` Resulting connectors: prefill `MultiConnector[Nixl, MooncakeStore(save_decode_cache)]`, decode `MultiConnector[Nixl, MooncakeStore(save_decode_cache)]`. Left out (not needed for this use case): GLM-5.3 example configs, monitoring dashboard and tooling, trainer DRAM contribution to the store, NUMA/UCX per-rank binding, router timeout/retry fields, lease TTL and staging-buffer fields, explicit `kv_transfer_config` passthrough. ## Checks run - `pytest -m "not gpu" tests/unit/test_configs.py tests/unit/inference`: 223 passed. Ruff: clean. Test: `test_pd_role_connectors`. - Dry run of `examples/advanced/glm-5.3/swe.toml` + router replay + `top_p = 0.95` (sampling masks) + prefill Mooncake (1 TB + disk) + decode Mooncake (`cpu.num_bytes = 0`): per-role JSON configs build `[Nixl, MooncakeStore]` for both roles; `bash -n` on the generated `rl.sbatch` passes; only prefill nodes start a storage client. ## GPU validation (2026-10-06, vLLM 0.31.0, 1x 8xH200, configs written by `write_pd_configs`, router v0.2.2) - Qwen3-30B-A3B TP1, `VLLM_BATCH_INVARIANT=1`, 1P1D NIXL, prefill Mooncake (RDMA, 40 GB) + decode Mooncake `cpu.num_bytes = 0` with `save_decode_cache`, router replay + sampling masks, 6 two-turn conversations (turn 2 start = `len(p1 + c1) - 1`) after evicting prefill's GPU cache, vs a colocated no-offload reference: **6/6 identical** (tokens, 264 rows, start, masks). Prefill external hits **16880** = sum of block_floor(start), vs 16512 when only prefill writes (prompt-only maximum, earlier run with the same conversations): the turn-1 completion blocks were loaded from decode-saved KV. No errors. - 1P1D without offload, same configs path: 10/10 and 20/20 (concurrency 10, fresh salt) identical to colocated. Not validated: multi-node Mooncake role launch on SLURM (template rendered and `bash -n` only). <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes disaggregated inference launch paths, KV connector composition, and Mooncake SLURM wiring; misconfiguration could break multi-turn prefix reuse, though config validation and unit tests cover the main connector shapes. > > **Overview** > Adds **per-role KV cache offload** for disaggregated prefill/decode via `inference.deployment.prefill_kv_cache_offload` and `decode_kv_cache_offload` (`inherit`, disabled, or a full offload block). Resolvers expose `kv_cache_offload_by_role`, `for_pd_role()`, and emit **`inference-{prefill,decode}.json`** so each node loads the right connector set. > > **Mooncake** now allows `cpu.num_bytes = 0` to join the shared store without starting a storage client (native offload still requires positive CPU bytes). Under P/D, the Mooncake connector gets **`save_decode_cache`** so decode-written KV is available on the next prefill turn; validation requires a **single disk path** across Mooncake roles sharing one master. > > SLURM templates take a **`mooncake_offloads`** map (replacing global offload flags), role-aware `_mooncake_store.sh.j2`, and **`inference-$ROLE.json`** in P/D jobs; inference nodes always run **`ulimit -l unlimited`** for NIXL/RDMA pinned memory. Docs cover the typical prefill-storage / decode-connect pattern and KV layout constraints across roles. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit ee5dcc6. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY -->
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
P/D deployments need different KV offload capacities: prefill contributes Mooncake storage while decode stays connected to the same pool and reserves host RAM for HiSparse. Add per-role offload configuration and optional trainer storage contributions, with a GLM-5.3 B300 RL example.
deployment.prefill_kv_cache_offloadanddeployment.decode_kv_cache_offloadsupport native and Mooncake backends. Overrides inherit the common config; explicitNonedisables offload for a role. Mooncake acceptscpu.num_bytes=0to retain its connector without starting a storage client. Native offload requires positive capacity. Explicitvllm.kv_transfer_configchains retain precedence over automatic connector composition.Multi-node RL can contribute independently sized trainer DRAM using
[deployment.trainer_mooncake], withnum_bytesand optionaldevice_name. Job-scoped clients join the inference pool, excluding the orchestrator host; they add no trainer KV connector or disk tier. Mooncake lease settings must agree across roles sharing a master. Router request timeout and log level are configurable in standalone and Slurm launches.The base
examples/advanced/glm-5.3/infer/pd-nvfp4-b300.tomlruns four trainer nodes plus 1P+1D inference. It fixes the GLM-5.3 BF16 checkpoint revision, uses online per-token NVFP4 4over6 for the first 85% of layers on trainer and inference, and transfers weights over NCCL. Inference uses EP8, FP8 KV, FlashInfer NVLink one-sided, decode HiSparse/index cache, sticky least-loaded routing and no EPLB. Prefill contributes 2,560 GiB Mooncake DRAM per node; decode contributes zero and reserves 300 GiB per GPU for HiSparse. The base trainer uses EP8/CP4, AdamW with optimizer-state CPU offload, selective activation checkpointing and activation offloading.The Scale-SWE quarter-zero training source uses Bash without an agentic judge, group size 16, top-p 0.95/top-k 512 sampling with sampling-mask replay, and the default GRPO algorithm. The debug algorithm override is removed. W&B logging, 1,000 maximum steps, fresh-run cleanup and disabled automatic dashboard startup are configured; credentials remain outside the repository.
The current
pd-nvfp4-b300-4p2d-ram.tomloverlay adds:glm53-pd-trainandint4-syn-gen-bash.TORCHINDUCTOR_COMPILE_THREADS=1, retaining compilation and caching without a Torch patch.Launch the current configuration with:
The existing
pd-nvfp4-b300-2p1d-disk.tomloverlay remains available for 2P+1D with filesystem offload, CP8/SignSGD, batch size 128 and concurrency 2,048. The optionaltools/inference_monitoring.pystarts job-scoped Prometheus/Grafana with 15-second collection, discovers inference roles from the launch log and exits with the Slurm job.Validation: the trainer-storage addition passed all 214 config tests, Ruff, generated shell syntax checks and an isolated launcher exercise covering orchestrator exclusion, storage capacity, NIC selection and shared metadata addresses. Subsequent targeted role-offload checks passed. For the current Bash config update, the linear penalty retains all three weights at 0.2; runtime labels, GRPO, batch size 256 and concurrency 3,600 are unchanged. The previous composed configuration passed launcher dry-run validation, and this update changes only the harness ID back to Bash. Local verifiers patches restore streaming with a 60-minute Bash client timeout, apply #2650 and add bounded retries during stream consumption using the buffered training replay guard. These dependency patches are not included in this config commit. All 13 focused transport/replay checks passed; production recovery remains to be measured.
Stacked on #3511 (on #3513). Related to #2949; this provides general per-role offload configs and zero-storage Mooncake connections.