Feat(megatron): EAGLE3 drafter training on Megatron backend

[ys2025-AI]

Basic information

Scope: Add Pipeline Parallelism (PP>1) and Sequence Parallelism (SP=True) support to the Megatron old-logprob hidden-state collection path, enabling EAGLE3 drafter co-training under TP+SP+PP topologies (e.g. Qwen3-4B TP4PP2, Qwen3-8B TP8PP1) on NPU.

Background: verl-SpeCo's EAGLE3 drafter online training captures specific-layer hidden states (aux layers + final-norm output) from the target model's compute_old_log_prob forward via register_forward_hook, then selects/merges/distributes them to drafter training workers. The previously merged efad0aa only supports Megatron TP; this PR adds PP>1 cross-stage exchange and fixes SP=True under PP async actors.

Problem solved: When PP>1, aux layers are split across pipeline stages — each stage's forward hook can only capture its local layers, so captures must be exchanged across stages and merged at the last stage. SP=True additionally shards hidden states across TP ranks, requiring all TP ranks to participate in the exchange.

Key changes

1. Temp-file cross-stage exchange (oldlogprob_runtime.py). Non-last stages pickle captures to /dev/shm/speco_pp_exchange/pp_{rank}_fbb_tp{N}.pkl (atomic .tmp+rename); the last stage polls (30s timeout), loads, and merges. This avoids HCCL P2P (deadlocks on NPU isend/irecv) and Ray object-store blocking (which freezes the async-actor event loop and deadlocks pipeline HCCL collectives).
2. Single ray.put of the final concatenated tensor at the last stage (owner is correct), instead of 132 per-microbatch ray.puts that would pin object store → spill → disk pressure. Task runner ray.gets once and immediately dels the ref.
3. index_copy_ reorder from micro-batch to full-batch order using indices saved by speco_aggregate_output — fixes the sample-misalignment regression where use_dynamic_bsz=True caused acceptance to drop (2.03→1.98).
4. SP=True hardening (c8fdc53): per-tp_rank file tag (fbb_tp{N}) so every TP rank exchanges its distinct SP shard; sparse-aware device inference + CPU→device remap after pickle round-trip; dense conversion; tp_rank=0-only attach guard; expected-warning downgrade to debug.
5. Context cleanup after exchange so the subsequent ref log_prob (no collect mask) does not enter exchange and write an empty temp file — fixes no_trainable_batch on step 2+.
6. Bug fixes (§2.5): wrong load_format key (acceptance=1.0), nbytes() TypeError, empty-file skip, context clear.

Design notes

• PP>1/PP=1 isolation: three branch points (speco_megatron_postprocess, _megatron_post_stage_exchange_and_consume, _megatron_save_stage_captures) fully isolate the two paths; PP=1 is byte-for-byte identical to efad0aa and unaffected.
• Upstream zero-modification: all TP/SP/PP logic lives in verl_speco/ via monkey-patch; Megatron-LM / MindSpeed / vllm-ascend / verl upstream core source is untouched.
• NPU/GPU not hard-bound: oldlogprob_runtime.py has 0 NPU code lines — device is inferred from saved-capture tensors; speco_mem_watchdog.py is 100% platform-agnostic (/proc/meminfo + disk_usage + OS signal).

Test environment

Item	Value
Hardware	Ascend NPU, 8 cards per node
Accelerators	ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 (8 NPUs)
Actor backend	Megatron (TP=4, PP=2, SP=True)
Rollout	vLLM TP=2, enforce_eager=True, gpu_memory_utilization=0.5
Model	Qwen3-4B (Qwen/Qwen3-4B)
Drafter	Qwen3-4B EAGLE3 (AngelSlim/Qwen3-4B_eagle3_AngelSlim)
Dataset	DAPO-Math-17k (train 10k / val)
Algorithm	GRPO, use_kl_in_reward=False, kl_loss_coef=0.001 (low_var_kl)
Train batch	64; ppo_mini_batch=16; ppo_micro_batch_per_gpu=10; dynamic_bsz=True
Drafter training	EAGLE3, spec_steps=3, spec_topk=1, spec_verify_tokens=4, step=20, collect/train_interval=1, publish_async=True, publish_dtype=bf16
Ray	object_store_memory=200GB hard cap, temp_dir=/dev/shm/ray, spill_threshold=0.9, num_cpus=64, worker_soft_limit=8
Watchdog	min_free_gb=250 (RAM), fs_min_free_gb=50 (/dev/shm + /), poll=10s
Max prompt / response	512 / 8192
Save/test freq	10 / 5; total_epochs=1

Test results

A. Qwen3-4B TP4PP2 SP=True — primary run

Drafter acceptance length (every 5 steps):

step	acceptance	step	acceptance
1	2.031	50	2.310
5	2.206	65	2.305
10	2.236	80	2.269
15	2.273	95	2.275
30	2.323	100	2.278
45 (peak)	2.355	102 (last)	2.277

Stats: min=2.031, max=2.355, mean=2.284. Acceptance rose +12% in the first 15 steps, peaked at step 45, then stabilized around 2.28.

Validation accuracy (val-core/math_dapo/acc/mean, every 5 steps):

step	val acc	step	val acc
5	0.125	55	0.469
10	0.289	65	0.523
15	0.336	75	0.484
25	0.375	85	0.563
35	0.461	95	0.547
45	0.461	100	0.523

Val accuracy rose from 0.125 → ~0.52–0.56 (4× improvement).

Reward (critic/score/mean): −0.762 (step 1) → +0.050 (step 10) → oscillates near 0 as training progresses.

Step-5 timing breakdown (SP=True vs SP=False, seconds):

phase	SP=False	SP=True
gen	335	286
old_log_prob	265	276
ref	38	37
update_actor	94	88
drafter	12	17
step total	750	709

SP=True old_log_prob slightly higher (+11s, SP all-reduce) but step total slightly lower (−41s) — equivalent.

B. SP=True vs SP=False equivalence (step 10)

metric	SP=False	SP=True
acceptance	2.235	2.236
reward	+0.044	+0.050

Equivalent — confirms the SP all-reduce deadlock is resolved; SP no longer needs to be disabled.

C. FSDP2 vs Megatron-TP8PP1 equivalence (Qwen3-8B, step 1-10)

step	FSDP2 acc	Megatron acc	FSDP2 reward	Megatron reward
1	2.034	2.034	−0.706	−0.644
5	2.310	2.309	−0.369	−0.400
10	2.328	2.337	+0.125	+0.150

Per-step acceptance and reward curves overlap to within <0.01 — verifies Megatron-TP drafter training is equivalent to FSDP2.

D. Stability

	FSDP2 8B	Megatron 8B (TP8PP1)	Megatron 4B (TP4PP2)
steps run	10	86 (first 10 compared)	102
trained=1.0 steps	10/10	10/10	102/102
crash/OOM	none	none	none
no_trainable_batch	none	none	none