Repository navigation

[PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP - #3545

Open

xrennvidia wants to merge 8 commits into

NVIDIA:mainfrom

xrennvidia:xren/dist-weight-saved-tensor-hooks

Open

[PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP#3545
xrennvidia wants to merge 8 commits into
NVIDIA:mainfrom
xrennvidia:xren/dist-weight-saved-tensor-hooks

Commits on Oct 23, 2026
- [PyTorch] Quantize distributed BF16 weights in the fused grouped MLP
  
  xrennvidia
  and
  claude
  committed
  06ab7f3
- [PyTorch] Drop NVFP4 rowwise input copies from the fused grouped MLP saved state
  
  xrennvidia
  and
  claude
  committed
  7595b1c
- Update transformer_engine/pytorch/ops/fused/grouped_mlp.py
  
  xrennvidia
  and
  timmoon10
  committed
  9a8faad
- [PyTorch] Move the NVFP4 saved-state test into test_grouped_mlp.py
  
  xrennvidia
  and
  claude
  committed
  13096d5
- [PyTorch] Move the fused grouped MLP distributed-weight tests into test_grouped_mlp.py
  
  xrennvidia
  and
  claude
  committed
  5d45d19
- [PyTorch] Build fused grouped MLP weight quantizer lists lazily
  
  xrennvidia
  and
  claude
  committed
  8ca9d08
- Apply suggestion from @ksivaman
  
  xrennvidia
  and
  ksivaman
  authored
  3faf7bf
- [PyTorch] Skip columnwise quantization of gathered distributed FC1 weights in forward
  
  xrennvidia
  and
  claude
  committed
  274706e