$ the-wire · showcase
NVFP4 GroupedExperts lands on Blackwell
By RepoJournal · Filed · About PyTorch · Composed from the cited sources · methodology
14 people shipped this
Ivy Zhou
@ivy-zhou
2 cited
iamzainhuda
@iamzainhuda
2 cited
rdspring1
@rdspring1
2 cited
Bin Bao
@desertfire
1 cited
Steven Chung
@schung-amd
1 cited
cherry77-cloud
@taking-lying-flat
Chien-Chin Huang
@fegin
felipemello1
@felipemello1
Anthony Shoumikhin
@shoumikhin
xingguo01
@xingguo01
Nikita Shulga
@malfet
rascani
@rascani
Warren Deng
@warrendeng
albanD
@albanD
TorchAO's new differentiable NVFP4 grouped GEMM stack brings high-throughput grouped experts to Blackwell, and it's the headline act.
TorchAO shipped a differentiable NVFP4 grouped GEMM stack for Mixture-of-Experts training on Blackwell, fusing randomized Hadamard transforms with NVFP4 quantization and feeding scaled_grouped_mm [1]. The same PR also fixes the NVFP4 quantize path to match TransformerEngine bitwise, using correctly-rounded division instead of the reciprocal path [2]. Meanwhile, torchtitan cleared out legacy import shims, with checkpoint_utils and checkpoint now fully removed, so "there are no compatibility shims left under torchtitan/components/" [3][4]. TorchAO also deleted two unused prototypes, ParetoQ and quant_logger, cutting 310 package LOC and 787 lines repository-wide [5][6]. In pytorch core, a lowercase 'exhaustive' compared against a Literal['DEFAULT', 'EXHAUSTIVE'] meant the EXHAUSTIVE escape hatch never fired in tuned_mm; that's fixed [7]. And the ROCm team fixed a GRU training stall caused by MIOpen RNN dropout mask reuse [8], which is the kind of silent killer that eats CI cycles for days.
Action items
- → Review the new NVFP4 GroupedExperts API if you train MoE on Blackwell pytorch/ao [plan]
- → Update torchtitan imports to drop legacy checkpoint shims pytorch/torchtitan [immediate]
- → Monitor GRU training on ROCm if you use dropout with num_layers > 1 pytorch/pytorch [monitor]
References
- [1] [nvfp4_training] Support NVFP4 GroupedExperts on Blackwell ↗ pytorch/ao
- [2] [nvfp4_training] Support NVFP4 GroupedExperts on Blackwell (#4598) ↗ pytorch/ao
- [3] Remove the legacy torchtitan.components.checkpoint_utils import path (#4241) ↗ pytorch/torchtitan
- [4] Remove the legacy torchtitan.components.checkpoint import path (#4184) ↗ pytorch/torchtitan
- [5] Delete unused ParetoQ prototype ↗ pytorch/ao
- [6] Delete unused quantization logger prototype ↗ pytorch/ao
- [7] [inductor] Fix EXHAUSTIVE search space check in tuned_mm (#192285) pytorch/pytorch ↗
- [8] [ROCm] Fix MIOpen RNN dropout mask reuse and gradient stall (#189285) pytorch/pytorch ↗
