[Bugfix] Restore DSpark draft param views after TMS resume

[CalvinXKY]

Summary

• Restore DSpark draft param views after TMS resume (commit 8fbfd891): Fixes param view corruption when TorchMemorySaver pause/resume interacts with DSpark draft model weight sync.
• Patch vLLM DFlash/DSpark meta-device bug (commit 934a91ff): Runtime monkeypatch for vllm-project/vllm#55076. IPC engine weight sync leaves some draft attention params on meta device; _build_context_kv_buffers then crashes on torch.cat with mixed CUDA/meta tensors. The patch skips rebuild when previous CUDA buffers exist (incremental IPC update), raises on first-load failures, and recomputes RoPE cos_sin_cache from meta.

Dependencies

• Requires vLLM with IPC engine support (vLLM 0.27.2+)
• vLLM meta-device fix: vllm-project/vllm#55076 (temporary runtime patch included; will be removed after upstream merge)

What was removed

• DSpark smoke test (tests/test_qwen3_4B_dspark_short.py) — blocked by TE 2.16.1 sm80 incompatibility on A800; will re-add after vLLM upstream fix lands
• CI assertions for spec_accept_rate in train_metric_utils.py

Test Plan

• DSpark 4B training verified on A800 (non-colocate, 100 steps)
• vLLM meta-device patch verified: 5/5 unit tests pass on .62
• vLLM upstream PR #55076 merged
• Remove runtime patch after upstream fix is in base image
• Re-add DSpark smoke test

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34368796 | 📁 Comparing 0d6921d against latest (85ac426)

  🔍 Preview build  

45 files changed · ± 45 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...

[gemini-code-assist]

Code Review

This pull request introduces synchronization of DSpark draft model parameters after the optimizer step in vime/backends/megatron_utils/model.py to fix broken view relationships, and removes the normalization argument from the model provider in vime/backends/megatron_utils/model_provider.py. Feedback on the changes points out potential TypeError issues when unpacking ParamIndexSpec and param_to_index entries due to differences in Megatron-LM versions, and suggests a robust fallback to handle both dataclass attributes and slice objects.

[gemini-code-assist]

Unpacking spec directly as (_start, _end, bucket_id) can cause a TypeError because ParamIndexSpec in standard Megatron-LM is a standard Python dataclass and is not iterable by default. Additionally, bucket.param_to_index[param_obj] can be a slice object rather than a tuple in some Megatron-LM/NeMo versions, which would also raise a TypeError when unpacked directly as local_start, local_end.

To ensure compatibility and robustness across different Megatron-LM versions, we should check if spec has a bucket_id attribute and handle entry as either a slice or a tuple/list.

Suggested change

	for param_obj, (_start, _end, bucket_id) in pim.items():
	if id(param_obj) not in draft_param_ids:
	continue
	bucket = buffer.buckets[bucket_id]
	if hasattr(bucket, "param_to_index") and param_obj in bucket.param_to_index:
	local_start, local_end = bucket.param_to_index[param_obj]
	if isinstance(local_start, int):
	view = bucket.param_data.view(-1)[local_start:local_end].view(param_obj.data.shape)
	param_obj.data.copy_(view)
	for param_obj, spec in pim.items():
	if id(param_obj) not in draft_param_ids:
	continue
	bucket_id = spec.bucket_id if hasattr(spec, "bucket_id") else spec[2]
	bucket = buffer.buckets[bucket_id]
	if hasattr(bucket, "param_to_index") and param_obj in bucket.param_to_index:
	entry = bucket.param_to_index[param_obj]
	if isinstance(entry, slice):
	local_start, local_end = entry.start, entry.stop
	else:
	local_start, local_end = entry[0], entry[1]
	if isinstance(local_start, int):
	view = bucket.param_data.view(-1)[local_start:local_end].view(param_obj.data.shape)
	param_obj.data.copy_(view)

[CalvinXKY]

Tests

#!/bin/bash

# DSpark 4-GPU 100-step: DSpark ON, policy NOT frozen
# Goal: Fair comparison with baseline - main model trains via GRPO, draft trains via 3-loss
# GPUs: 0-3 | TP=2 | colocate | 4 GPUs
#
# Key change: --dspark-freeze-policy REMOVED (policy model trains normally)
# DSpark config:
#   - Pretrained draft model (Qwen3-4B-dspark-pretrained)
#   - 5 draft layers, block_size=7, num_speculative_tokens=7
#   - 3-loss training: CE + L1 + confidence
#   - Policy model trains via GRPO policy_loss (NOT frozen)

set -ex

# Cleanup residual processes
pkill -9 -f '[v]llm serve|VLL[M]::' || true
sleep 3
ray stop --force || true
pkill -9 ray || true
pkill -9 python || true
sleep 3
pkill -9 ray || true
pkill -9 python || true

export PYTHONUNBUFFERED=1

# Detect NVLink
NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
    HAS_NVLINK=1
else
    HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"

NUM_GPUS=4
echo "NUM_GPUS: $NUM_GPUS"

# Model config (Qwen3-4B) - inlined from models/qwen3-4B.sh
MODEL_ARGS=(
   --swiglu
   --num-layers 36
   --hidden-size 2560
   --ffn-hidden-size 9728
   --num-attention-heads 32
   --group-query-attention
   --num-query-groups 8
   --use-rotary-position-embeddings
   --disable-bias-linear
   --normalization "RMSNorm"
   --norm-epsilon 1e-6
   --rotary-base 1000000
   --vocab-size 151936
   --kv-channels 128
   --qk-layernorm
)

CKPT_ARGS=(
   --hf-checkpoint /data/nfs_87/xky/models/Qwen3-4B
   --ref-load /data/nfs_87/xky/models/Qwen3-4B_torch_dist
   --save /data/nfs_87/xky/dspark_test/save_dspark_4gpu_100step_nofreeze
   --save-interval 50
)

ROLLOUT_ARGS=(
   --prompt-data /data/nfs_87/datasets/dapo-math-17k/dapo-math-17k.jsonl
   --input-key prompt
   --label-key label
   --apply-chat-template
   --rollout-shuffle
   --rm-type deepscaler
   --num-rollout 100
   --rollout-batch-size 32
   --n-samples-per-prompt 8
   --rollout-max-response-len 4096
   --rollout-temperature 1

   --global-batch-size 256
   --balance-data
)

EVAL_ARGS=(
   --eval-interval 20
   --eval-prompt-data aime /data/nfs_87/datasets/aime-2024/aime-2024.jsonl
   --n-samples-per-eval-prompt 16
   --eval-max-response-len 16384
   --eval-top-p 1
)

PERF_ARGS=(
   --tensor-model-parallel-size 2
   --sequence-parallel
   --pipeline-model-parallel-size 1
   --context-parallel-size 1
   --expert-model-parallel-size 1
   --expert-tensor-parallel-size 1

   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   --use-dynamic-batch-size
   --max-tokens-per-gpu 9216
)

GRPO_ARGS=(
   --advantage-estimator grpo
   --use-kl-loss
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl
   --entropy-coef 0.00
   --eps-clip 0.2
   --eps-clip-high 0.28
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-6
   --lr-decay-style constant
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98
)

VLLM_ARGS=(
   --rollout-num-gpus-per-engine 2
   --vllm-gpu-memory-utilization 0.5
   --vllm-enforce-eager
   --vllm-speculative-config '{"method":"dspark","model":"/data/nfs_87/xky/models/Qwen3-4B-dspark-pretrained","num_speculative_tokens":7}'
)

MISC_ARGS=(
   --attention-dropout 0.0
   --hidden-dropout 0.0
   --accumulate-allreduce-grads-in-fp32
   --attention-softmax-in-fp32
   --attention-backend flash
)

DSPARK_ARGS=(
   --dspark-enabled
   --dspark-block-size 7
   --dspark-num-draft-layers 5
   --dspark-target-layer-ids 1,9,17,25,33
   --dspark-markov-rank 256
   --dspark-markov-head-type vanilla
   --dspark-num-anchors 512
   --dspark-mask-token-id 151669
   --dspark-ce-loss-alpha 0.5
   --dspark-l1-loss-alpha 0.5
   --dspark-confidence-head-alpha 0.1
   --dspark-loss-decay-gamma 4.0
   --dspark-draft-loss-weight 1.0
   --dspark-intermediate-size 9728
   --dspark-pretrained-model /data/nfs_87/xky/models/Qwen3-4B-dspark-pretrained/model.safetensors
)

# Launch ray
export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265

VIME_DIR="/data/nfs_87/xky/dspark_new/vime"

RUNTIME_ENV_JSON="{
  \"working_dir\": \"${VIME_DIR}\",
  \"env_vars\": {
    \"PYTHONPATH\": \"/root/Megatron-LM/\",
    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
  }
}"

ray job submit --address="http://127.0.0.1:8265" \
   --runtime-env-json="${RUNTIME_ENV_JSON}" \
   -- python3 train.py \
   --actor-num-nodes 1 \
   --actor-num-gpus-per-node ${NUM_GPUS} \
   --colocate \
   ${MODEL_ARGS[@]} \
   ${CKPT_ARGS[@]} \
   ${ROLLOUT_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${GRPO_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${EVAL_ARGS[@]} \
   ${VLLM_ARGS[@]} \
   ${MISC_ARGS[@]} \
   ${DSPARK_ARGS[@]}

[aoshen02]

Thanks for tracking this down. I think the current patch fixes the symptom, but the synchronization is placed too late and on a very hot path.

The root problem described here is that torch_memory_saver.resume() breaks the alias between each DSpark param.data and its Megatron DDP contiguous parameter buffer. Copying values after every optimizer.step() compensates for that broken alias indefinitely:

• it scans every chunk, buffer, and parameter on every training step;
• it puts a DSpark/TMS lifecycle workaround in the generic forward_step();
• it depends on private/version-sensitive Megatron structures (param_index_map, buckets, and param_to_index);
• it may silently skip synchronization, and direct ParamIndexSpec unpacking is not valid in every Megatron version.

A more fundamental and smaller fix would be to restore the view once, immediately after TMS resume:

torch_memory_saver.resume()

if self.args.dspark_enabled:
    restore_dspark_param_views(self.model)

The helper should live with the DSpark model integration (for example dspark/modeling.py), and rebind each draft parameter to its slice in the DDP contiguous parameter buffer. Once the alias is restored, the normal optimizer path updates draft parameters without an extra per-step copy.

Conceptually:

def restore_dspark_param_views(model):
    for chunk in model:
        draft_params = set(unwrap_model(chunk).draft_model.parameters())
        for buffer in chunk.buffers:
            for param, spec in buffer.param_index_map.items():
                if param in draft_params:
                    param.data = buffer.param_data[
                        spec.start : spec.end
                    ].view_as(param)

The exact ParamIndexSpec fields should be taken from the pinned Megatron version rather than adding multi-version hasattr fallbacks.

The normalization correction looks right. Using getattr(args, "dspark_enabled", False) at the model-provider boundary is also reasonable because the provider is reused by Megatron conversion tooling with a different argument namespace.

I would also add a focused regression test instead of relying only on the 100-step convergence run:

1. capture the draft parameter and DDP-buffer relationship;
2. execute TMS pause/resume;
3. restore the views;
4. assert the draft parameter points at the expected buffer storage;
5. run one optimizer step and assert the draft parameter changes without any post-step copy.

That directly tests the reported failure mode and would make the fix both smaller and more first-principled.