Support row-only MXFP8 distributed master-weight casts by xiuhu17 · Pull Request #3488 · NVIDIA/TransformerEngine

xiuhu17 · 2026-09-06T06:28:18Z

Description

Support distributed master-weight writeback into MXFP8 primary weights whose columnwise data and scales are absent, while preserving the bidirectional path.

Memory motivation: radixark/miles#615
Related initialization change: #3468
Related Megatron integration: NVIDIA/Megatron-LM#7095

Changes

Pack/reduce amax only for directions that exist, including mixed row-only/bidirectional batches. Retain one packed-amax MAX all-reduce; no high-precision reconstruction collective.
Skip absent columnwise scale updates and output writes. Preserve rowwise storage and rounding through the model dtype. Support in-place writes and the helper's FSDP fragment-output interface.
Keep public signatures unchanged. Native omission uses empty NVTETensor handles (has_data() == false). The PyTorch extension normalizes zero-element tensors, including views with backing storage, into that native representation. Common code no longer depends on a magic empty 2D shape.
Specialize the two existing kernels on a boolean columnwise template parameter, using the existing dispatch macro. There is no duplicated kernel implementation or new user-facing configuration.
Move the standalone regression tests into existing suites: partial kernels in test_partial_cast.py (L0), distributed updates in distributed/test_cast_master_weights_to_fp8.py (L1). The existing two-rank launch now covers mixed layouts, fragment outputs, a boundary inside a 32-value block, and an empty master shard on a participating rank.

Scope

MXFP8 distributed cast only. No initialization policy, TMS/backup, NVFP4, or Megatron fallback changes. This is not complete FSDP or full-training integration coverage. Megatron's separate gather copy-back layout handling remains relevant.

Validation of the revised implementation

Source-built TE on NVIDIA GB200: CUDA compiler 13.3, Torch 2.14.0+cu130, Python 3.14.6. Python/native import paths point to the isolated new build. The benchmark library and installed test library have the same ELF build ID.

test_partial_cast.py plus mxfp8/test_mxfp8_master_weight_empty_shard.py: 5 passed.
Entire distributed/test_cast_master_weights_to_fp8.py, including its real two-rank subprocess: 4 passed.
Additional one-rank and four-rank runs of the same distributed worker: all four row-only/mixed/fragment combinations and the existing empty-shard regression passed on every rank.
Both public update entry points, consecutive updates, exact native data/active scales, storage-pointer stability, and a single packed reduction are checked.
Black, clang-format and git diff --check: pass.

The environment emits Python 3.14 TorchScript/opaque-type deprecation warnings. Full Megatron training, checkpoint/resharding and end-to-end FSDP are not covered here.

Bidirectional kernel performance

The initial runtime-pointer-guard implementation did regress amax by 21–35% in this matrix. Following the review suggestion, both kernels now use compile-time boolean specialization.

Baseline: unmodified kernel source at merge base 5f6105b9; revised source: 704b1707. Both are built with the same CUDA 13.3 compiler, -O3, and GB200 sm_100a target. Both columnwise and rowwise outputs are enabled in every measurement. Outputs are checked bitwise before timing.

Method: CUDA-event timing of CUDA-graph replays, seven alternating old/new rounds; median time per kernel call, excluding Python dispatch and allocation. Each graph captures 50 calls for input sizes up to 4M elements, otherwise 10 calls, with 20 replays per measured round. Four shapes × BF16/FP32 × full/partial input × two kernels = 32 comparisons. Partial inputs start at offset 17 and contain one quarter of the full matrix.

Full-input results, microseconds per call:

Shape	Input	Amax old → revised	Cast old → revised
160 × 96	BF16	3.278 → 3.278	7.577 → 7.578
1024 × 1024	BF16	3.693 → 3.690	7.884 → 7.894
4096 × 4096	BF16	20.109 → 20.099	24.656 → 24.784
16384 × 6144	BF16	113.365 → 113.376	202.489 → 202.658
160 × 96	FP32	2.953 → 2.953	7.413 → 7.378
1024 × 1024	FP32	3.359 → 3.362	7.890 → 7.870
4096 × 4096	FP32	19.516 → 19.556	34.558 → 35.254
16384 × 6144	FP32	110.396 → 110.395	210.092 → 210.154

Across all 32 full/partial comparisons, revised/baseline time is 0.994–1.020×. The large amax regression is eliminated; these microbenchmarks do not establish an end-to-end training speedup. The standalone benchmark harness is not added to repository runtime or CI code.

Size and checklist

This review-fix commit changes six files, +150 / -169 lines (net -19), including removing the separate test file and reusing the existing suites.

Bug fix; contributing guidelines followed
Native omission contract documented
Regression coverage wired into existing CI
Relevant existing and new GPU tests pass
Bidirectional old/new kernel benchmark included

Signed-off-by: xiuhu17 <zhihao.wang@perplexity.ai>

for more information, see https://pre-commit.ci

Signed-off-by: xiuhu17 <zhihao.wang@perplexity.ai>

greptile-apps · 2026-09-06T06:38:19Z

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable new issues were found.

Summary

Supports distributed master-weight updates for row-only MXFP8 weights while keeping the bidirectional path.

Distributed MXFP8 updates now write row-only and mixed-direction weights.

The previous, unnumbered CI-coverage finding is addressed: the existing L1 suite launches two ranks and runs the row-only, mixed-layout, fragment-output, and empty-shard cases. Review checks were source-only; no GPU tests were run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Master-weight shards] --> B[Compute amax for present directions]
    B --> C[One packed MAX all-reduce]
    C --> D[Update scales for present directions]
    D --> E[Write rowwise output]
    D --> F{Columnwise output present?}
    F -->|Yes| G[Require scales and write columnwise output]
    F -->|No| H[Skip columnwise reads and writes]

_{Reviews (5) · Last reviewed commit: "Select MXFP8 partial-cast directions fro..." · Reviewed by Greptile}

ptrendx · 2026-09-29T23:29:25Z

 *  \param[in]     input            Input tensor (continuous segment of flattened original tensor).
 *  \param[in,out] amax_rowwise     Output tensor for row-wise maximum absolute values.
 *  \param[in,out] amax_colwise     Output tensor for column-wise maximum absolute values.
+ *                                Shape [0, 0] skips column-wise computation.


Checking for empty 2D shape is not great. Just check for has_data being false (empty NVTETensor).

ptrendx · 2026-09-29T23:29:43Z

+ *  original tensor. To skip column-wise computation, pass output_colwise with shape [0]
+ *  and scale_inv_colwise with shape [0, 0]. Row-wise outputs remain required.


Same comment as above.

ptrendx · 2026-09-29T23:35:41Z

-             "Wrong padding of amax_colwise's cols");
-  NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols");
-  NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise");
+  // Empty views can have a non-null data pointer: use shape to detect omission.


If there is an issue of the pointer passed from the PyTorch side, then that should be handled in the PyTorch extension rather than WARing it here.

Addressed in 704b170.

ptrendx · 2026-09-29T23:38:00Z


  // Load scales_colwise
-  {
+  if (output_colwise != nullptr) {


Please include some benchmark results on the old vs new time for the case still using columnwise output to see whether those additional conditions didn't introduce overhead. If they did then they would need to be handled via templated kernels.

Fixed the amax regression with boolean templates in 704b170; benchmark results are in the PR description.

ptrendx · 2026-09-29T23:39:28Z

+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# See LICENSE for license information.
+
+"""Run with pytest on one GPU or torchrun -m pytest on multiple GPUs."""


Why do we need a separate test file for that? It should be part of one of the larger suites like test_sanity or test_numerics. Also, for the separate test files, they would not actually be run unless they are part of the test.sh in one of the directories in qa/.

Moved the tests into the existing L0 and L1 suites in 704b170.

Signed-off-by: Zhihao Wang <zhihao.wang@perplexity.ai>

xiuhu17 · 2026-10-14T19:50:35Z

@ptrendx would you like to review on this again? Thanks a lot!

Also would you like to take a look at the initial memory optimization pr: #3468? Thanks!

Use columnwise output presence to select partial-cast computation in the common API and PyTorch binding. Require scales for requested columnwise output and ignore them when that output is omitted. Document the contract and cover unused scales, empty shards, and missing required scales in regression tests. Signed-off-by: Codex <codex@openai.com>

ptrendx · 2026-10-26T23:13:07Z

/te-ci

xiuhu17 · 2026-10-27T02:46:07Z

Hi @ptrendx, seems ci is passed. Would you like help to merge? Thanks!

Support row-only MXFP8 distributed master-weight casts

1f306a1

Signed-off-by: xiuhu17 <zhihao.wang@perplexity.ai>

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Sep 6, 2026

xiuhu17 mentioned this pull request Sep 6, 2026

Preserve row-only MXFP8 primaries with native TE writeback NVIDIA/Megatron-LM#7095

Open

2 of 4 tasks

[pre-commit.ci] auto fixes from pre-commit.com hooks

9c04c1a

for more information, see https://pre-commit.ci

xiuhu17 marked this pull request as ready for review September 6, 2026 06:34

xiuhu17 requested review from ksivaman and ptrendx as code owners September 6, 2026 06:34

Share the NCCL process group across row-only cast tests

497dd39

Signed-off-by: xiuhu17 <zhihao.wang@perplexity.ai>

greptile-apps Bot reviewed Sep 6, 2026

View reviewed changes

Comment thread tests/pytorch/mxfp8/test_rowwise_master_cast.py Outdated

xiuhu17 mentioned this pull request Sep 10, 2026

Automatically omit unused columnwise primary weights for backward overrides #3468

Open

zianglih mentioned this pull request Sep 14, 2026

[Roadmap] Blackwell MXFP8 and NVFP4 RL training radixark/miles#615

Open

63 of 70 tasks

ptrendx reviewed Sep 29, 2026

View reviewed changes

xiuhu17 and others added 2 commits October 7, 2026 22:58

Address row-only MXFP8 review and preserve bidirectional performance

704b170

Signed-off-by: Zhihao Wang <zhihao.wang@perplexity.ai>

Merge branch 'main' into fix-mxfp8-rowwise-master-cast

41f6cbf

xiuhu17 and others added 2 commits October 18, 2026 16:43

Merge branch 'main' into fix-mxfp8-rowwise-master-cast

dc3a5e2

ptrendx approved these changes Oct 26, 2026

View reviewed changes

ptrendx merged commit 3850986 into NVIDIA:main Oct 27, 2026
32 of 37 checks passed

		* original tensor. To skip column-wise computation, pass output_colwise with shape [0]
		* and scale_inv_colwise with shape [0, 0]. Row-wise outputs remain required.

Conversation

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

Uh oh!

Description

Changes

Scope

Validation of the revised implementation

Bidirectional kernel performance

Size and checklist

Uh oh!

greptile-apps Bot commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

ptrendx Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

ptrendx Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

ptrendx Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

xiuhu17 Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ptrendx Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

xiuhu17 Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ptrendx Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

xiuhu17 Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xiuhu17 commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ptrendx commented Oct 26, 2026

Uh oh!

xiuhu17 commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

xiuhu17 commented Sep 6, 2026 •

edited

Loading

greptile-apps Bot commented Sep 6, 2026 •

edited

Loading

xiuhu17 Oct 8, 2026 •

edited

Loading

xiuhu17 Oct 8, 2026 •

edited

Loading

xiuhu17 Oct 8, 2026 •

edited

Loading

xiuhu17 commented Oct 14, 2026 •

edited

Loading