[xiuhu17]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Preserve row-only MXFP8 primary storage during BF16 parameter-gather copy-back, and validate distributed writeback through TE's native cast helper.

Related TE changes:

• Automatically omit unused columnwise primary weights for backward overrides TransformerEngine#3468 — recipe-driven row-only primary initialization for HP/dequantized backward overrides.
• Support row-only MXFP8 distributed master-weight casts TransformerEngine#3488 — native distributed MXFP8 cast support for absent columnwise data/scales. Required for row-only distributed writeback.

Implementation

• Preserve the existing MXFP8 storage directions when copying gathered BF16 weights back: do not force a columnwise allocation for a row-only primary.
• Use the normal cast_master_weights_to_fp8 path for row-only and bidirectional parameters. TE performs packed-amax reduction and writes each local shard; parameter synchronization then gathers quantized data.
• Remove this PR's earlier _quantize_rowwise_mxfp8_param_shard compatibility fallback, its high-precision tile reconstruction/all-reduces, and its special dispatch/FSDP rejection.
• Keep existing post-all-gather processing and non-MXFP8 paths unchanged.

There is no fallback for TE builds lacking #3488. That change is not yet associated with a released TE version, so this PR does not invent a minimum version or infer capability from the presence of an older public helper symbol. Use a TE build containing #3488 for row-only writeback.

Based directly on NVIDIA/Megatron-LM dev. No TMS allocation hook, quantized_param_init_memory_context, new backward-override config, initialization API change, or NVFP4 support is included.

Validation

All testing for this update ran on training_gb200_dev, NVIDIA GB200, with the source-built TE native extension from #3488 and the synced Megatron worktree. Import paths were checked.

• Single GPU: 6 passed.
• Four NCCL ranks: 6 passed on each rank, exit code 0.
• Tests exercise Megatron quantize_param_shard -> native TE cast -> simulated quantized-byte parameter synchronization -> Megatron post-all-gather processing.
• Exact native data/active-scale comparisons over two updates, row-only and bidirectional layouts, unaligned/empty/replicated master shards, and stable primary storage pointers.
• Assert cast performs exactly one packed-amax reduction, with no high-precision parameter-tile collective.
• BF16 copy-back preserves row-only versus bidirectional storage.

Commands used pytest --noconftest and torchrun --standalone --nproc_per_node=4 -m pytest --noconftest on tests/unit_tests/test_mxfp8_rowwise_writeback.py. The tests own their NCCL fixture; unrelated suite-wide dataset-download fixtures were deliberately not run. Environment warnings include optional FlashAttention/CUTLASS, TorchScript/Python 3.14, and optional dependencies.

This is targeted helper/copy-back integration coverage, not a full distributed-optimizer training loop, end-to-end FSDP, checkpoint/resharding, or performance validation. Those remain unverified.

Contribution process

• I have added relevant unit tests
• I have added relevant full functional tests
• I have run the full autoformatter script

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

[xiuhu17]

Native TE follow-up: NVIDIA/TransformerEngine#3488

That PR extends the existing MXFP8 distributed cast helper and partial-amax/partial-cast CUDA kernels to handle absent columnwise storage. It preserves the single packed-amax MAX reduction and skips columnwise work, instead of reconstructing high-precision tiles across ranks as this PR's compatibility fallback does.

It complements the row-only initialization change in NVIDIA/TransformerEngine#3468. Once the native TE change is available, the row-only writeback fallback here can be replaced with the normal TE cast path; the separate BF16 gather copy-back layout fix remains relevant.

The new TE PR has now passed GPU validation on training_gb200_dev: 9 single-GPU regression tests, plus all 7 new cases on each of four NCCL ranks (exit code 0), using the source-built native extension. Coverage includes mixed row-only/bidirectional weights, unaligned and empty master shards, fragment outputs, consecutive updates, and exact native data/active-scale equality. This does not claim full Megatron training/checkpoint integration validation. No Megatron source changes are included in this cross-reference update.