feat(automodel): speculative decoding draft co-training

[yuekaizhang]

Summary

Ports speculative-decoding draft co-training (DSpark and DFlash block
drafters, and an EAGLE3 TTT drafter) from RL_dspark onto NeMo-RL's
automodel (DTensor v2) backend, so a draft model can be trained jointly
with the policy and refit into vLLM (draft.* weight stream) for
speculative-decoding generation.

• New nemo_rl/models/automodel/draft/ module: EAGLE3 loss/integration,
  block-draft (dspark/dflash) loss, and a shared markov head, plus CP
  draft-loss rescale, long-context chunking (TTT lm_head/KL), embed-sharing
  and weights/model resume fixes, and support for native/flat (SGLang
  SpecForge-style) EAGLE3 checkpoints.
• New recipes: Qwen3-8B (GRPO + DAPO, 1n8g) DSpark/DFlash/EAGLE3
  co-training siblings, and Qwen3-30B-A3B-Instruct-2507 (4n8g, EP32)
  EAGLE3 co-training with the lmsys SpecForge drafter (published as a
  config-patched fork at
  yuekai/SGLang-EAGLE3-Qwen3-30B-A3B-Instruct-2507-SpecForge-Nex —
  only max_position_embeddings changed, 2048 -> 40960, to stop vLLM
  clamping the drafter's RoPE).
• Nightly test coverage for dspark/dflash/eagle3 co-training on the
  automodel backend.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuekaizhang]

Merge conflict note: gh pr view reports mergeable=CONFLICTING, mergeStateStatus=DIRTY against main — please rebase and resolve conflicts.

Reviewed by a 5-agent team (rl-expert, comment-reviewer, test-agent, bug-finder, design-reviewer) plus an adversarial devil's-advocate pass. This is a well-guarded feature port — no functional bugs found, extensive setup-time asserts/raises, strong existing unit-test coverage for the tricky numerics. Affirmations: grpo.py/grpo_sync.py have zero dspark/dflash/eagle3 string references (clean _DraftRuntimeBase polymorphic seam, single dispatch factory in dtensor_policy_worker_v2.py); dspark/dflash unification doesn't leak per-variant conditionals into shared runtime code; common.py's header comment cleanly documents the vendored-vs-local boundary against upstream Automodel.

1 critical item (this PR fails its own configs-minimize-check CI hook on 2 recipes), 9 suggestions, 2 informational notes — see inline comments.

Generated by Claude Code

[yuekaizhang]

grpo-qwen3-8b-4n8g-automodel-eagle3.yaml:20

1 action item.

PR-introduced. algo: "eagle3" is byte-identical to DEFAULT_DRAFT_ALGO (nemo_rl/models/policy/__init__.py:532). Reproduced directly: uv run tools/config_cli.py minimize-check examples/configs/recipes/llm/grpo-qwen3-8b-4n8g-automodel-eagle3.yaml reports "is not minimized" — this PR fails its own configs-minimize-check pre-commit hook (same for the 30B sibling recipe, commented separately).

Root cause: the default is scattered via .get("algo", DEFAULT_DRAFT_ALGO) across 3 call sites (setup.py:764, lm_policy.py:138, dtensor_policy_worker_v2.py:391) instead of one schema default, so an explicit algo: "eagle3" line reads as meaningful to a human but is silently treated as redundant by the minimizer.

AI-1

