[seonjinn]

Summary

• Adds a bounded multi-pass EAGLE-3 TTT provider.
• Uses structured MCore attention without retaining dense pass logits or sequence-square masks.
• Validates rotary-position support when the session is created.

Scope

Covers the provider, pass state, attention, loss, and split-step training primitives. Higher-level configuration and recipes remain out of scope.

Validation

• Exact head: fdf2a971be0b396ed7282a9a99c79a6af24a5e7b
• Ruff and Pyrefly pass.
• Focused CPU suites: 50 passed / 11 CUDA-skipped; TP suite: 2 passed / 1 CUDA-skipped.
• All completed GitHub checks pass.

Prior GPU evidence

• Prior 4×GB200 focused suite: 139 passed with 1 optional-dependency skip.

Dependencies / Remaining

Depends on #3703 and #3706. CUDA memory cases remain GPU-only validation.

[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.

[seonjinn]

/review-pr-team

[seonjinn]

Self-review before requesting reviewers.

Method: six review agents over the PR8a-only diff (8 files vs the #3703+#3706 base), then an adversarial pass that re-verified citations, ran two empirical checks on CPU torch, and re-checked CI/PR state live.

The core math held up well under adversarial verification: the multi-pass staircase visibility and the RoPE-relative shift alignment in _merge_branch were re-derived against an independently written rotate-half oracle and match to machine precision (4.4e-16, fp64); label alignment (teacher_offset = pass_index + 1) was re-worked by hand; the cross-pass detach asymmetry is pinned by an explicit test; and the TP/SP integration follows the clean pattern -- EagleTTTCoreAttention owns zero parameters, so the sharded-checkpoint bug classes from earlier PRs in this series are structurally inapplicable. The SP test (real NCCL SelfAttention with sequence_parallel=True, output + input grads + parameter grads vs a global dense oracle) is the strongest distributed test in the series so far.

The one substantive finding is the activation-recompute incompatibility (inline, empirically reproduced): with recompute_granularity=full or selective, backward re-enters the armed forward and crashes on step 1.

Process, before merge:

1. No /ok to test has ever been issued on this PR -- zero substantive CI on any commit. Needs a fresh /ok to test 292e29d31a6276c0e4823c37d94f7c6e917f2fc1 (after the rebase below, on the new head).
2. Rebase required: both stacked dependencies have diverged (not just advanced) -- #3706's head (40badeea) and #3703's head (5133c28d) no longer contain this PR's base 660ee3c6, so the rebase may not be clean.
3. Body accuracy: the Ruff/pyrefly/4xGB200 numbers are self-reported (no check-runs exist); "139 passed" spans the full 13-file stacked suite, not this PR's 5 test files (~91 collected) -- worth stating the basis.
4. Below inline threshold: MCoreEagleTTTSession is never driven past pass_count=2 off-GPU, and the "pinned_modelopt" test name pins an internal formula, not an external ModelOpt oracle -- consider a rename or a small modelopt-driven geometry diff test later.

[seonjinn]

/ok to test 95e639b

[seonjinn]

/ok to test be80ab6

[seonjinn]

Merged latest main (b03da0f469) and re-fired CI on be80ab656b.

main now carries the typed speculator contract from #3701, so this branch had to be adapted rather than merged verbatim. Three changes, one of which git merged clean but broke at runtime:

1. nemo_rl/models/policy/__init__.py — took main's deletion.
#3701 removed the DraftConfig / DraftConfigDisabled TypedDicts and replaced them with Eagle3DraftConfig(BaseModel, extra="allow") in nemo_rl/models/policy/draft_config.py. This branch's only addition to that file was the token_chunk_size: NotRequired[int] field, so the whole conflicted block is dropped in favour of the new contract. policy.draft.token_chunk_size still validates and still reaches the loss wrapper, as a pydantic extra.

2. pyrefly.toml — sorted union.
Both sides inserted into the same sorted list under nemo_rl/models/megatron/draft/: eagle_provider.py / eagle_ttt.py / step_state.py from here, training.py / utils.py from main.

3. nemo_rl/models/megatron/train.py — .get() no longer exists on the draft config.
This one is worth calling out because it was not a merge conflict. policy_cfg["draft"] used to be a TypedDict; after #3701 it is a pydantic model. Two sites here still read

self.cfg["draft"].get("token_chunk_size", DEFAULT_DRAFT_TOKEN_CHUNK_SIZE)

which git merged clean, passes lint, and raises AttributeError only on a draft-enabled run. Rewritten as getattr(self.cfg["draft"], "token_chunk_size", DEFAULT_DRAFT_TOKEN_CHUNK_SIZE), which preserves the exact prior semantics.

Known gap, deliberately not fixed here. token_chunk_size is left undeclared on Eagle3DraftConfig, so #3701's _reject_near_miss_extra_keys validator cannot catch a typo of it - token_chunk_sze would silently fall back to 4096 instead of raising. Declaring the field properly would mean importing DEFAULT_DRAFT_TOKEN_CHUNK_SIZE from nemo_rl/algorithms/loss/draft.py, which pulls torch into draft_config.py - a module #3701 deliberately kept to pydantic and stdlib only. Happy to declare it with a literal default here instead if reviewers prefer the typo protection over the duplicated constant.