Skip to content

feat(dflash): add fused CE and total anchor budget - #3673

Draft
Slyne wants to merge 4 commits into
NVIDIA-NeMo:mainfrom
Slyne:slyne/feat/dflash-fused-anchor-budget
Draft

feat(dflash): add fused CE and total anchor budget#3673
Slyne wants to merge 4 commits into
NVIDIA-NeMo:mainfrom
Slyne:slyne/feat/dflash-fused-anchor-budget

Conversation

@Slyne

@Slyne Slyne commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

What

  • add a local-microbatch max_total_anchors budget so anchor construction has a predictable memory ceiling
  • add chunked fused LM-head + cross-entropy execution that avoids retaining full-vocabulary logits
  • preserve dense/fused metric and gradient parity, including TP and ancestor-owned FSDP2 target heads
  • normalize weighted DFlash-family objectives over the draft data-parallel group, with separate DFlash2 base and selector denominators
  • thread the anchor budget through DFlash2, Domino, and JetSpec and document the new controls

The shipped Qwen3 DFlash examples enable fused CE to reduce peak training memory. The recipe API still defaults it to disabled. Fused CE supports loss_type: dflash; variable-prefix training must disable it.

Testing

  • 243 passed across the focused DFlash, DFlash2, Domino, JetSpec, loss, draft, TP, FSDP, and recipe suites
  • real two-rank DDP/no_sync gradient parity with uneven per-rank base and selector counts, including a rank with no selector examples
  • Ruff formatting/lint and git diff --check

Review

  • independent review-agent loop: no actionable findings
  • Claude Code Opus, maximum effort, repeated after fixes: no actionable findings

Slyne added 3 commits August 25, 2026 11:45
Signed-off-by: SlyneD <slyned@nvidia.com>
Signed-off-by: SlyneD <slyned@nvidia.com>
Signed-off-by: SlyneD <slyned@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 25, 2026

Copy link
Copy Markdown

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.

@github-actions

github-actions Bot commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Signed-off-by: SlyneD <slyned@nvidia.com>
@Slyne
Slyne force-pushed the slyne/feat/dflash-fused-anchor-budget branch from 2924c8d to ee03a53 Compare August 26, 2026 21:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant