No.41 SDPO实现与复现 - #237
Open
ZiyiTsang wants to merge 274 commits into
Open
Conversation
# 🐛 Bug Fix ## Fix torch.cuda patch broken by device abstraction refactor - The device abstraction commit (632b29c) replaced hardcoded `torch.cuda.get_device_properties` / `torch.cuda.get_device_capability` patch targets with `torch.{device_utils.get_device_name()}.*` - When no accelerator is available, `get_device_name()` returns `"cpu"`, so patches targeted `torch.cpu.*` instead of `torch.cuda.*` - Megatron validate_args internally always calls `torch.cuda.*`, so the patches must target `torch.cuda` regardless of device abstraction - Restore hardcoded `torch.cuda.*` patch targets with explanatory comment
# 🔩 Chore ## Remove dead helpers from data module - Delete the unreferenced `filter_long_prompt` helper from `relax/utils/data/data.py` - Delete the dead `_build_messages` helper that was shadowed by `relax/utils/data/data_utils.py` - Delete the unused `process_rollout_data` helper and the imports it required
# 🐛 Bug Fix ## Keep multimodal prompt building non-destructive - Build multimodal message content without mutating cached prompt rows - Prevent reused raw samples from carrying expanded message content into later reads ## Support sliced eager dataset paths - Parse per-file generalized slice syntax in eager file readers - Keep multi-file eager path behavior aligned with streaming path semantics
The rollout component exits its main loop on the final training step, leaving the eval handler un-awaited. This caused a race condition where the controller's atexit shutdown tore down SGLang engines mid-flight. This fix blocks until the evaluation finishes at the end of training.
# ⭐ Feature ## Migrate from Megatron-LM to Megatron-Bridge - Replace direct Megatron-LM checkout with Megatron-Bridge (commit 2faedbf6) in Dockerfile - Upgrade transformer_engine from 2.10.0 to 2.14.1 - Archive old megatron patch (3714d81d) and add new patch for 20260506-85bced0ae ## Adapt Relax backend to Megatron-Bridge API changes - Update vocab_size_with_padding import with fallback for new module path - Rename enable_gloo_process_groups to use_gloo_process_groups - Rename norm_epsilon to layernorm_epsilon in HF config validation - Accept **kwargs in wrapped_provider for new model_provider signature - Relax partition_stride assertion for GLU/SwiGLU linear_fc1 layers (stride=2) - Guard checkpoint_write_patch against removed write_preloaded_data_multiproc
# ⭐ Feature ## Add Qwen3.6 model support with automatic expert format detection - Add Qwen3.6-35B-A3B model configuration script with MoE parameters (256 experts, 8-way routing) - Implement MTP MoE expert weight format detection in Qwen35VL bridge - Qwen3.5: per-expert storage (gate_proj/up_proj/down_proj per expert) - Qwen3.6: packed format (gate_up_proj/down_proj shared tensor) - Add training script for Qwen3.6-35B-A3B 8xGPU colocate mode with multimodal support - Extend Megatron bridge patch with format-aware weight mappings --- # 🐛 Bug Fix ## Fix multimodal data counting and training script paths - Fix remain_data counter for pre-structured multimodal content (was skipping already-processed items) - Remove invalid dataset slice notation (@[0:1000]) from PROMPT_SET path in training script
# ⭐ Feature ## Add rollout reward field metrics - Aggregate numeric fields from reward dictionaries during rollout logging - Skip the primary reward key and raw_reward to preserve existing reward metrics - Reuse the shared helper from the SGLang rollout metrics path
# 🐛 Bug Fix ## Strip consecutive `<|image_pad|>` tokens in pre-tokenized prompts - Add `QwenVLImageProcessor._strip_image_token` static helper that collapses runs of `<|image_pad|>` (token id 151655) into a single placeholder while leaving the surrounding `<|vision_start|>`/`<|vision_end|>` markers intact. - Apply the helper to `prompt` before calling `load_mm_data` in `process_mm_data_async`, so pre-tokenized `input_ids` (where each image is already expanded to N image-pad tokens, one per visual patch) no longer collide with `load_mm_data` re-expanding the placeholder itself. Without the collapse, the pipeline saw `N x M` image-pad tokens and miscounted positions, breaking mrope bookkeeping. - Raw text (`str`) prompts are passed through unchanged.
# 🐛 Bug Fix ## Disable torch.compile around GatedDeltaNet QKV prep - Wrap `_prepare_qkv_for_gated_delta_rule` with `torch._dynamo.config.patch(disable=True)` in the Megatron patch - Avoids torch.compile failure on the Qwen3.6 GatedDeltaNet path --- # ⭐ Feature ## Generalize unsplit-forward path to text-only Qwen3.6 / Qwen3.5 - Detect `Qwen3VLModel` at model build and set `args.uses_unsplit_forward`; the bridge model does CP+SP splitting internally for both VL and text-only Qwen3.5/3.6 sharing the same architecture - Route unsplit tokens + tp*cp*2-aligned `cu_seqlens` through `forward_only` / `train_one_step` whenever the flag is on, not just for VL inputs - Propagate the flag through `data.get_batch`, `loss.compute_advantages_and_returns`, `log_rollout_data`, and `stream_dataloader.post_process_rollout_data` so padding stays consistent ## Add Qwen3.6-35B-A3B 8xGPU DAPO-math training script - New `scripts/training/text/run-qwen36-35B-A3B-8xgpu.sh` for sync GRPO training with TP=2/PP=2/CP=2/EP=4 and partial rollout
# 🐛 Bug Fix ## Propagate fp16 to SGLang Mamba conv dtype - `relax/distributed/ray/genrm.py` and `relax/distributed/ray/rollout.py`: pass `SGLANG_MAMBA_CONV_DTYPE=float16` to the engine env when `--fp16` is set, so Qwen3.6 hybrid-Mamba layers use the matching dtype in rollout/GenRM - `scripts/training/multimodal/run-qwen3-vl-30B-A3B-8xgpu.sh`: add `--fp16 --use-rollout-routing-replay --use-slime-router` - `scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu.sh`: add `--use-rollout-routing-replay --use-slime-router`; document why fp16 is intentionally left disabled for Qwen3.5 ## Skip routing replay for MTP layers - `relax/utils/training/routing_replay.py`: MTP routers exist in training but rollout (sglang) does not run MTP, so there is nothing to record or replay against. Install a pre-hook that clears the global `ROUTING_REPLAY` (so `compute_topk` falls through to the original impl) and skip registration in `all_routing_replays` to keep the per-layer accounting consistent - Guard `compute_topk` against `ROUTING_REPLAY is None` ## Detect Ray 2.x head node by internal resource - `relax/utils/utils.py::get_serve_url`: Ray 2.x auto-registers `node:__internal_head__` on the head node; legacy setups also tag it with a custom `head` resource. Accept either when scanning `ray.nodes()` so head IP discovery works on both --- # 📝 Documentation ## Announce Qwen3.6 support in README - `README.md` / `README_zh.md`: add 05/11/2026 news entry noting Qwen3.6 series (text + VLM) support
# ⭐ Feature ## Auto-detect shared-GPU colocate sub-mode for GenRM - Pick mode from GPU allocation: `R+G==A` keeps the existing split layout; `R==G==A` activates the new shared layout where rollout and genrm overlap on the same bundles. Other combinations are now rejected at startup with a clear error. - Drop the bundle offset for genrm in shared mode so both engines schedule on the same `[0, A)` bundles, and lower genrm Ray fractional `num_gpus` default from 0.2 to 0.1 to leave room alongside rollout. - Plumb `mem_fraction_static` through `--genrm-engine-config`; rollout keeps using `--sglang-mem-fraction-static`. Two engines can now split each GPU independently. - Onload rollout weights and genrm KV in parallel inside `update_weights()` so both engines come back together before the next rollout step. --- # 📝 Documentation ## Document the new GenRM colocate sub-mode (en + zh) - Add a second ASCII architecture diagram for the shared layout and a sub-mode auto-detection table. - Introduce a Shared-mode launch example with `mem_fraction_static` settings and a warning to keep the per-GPU sum < 1.0. - Update Best Practices with sub-mode selection guidance and OOM troubleshooting for shared mode. - Fix stale defaults in the sampling-config table (temperature 0.2 -> 0.1, max_response_len 1024 -> 4096) and add `ep_size` / `mem_fraction_static` to the engine-config table. - Update the example launch script to demonstrate shared mode (rollout 0.5 + genrm 0.3, both on 8 GPU). --- # 🔩 Chore ## Add py-spy multi-PID dump helper - `scripts/tools/_pyspy_dump.sh` runs `py-spy dump` over a list of PIDs in one ray-job submission, used by the debug-hang skill to avoid per-PID submission overhead.
This commit integrates the glm_moe_dsa model, updates Megatron backend, sets up corresponding training scripts, and modifies entrypoint scripts (local.sh, ray-job.sh, spmd-multinode.sh) to support environment variable overrides, keeping the cleanup logic intact.
# 🔩 Chore ## Update torch_memory_saver dependency - Switch source repo from `fzyzcjy/torch_memory_saver` to `redai-infra/torch_memory_saver` fork - Pin to commit `afc13785c50119048e2dd8ac497cc9e29ec75bd4` - Set `TMS_CUDA_MAJOR=12` build-time env var for CUDA 12 compatibility
# ♻️ Refactor
## Split path env vars in launcher scripts
- Introduce `MODEL_DIR` (HF weights / `--ref-load`) and `DATA_DIR`
(`PROMPT_SET` / `--eval-prompt-data`) alongside `EXP_DIR`
(`--load` / `--save`) across 31 training and example scripts
- Each variable is overridable independently; `MODEL_DIR` and
`DATA_DIR` fall back to `EXP_DIR`, while `EXP_DIR` falls back to
`MODEL_DIR` to preserve the legacy `export MODEL_DIR=/root` flow
- Wire `omni-16xgpu-async` defaults (`HF_CHECKPOINT`, `PROMPT_SET`,
`EVAL_PROMPT_DATA`) through the new vars instead of `/path/to/...`
placeholders so the convention is uniform
---
# 📝 Documentation
## Document the path variable convention
- Add a tip block in `docs/{en,zh}/guide/customize-training.md`
describing the three directories and the fallback chain
- Update the example `--hf-checkpoint` / `--ref-load` snippet to use
`${MODEL_DIR}` instead of `${EXP_DIR}`
…eyes script # 🔩 Chore ## Align launch argument order in deepeyes run script - Remove duplicate `--rollout-num-gpus-per-engine` definition in `examples/deepeyes/run_deepeyes.sh`. - Keep the effective `--rollout-num-gpus-per-engine 2` from the SGLang config section and remove the earlier conflicting value.
# 🐛 Bug Fix ## Preserve DeepEyes multimodal state across partial resume - Keep `sample.multimodal_inputs` as the read-only dataset input instead of appending observation images into the shared shallow-copy reference - Make `_prepare_initial_inputs()` return the initial processor output as local `init_mm_train` without overwriting `sample.multimodal_train_inputs` - Keep observation images in `current_image_data` and append observation processor chunks to `multimodal_train_inputs_buffer` - Route the budget-exhausted early return through `_finalize_sample()` so resumed samples leave `generate()` with merged multimodal train inputs
# ⭐ Feature ## Auto-enable true-on-policy mode in fully-async - Auto-enable `--true-on-policy-mode` in `slime_validate_args` when `fully_async` and `rollout_batch_size * n_samples_per_prompt == global_batch_size`, since the train forward log_probs equal what actor_fwd would produce - Add `ROLES_FULLY_ASYNC_ON_POLICY` (no actor_fwd) and route to it from `process_role` when the mode is on - Recompute `old_log_probs` inline in `policy_loss_function` via `log_probs.detach()`, recovering vanilla PG (ratio ≡ 1) while keeping TIS valid against rollout_log_probs - Drop `log_probs` from required data fields in `MegatronTrainRayActor` and `Advantages` when actor_fwd is absent; treat `rollout_log_probs` as kl-zero template in advantages compute - Make `Controller` fully-async weight-sync skip `actor_fwd` recv when the role is not registered, and skip the actor_fwd HTTP probe in the actor health check - Fall back to `rollout_log_probs` for entropy logging in `log_rollout_data` when `log_probs` is unavailable --- # 🔩 Chore ## Tune training scripts and runtime env - Add `NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME` default in `scripts/entrypoint/local.sh` - Drop `actor_fwd`/`reference` from resource specs in async scripts (35B/9B text, 9B openr1mm-mm) now that true-on-policy mode handles them - Add `--log-probs-max-tokens-per-gpu` and adjust parallelism/recompute/MoE-dispatcher knobs across qwen3/qwen35/qwen36 scripts to fit larger micro-batches
# 🐛 Bug Fix ## PR #1889: validate MoE HF config when dense layers exist - Add `_has_dense_moe_layers` / `_is_moe_config` helpers in `relax/backends/megatron/arguments.py` - Validate `moe_intermediate_size` and `shared_expert_intermediate_size` - Skip `intermediate_size` check when model is pure MoE with no dense layer - Ref: THUDM/slime#1889 ## PR #1880: fix per-actor HTTP POST concurrency split - Replace node-count-only divisor with `len(nodes) * num_gpus_per_node` - Use ceiling division `(c + n - 1) // n` to avoid losing concurrency budget - File: `relax/utils/http_utils.py` - Ref: THUDM/slime#1880 ## PR #1873: avoid blocking the asyncio event loop on ray.get - Replace `asyncio.to_thread(ray.get, obj_ref)` with direct `await obj_ref` - Removes the hard ThreadPoolExecutor cap on parallel POSTs - File: `relax/utils/http_utils.py` - Ref: THUDM/slime#1873 ## PR #1888: actor save_model must wake_up / sleep, not just reload PG - `save_model` was calling `reload_process_groups()` without resuming `torch_memory_saver`, leaving GPU tensors unreachable when NCCL is triggered during checkpoint save - Use `self.wake_up()` / `self.sleep()` instead so PG and TMS stay in sync - File: `relax/backends/megatron/actor.py` - Ref: THUDM/slime#1888 ## PR #1882: disaggregate PPO must reconnect rollout NCCL across sleep - Add `disconnect_rollout_engines` to `UpdateWeightFromDistributed` - In `actor.sleep()`, tear down weight-sync NCCL group for disaggregate PPO (`use_critic` + not `colocate`) before `destroy_process_groups()` - In `actor.update_weights()`, force `wake_up` + `connect_rollout_engines` + `sleep` when the disconnect path was taken - Drop the buggy `args.use_critic` GPU-offset branch in `sglang_engine.get_base_gpu_id` - Auto-enable `offload_train` when `use_critic` is on - Files: `relax/backends/megatron/actor.py`, `relax/backends/megatron/weight_update/update_weight_from_distributed.py`, `relax/backends/sglang/sglang_engine.py`, `relax/utils/arguments.py` - Ref: THUDM/slime#1882 ## PR #1878: reinitialize critic output_layer when ckpt shape mismatches - Detect missing or shape-mismatched `output_layer.{weight,bias}` in the critic checkpoint metadata before / after `load_checkpoint` - Reinitialize with `normal_(0, 0.02)` for weight and zero for bias, plus `optimizer.reload_model_params()` for fp16/bf16 master sync - Gated on `role == "critic"`, no effect on actor or non-PPO algorithms - File: `relax/backends/megatron/model.py` - Ref: THUDM/slime#1878 --- # ⭐ Feature ## PR #1890: add missing spec / prefix-cache rollout metrics - Wire `_compute_spec_metrics` and `_compute_prefix_cache_metrics` into `compute_metrics_from_samples` - File: `relax/distributed/ray/rollout.py` - Ref: THUDM/slime#1890 --- # 🔩 Chore ## PR #1862: improve `slice_log_prob_with_cp` assert message - Include `len(log_prob)`, `response_length`, `total_length` in the assert so failures are diagnosable - File: `relax/backends/megatron/cp_utils.py` - Ref: THUDM/slime#1862 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
# ⭐ Feature ## Add --custom-prompt-path for prompt transformation hook - Add CLI argument in add_data_arguments (arguments.py) - Load custom function via load_function in data_source.py - Thread custom_prompt_func through build_messages, process_raw_sample, BaseDataset, Dataset, and StreamingDataset - Custom function is called after prompt extraction, before conversation/multimodal processing ## Add --image-resize-scale-factor for image dimension alignment control - Add CLI argument in add_data_arguments (arguments.py) - Add image_resize_scale_factor field to MultimodalConfig with from_args propagation and getter function - Update fetch_image with 3-way logic: None uses default patch_factor, 0 disables alignment, positive int uses custom value --- # 📝 Documentation ## Update configuration reference docs - Add --custom-prompt-path to Dataset table (EN + ZH) - Add --image-resize-scale-factor to Multimodal Data table (EN + ZH) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
# ⭐ Feature ## Add hybrid training mode combining async data pipeline with colocate weight sharing - Introduce `--hybrid` flag that sets `fully_async=True` and `colocate=True` so actor/ref/actor_fwd share GPUs via TensorBackuper+_switch_model while rollout runs on a separate GPU placement group with streaming transfer queue - Add `train_hybrid()` method in MegatronTrainRayActor: collects sub-batches from transfer queue, runs ref/teacher/actor forward per sub-batch, merges all sub-batches, computes advantages with correct global normalization, then trains on the full merged batch - Register hybrid mode in `process_role()` to use ROLES_COLOCATE (actor + rollout only, no separate reference/actor_fwd services) - Update controller to skip shared placement groups and fully-async DCS weight sync setup when hybrid is active - Skip actor_fwd health probe in `_check_services_health` for hybrid mode to avoid spurious warnings - Add `train_hybrid()` dispatch in RayTrainGroup and Actor component - Validate argument combinations: `--hybrid` is the supported way to combine async pipeline with colocate weight sharing; bare `--fully-async --colocate` now raises ValueError - Set `offload_train=False`, `offload_rollout=False`, and `compute_advantages_and_returns=True` for hybrid mode - Add Qwen3-4B 8xGPU hybrid-async training launch script
# ⭐ Feature
## Add `relax/utils/visualize` rollout result viewer
- Add web viewer adapted from rlsp/utils/visualize: FastAPI + single-page UI
for browsing `<save>/rollout_result/{train,eval}/{step}.jsonl`, with
step dropdown, sample nav, sort by reward / response_length, sample-info
card, and prompt / response / label rendering with chat-template / tool-call
/ `<think>` highlighting
- Auto-discover `train/` and `eval/` subdirs and render a tab toggle when both
exist; fall back to a single anonymous bucket for flat dirs
- Add terminal UI mode (`--tui`) adapted from redaccel/verl reward_viewer_v2:
sync-load the first step then stream remaining steps via a daemon thread
(newest first), with step / sample / dataset / sort dropdowns, field
filter, fuzzy search (`f`/`enter`/`esc`), vim-style page nav, and
text/table render toggle
- Default theme switched to dark; header shows the Relax wordmark linking
to the GitHub repo plus a GitHub icon
- Mask multimodal pad tokens (`<|image_pad|>`, etc.) by default in the TUI
via `--mask-str`
## Add `relax/entrypoints/visualize` thin wrapper
- One command for both modes: `python -m relax.entrypoints.visualize <dir>`
(web) and `... --tui` (terminal)
- `DATA_DIR` is a required positional argument
- TUI dependencies (`textual`, `rich`) are lazy-imported; clear error
message if missing
---
# 📝 Documentation
## Add bilingual rollout result viewer guide
- Add `docs/{en,zh}/guide/rollout-result-viewer.md` covering data layout,
launch command, flags, page features, terminal UI key bindings, and
reverse-proxy notes
- Register the new pages under the existing "Operations & Debugging" /
"运维与调试" sidebar group in `docs/.vitepress/config.mts`
- Add `docs/public/relax-viewer.png` screenshot (palette-compressed PNG
to stay under the 500 KB pre-commit limit)
# 🐛 Bug Fix
## Fix MODEL_DIR/EXP_DIR initialization in qwen35-9B hybrid-async script
- Replace buggy `EXP_DIR="${MODEL_DIR:=...}"` side-effect assignment with separate `EXP_DIR`/`MODEL_DIR`/`DATA_DIR` defaults, matching `run-qwen35-9B-8xgpu-openr1mm-async.sh`
- Point `--hf-checkpoint` / `--ref-load` at `${MODEL_DIR}` and `PROMPT_SET` at `${DATA_DIR}` so model and dataset roots can be overridden independently of the experiment output dir
(cherry picked from commit 466c779)
(cherry picked from commit 8034d15)
# 🔩 Chore ## Align DeepEyes dataset inputs - Align fp16, GenRM, and partial-rollout scripts with examples/deepeyes/run_deepeyes.sh - Use the same Deepeyes v1 training shards as the main script - Use the same thinklite reasoning accuracy eval slice as the main script (cherry picked from commit df613b4)
# 📝 Documentation ## Add bilingual Hybrid training mode guide - Add `docs/en/guide/hybrid-training.md` and `docs/zh/guide/hybrid-training.md` describing the hybrid execution mode (streaming TransferQueue + in-process TensorBackuper weight sharing) - Cover mode comparison vs Colocate / Fully Async, role layout (`ROLES_COLOCATE` with disjoint actor/rollout placement groups), `--hybrid` flag resolution, and the three-phase `train_hybrid` loop - Document required and optional flags (`--hybrid`, `--num-iters-per-train-update`, `--max-staleness`, `--balance-data`) and the default overrides applied in `relax/utils/arguments.py` - Include the 8-GPU multimodal reference launch from `scripts/training/multimodal/run-qwen35-9B-8xgpu-openr1mm-hybrid-async.sh` and troubleshooting tips (stalled sub-batches, balance-data rejection) - List planned next steps: integrate DCS for weight sync, split `train_actor` by `num_iters_per_train_update` ## Register pages in VitePress sidebar - Add Hybrid Training Mode under the Advanced group in both `en` and `zh` sidebars in `docs/.vitepress/config.mts` (cherry picked from commit 7f78ff7)
# ⭐ Feature ## Add JSON provider config dump - Keep the existing transformer_config.pkl dump for compatibility - Also write transformer_config.json next to it for easier inspection - Convert non-JSON-safe values recursively and fall back to str() when needed (cherry picked from commit 598c15c)
- restore configs/env.yaml (clean upstream template) so the PR does not delete it; local secret-bearing copy stays on disk untracked-by-intent - drop sdpo-only .gitignore entries (configs/env.yaml, wandb/, scripts/training/sdpo/) to match upstream/main - remove tests/utils/opd/test_opd_legacy_regression.py (not needed)
- ruff, ruff-format, mdformat, docformatter, clang-format fixes - covers sdpo README, megatron data.py, opd/metrics/tracking utils
# 🐛 Bug Fix ## Score tool-use trajectories as ordered (Action, Input) pairs - _extract_tool_calls parses (Action, Action Input) pairs in document order instead of merging every input into one flat dict; _golden_tool_call returns the same ordered pair structure - Correctness is now format_ok and predicted pairs == golden pairs, so step order and per-step Action<->Input pairing matter: 37% of the tooluse rows are multi-step and 12% carry the same input key with different values across steps, which the old merged-dict comparison could not distinguish from a correct trajectory # ✅ Tests - Add order-swap, cross-step value-swap and exact multi-step cases
# ♻️ Refactor ## Route every algorithm difference through EnvironmentFeedback hooks - Split feedback implementations per algorithm: feedback.py (base + OPD), opsd_feedback.py (OPSD), sdpo/feedback.py (SDPO); base-class defaults reproduce plain OPD so OPD/MOPD need zero overrides - Delete all 21 is_sdpo branches from OpdManager together with the is_sdpo_feedback / is_sdpo_prompt_routing_enabled reflection; --opd-feedback-class now defaults to OPDFeedback - SDPO preflight (text-only check, stale-payload clearing, teacher prompt presence) moves into SDPOFeedback.prepare_teacher_prompts; teacher fetch/response failures degrade like plain OPD and surface via the assembly-time check_transfer_channels validation - GenerateState takes its feedback from OpdManager; the generate_and_rm_group call sites are unchanged - Rebind example launchers: vision_opd/mopd -> opsd_feedback. OPSDFeedback, sdpo -> sdpo.feedback.*, math_opd uses the default ## Drop the OPD top-K context-parallel slicing - Delete slice_opd_topk_rollout_fields and its call sites; launch validation rejects top-K token selection combined with CP > 1 / dynamic CP / allgather-cp - Restore cp_utils get_cp_local_num_tokens to upstream shape (the int cast lives on fix/cp-num-tokens-int) and restore the stream_dataloader / metrics / tracking churn to upstream - Gate _apply_opd_sample_mask behind use_opd so non-OPD runs skip it # ✅ Tests - Update payload / feedback / arguments tests to the hook contract - Add OPD no-op contract, SDPO escalation hooks, default-class resolution, top-K x CP rejection and SDPO launch validation coverage
# ♻️ Refactor ## Make EnvironmentFeedback concrete with behavioral defaults - Replace the abstractmethod pair with defaults: record_sample_feedback records reward-dict env feedback and prepare_teacher_prompts clears sample.teacher_prompt/opd_sample_mask; OPDFeedback becomes an explicit alias subclass ## Route the OPSD dataset privilege through OPSDFeedback - process_raw_sample stores the rendered teacher prompt under metadata["opd_teacher_prompt"] instead of Sample(teacher_prompt=...) - OPSDFeedback.prepare_teacher_prompts assigns metadata["opd_teacher_prompt"] -> sample.teacher_prompt at rollout time; samples without it keep the OpsdWorker student-prompt fallback ## Run both feedback hooks on every OPD rollout path - Add _record_feedback_and_prefill in sglang_rollout; the group_rm branch and both non-group branches (multi-sample / single-sample) now share record -> prepare -> prefill - Delete SDPOFeedback.record_sample_feedback override and its module helper (base default is equivalent) --- # ✅ Tests ## Pin the new feedback contract - Base recording applies to OPD/OPSD/SDPO subclasses alike - OPSD assigns privilege from metadata; missing key leaves teacher_prompt None - process_raw_sample surfaces teacher prompt as metadata instead of the Sample field
# Conflicts: # relax/backends/megatron/model.py
# 🐛 Bug Fix ## Restore Megatron/sglang-free import isolation - Stub sglang_router and relax.backends.sglang.arguments before importing relax.utils.arguments so the CPU CI runners (no sglang installed) can collect the opd teacher colocate test suite - Keep the sdpo-specific assertions (per-token loss, feedback class, allgather-cp rejection) on top of the restored fixture
# ♻️ Refactor ## Remove unused TopkWorker.TRANSFER_FIELDS - Superseded by topk_transfer_fields(); no consumers remain # ✅ Tests ## Remove ghost attrs from colocate arguments test - Drop opd_mask_on_success and opd_log_prob_dump_dir kwargs that the argument parser no longer defines
# ⭐ Feature ## Unify the sdpo example launchers - Source examples/on_policy_distillation/sdpo/env.sh in every 4xgpu launcher - Add --teacher-sglang-enable-weights-cpu-backup to every launcher - Align training and eval params with the biology script (5000 rollouts, batch 32, 8 samples per prompt, global batch 256, eval every 5 iters with 16 samples per eval prompt) - Wire the eval split into the toolalpaca launcher --- # ♻️ Refactor ## Stop writing unused choices metadata - prepare_data.py no longer emits a never-consumed choices field --- # 📝 Documentation ## Update the sdpo README - Correct model size, resource layout, config table and eval flow - Document the env.sh contract sourced by every launcher
kkyyxhll
reviewed
Aug 30, 2026
|
|
||
| # OPD manager (singleton — one OpdManager per GenerateState) | ||
| self.opd_manager = opd.OpdManager(args) if opd.is_opd_enabled(args) else None | ||
| self.feedback = self.opd_manager.feedback if self.opd_manager is not None else None |
Contributor
There was a problem hiding this comment.
这个是否可以直接opd_manager.feedback？
|
|
||
| OPD_ARGS=( | ||
| --use-opd | ||
| --opd-feedback-class relax.utils.opd.opsd.feedback.OPSDFeedback |
Contributor
There was a problem hiding this comment.
这个是否可以保持其他的recipe没有变动？且mopd也不是opsd
# ♻️ Refactor ## Remove GenerateState.feedback alias - Delete the `self.feedback` mirror of `opd_manager.feedback` in GenerateState - Access `state.opd_manager.feedback` directly inside `_record_feedback_and_prefill` - Rollout behavior is unchanged; the OPD injection point stays between rm and prefill
# ♻️ Refactor ## Revert feedback-class flags in existing recipes - Drop `--opd-feedback-class OPSDFeedback` from the three mopd and three vision_opd launchers: none passes `--opd-teacher-prompt-key`, so the flag never had an effect there and MOPD/vision recipes are not OPSD - Restore `--sglang-disable-cuda-graph` in the three mopd launchers - The 2B mopd launcher now matches main byte-for-byte
# ⭐ Feature ## Interface-level feedback parameter schema - Declare the constructor schema once on EnvironmentFeedback (teacher_prompt_key, success_reward_threshold); the selected subclass binds the same --opd-feedback-kwargs dict at construction time - Add load_feedback(path, kwargs) with loud TypeError on unknown fields; OPD default stays OPDFeedback with empty kwargs - OPSDFeedback now requires teacher_prompt_key and validates the sglang teacher path in validate_launch_args - SDPOFeedback consumes success_reward_threshold (default 1.0) instead of a hardcoded score >= 1.0 boundary --- # 🐛 Bug Fix ## Drop the --opd-teacher-prompt-key engine flag - The flag silently lost its effect when the teacher prompt moved to the feedback strategy; the key now enters only via --opd-feedback-kwargs - OpsdWorker.from_args activates OPSD from image keys alone; prompt-key routing is owned by OPSDFeedback --- # ♻️ Refactor ## Collapse empty SDPO domain subclasses - Merge SciKnowEvalSDPOFeedback and ToolUseSDPOFeedback (both were `pass`) into GoldenAnswerSDPOFeedback for static golden-answer text datasets; keep CodeSDPOFeedback as the raising placeholder - Point the six SDPO launchers and the README at the merged class - Ingestion reads teacher_prompt_key from --opd-feedback-kwargs in data_source; rendering and metadata hand-off are unchanged --- # ✅ Tests ## Cover kwargs binding and the merged classes - Add load_feedback binding/error tests and a success-threshold test - Update payload, feedback, and arguments tests for the new flag and class
# 🎨 Style ## Make mdformat and docformatter hooks idempotent - Re-wrap GoldenAnswerSDPOFeedback docstring to satisfy docformatter 1.3.1 (--wrap-descriptions 79) - Shrink the Feedback column of the launcher table by one cell width to satisfy mdformat 0.7.9
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.
总结
本 PR 在 OPD/OPSD teacher-prefill 通路上增加一个低侵入的 Relax-SDPO 路径：plain OPD
直接用 rollout token；ordinary OPSD 经
--opd-feedback-class OPSDFeedback消费数据集teacher prompt 列（ingestion 渲染进
metadata["opd_teacher_prompt"]，rollout 时由feedback 类赋给
Sample.teacher_prompt）；SDPO 在 group rollout 完成并获得 reward 后，统一经
EnvironmentFeedback.record_sample_feedback()与prepare_teacher_prompts(group, rewards)，由 feedback 类动态构造 teacher prompt。teacher 仅对原 response 在 studentTop-K token ids 上重打分。纯蒸馏复用现有 OPD loss：
SDPO 仅支持
student_topk，仅走文本、TP/CP、无 PP 路径。普通 OPD/OPSD 继续复用已有teacher prefill、Top-K selection、KL/JSD 计算和 loss reducer。本 PR 另增
opd_sample_mask过滤无动态 teacher context 的 sample；teacher 为冻结 snapshot（EMA 更新为后续阶段）。
OPD 与 SDPO 的差异
sample.rollout_tokens or sample.tokens；ordinary OPSD 用数据集列渲染的 privileged prompt（OPSDFeedback赋值）Sample.teacher_prompt，可附同 group 成功 responsegenerate_and_rm_group()的 reward 后、OpdManager.prefill()前注入static：保持初始 teacher snapshotstudent_sampled/student_topk/teacher_topk/unionstudent_topkTopkWorkerStrategy / SpecificationTopkWorker.build_teacher_payload()compute_policy_opd_loss配置为纯蒸馏--opd-loss-coef=1 --opd-kl-coef=0Sample.opd_sample_mask；get_batch 折叠进 loss_masks，compute_policy_opd_loss 据其存在选 CP 感知归约validate_sdpo_text_only、validate_sdpo_topk_payloadvalidate_sdpo_text_onlyGuardflowchart TD classDef existing fill:#eef3f8,stroke:#64748b,stroke-width:1px,color:#1e293b classDef changed fill:#fee2e2,stroke:#dc2626,stroke-width:2px,color:#991b1b classDef shared fill:#eef3f8,stroke:#2563eb,stroke-width:2px,color:#1e293b DATA["RolloutDataSource"]:::existing --> ROLLOUT["SGLangRollout<br/>student group rollout"]:::existing ROLLOUT --> RM["batched_async_rm(args, group)<br/>one reward per sample"]:::existing subgraph FEEDBACK_API["EnvironmentFeedback"] direction TB subgraph RECORD_STAGE["record_sample_feedback(sample, reward)"] direction LR OPD_RECORD["OPDFeedback<br/>no-op"]:::existing OPSD_RECORD["OPSDFeedback<br/>no-op"]:::existing SDPO_RECORD["GoldenAnswer / Code<br/>record sample feedback metadata"]:::changed end subgraph PREPARE_STAGE["prepare_teacher_prompts(group, rewards)"] direction LR OPD_PREPARE["OPDFeedback<br/>plain OPD: rollout tokens"]:::existing OPSD_PREPARE["OPSDFeedback<br/>ordinary OPSD: Sample.teacher_prompt"]:::existing SDPO_PREPARE["GoldenAnswer / Code<br/>dynamic Sample.teacher_prompt<br/>+ Sample.opd_sample_mask"]:::changed end OPD_RECORD --> OPD_PREPARE OPSD_RECORD --> OPSD_PREPARE SDPO_RECORD --> SDPO_PREPARE end RM --> OPD_RECORD RM --> OPSD_RECORD RM --> SDPO_RECORD subgraph LIFECYCLE["Existing OpdManager teacher-prefill lifecycle"] MANAGER["OpdManager.prefill(samples)"]:::existing PREPARE_INPUT["teacher input prep<br/>OPSD/SDPO: OpsdWorker 读 Sample.teacher_prompt"]:::existing REQUEST["_teacher_prefill()<br/>teacher request + response offset"]:::existing TRANSFER["_assemble_transfer()<br/>Top-K / teacher log-prob fields"]:::shared MANAGER --> PREPARE_INPUT --> REQUEST --> TRANSFER end OPD_PREPARE --> MANAGER OPSD_PREPARE --> MANAGER SDPO_PREPARE --> MANAGER TRANSFER --> LOSS_INPUT["OPD transfer to training"]:::shared LOSS_INPUT --> LAYOUT["existing OPD TP/CP layout<br/>+ SDPO sample gate"]:::changed LAYOUT --> FORWARD["Megatron student forward"]:::shared FORWARD --> LOSS["compute_policy_opd_loss<br/>existing OPD loss + sample mask"]:::existing RM --> RL_MASK["advantage + loss_masks"]:::existing RL_MASK --> LOSS LOSS --> OPT["Megatron optimizer update"]:::shared class DATA,ROLLOUT,RM,MANAGER,PREPARE_INPUT,REQUEST,LOSS existing class SDPO_RECORD,SDPO_PREPARE,LAYOUT changed class TRANSFER,LOSS_INPUT,FORWARD,OPT,RL_MASK shared红色节点为本 PR 新增/改变的边界逻辑；蓝色边框为复用组件；OPD/OPSD 分支保持灰色。三者经同一
调用路径进入
record_sample_feedback()与prepare_teacher_prompts()，但实现不同：OPD/OPSD为空实现（OPSD 在
prepare_teacher_prompts中赋值数据集 privileged prompt），GoldenAnswer SDPOfeedback 记录当前 sample feedback 并由 batch-level 实现决定是否读取同组成功 response。特权信息只进入
Sample.teacher_prompt与 teacher input，不改 studentprompt/response 或 rollout token ids。
opd_sample_mask随 transfer 进入训练侧，在get_batch的_apply_opd_sample_mask()折叠进
loss_masks、并由compute_policy_opd_loss选择 CP 感知归约来屏蔽无有效 teachertarget 的 sample；普通 OPD/OPSD 不设该字段，保持原数值路径。
组件与设计模式
设计原则：低侵入地增加 dynamic feedback prompt、sample-level OPD mask 与
staticteacher 更新。compute_policy_opd_loss()；--opd-loss-coef 1.0 --opd-kl-coef 0.0record_sample_feedback(sample, reward)与prepare_teacher_prompts(group, rewards)；generate_and_rm_group()在 reward 后依次调用--opd-feedback-class选择GoldenAnswerSDPOFeedback（静态 golden-answer 文本任务）、CodeSDPOFeedback（占位）Sample.opd_sample_maskOpdManager.produce_opd_transfer_data()写入train_data；get_batch的_apply_opd_sample_mask()折叠进loss_masks，compute_policy_opd_loss()据其存在选 CP 感知归约relax/utils/opd/sdpo/validation.py；OpdManager.prefill()/_assemble_transfer()调用OpdManager、TopkWorker及现有 OPD loss/reducerEnvironmentFeedback所有 OPD/OPSD/SDPO 路径经同一组接口，行为由
--opd-feedback-class选定的 feedback 类决定。参数 schema 声明在接口层，经单一 CLI 入口
--opd-feedback-kwargs（JSON）传入，运行时cls(**kwargs)绑定到具体构造函数，各子类只消费自己用到的字段：OPDFeedback+ 空 kwargs（teacher 复用 student prompt）OPSDFeedback+{"teacher_prompt_key": "<数据列名>"}（必填， ingestion 用该列渲染privileged prompt）
SDPOFeedback系 + 可选{"success_reward_threshold": 0.8}（成功分数线，默认 1.0）注入点在
relax/engine/rollout/sglang_rollout.py的generate_and_rm_group()：studentgroup 完成 rollout 后，先由
batched_async_rm()产生与 group 对应的 rewards，再依次record_sample_feedback()、prepare_teacher_prompts()，最后进现有OpdManager.prefill()。record_sample_feedback()只处理当前 sample：OPD/OPSD 空实现，SDPO 将 reward payload 中feedback 写入
sample.metadata["env_feedback"]。prepare_teacher_prompts()处理完整 group，按决策矩阵逐样本决定注入内容（见「METH 决策矩阵」）。无有效 solution/feedback 时复制原 prompt
并将
opd_sample_mask设为False。student 的 prompt/response/rollout token ids/studentTop-K ids 均不变。
特权信息使用
GoldenAnswerSDPOFeedback与占位的CodeSDPOFeedback共用一套矩阵（reward 无关），各数据集反馈文本由示例reward.py生成。所有注入内容均不含金标（正确答案 / expected 动作参数）。
group_index(或metadata.uid) 内reward ≥ success_reward_threshold（默认 1.0）的成功回答，包装为<successful_attempt>。opd_sample_mask=False，因折叠进共享loss_masks而不贡献任何 loss 梯度（纯蒸馏配置下基础 RL 梯度本就≈0，可观察效果集中在 OPD）。变更
CLI 与兼容性
--opd-feedback-classEnvironmentFeedback子类；不传默认relax.utils.opd.feedback.OPDFeedback（plain OPD/MOPD：teacher 复用 student prompt）。--opd-feedback-kwargs（新增）teacher_prompt_key（OPSD 必填，数据集 teacher prompt 列名）、success_reward_threshold（SDPO 成功分数线，默认 1.0）。传未知字段在构造时报TypeError。--opd-teacher-prompt-key（删除）--opd-feedback-kwargs {"teacher_prompt_key": ...}，单一入口。--opd-loss-coef/--opd-kl-coef--opd-loss-coef 1.0 --opd-kl-coef 0.0表达纯蒸馏。--opd-token-selection/--opd-jsd-alphastudent_topk；对称 JSD 配--opd-kl-type jsd --opd-jsd-alpha 0.5。--group-rm、student_topk、正opd_loss_coef、per-token loss；拒绝 PP/MTP/多模态。文件变更
relax/engine/rollout/on_policy_distillation.pyOpdManager经load_feedback(path, kwargs)构造 feedback；将opd_sample_mask纳入 SDPO transfer schema，在 teacher prefill/transfer 边界执行 SDPO payload 校验。relax/utils/opd/feedback.pyload_feedback；relax/utils/opd/opsd/feedback.py的OPSDFeedback赋值数据集 privileged prompt；relax/utils/opd/sdpo/feedback.py依 reward/group context 生成 teacher prompt 与 sample mask（GoldenAnswerSDPOFeedback/CodeSDPOFeedback）。relax/utils/opd/opd_utils.py--opd-feedback-kwargs、删除--opd-teacher-prompt-key；在compute_policy_opd_loss增加 CP 感知归约分支（get_sum_of_sample_mean）；普通 OPD 走原reduce_opd_loss，数值路径不变。relax/utils/types.py/relax/engine/rollout/data_source.pySample增加opd_sample_mask字段；ingestion 从--opd-feedback-kwargs读teacher_prompt_key渲染数据集 teacher prompt 列。examples/on_policy_distillation/sdpo/验证
可复制命令
聚焦单元测试：
Ray / SGLang teacher prefill 集成（需空闲训练 pod）：
语法与静态检查：
python -m compileall -q relax examples/on_policy_distillation/sdpo tests bash -n examples/on_policy_distillation/sdpo/*.sh git diff --check pre-commit run --all-files新增测试
tests/utils/opd/test_feedback.pytests/utils/opd/test_opd_legacy_regression.pyTrue与普通 OPD 数值一致；Falsesample 不产生 OPD 梯度；全False时 loss/梯度为零；普通路径保持原 oracle。tests/engine/rollout/test_on_policy_distillation_payload.pyopd_sample_mask进 transfer；student Top-K 与 teacher payload 的 shape/内容边界正确。tests/utils/opd/test_opd_topk_log_probs.py0、负 sentinel、越界输入处理正确。tests/utils/test_arguments_opd_teacher_colocate.pytests/backends/megatron/test_opd_loss_aggregation.pytests/backends/megatron/weight_update/test_lora_weight_sync.py、tests/distributed/ray/test_weight_sync.pytests/examples/sdpo/test_prepare_data.py、tests/examples/sdpo/test_data_reward.py单元、集成与端到端测试结果
环境：
python -m compileall与bash -n全部通过（exit 0）。聚焦测试套件结果如下：test_feedback）test_opd_legacy_regression）test_on_policy_distillation_payload）test_opd_topk_log_probs）test_arguments_opd_teacher_colocate）test_prepare_data+test_data_reward）test_opd_loss_aggregation）test_lora_weight_sync）test_weight_sync）test_opd_teacher_controller_colocate）*
test_opd_loss_token_mean在其 fixture 调用pytest.importorskip("torch", exc_type=...)，该 kwarg 在 pytest 9 已移除，属测试脚手架与 pytest 版本不兼容，非产品逻辑回归；同文件其余
用例因共享该 fixture 未能执行。
实际训练结果
仅汇报 SciKnowEval Biology 子集的 SDPO 训练（run
ZiyiTsang/relax-sdpo/7f3n0ylw，wandb 记录至 step 400）。
训练环境介绍
[1,4]，rollout SGLang[1,3]，teacher SGLang[1,1]--opd-loss-coef 1.0 --opd-kl-coef 0.0（纯蒸馏），student_topkk=16，jsd α=0.5，--opd-norm-mode tail，teacherstaticsciknoweval/biologytrain.jsonl；eval-interval=5，n=16 eval samples/prompt结果汇报（Biology · SDPO · ≤ step 400）
opd_sample_mask）风险与回退
已知限制
student_topk、无 PP/MTP 路径；多模态 teacher prompt与
allgather_cp=TrueTop-K 路径不在本 PR 范围。经
opd_sample_mask=False不参与 OPD loss。风险
test_on_policy_distillation_payload.py与真实 teacher prefill 验收。layout / global-id contract 验证。
关闭开关或回退方式
examples/on_policy_distillation/sdpo/run-grpo.sh，不传 SDPOfeedback 与 teacher 参数。
其他
复现过程中发现若干bug，已经提交PR解决：#287
检查清单
opd_sample_mask作为 sample-level transfer field 在 CP=1/CP>1 reduction 生效；fallback sample 不贡献 OPD 梯度。static模式（actor update 后不更新 teacher 权重）。EVAL_ARGS/ROLLOUT_ARGS/OPD_ARGS/GRPO_ARGS/OPTIMIZER_ARGS/PERF_ARGS分组组织。More info: issue #86