Repository navigation

[PyTorch] Fix DistributedWeight handling under saved-tensor hooks and in the fused grouped MLP - #3545

Open

xrennvidia wants to merge 9 commits into

NVIDIA:mainfrom

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

Open

[PyTorch] Fix DistributedWeight handling under saved-tensor hooks and in the fused grouped MLP#3545
xrennvidia wants to merge 9 commits into
NVIDIA:mainfrom
xrennvidia:xren/dist-weight-saved-tensor-hooks

Commits on Oct 18, 2026
- [PyTorch] Keep DistributedWeight objects on ctx across saved-tensor hooks
  
  xrennvidia
  and
  claude
  committed
  7638a3c
- [PyTorch] Quantize distributed BF16 weights in the fused grouped MLP
  
  xrennvidia
  and
  claude
  committed
  a9a8420
- [PyTorch] Drop NVFP4 rowwise input copies from the fused grouped MLP saved state
  
  xrennvidia
  and
  claude
  committed
  250974f
Commits on Oct 21, 2026
- Update transformer_engine/pytorch/ops/fused/grouped_mlp.py
  
  xrennvidia
  and
  timmoon10
  authored
  66705ee
- [PyTorch] Move the NVFP4 saved-state test into test_grouped_mlp.py
  
  xrennvidia
  and
  claude
  committed
  a085179
- [PyTorch] Move the fused grouped MLP distributed-weight tests into test_grouped_mlp.py
  
  xrennvidia
  and
  claude
  committed
  14f5a3c
- [PyTorch] Fold the saved-tensor-hooks DistributedWeight tests into test_distributed_weight.py
  
  xrennvidia
  and
  claude
  committed
  8ff5384
- [PyTorch] Build fused grouped MLP weight quantizer lists lazily
  
  xrennvidia
  and
  claude
  committed
  f8ec19d
- Merge remote-tracking branch 'origin/main' into xren/dist-weight-saved-tensor-hooks
  xrennvidia
  committed
  d0bff45