Skip to content

feat: add orchestrator debug mode (no-inference, no-trainer)#2867

Open
samsja wants to merge 1 commit into
mainfrom
feat/sane-memory-defaults
Open

feat: add orchestrator debug mode (no-inference, no-trainer)#2867
samsja wants to merge 1 commit into
mainfrom
feat/sane-memory-defaults

Conversation

@samsja

@samsja samsja commented Jun 24, 2026

Copy link
Copy Markdown
Member

Summary

Adds an orchestrator debug mode for running without inference or a trainer, plus a deterministic test environment for memory/pipeline ablations.

Changes

Orchestrator debug mode (orchestrator.py, utils.py, watcher.py, client.py)

Adds [orchestrator.debug] with two flags:

  • no_inference: Uses a NoOpInferencePool (no model calls), clears the renderer, skips renderer/VLM/pool-size validators, and skips NCCL weight broadcast init and teacher inference pool setup.
  • no_trainer: Uses filesystem rollout transport and a DebugWeightWatcher that advances policy.version locally so dispatch gating and off-policy cancellation still run.

update_dispatch_gate is now async and drives local policy updates in no-trainer mode. RSS memory logging via psutil is added after each step when debug mode is enabled.

no-inference-env (tests/debug_envs/no_inference_env/)

A deterministic test environment that synthesizes heavy multi-turn trajectories (configurable sequence length, turns, optional routed experts, rollout delay) without requiring vLLM. Useful for memory/pipeline ablations.

Docs (skills/training/start-run/SKILL.md)

Documents uv run orchestrator ... --debug.no-inference --debug.no-trainer and the no-inference env configuration.


Note

Medium Risk
Touches orchestrator startup, weight watching, and dispatch gating on the training hot path, but debug flags default off; production behavior should be unchanged unless debug is enabled.

Overview
Adds orchestrator.debug with no_inference and no_trainer so you can run the orchestrator alone for pipeline/memory ablations without vLLM or a trainer.

With no_inference, setup uses a NoOpInferencePool, clears the renderer, skips teacher/NCCL init, and relaxes renderer/VLM/pool-size validators. With no_trainer, rollout transport is forced to filesystem and a DebugWeightWatcher bumps policy.version locally after each shipped batch (via async update_dispatch_gate) so dispatch gating and off-policy cancellation still behave like production. Debug runs also log process RSS (orchestrator + children) with psutil after each step.

Ships no-inference-env: a deterministic env that builds heavy multi-turn trajectories (long seqs, tools, optional routed-expert payloads, configurable rollout delay) without calling the model. Wired into the envs extra and documented for uv run orchestrator ... --debug.no-inference --debug.no-trainer.

Reviewed by Cursor Bugbot for commit f2c79f7. Bugbot is set up for automated code reviews on this repo. Configure here.

@samsja samsja force-pushed the feat/sane-memory-defaults branch from 0f3d5d1 to df0f3d1 Compare June 24, 2026 18:52
@samsja

samsja commented Jun 24, 2026

Copy link
Copy Markdown
Member Author

Closing in favor of #2868 which now includes both the memory-saving defaults and the EP auto changes.

@samsja samsja closed this Jun 24, 2026
@samsja samsja reopened this Jun 24, 2026
Comment thread src/prime_rl/orchestrator/watcher.py
Comment thread packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py Outdated
Comment thread tests/debug_envs/no_inference_env/no_inference_env.py
mikasenghaas added a commit that referenced this pull request Jun 25, 2026
- Drop router replay (trainer.enable_router_replay + inference.enable_return_routed_experts):
  mutually exclusive with inference.kv_cache_offload (rl.py validator — external KV cache hits
  don't carry routed-expert decisions).
- Enable native KV-cache offloading with a 128GB CPU tier (extends the prefix cache).
- FP8 trainer (DeepGEMM blockwise linear/MoE) — impl=custom is already set.
- Use the bash_edit harness (bash + local edit tool) for train + eval, replacing pure bash.
- Bump max_inflight_rollouts 384 -> 512.
- Keep LM-head token chunking + activation checkpointing explicit at the values that become
  trainer defaults in #2867 (not merged yet, so removing them would disable the features).
- Drop the -lp length-penalty variant.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
mikasenghaas added a commit that referenced this pull request Jun 25, 2026
- Drop router replay (trainer.enable_router_replay + inference.enable_return_routed_experts):
  mutually exclusive with inference.kv_cache_offload (rl.py validator — external KV cache hits
  don't carry routed-expert decisions).
- Enable native KV-cache offloading with a 128GB CPU tier (extends the prefix cache).
- FP8 trainer (DeepGEMM blockwise linear/MoE) — impl=custom is already set.
- Use the bash_edit harness (bash + local edit tool) for train + eval, replacing pure bash.
- Bump max_inflight_rollouts 384 -> 512.
- Keep LM-head token chunking + activation checkpointing explicit at the values that become
  trainer defaults in #2867 (not merged yet, so removing them would disable the features).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@samsja samsja force-pushed the feat/sane-memory-defaults branch from 9932d00 to 635fe1c Compare June 25, 2026 17:19
Comment thread src/prime_rl/orchestrator/orchestrator.py Outdated
samsja added a commit that referenced this pull request Jun 26, 2026
- DebugWeightWatcher.apply_policy_update now calls on_version_pending
  before advancing policy.version, mirroring production WeightWatcher
  so debug runs exercise off-policy cancellation
- validate_training_mode no longer skips when no_inference is set,
  so invalid training_mode/teacher combos are caught in debug mode
- no_inference_env only adds tool-role messages on turns with actual
  tool_calls, and _token_plan allocates env deltas only for tool-call
  turns to keep the seq_len invariant
- NCCL weight broadcast init is now skipped when no_inference is set,
  preventing a crash from empty admin_clients on NoOpInferencePool

fixes #2867

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit bcb4c2a. Configure here.

Comment thread tests/debug_envs/no_inference_env/no_inference_env.py
@samsja samsja changed the title feat: enable memory-saving defaults for trainer feat: add orchestrator debug mode (no-inference, no-trainer) Jun 26, 2026
Adds an orchestrator debug mode for running without inference or a trainer,
plus a deterministic test environment for memory/pipeline ablations.

Orchestrator debug mode (orchestrator.py, utils.py, watcher.py, client.py):
- no_inference: Uses NoOpInferencePool (no model calls), clears the renderer,
  skips renderer/VLM/pool-size validators, and skips NCCL weight broadcast
  init and teacher inference pool setup.
- no_trainer: Uses filesystem rollout transport and DebugWeightWatcher that
  advances policy.version locally so dispatch gating and off-policy
  cancellation still run.
- update_dispatch_gate is now async and drives local policy updates in
  no-trainer mode. RSS memory logging via psutil is added after each step
  when debug mode is enabled.

no-inference-env (tests/debug_envs/no_inference_env/):
A deterministic test environment that synthesizes heavy multi-turn
trajectories (configurable sequence length, turns, optional routed experts,
rollout delay) without requiring vLLM.

Docs (skills/training/start-run/SKILL.md):
Documents uv run orchestrator ... --debug.no-inference --debug.no-trainer
and the no-inference env configuration.

fixes #2867
@samsja samsja force-pushed the feat/sane-memory-defaults branch from 0edbd7c to f2c79f7 Compare June 26, 2026 22:52
mikasenghaas added a commit that referenced this pull request Jun 29, 2026
* feat(v1): GLM-4.5-Air scaleswe SWE ablation config (router replay)

v1 RL config: GLM-4.5-Air (zai-org/GLM-4.5-Air, 100B MoE) on scaleswe-v1 (train)
+ swebench-verified (eval), bash harness on prime sandboxes. Router replay on
(trainer.enable_router_replay + inference.enable_return_routed_experts).

2 train + 2 infer nodes. Trainer: cp=8 ulysses, muon, + GLM-5.1 prod-run memory
improvements (LM-head chunking, AC + activation offload, optimizer CPU offload,
skip-gather/skip-optimizer ckpt). Inference: 2x tp=8 replicas, NO expert
parallelism (inference EP + router-replay capture deadlocked the engine via
cross-node EP all-to-all). Renderer/parsers auto-resolve from the official slug.

Relies on fixes already in main: the glm4_moe routed_experts .contiguous() slice
(torch.compile stride assert) and the verifiers always-install-uv bootstrap.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(slurm): VLLM_CACHE_ROOT=/tmp in multi_node_rl template

vLLM's compile cache defaulted to NFS (~/.cache/vllm), which hung inference
startup on slow shared FS. Point it at node-local /tmp (matching inference.sbatch.j2).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat(swe-abl): online fp8 inference + 384 rollout concurrency for glm45air

Inference runs vLLM online fp8 quant (vllm_extra={quantization="fp8"}) over the
bf16 policy for faster generation; trainer/inference/orchestrator all use the
bf16 zai-org/GLM-4.5-Air. The per-channel GLM-4.5-Air-FP8 checkpoint is
incompatible with prime-rl's block-wise fp8 path (use_deep_gemm /
quantize_in_weight_transfer), so we use online quant instead. max_inflight_rollouts=384.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* feat(swe-abl): glm45air-scaleswe -lp variant (default length penalty)

Sibling of glm45air_scaleswe.toml with orchestrator.advantage.length_penalty
enabled at defaults (coef=0.25, gate_by_correctness=false). Distinct slurm
job_name + sandbox labels (glm45air-swe-lp) so it runs alongside the no-penalty
run without sharing prime sandboxes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(swe-abl): bump weight_broadcast timeout to 3600s for cold NFS loads

Cold-cache nodes read the 206GB bf16 model from NFS at ~50s/shard (~46min
total). The weight-broadcast store rendezvous default (1200s/20min) times out
before inference finishes loading (DistStoreError: 1/17 clients joined), killing
the trainer. 3600s covers the cold-load worst case with margin. Applied to both
the base and -lp ablation configs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): kv-offload + fp8, drop router replay, bash_edit

- Drop router replay (trainer.enable_router_replay + inference.enable_return_routed_experts):
  mutually exclusive with inference.kv_cache_offload (rl.py validator — external KV cache hits
  don't carry routed-expert decisions).
- Enable native KV-cache offloading with a 128GB CPU tier (extends the prefix cache).
- FP8 trainer (DeepGEMM blockwise linear/MoE) — impl=custom is already set.
- Use the bash_edit harness (bash + local edit tool) for train + eval, replacing pure bash.
- Bump max_inflight_rollouts 384 -> 512.
- Keep LM-head token chunking + activation checkpointing explicit at the values that become
  trainer defaults in #2867 (not merged yet, so removing them would disable the features).
- Drop the -lp length-penalty variant.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): kv-offload + fp8, drop router replay, bash_edit

- Drop router replay (trainer.enable_router_replay + inference.enable_return_routed_experts):
  mutually exclusive with inference.kv_cache_offload (rl.py validator — external KV cache hits
  don't carry routed-expert decisions).
- Enable native KV-cache offloading with a 128GB CPU tier (extends the prefix cache).
- FP8 trainer (DeepGEMM blockwise linear/MoE) — impl=custom is already set.
- Use the bash_edit harness (bash + local edit tool) for train + eval, replacing pure bash.
- Bump max_inflight_rollouts 384 -> 512.
- Keep LM-head token chunking + activation checkpointing explicit at the values that become
  trainer defaults in #2867 (not merged yet, so removing them would disable the features).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): drop inference fp8 quant to lower mismatch

Run bf16 inference (remove vllm_extra quantization=fp8). fp8 inference added
~10x mismatch KL (~0.002 vs ~0.0002); bf16 inference lowers it. Trainer fp8 +
native KV-cache offload (and the now-disabled router replay) unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): bump context to 131072

seq_len + inference.model.max_model_len 65536 -> 131072.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): revert to 65k ctx, drop kv offload, 384 inflight

- seq_len + inference.model.max_model_len 131072 -> 65536
- Remove inference.kv_cache_offload (native CPU tier)
- max_inflight_rollouts 512 -> 384

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): eval at step 0 (explicit skip_first_step=false)

Make the startup eval explicit so SWE-Bench Verified runs at step 0 before
any train rollouts (already the default; pinned for clarity).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): drop fp8 trainer

Remove trainer.model.fp8 (back to bf16 training).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): bring back router replay

Re-enable trainer.enable_router_replay + inference.enable_return_routed_experts.
Compatible again now that kv-cache offload and fp8 (both mutually exclusive with
router replay) have been dropped; replaying inference's routed-expert decisions in
the trainer cuts the train/inference mismatch by ~an order of magnitude.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): scale inference to 4 nodes / 4 replicas

num_infer_replicas 2 -> 4 (num_infer_nodes is per-replica, so total inference
nodes = num_infer_nodes * num_infer_replicas = 4). Total job = 2 train + 4 infer.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): bump max_inflight_rollouts 384 -> 512

Increase orchestrator rollout concurrency to better saturate the 4 inference replicas.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(slurm): node-local FlashInfer JIT cache to avoid weka lock deadlock

FLASHINFER_WORKSPACE_BASE defaulted to $HOME/.cache/flashinfer on shared weka.
With concurrent GLM-4.5-Air/MoE runs every TP worker contends on the same
fused_moe_*.lock there and deadlocks in uninterruptible (D-state) filesystem
I/O during the CUTLASS fused-MoE JIT build, so inference never serves. Pin it
to node-local /tmp like the Triton/vLLM caches.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* glm45air-scaleswe: bump max_steps 400 -> 1000

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): set node-local cache dirs via config env_vars

Move the cache-dir env vars out of the SLURM templates and into the config
so the PR only adds the TOML (the merged env-var feature, #2863, makes this
possible):
- [env_vars] TRITON_CACHE_DIR (trainer + inference)
- [inference.env_vars] VLLM_CACHE_ROOT, FLASHINFER_WORKSPACE_BASE

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* exp(glm45air-scaleswe): migrate renderer preserve_all_thinking -> thinking_retention

Main #2900 replaced the renderer `preserve_all_thinking` bool with
`thinking_retention`; the merge left the config on the removed field, so it
failed the config-load unit test ("No config class could be parsed"). Switch
to `thinking_retention = "all"` (the documented equivalent).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Mika Senghaas <mail@mikasenghaas.de>
Co-authored-by: Mika Senghaas <mika@primeintellect.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant