[Pytorch][Common] Hybrid quantization - #2817
Merged
Merged
Conversation
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Contributor
Greptile SummaryThis PR introduces hybrid (per-direction) quantization for TransformerEngine, allowing rowwise and columnwise representations to use different quantization formats (e.g. MXFP8 forward + NVFP4 backward). It adds
Confidence Score: 4/5Safe to merge for the experimental hybrid quantization feature; the major correctness bugs from previous review rounds are confirmed fixed, and the remaining known gaps are well-documented TODOs (NVFP4 FSDP2, per-block distopt). All P1 findings from previous review iterations — MX scale truncation floor-vs-ceil, columnwise_source='rowwise_dequantized' ignored in grouped path, make_empty internal flag leak, fsdp_buffer_fields not direction-aware on Hopper, update_usage missing after post_all_gather, None entry crash in _hybrid_split_quantize — are confirmed fixed. The Hopper columnwise-only distopt guard now covers both Float8Quantizer and Float8CurrentScalingQuantizer. The only new finding is a P2 CPU–GPU sync in IdentityQuantizer.update_quantized when noop_flag is active and buffer metadata changes, a rare code path. Known limitations (NVFP4 FSDP2 sub-storage, per-block hybrid distopt, delayed-scaling requests inside HybridQuantizer) are clearly documented as in-code TODOs with tracking issues. Files Needing Attention: tensor/utils.py (_route_hybrid_to_buckets and _validate_per_tensor_fp8_fsdp_hopper_policy) and module/grouped_linear.py (_split_quantize_hybrid) are the most operationally critical paths and would benefit from careful integration testing on Hopper/L40 hardware before wide adoption. Important Files Changed
Class Diagram%%{init: {'theme': 'neutral'}}%%
classDiagram
class Quantizer {
+rowwise_usage: bool
+columnwise_usage: bool
+quantize(tensor) QuantizedTensorStorage
+make_empty(shape) QuantizedTensorStorage
+update_quantized(src, dst) QuantizedTensorStorage
}
class HybridQuantizer {
+rowwise_quantizer: Quantizer
+columnwise_quantizer: Quantizer
+columnwise_source: str
+quantize_impl(tensor) HybridQuantizedTensor
+update_quantized(src, dst)
+supports_only_rowwise_all_gather() bool
+is_requantization_safe() bool
}
class IdentityQuantizer {
+dtype: Optional[torch.dtype]
+quantize_impl(tensor) IdentityTensor
+update_quantized(src, dst)
+is_requantization_safe() bool
}
class QuantizedTensorStorage {
+fsdp_buffer_fields() Tuple
+fsdp_extract_buffers()
+fsdp_assign_gathered(gathered, meta)
+prepare_for_saving()
+restore_from_saved(tensors)
}
class HybridQuantizedTensorStorage {
+_rowwise_storage: Optional[QuantizedTensorStorage]
+_columnwise_storage: Optional[QuantizedTensorStorage]
+_quantizer: HybridQuantizer
+dequantize(dtype) Tensor
+update_usage(rowwise, columnwise)
+clear()
}
class HybridQuantizedTensor {
+fsdp_pre_all_gather(mesh)
+fsdp_post_all_gather(outputs, meta)
+__torch_dispatch__(func)
+detach() HybridQuantizedTensor
}
class IdentityTensorStorage {
+_hp_data: Optional[Tensor]
+fsdp_buffer_fields()
+dequantize(dtype) Tensor
}
class IdentityTensor {
+fsdp_pre_all_gather(mesh)
+fsdp_post_all_gather(outputs, meta)
+__torch_dispatch__(func)
}
Quantizer <|-- HybridQuantizer
Quantizer <|-- IdentityQuantizer
QuantizedTensorStorage <|-- HybridQuantizedTensorStorage
HybridQuantizedTensorStorage <|-- HybridQuantizedTensor
QuantizedTensorStorage <|-- IdentityTensorStorage
IdentityTensorStorage <|-- IdentityTensor
HybridQuantizer --> Quantizer : rowwise_quantizer
HybridQuantizer --> Quantizer : columnwise_quantizer
HybridQuantizedTensorStorage --> QuantizedTensorStorage : _rowwise_storage
HybridQuantizedTensorStorage --> QuantizedTensorStorage : _columnwise_storage
Reviews (37): Last reviewed commit: "Fix CI" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
timmoon10
reviewed
Apr 2, 2026
timmoon10
left a comment
timmoon10
left a comment
Member
There was a problem hiding this comment.
Overall I think this moves us in a good direction. I see some minor bugs, as well as bugs reported by @greptile-apps.
Comment on lines
+52
to
+53
| rowwise_result = self.rowwise_quantizer.quantize(tensor) | ||
| columnwise_result = self.columnwise_quantizer.quantize(tensor) |
Member
There was a problem hiding this comment.
Do we handle the case where not all usages are needed? I'd expect something like:
Suggested change
| rowwise_result = self.rowwise_quantizer.quantize(tensor) | |
| columnwise_result = self.columnwise_quantizer.quantize(tensor) | |
| rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None | |
| columnwise_result = self.columnwise_quantizer.quantize(tensor) if self.columnwise_usage else None |
| requires_grad: bool = False, | ||
| pin_memory: bool = False, | ||
| ) -> HybridQuantizedTensor: | ||
| self.rowwise_quantizer.internal = True |
Member
There was a problem hiding this comment.
Could we just set internal=True in the constructor? I don't think we ever need PyTorch tensor functionality in the per-usage data.
Collaborator
Author
There was a problem hiding this comment.
This would not work under FSDP2.
Comment on lines
+1339
to
+1355
| def factory(role): | ||
| if role == "linear_weight": | ||
| return HybridQuantizer( | ||
| rowwise_quantizer=_make_fp8_quantizer(), | ||
| columnwise_quantizer=_make_mxfp8_quantizer(), | ||
| ) | ||
| if role == "linear_input": | ||
| return HybridQuantizer( | ||
| rowwise_quantizer=_make_fp8_quantizer(), | ||
| columnwise_quantizer=_make_nvfp4_quantizer(), | ||
| ) | ||
| if role in ("linear_grad_output", "linear_grad_input"): | ||
| return HybridQuantizer( | ||
| rowwise_quantizer=_make_mxfp8_quantizer(), | ||
| columnwise_quantizer=_make_nvfp4_quantizer(), | ||
| ) | ||
| return None |
Member
There was a problem hiding this comment.
This is horrifying. Good test.
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
13 tasks
ptrendx
reviewed
May 29, 2026
Comment on lines
+27
to
+30
| # DCP serializes ``CustomRecipe`` via ``pickle``; closure-based qfactories | ||
| # (lambdas, inner functions referencing captured state) are not picklable, | ||
| # so the qfactory must live at module scope. See | ||
| # ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. |
Member
There was a problem hiding this comment.
This comment is potentially useful, but I don't think it is in the right place - shouldn't it be closer to the actual implementation?
ptrendx
reviewed
May 29, 2026
Comment on lines
+1177
to
+1184
| for param in model.parameters(): | ||
| state = optimizer.state[param] | ||
| assert state["exp_avg"].dtype == torch.float32 | ||
| assert state["exp_avg_sq"].dtype == torch.float32 | ||
| if "master_param" in state: | ||
| assert state["master_param"].dtype == torch.float32 | ||
|
|
||
| assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" |
Member
There was a problem hiding this comment.
That's not a very strict test, is there a way for us to do some numerical correctness comparisons?
Collaborator
Author
There was a problem hiding this comment.
Enabled check for the monotonic loss decrease (still mostly sanity), and also enabled hybrid vs vanilla bitwise recipe comparizon, see e.g. test_fused_adam_hybrid_vs_base_recipe_parity.
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Collaborator
Author
|
/te-ci L0 L1 |
timmoon10
reviewed
Jul 22, 2026
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Collaborator
Author
|
/te-ci L0 L1 |
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Collaborator
Author
|
/te-ci L0 L1 |
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Collaborator
Author
|
/te-ci L0 L1 |
timmoon10
previously approved these changes
Aug 3, 2026
Collaborator
Author
|
/te-ci L0 L1 |
for more information, see https://pre-commit.ci
Collaborator
Author
|
/te-ci L0 L1 |
ptrendx
approved these changes
Aug 4, 2026
pggPL
added a commit
to pggPL/TransformerEngine
that referenced
this pull request
Aug 4, 2026
Brings main's 12 commits, including hybrid quantization (NVIDIA#2817), in through the rebased base branch. Four conflicts: - float8_tensor.py / storage/float8_tensor_storage.py: import lists, unioned (_resolve_view_shape from main, tensor_can_be_materialized from here). - linear.py, tp_world_size block: kept this branch's backward_needs_input, which reads the args bundle. - linear.py, `if fp8:` guard: both sides deleted a different statement from it -- this branch dropped assert_dim_for_fp8_exec (host-side shape check that does not trace), main dropped the save_original_input assert. Applying both deletions empties the block, so the `if fp8:` goes away. The save_original_input case is already handled earlier on this branch, which warns and disables it for delayed-scaling quantizers instead of asserting. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
This was referenced Aug 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Hybrid (per-direction) quantization. Hybrid means rowwise/colwise can use different formats via CustomRecipe(qfactory).
This is an experimental feature.
The main problem that it tries to solve is that precision requirements are non-uniform.
Current recipes set one format for both rowwise and colwise directions.
Hybrid quantization enables, e.g. MXFP8 fwd and NVFP4 bwd (or vice versa) or any other valid combination. No need for a hardcoded recipe for every combination.
Composer-style (Composer 2 paper) grouped GEMM recipe, e.g. row-scaled NVFP4 fwd + MXFP8 bwd:
By default, the above factory uses
columnwise_source="original", so MXFP8 backward operands are quantized from the original high-precision tensor. Usecolumnwise_source="rowwise_dequantized"when the backward operand should be derived from the dequantized rowwise NVFP4 forward value.C++ optimizations (fusions, etc.) will come as standalone PRs. cc @kainzhong
TODO:
Follow-up issue tracker #3158.
Integration
Ecosystem integration (all functional, unit-tested):
Megatron-LM integration status:
--fp{4,8}-param-gather+ dist opt (persistent low-precision params viaquantized_model_init+ sharded-master FP32 → quantized cast viaquantize_master_weights.)- [Done] Per-tensor Float8 hybrid (delayed and/or current, any per-direction combination
including same-format, cross-format Float8, single-direction)
- [TODO] Per-block hybrid sub-quantizers (MXFP8, NVFP4, Float8Blockwise) — each rejected per-direction by
quantize_master_weights; unblocker is TE-side cast-helper / kernel.--fp{4,8}-param-gather(fix private attribute access)--fp{4,8}-param-gather- [Done] TE-side hybrid FSDP2 path works end-to-end for Float8 / MXFP8 / Float8Blockwise sub-storages (TODO: need some minor MLM update)
- [TODO] NVFP4 sub-storage FSDP2 hooks
_hybrid_split_quantizeunder Megatron MoE)Review
Total diff +14000
New hybrid source (
hybrid_tensor.py,hybrid_tensor_storage.py,identity_tensor.py,identity_tensor_storage.py) ~1800Adjacent modifications ~1500
Tests are the rest (~10K)
Suggested reading order
-columnwise_source controls whether columnwise quantization uses the original input or the rowwise-dequantized value.
1.1 Identity passthrough — b99277a
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: