Skip to content

Complete main DSv4 determinism port for packed THD/CP - #8194

Draft
FDecaYed wants to merge 9 commits into
NVIDIA:mainfrom
FDecaYed:codex/dsv4-determinism-main-port-20261013
Draft

Complete main DSv4 determinism port for packed THD/CP#8194
FDecaYed wants to merge 9 commits into
NVIDIA:mainfrom
FDecaYed:codex/dsv4-determinism-main-port-20261013

Conversation

@FDecaYed

@FDecaYed FDecaYed commented Oct 13, 2026

Copy link
Copy Markdown
Contributor

Purpose and relationship to the existing ports

Draft main-based companion for #7441, completing its packed THD/static-CP attention/indexer determinism coverage. The automated Mirror Dev to Main run for #7441 did not create a mirror: the source PR has no mirror-to-main label. The previous packed integration #7921 is closed without merging.

This is intentionally not a second monolithic copy of every #7441 change:

Area Main landing path
Stable Top-K, SBHD forward ordering, attention/indexer backward #7708 (shared prerequisite, temporarily included here)
Packed THD/static CP, padding and forward/replay ordering This draft, adapted from the two packed commits of #7921
TE operation-fuser grouped-MLP probability gradients Existing independent #7769; not duplicated here
Deterministic router auxiliary-loss guard Already present on main through its dedicated auxiliary-loss-fusion option

The branch starts from main 3bb482c487cd447a2b40a3bf076ab14cad514e13 and temporarily carries the three shared #7708 commits. Please land/reconcile #7708 first, then rebase this draft to keep the remaining packed-only diff. Its current scope includes the shared prerequisite so the branch is buildable against main; it does not target a bot pull-request/* branch.

Changes

  • Pass explicit deterministic mode through packed indexer Top-K and both packed attention backward paths, preserving default-mode dispatch and static CP process-group plumbing.
  • Preserve compressed-first compact key order between no-grad and training forwards, with sink-only padded rows and no mutation of caller-owned index tensors.
  • Keep clipped-KL/indexer gradients, global real-token normalization, and deferred CP reductions from the existing port.
  • Remove main's now-obsolete blanket rejection of packed deterministic mode. Keep backend/layout/static-CP validation, and test both mode settings against valid and invalid inputs. Unsupported deterministic runtimes and dense indexer backward still fail at their respective wrappers.
  • Mark the five parameterized SM100 attention/indexer regression cases for the GB200 lane. On Make DSv4 attention and indexer backward deterministic #7708's current CI, H100 skips those cases and the GB200 bucket excludes their unmarked file; this draft explicitly selects them.
  • Reconcile the single test-import conflict by retaining both main's cp_utils use and the packed port's packed_sparse_attention use.

Scratch allocation follows #7708's per-call cuDNN allocation strategy. This does not copy #7441's shared device-global workspace cache or resolve the outstanding review on that dev implementation.

Attribution

Original deterministic DSv4 work: Haochen Yuan (@HaochenYuan), #7441. Stable Top-K/forward-ordering approach: Bo Li (@bobboli), #7169. Main shared and packed adaptations, batch-local indexing and clipped-KL corrections: Yan Xu (@Connor-XY), #7708 / #7706 / #7921. Cherry-picked commits retain original authors, co-author trailers, source SHAs and sign-offs, with signed Deyu Fu adaptation commits.

Validation on this branch

  • BASE_REF=main CHECK_ONLY=true bash tools/autoformat.sh: mandatory Black/isort/pylint/ruff checks passed for all nine changed files.
  • Kernel determinism coverage checker: passed for all nine changed files; no exemption label or test disabling.
  • 39 CPU-adapted exact-source numerical/replay/dispatch cases passed for the shared helpers.
  • 10 exact-source CPU packed-wrapper/layout cases passed, including padding preservation, deterministic compaction and the new layout-validation regression.
  • Structural check confirms explicit GB200 selection for the five SM100 regression variants.
  • git diff --check and local signature/sign-off verification passed.

Validation boundary: Osaka has no GPU. These isolated CPU probes bypass GPU package/conftest initialization and stub kernel/collective work; they do not certify CUDA, CP1/2/4, graph replay, or the rebased branch's end-to-end numerics. Optional mypy still reports missing training dependencies/existing typing issues and is not certified clean. The earlier GB300 results documented in #7708/#7921 are provenance, not a claim that this new head passed GPU tests.

