Hybrid quantization follow-ups
维护者通常 2 天内回复
@negvet 已经在做这个了。
开始于 2026年6月30日。
评估
这个 Issue 还没有评估数据。
描述
Tracking follow-ups from PR #2817. Detailed implementation notes are already in code comments / xfails.
Recipe / validation
- Validate per-GEMM scaling-mode compatibility for hybrid qfactories.
- See
HybridQuantizer._get_compatible_recipe()intransformer_engine/pytorch/tensor/hybrid_tensor.py.
- See
- Support delayed-scaling requests inside
HybridQuantizer.- See
HybridQuantizer.__init__()in `transformer_engine/pytorch/tensor/hybrid_tensor.py
- See
- Support producing both directional representations at the fused normalization boundary (transformer_engine/pytorch/ops/basic/layer_norm.py)
Performance / kernels
- Fused kernels: one read -> two writes.
- Introduce a policy for when to materialize a tensor for
columnwise_source="rowwise_dequantized". - Support module-/role-specific CustomRecipe quantization alignment. The current recipe-global quantization_alignment must use the maximum requirement across all factory outputs, which can overpad lower-alignment MXFP8/FP8 layers in mixed-format models. Derive alignment from canonical cached module quantizers and propagate it to the corresponding Fp8Padding without speculative qfactory calls.
- Support HybridQuantizer with standard Userbuffers.
- Add native columnwise-only per-tensor FP8 quantization for Hopper, covering both CurrentScaling and DelayedScaling. The legacy per-tensor cast-transpose kernel currently requires both rowwise and columnwise output buffers.
TP/SP
- Add native
HybridQuantizerdispatch togather_along_first_dim. - Preserve existing SP amax-reduction semantics and cover them in distributed tests.
FSDP2
- Optimize hybrid FSDP2 communication buffers.
- See
HybridQuantizedTensor.fsdp_pre_all_gather()intransformer_engine/pytorch/tensor/hybrid_tensor.py.
- See
- Fix
HybridFloat8BlockScalingFSDP2 xfail.- See
_HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASONintests/pytorch/distributed/fsdp2_tests/conftest.py.
- See
- Add NVFP4 hybrid sub-storage FSDP2 hooks.
- See
TestHybridFsdpPreAllGatherProtocol.test_nvfp4_sub_storage_raises_on_pre_all_gather()intests/pytorch/test_hybrid_quantization.py. - Non-hybrid FSDP alignment: reuse fsdp extract buffers and fsdp_assign_gather for the current base tensor class (non-hybrid)
- See
- Support Hybrid sub-storages that fall back to a high-precision tensor for an unquantizable local shard (TransformerEngine/transformer_engine/pytorch/tensor/hybrid_tensor.py).
GEMM / quantization
- Support
HybridQuantizer/IdentityQuantizeras GEMM output quantizers.- See
_reject_unsupported_output_quantizer()intransformer_engine/pytorch/cpp_extensions/gemm.py.
- See
GroupedLinear / grouped storage
- Support
IdentityQuantizerandHybridQuantizerwithGroupedLinear(single_grouped_weight=True).
Distributed optimizer / Megatron
- Support per-block hybrid sub-quantizers in
quantize_master_weights. - Hybrid distributed optimizer: partial-master support with columnwise_source="original". The current one-payload distributed-optimizer path cannot preserve both independently quantized Hybrid directions when each rank owns only a partial master weight. A safety guard rejects this configuration unless full-master data or an FSDP Hybrid shard containing both directions is provided. Unblock by adding a two-payload sharding/all-gather contract for Hybrid rowwise and columnwise storage, including their scale/amax metadata, then relax _validate_hybrid_partial_master_policy.
- Hybrid distributed optimizer with columnwise_source="rowwise_dequantized": reconstruct the column only after the rowwise update/all-gather.
- Complete Megatron-LM
quantized_model_init + --fp{4,8}-param-gather + dist opt.- See PR #2817 integration notes.
- Complete Megatron-FSDP +
--fp{4,8}-param-gather.- See PR #2817 integration notes.
- Complete Torch FSDP2 +
--fp{4,8}-param-gather.- See PR #2817 integration notes and
tests/pytorch/distributed/fsdp2_tests/.
- See PR #2817 integration notes and
Activation recompute
- Investigate vanilla
torch.utils.checkpoint(use_reentrant=False)with TE weight-workspace cache.- See xfails in
TestHybridActivationRecomputeintests/pytorch/test_hybrid_quantization.py. te.checkpointpath is already covered and works.
- See xfails in
Validation
- Convergence validation of base non-hybrid recipes.
- 主要语言
- Python
- 星标
- 3.6k
- 派生
- 862
- 平均合并
- 4 天 9 小时
- 30 天内合并 PR
- 61
环境准备
- 没有 Dockerfile 或 Docker Compose 文件
- 有 Pull Request 模板
- 阅读贡献指南
从这里开始
- 先读完整个 Issue，再读项目的贡献指南。
- 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
- Fork 仓库，在一个分支上完成修改。
- 提交 Pull Request，并在描述里引用这个 Issue 编号。
NVIDIA/TransformerEngine 的其他 Issue
-
[PyTorch] fp8_cs_quantize fake implementation returns a vector inverse scale instead of a scalar可能重新可做 @sanjana658 于 12 天前认领，目前没有进行中的 PR。 未关闭
难度 2/5 1-3 小时 新手友好度 82/100
NVIDIA/TransformerEngine#3636 · 2 条评论 ·
维护者通常 2 天内回复
-
[Bug] Backend selection picks FA3 for training with head_dim_qk=192 / v_head_dim=128, but FA3 backward cannot run it可能重新可做 @modelpath-dev 于 46 天前认领，目前没有进行中的 PR。 未关闭attention
难度 2/5 1-3 小时 新手友好度 85/100
NVIDIA/TransformerEngine#3481 · 4 条评论 ·
维护者通常 2 天内回复
-
bug
难度 2/5 1-3 小时 新手友好度 68/100
NVIDIA/TransformerEngine#2189 · 7 条评论 · 5 个 reaction ·
维护者通常 2 天内回复
-
[PyTorch] torch.compile fails on FusedAttnBackend IntEnum comparisons可能已有人在做 @pggPL 今天认领。 未关闭bug
NVIDIA/TransformerEngine#3716 · 已指派 1 人 ·
维护者通常 2 天内回复
-
[PyTorch] FA2 is selected for head_dim > 192 with dropout on sm120, where flash-attn's backward launches no kernel可能已有人在做 @0z5a 今天认领。 未关闭
难度 3/5 1-2 天 新手友好度 69/100
NVIDIA/TransformerEngine#3711 · 2 条评论 ·
维护者通常 2 天内回复
查看 NVIDIA/TransformerEngine 的全部 Issue
相似的 Issue
-
documentation
难度 2/5 1-3 小时 新手友好度 67/100
维护者通常 1 天内回复
-
good first issue module-baseline needs triage ux
难度 2/5 1-3 小时 新手友好度 82/100
darnitdevorg/darnit#650 ·
维护者通常 2 天内回复
-
难度 2/5 1-3 小时 新手友好度 65/100
nyjc-computing/campus#938 ·
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 70/100
nf-osi/nf-metadata-dictionary#1027 ·
维护者通常 1 天内回复
-
bug documentation
难度 1/5 1 小时以内 新手友好度 92/100
qatrackplus/qatrackplus#913 ·
维护者通常 3 天内回复