Conversation
4 tasks
aoshen02
added a commit
that referenced
this pull request
Sep 1, 2026
… acceleration (vllm-project#397) * feat: add DSpark speculative decoding draft model training for RL rollout acceleration Integrates DSpark speculative decoding into vime RL training pipeline to accelerate rollout generation. DSpark trains a lightweight draft model (Markov heads + attention layers) alongside the policy model, enabling speculative decoding during vLLM rollout with ~2.5x throughput improvement. Key changes: - New module vime/backends/megatron_utils/dspark/ with draft model architecture (Markov heads, attention, hidden state capture), training losses (CE + L1 + confidence), and weight export to vLLM - Actor integration: draft weight cache/restore lifecycle around sleep/wake_up, synchronize()+empty_cache() before pause() to prevent memory fragmentation crash - Weight sync: vLLM 0.26/0.27 compatibility shims via NCCLTrainerAdapter, create_nccl_trainer with try/except fallback, send_custom_weights for metadata via Ray before NCCL broadcast - vLLM engine: router retry + extended timeouts for DSpark weight sync, spec_accept_rate extraction from /inference/v1/generate endpoint - TP>1 support with vocab padding/unpadding (all-gather, zero-pad, strip) - 14 new --dspark-* CLI flags for draft model configuration Signed-off-by: kaiyuanxie <kyxiezju@163.com> * feat: add DSpark speculative decoding draft model training for RL rollout acceleration - New module: vime/backends/megatron_utils/dspark/ (8 files: modeling, loss, attention, markov_head, hidden_capture, export, common, __init__) - Actor integration: draft weight cache/restore lifecycle, freeze-policy - vLLM engine: load_draft_weights_from_file via HTTP /collective_rpc endpoint - L1 loss OOM fix: chunked logsumexp + gradient checkpointing - TP>1 vocab padding 3-stage fix: all-gather (model build), zero-pad (pretrained load), strip (vLLM export) - 14 new --dspark-* CLI arguments - Restored origin/main weight transfer API (common.py, update_weight_from_distributed.py) - DSpark draft weight sync methods added on top of origin/main update_weight_from_tensor.py - Removed unrelated changes (gemma4, gpt_oss, fp8_helpers, ppo_utils, mask_utils, etc.) that were accidentally included in previous commit Verified: all CI tests pass (27 CPU test files + 150 utils tests + 25 upstream-sync tests), pre-commit clean. Signed-off-by: kaiyuanxie <kyxiezju@163.com> * refactor: include DSpark draft params in weights_backuper Replace manual _save_dspark_draft_to_cpu() / _restore_dspark_draft_from_cpu() with automatic draft param management via weights_backuper source_getter. Changes: - Add _iter_dspark_draft_params() to iterate draft model params - Chain draft params into weights_backuper source_getter - Delete _save_dspark_draft_to_cpu() and _restore_dspark_draft_from_cpu() - Add _switch_model(actor) in update_weights() for colocate mode (disable() does not restore GPU memory from TMS backup) - Update stale comments This simplifies the pause/resume cycle: weights_backuper now automatically saves/restores draft params alongside policy params. Signed-off-by: kaiyuanxie <kyxiezju@163.com> * fix: DSpark draft weight sync and local spec compatibility Fix 5 bugs discovered during end-to-end validation on A800: 1. Local spec layernorm names (qwen2.py): Add input_layernorm.weight and pre_mlp_layernorm.weight mappings for --transformer-impl local (TE spec uses different param names). 2. Filter draft params from weight sync (actor.py, hf_weight_iterator_direct.py): weights_getter and _get_megatron_local_param_infos now exclude .draft_model. params to prevent KeyError in HF weight conversion. 3. Non-MoE _ipc_engine setup (update_weight_from_tensor.py): connect_rollout_engines returned early for non-MoE models before setting _ipc_engine/_ipc_gather_src, causing draft weights to never sync to vLLM (spec_accept_rate=0%). Now sets these for all models. 4. Disable RoPE fusion when TE is broken (model_provider.py): When DSPARK_DISABLE_TE is set, also disable apply_rope_fusion to avoid TE fused RoPE kernel errors on A800 (SM 8.0). 5. Safe draft param restore (actor.py): Replace _switch_model(actor) with _restore_dspark_draft_params_safe() which assigns fresh GPU tensors instead of copy_(), working even when TMS pause() freed the original GPU storage. Also adapt to origin/main API changes: - Unpack engine_parallel_configs from get_updatable_engines_and_lock - Remove normalization kwarg from get_gpt_layer_with_transformer_engine_spec - Use convert_to_global_name=True for consistent draft param naming Verified: spec_accept_rate 32.5% (vime rollout), 31-40% (vLLM SpecDecoding) on Qwen3-4B DSpark colocate 8-GPU. Signed-off-by: kaiyuanxie <kyxiezju@163.com> * refactor: simplify DSpark weight sync and validate TP paths (#2) * refactor: simplify DSpark weight sync and TP support * test: remove redundant DSpark hidden capture unit test Signed-off-by: aoshen02 <aoshen@inferact.ai> * fix: add normalization parameter to get_gpt_layer_local_spec The get_gpt_layer_local_spec function requires normalization parameter but it was missing, causing RMSNorm models (e.g. Qwen3) to use incorrect normalization. Pass args.normalization to both MoE and dense model paths. Signed-off-by: CalvinXKY <kyxiezju@163.com> * refactor: simplify DSpark weight export integration Signed-off-by: aoshen02 <aoshen@inferact.ai> --------- Signed-off-by: kaiyuanxie <kyxiezju@163.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: CalvinXKY <kyxiezju@163.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>
aoshen02
added a commit
that referenced
this pull request
Sep 14, 2026
… acceleration (vllm-project#397) * feat: add DSpark speculative decoding draft model training for RL rollout acceleration Integrates DSpark speculative decoding into vime RL training pipeline to accelerate rollout generation. DSpark trains a lightweight draft model (Markov heads + attention layers) alongside the policy model, enabling speculative decoding during vLLM rollout with ~2.5x throughput improvement. Key changes: - New module vime/backends/megatron_utils/dspark/ with draft model architecture (Markov heads, attention, hidden state capture), training losses (CE + L1 + confidence), and weight export to vLLM - Actor integration: draft weight cache/restore lifecycle around sleep/wake_up, synchronize()+empty_cache() before pause() to prevent memory fragmentation crash - Weight sync: vLLM 0.26/0.27 compatibility shims via NCCLTrainerAdapter, create_nccl_trainer with try/except fallback, send_custom_weights for metadata via Ray before NCCL broadcast - vLLM engine: router retry + extended timeouts for DSpark weight sync, spec_accept_rate extraction from /inference/v1/generate endpoint - TP>1 support with vocab padding/unpadding (all-gather, zero-pad, strip) - 14 new --dspark-* CLI flags for draft model configuration Signed-off-by: kaiyuanxie <kyxiezju@163.com> * feat: add DSpark speculative decoding draft model training for RL rollout acceleration - New module: vime/backends/megatron_utils/dspark/ (8 files: modeling, loss, attention, markov_head, hidden_capture, export, common, __init__) - Actor integration: draft weight cache/restore lifecycle, freeze-policy - vLLM engine: load_draft_weights_from_file via HTTP /collective_rpc endpoint - L1 loss OOM fix: chunked logsumexp + gradient checkpointing - TP>1 vocab padding 3-stage fix: all-gather (model build), zero-pad (pretrained load), strip (vLLM export) - 14 new --dspark-* CLI arguments - Restored origin/main weight transfer API (common.py, update_weight_from_distributed.py) - DSpark draft weight sync methods added on top of origin/main update_weight_from_tensor.py - Removed unrelated changes (gemma4, gpt_oss, fp8_helpers, ppo_utils, mask_utils, etc.) that were accidentally included in previous commit Verified: all CI tests pass (27 CPU test files + 150 utils tests + 25 upstream-sync tests), pre-commit clean. Signed-off-by: kaiyuanxie <kyxiezju@163.com> * refactor: include DSpark draft params in weights_backuper Replace manual _save_dspark_draft_to_cpu() / _restore_dspark_draft_from_cpu() with automatic draft param management via weights_backuper source_getter. Changes: - Add _iter_dspark_draft_params() to iterate draft model params - Chain draft params into weights_backuper source_getter - Delete _save_dspark_draft_to_cpu() and _restore_dspark_draft_from_cpu() - Add _switch_model(actor) in update_weights() for colocate mode (disable() does not restore GPU memory from TMS backup) - Update stale comments This simplifies the pause/resume cycle: weights_backuper now automatically saves/restores draft params alongside policy params. Signed-off-by: kaiyuanxie <kyxiezju@163.com> * fix: DSpark draft weight sync and local spec compatibility Fix 5 bugs discovered during end-to-end validation on A800: 1. Local spec layernorm names (qwen2.py): Add input_layernorm.weight and pre_mlp_layernorm.weight mappings for --transformer-impl local (TE spec uses different param names). 2. Filter draft params from weight sync (actor.py, hf_weight_iterator_direct.py): weights_getter and _get_megatron_local_param_infos now exclude .draft_model. params to prevent KeyError in HF weight conversion. 3. Non-MoE _ipc_engine setup (update_weight_from_tensor.py): connect_rollout_engines returned early for non-MoE models before setting _ipc_engine/_ipc_gather_src, causing draft weights to never sync to vLLM (spec_accept_rate=0%). Now sets these for all models. 4. Disable RoPE fusion when TE is broken (model_provider.py): When DSPARK_DISABLE_TE is set, also disable apply_rope_fusion to avoid TE fused RoPE kernel errors on A800 (SM 8.0). 5. Safe draft param restore (actor.py): Replace _switch_model(actor) with _restore_dspark_draft_params_safe() which assigns fresh GPU tensors instead of copy_(), working even when TMS pause() freed the original GPU storage. Also adapt to origin/main API changes: - Unpack engine_parallel_configs from get_updatable_engines_and_lock - Remove normalization kwarg from get_gpt_layer_with_transformer_engine_spec - Use convert_to_global_name=True for consistent draft param naming Verified: spec_accept_rate 32.5% (vime rollout), 31-40% (vLLM SpecDecoding) on Qwen3-4B DSpark colocate 8-GPU. Signed-off-by: kaiyuanxie <kyxiezju@163.com> * refactor: simplify DSpark weight sync and validate TP paths (#2) * refactor: simplify DSpark weight sync and TP support * test: remove redundant DSpark hidden capture unit test Signed-off-by: aoshen02 <aoshen@inferact.ai> * fix: add normalization parameter to get_gpt_layer_local_spec The get_gpt_layer_local_spec function requires normalization parameter but it was missing, causing RMSNorm models (e.g. Qwen3) to use incorrect normalization. Pass args.normalization to both MoE and dense model paths. Signed-off-by: CalvinXKY <kyxiezju@163.com> * refactor: simplify DSpark weight export integration Signed-off-by: aoshen02 <aoshen@inferact.ai> --------- Signed-off-by: kaiyuanxie <kyxiezju@163.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: CalvinXKY <kyxiezju@163.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: aoshen02 <aoshen@inferact.ai>
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
1 participant
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.
Summary
This PR targets
codex/slime-sync-v0271directly, the head branch of vllm-project#386, so it does not stack the sync diff again.Issue vllm-project#388
The issue revision (
6cefd846) predates vllm-project#340. Its raw iterator could map a global owner rank to the wrong EP-subgroup root when PP and EP were both greater than one. The reporter's repeated experts at boundaries 0/32 and 31/127 match that defect; vllm-project#340 already fixed the generic raw path.PR vllm-project#386 then removed the old Gemma4 implementation entirely. This PR restores Gemma4 through Bridge and validates that replacement against the issue's failing conversion boundary.
Validation
Gemma4-26B-A4B-it, 30 layers / 128 experts, two 8xH200 nodes:
local=0 remote=0688d34b85c6f3785d4a542ab202c35ab077a904bgit diff --check: passedThe exact 4-node actor/rollout deployment from vllm-project#388 was not rerun; its independently isolated 16-rank HF round-trip boundary was reproduced with the same TP2/PP2/EP4 topology.