Keep this PR draft pending exact-head GPU CI (including the real packed CP and SM100 cases), reconciliation with #7708, and the required reviews. No performance claim and no whole-recipe determinism certification are made.

Connor-XY and others added 6 commits October 13, 2026 17:49
@Connor-XY @FDecaYed
Make DSv4 attention and indexer backward deterministic
f5378e8 
Signed-off-by: Yan Xu <yxu1@nvidia.com>
(cherry picked from commit 915d73c)
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@Connor-XY @bobboli @FDecaYed
Fix deterministic CSA forward key order across grad modes
7fa3f13 
Use the training compressed-first order for deterministic no-grad forwards,
preserving the existing default-mode order. Add dispatch coverage and real
learned-indexer forward parity tests for zero and positive sparse KL loss.

Adapted from Bo Li (@bobboli)'s Megatron-LM PR NVIDIA#7169. The existing stable
Top-K helper also originates in that PR. Deterministic sparse-attention and
indexer backward in this series builds on Haochen Yuan (@HaochenYuan)'s
PR NVIDIA#7441, with our batch-indexing and clipped-KL corrections in NVIDIA#7706.

Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
(cherry picked from commit 3989b77)
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@Connor-XY @FDecaYed
Align CSA parity test TopK with cuDNN score-recompute tiles
c94c34d 
Use 128 TopK slots in the learned-indexer forward-parity fixture. The
positive sparse-loss path requires a multiple of the cuDNN attention
score-recompute kernel's 128-key tile. The 32 compressed keys are padded
by the existing stable selector; production behavior is unchanged.

Signed-off-by: Yan Xu <yxu1@nvidia.com>
(cherry picked from commit cc6e840)
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@Connor-XY @FDecaYed
Make packed DSv4 context-parallel attention deterministic
8cdb628 
Signed-off-by: Yan Xu <yxu1@nvidia.com>
(cherry picked from commit a64d341)
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@Connor-XY @bobboli @FDecaYed
Preserve packed CSA forward ordering and sink-only padding
4305000 
Use compressed-first indices in deterministic packed forwards regardless of
grad mode, compacting no-grad indices as in training. Enforce sink-only
padding at both attention wrapper boundaries without mutating caller data.
Add real CP1/2/4 forward parity and focused wrapper regressions.

Extends Bo Li (@bobboli)'s key-ordering approach from PR NVIDIA#7169 to main's
packed CP implementation; that PR explicitly left CP ordering for follow-up.
The wrapper padding guard is adapted from the same PR. Existing deterministic
backward builds on Haochen Yuan (@HaochenYuan)'s PR NVIDIA#7441 and the batch/KL
corrections in NVIDIA#7706; this commit adds the packed forward contract.

Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
(cherry picked from commit 8e2d0fe)
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@FDecaYed
Enable main packed determinism validation and select SM100 CI coverage
4ee3ca5 
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Oct 13, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@FDecaYed

FDecaYed commented Oct 13, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4ee3ca5

@FDecaYed
Preserve query norm rounding in deterministic forwards
38463cb 
Keep BF16 intermediate rounding consistent between compiled inference and AOTAutograd forwards when deterministic_mode is enabled. Preserve the default compiled path and the existing packed end-to-end bitwise assertions. Add registered grad-mode parity and gradient replay coverage for BF16/FP32 and contiguous/transposed inputs.

Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@FDecaYed

FDecaYed commented Oct 15, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 38463cb

@FDecaYed FDecaYed mentioned this pull request Oct 15, 2026
Open
@balasaajay

balasaajay commented Oct 15, 2026

Copy link
Copy Markdown
Contributor

/ok to test 5c348b0

@FDecaYed
Use eager FP32 query norm in deterministic mode
4185405 
Preserve the compiled low-precision path and the default query norm while avoiding grad-mode compiler differences for FP32. Keep strict parity checks and add a regression for the FP32 dispatch boundary.

Signed-off-by: Deyu Fu <deyuf@nvidia.com>
@FDecaYed

FDecaYed commented Oct 16, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4185405

@lhb8125 lhb8125 mentioned this pull request Oct 26, 2026
Open

This branch was previously deployed

2 inactive deployments
public — 41854050 Deployed Oct 16, 2026 by copy-pr-bot[bot] via build-test-publish-wheels / publish-wheels (megatron-core, arm64) #12079
test — 41854050 Deployed Oct 16, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #24375
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

26.10.01

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@FDecaYed @balasaajay @Connor-XY