Action: drop this line (it's already the default), or better, make the default an explicit schema field so .get(key, MODULE_CONSTANT) doesn't need to be repeated at 3 call sites.

Suggested change

algo: "eagle3"

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

grpo-qwen3-30ba3b-2507-instruct-4n8g-automodel-eagle3.yaml:101

1 action item.

PR-introduced. Same issue as the Qwen3-8B eagle3 recipe (see that comment): algo: "eagle3" duplicates DEFAULT_DRAFT_ALGO and fails configs-minimize-check here too — verified by running the tool directly on this file.

AI-1

Action: drop this line, or fix the root cause (see the other recipe's comment).

Suggested change

algo: "eagle3"

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

integration.py:530

1 action item.

PR-introduced. Redundant in-function import os — already imported at module top (line 25).

AI-1

Suggested change

"""

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

common.py:81

1 action item.

PR-introduced. build_anchor_candidate_mask is a cross-file interface (imported by draft_qwen3.py) with no docstring.

AI-1

Action: add a Google-style docstring, e.g.:

Suggested change

	def build_anchor_candidate_mask(
	*,
	seq_len: int,
	loss_mask: torch.Tensor,
	doc_remaining: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""Mask of anchor positions whose first predicted token is supervised.
	Args:
	seq_len: source sequence length.
	loss_mask: [batch_size, seq_len] per-token supervision mask.
	doc_remaining: [batch_size, seq_len] tokens remaining in the current
	packed document, or None when unpacked.
	Returns:
	[batch_size, seq_len - 1] boolean mask; True at anchor position p
	when token p + 1 is supervised (and, when packing, stays inside the
	anchor's document).
	"""

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

common.py:106

1 action item.

PR-introduced. sample_anchor_positions — same gap, also cross-file (used by draft_qwen3.py).

AI-1

Action: add a Google-style docstring, e.g.:

Suggested change

	def sample_anchor_positions(
	*,
	seq_len: int,
	loss_mask: torch.Tensor,
	num_anchors: int,
	device: torch.device,
	doc_remaining: Optional[torch.Tensor] = None,
	generator: Optional[torch.Generator] = None,
	) -> tuple[torch.Tensor, torch.Tensor]:
	"""Sample up to num_anchors valid anchor positions per sample.
	Args:
	seq_len: source sequence length.
	loss_mask: [batch_size, seq_len] per-token supervision mask.
	num_anchors: maximum anchors to sample per sample.
	device: device to allocate outputs on.
	doc_remaining: packed-document remaining-token counts, or None.
	generator: optional RNG for a reproducible draw (must be seeded
	identically across TP peers sharing the same draft replica).
	Returns:
	(anchors, keep_mask): anchors is [batch_size, num_anchors] sorted
	anchor token positions (zero-padded past the valid count); keep_mask
	marks which of those are real vs. padding.
	"""

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

integration.py:1282

No action needed — low-severity design suggestion, not a blocker.

Context — no action. This file is 1431 lines with 4 near-independent clusters (config adaptation, model building, runtime/loss, checkpoint I/O). The checkpoint I/O cluster from here (draft_checkpoint_dir) through EOF (~150 lines) has zero references to the runtime/loss classes above it (_DraftRuntimeBase/DSparkRuntime/Eagle3Runtime/DSparkHiddenCapture) — verified by grep. A future engineer debugging checkpoint-metadata issues has to scroll past ~1100 lines of unrelated config/runtime code to reach it.

One option — of several: split this cluster into its own checkpoint.py module; the near-zero cross-references make it a low-risk, mechanical extraction.

[yuekaizhang]

vllm_worker.py:615-621

1 action item.

PR-introduced. _draft_module_sharing_disable_required(config) is defined in this same file, but this block re-implements the identical method in ("dspark","dflash") or (method=="eagle3" and draft_full_refit) predicate inline instead of calling it. A future 4th co-trained method (or a policy change to when eagle3 counts as full-stream) needs updating in both places by hand.

AI-1

Action: call the existing helper here instead of re-deriving the predicate.

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

test_draft_config_adapters.py:23

1 action item.

PR-introduced. This file (and 7 siblings: test_dspark_checkpoint_meta.py, test_draft_ratio_metrics.py, test_draft_refit_transports.py, test_draft_loss_decay.py, test_draft_vocab_maps.py, test_dspark_anchor_sampling.py, test_eagle3_trainer_grads.py) import from nemo_rl.models.automodel.draft.*, which unconditionally imports nemo_automodel at module scope — but none of the 8 have the try: import nemo_automodel / pytest.skip(allow_module_level=True) guard or the @pytest.mark.automodel marker that every pre-existing sibling test file in this directory has. Verified against conftest.py's marker-based include/exclude logic and the CI Docker build (single shared venv with all extras, so this won't crash any shard today) — but these 8 files are excluded from the --automodel-only shard and instead run unmarked in the generic Models shard, inconsistent with the established convention and fragile if a future change makes nemo_automodel non-default there.

AI-1

Action: add the skip guard + @pytest.mark.automodel (or module-level pytestmark) to all 8 files, matching test_automodel_checkpoint.py's pattern.

[yuekaizhang]

fixed in a18b69d

[yuekaizhang]

worker_groups.py:204-210

No action needed — informational scoping note.

Context — no action. This 185-line diff (per-pg_idx IsolatedWorkerInitializer pooling, batched address allocation, runtime_env_overrides passthrough) has zero draft/eagle3/dspark references (verified by grep) — general Ray worker-group infra bundled into this feature-named PR. Not a defect, just flagging it widens this PR's review surface/blast radius and makes a future revert of the draft feature alone harder; consider a separate PR next time for this class of change.

[yuekaizhang]

integration.py:951-977

No action needed — minor nit, not a blocker.

Context — no action. compute_dspark_loss's confidence-head diagnostic terms (confidence_abs_error_num, confidence_bias_num, confidence_cumprod_bias_num, confidence_diag_den) are computed whenever a confidence head is enabled, but DSparkRuntime._terms_to_metrics never forwards them into the metrics dict — verified genuinely dropped via repo-wide grep (no consumer anywhere). Not a correctness bug, just wasted compute and no calibration visibility for confidence-head tuning. Worth a follow-up to either wire these into metrics or drop the computation.