feat(draft): preserve split-step normalization parity by seonjinn · Pull Request #3706 · NVIDIA-NeMo/RL

seonjinn · 2026-08-19T00:08:28Z

Summary

Accumulates draft-loss statistics across split microbatches.
Normalizes draft gradients and metrics using the DP-global draft-token count.
Preserves pipeline-wide gradient-norm reporting and clears state on abort or retry.

Scope

Split-step state and lifecycle handling around the soft-CE primitive; no new draft model or runtime integration.

Validation

Exact head: 76c2648d75c76b9f107f5de3d28f48cfa879316b
27 current-head GitHub checks pass; remaining unit jobs are queued.

Prior distributed evidence

OCI 6327354, 6326981, 6327012, 6327013: split lifecycle, PP2 grad-norm propagation, and DP2/DP4 uneven/zero-count FP32/BF16 parity passed.

Dependencies / Remaining

Depends on #3703. The branch-up-to-date check is the only current failure.

Signed-off-by: seonjinn <sna@nvidia.com>

copy-pr-bot · 2026-08-19T00:08:33Z

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 · 2026-08-19T00:08:54Z

Self-review: review-pr-team

Reviewed the PR3-only range on top of #3703, including split-step loss semantics, DP reduction ordering, draft-only gradient rescaling, metrics, and tests.

Resolved during self-review:

Preserved the runtime policy-loss metric contract instead of applying an invalid combined-loss normalization.
Made the reduction expectation device-local for GPU execution.
Corrected the OCI selector so the TQ draft-grad-norm test is exercised.

No unresolved correctness, security, or design findings remain in the PR3-only range.

Validation: OCI-Hsg job 6299067 completed successfully with 6 base tests, 2 MCore split tests, DP=2 NCCL uneven-count gradient/metric parity, Ruff, and compile checks.

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-08-20T05:50:28Z

/review-pr-team

seonjinn · 2026-08-20T06:32:20Z

/ok to test 40badee

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn

Self-review before requesting reviewers.

Method: a focused review agent, then an adversarial pass with empirical CPU execution of the real modules.

The headline claim survived a serious challenge worth recording: a candidate CRITICAL argued the split path normalizes by the draft's own count while the monolithic path uses the policy's global_valid_toks, breaking parity. Refuted on four grounds: (1) roll_tensor(token_mask, shifts=-1) zeroes the vacated last slot, so the draft mask count equals sum(token_mask[:, 1:]) -- exactly the policy formula -- per sample, by construction; (2) both counts ride the same DP all-reduce; (3) the gradient scale composition (scale_gradients(1/n) then correct_main_grads x n x 1/draft_count) nets to raw/draft_count in both paths; (4) an empirical CPU run with a mask exercising the shifted first slot produced identical losses (relative diff 0.0). Also verified clean: zero-count guards, single-DP-collective accumulation, abort/retry state freshness, PP grad-norm collective symmetry for non-owner ranks.

Follow-up commit: the split-vs-monolithic parity regression test (inline) and the draft_loss shared constant (inline).

Notes: CI on this head is green except two non-code failures (an HF-CDN 429 download flake and the label-gated quality check). One pre-existing limitation worth a body line: draft + CP>1 fails the loss's mask/logits shape check identically in both paths -- not a parity regression, but worth stating. #3725 stacks on this PR and relies on the scale_gradients -> correct_main_grads -> optimizer.step() finalization order; flag any reordering there.

seonjinn · 2026-09-17T02:27:53Z

/ok to test fb138a7

Signed-off-by: seonjinn <sna@nvidia.com>

The merge-base changed after approval.

seonjinn · 2026-09-17T09:58:39Z

/ok to test 89a0a09

seonjinn · 2026-09-17T20:00:45Z

/ok to test 6d5c106

seonjinn · 2026-09-18T05:57:01Z

/ok to test 53a9e12

Copied pull-request branches use push events, so Codecov needs an explicit PR number. Pass commit_parent for upload-coverage instead of base_sha, which is only supported by pr-base-picking. Preserve the coverage threshold and all training behavior. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-09-18T08:53:19Z

/ok to test e94b052

Preserve split-step normalization alongside the merged draft optimizer support. Keep both type-check entries and inherit the shared CI workflow from main. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-09-19T00:57:04Z

/ok to test 18138e8

seonjinn · 2026-09-20T01:31:07Z

/ok to test 0cbe19c

seonjinn · 2026-09-20T01:31:42Z

/ok to test 0cbe19c

Preserve deferred draft normalization alongside the indexed teacher head contract. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-09-22T03:10:39Z

/ok to test b4fbcc4

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added 2 commits August 18, 2026 15:16

feat(draft): stream vocab-parallel soft CE

5239101

Signed-off-by: seonjinn <sna@nvidia.com>

test(draft): cover cached single-tile soft CE

1e3d8eb

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn requested review from a team as code owners August 19, 2026 00:08

seonjinn added 6 commits August 18, 2026 17:18

perf(draft): cache single-tile soft CE

9cdc84f

Signed-off-by: seonjinn <sna@nvidia.com>

feat(draft): stream vocab-parallel soft CE

5930782

Signed-off-by: seonjinn <sna@nvidia.com>

test(draft): cover cached single-tile soft CE

43e1b8c

Signed-off-by: seonjinn <sna@nvidia.com>

perf(draft): cache single-tile soft CE

f6af1ff

Signed-off-by: seonjinn <sna@nvidia.com>

feat(draft): preserve split-step normalization parity

681a524

Signed-off-by: seonjinn <sna@nvidia.com>

test(draft): keep reduction expectation on device

1392bb9

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn force-pushed the seonjinn/draft-cotrain-step-state-20260818 branch from 3f5cadd to 1392bb9 Compare August 19, 2026 01:26

seonjinn added 6 commits August 18, 2026 20:01

test(draft): expand split-step normalization coverage

5edddaa

Signed-off-by: seonjinn <sna@nvidia.com>

test(draft): require PP-safe grad norm reduction

8b4bcb5

Signed-off-by: seonjinn <sna@nvidia.com>

fix(draft): reduce grad norm across pipeline ranks

83545f9

Signed-off-by: seonjinn <sna@nvidia.com>

test(draft): preserve split grad norm metric

049ba59

Signed-off-by: seonjinn <sna@nvidia.com>

fix(draft): retain split grad norm metric

0090cd1

Signed-off-by: seonjinn <sna@nvidia.com>

style(draft): format metric propagation tests

40badee

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added the CI:L2 Run doctests, unit tests, functional tests, and convergence tests label Aug 20, 2026

seonjinn mentioned this pull request Aug 20, 2026

[Do not merge into Main] feat(draft): add bounded multi-pass Eagle TTT #3725

Open

seonjinn added 3 commits August 20, 2026 15:41

Merge branch 'main' into seonjinn/draft-cotrain-streaming-loss-20260818

5133c28

Signed-off-by: seonjinn <sna@nvidia.com>

test(draft): assert streaming loss normalization shape

8887579

Signed-off-by: seonjinn <sna@nvidia.com>

Validate caller-supplied bin_ids range in streaming soft CE

895e0c2

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn commented Aug 22, 2026

View reviewed changes

Comment thread tests/unit/models/megatron/test_draft_step_state.py

Comment thread nemo_rl/algorithms/loss/wrapper.py Outdated

terrykong previously approved these changes Sep 17, 2026

View reviewed changes

seonjinn added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:L2 Run doctests, unit tests, functional tests, and convergence tests labels Sep 17, 2026

Merge latest main and remove upstream soft-CE duplication from PR diff

89a0a09

Signed-off-by: seonjinn <sna@nvidia.com>

Merge branch 'main' into seonjinn/draft-cotrain-step-state-20260818

6d5c106

Merge branch 'main' into seonjinn/draft-cotrain-step-state-20260818

53a9e12

github-actions Bot added the CI Relating to CI label Sep 18, 2026

chore(draft): merge latest main and deduplicate PR changes

18138e8

Preserve split-step normalization alongside the merged draft optimizer support. Keep both type-check entries and inherit the shared CI workflow from main. Signed-off-by: seonjinn <sna@nvidia.com>

github-actions Bot removed the CI Relating to CI label Sep 19, 2026

Merge branch 'main' into seonjinn/draft-cotrain-step-state-20260818

0cbe19c

chore(draft): merge latest main for PR NVIDIA-NeMo#3706

b4fbcc4

Preserve deferred draft normalization alongside the indexed teacher head contract. Signed-off-by: seonjinn <sna@nvidia.com>

terrykong enabled auto-merge (squash) September 22, 2026 04:46

terrykong approved these changes Sep 22, 2026

View reviewed changes

terrykong merged commit a9c3178 into NVIDIA-NeMo:main Sep 22, 2026
84 checks passed

yfw mentioned this pull request Sep 22, 2026

build: bump vllm to 0.29.0 #4080

Merged

4 tasks done

seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 22, 2026

Merge latest main after NVIDIA-NeMo#3706 into PR NVIDIA-NeMo#3715

13aa90c

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 22, 2026

Merge latest main after NVIDIA-NeMo#3706 into PR NVIDIA-NeMo#3723

3186bd1

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 22, 2026

Merge latest main after NVIDIA-NeMo#3706 into PR NVIDIA-NeMo#3711

b97a682

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 22, 2026

Merge latest main after NVIDIA-NeMo#3706 into PR NVIDIA-NeMo#3726

9b0561b

Signed-off-by: seonjinn <sna@nvidia.com>

Conversation

seonjinn commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Scope

Validation

Prior distributed evidence

Dependencies / Remaining

Uh oh!

copy-pr-bot Bot commented Aug 19, 2026

Uh oh!

seonjinn commented Aug 19, 2026

Self-review: review-pr-team

Uh oh!

seonjinn commented Aug 20, 2026

Uh oh!

seonjinn commented Aug 20, 2026

Uh oh!

seonjinn left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

seonjinn commented Sep 17, 2026

Uh oh!

seonjinn commented Sep 17, 2026

Uh oh!

seonjinn commented Sep 17, 2026

Uh oh!

seonjinn commented Sep 18, 2026

Uh oh!

seonjinn commented Sep 18, 2026

Uh oh!

seonjinn commented Sep 19, 2026

Uh oh!

seonjinn commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

seonjinn commented Sep 20, 2026

Uh oh!

seonjinn commented Sep 22, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

seonjinn commented Aug 19, 2026 •

edited

Loading

seonjinn commented Sep 20, 2026 •

edited

Loading