[Pytorch][Common] Hybrid quantization - #2817
Conversation
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR introduces hybrid quantization for Transformer Engine, enabling weights to hold two independently-quantized representations (rowwise and columnwise) in different formats simultaneously — for example, FP8 rowwise paired with Identity (high-precision) columnwise, or FP8 rowwise paired with MXFP8 columnwise.
Confidence Score: 4/5Safe to merge with one P1 to resolve: _unwrap_tensor's silent None return should raise explicitly before reaching the GEMM kernel. The PR is a large, well-structured addition; FSDP2, distopt, and grouped-linear paths are correctly extended; known limitations are guarded with NotImplementedError. The one P1 (_unwrap_tensor None return) is low blast-radius since it only triggers on misconfigured hybrid tensors, not normal usage. Files Needing Attention: transformer_engine/pytorch/cpp_extensions/gemm.py (_unwrap_tensor None-return path) and transformer_engine/pytorch/tensor/utils.py (_cast_master_weights_to_identity duplicated bounds logic). Important Files Changed
|
timmoon10
left a comment
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.
| rowwise_result = self.rowwise_quantizer.quantize(tensor) | ||
| columnwise_result = self.columnwise_quantizer.quantize(tensor) |
There was a problem hiding this comment.
Do we handle the case where not all usages are needed? I'd expect something like:
| 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 |
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.
There was a problem hiding this comment.
This would not work under FSDP2.
| 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 |
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
| # 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``. |
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?
| 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}" |
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?
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>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
|
/te-ci pytorch L1 |
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
|
/te-ci pytorch L1 |
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
vthumbe1503
left a comment
There was a problem hiding this comment.
Have a few comments on the tests. Apart from that LGTM from FSDP2 side of things
Signed-off-by: Evgeny <etsykunov@nvidia.com>
|
/te-ci L0 L1 |
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
|
/te-ci L0 L1 |
| # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP | ||
| # activations are gathered in high precision and re-quantized whole, so | ||
| # every rank already sees the same global amax. | ||
| # TODO(#3158): once native quantized all-gather lands (see | ||
| # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path | ||
| # quantizes per-shard, needing a hybrid branch here that mirrors the | ||
| # current-scaling / NVFP4 SP reduction above: | ||
| # elif recipe.custom(): | ||
| # ... # enable SP amax reduction on the hybrid input/grad quantizer |
There was a problem hiding this comment.
These comments are adding a lot of noise and will very easily become stale.
I get that this is a guide for future TP support. However, in my experience Claude is extremely timid about deleting code and will prefer to let comments become wrong than to fix them. If we do keep these comments, we should make sure future PRs remove them as intended.
There was a problem hiding this comment.
I added this explicitly because it is easy to miss amax reduction when enabling low-precision comm, and this can be difficult to debug.
But I can see that this comment is stale since PR #3104. SP amax reduction is now configured at point of use through set_quantizer_amax_reduction_group(). HybridQuantizer already forwards with_amax_reduction and amax_reduction_group to its sub-quantizers. So no need for additional branches in set_meta_tensor(). So removed it completely from all modules.
| # Hybrid override: callers drop columnwise before AG, expecting to | ||
| # synthesize it post-AG via ``update_usage(columnwise_usage=True)`` | ||
| # (native FP8's ``_create_transpose``). Hybrid has no synthesis path — | ||
| # that update_usage is a no-op — so re-quantize with both directions, | ||
| # mirroring what the planned native hybrid AG dispatch would produce | ||
| # (see the TODO in | ||
| # :meth:`HybridQuantizer.supports_only_rowwise_all_gather`); once | ||
| # native AG lands, hybrid won't reach this fallback. | ||
| if isinstance(quantizer, HybridQuantizer): | ||
| prev_row, prev_col = quantizer.rowwise_usage, quantizer.columnwise_usage | ||
| quantizer.set_usage(rowwise=True, columnwise=True) | ||
| try: | ||
| out = quantizer(out) | ||
| finally: | ||
| quantizer.set_usage(rowwise=prev_row, columnwise=prev_col) | ||
| else: | ||
| out = quantizer(out) |
There was a problem hiding this comment.
This override is unnecessary:
- Native hybrid AG is not supported: The correct behavior is to AG in high precision and then call the quantizer directly.
- Native hybrid AG is supported: The correct behavior is to add a custom AG implementation above, the same pattern as FP8/MXFP8/NVFP4.
| # Hybrid override: callers drop columnwise before AG, expecting to | |
| # synthesize it post-AG via ``update_usage(columnwise_usage=True)`` | |
| # (native FP8's ``_create_transpose``). Hybrid has no synthesis path — | |
| # that update_usage is a no-op — so re-quantize with both directions, | |
| # mirroring what the planned native hybrid AG dispatch would produce | |
| # (see the TODO in | |
| # :meth:`HybridQuantizer.supports_only_rowwise_all_gather`); once | |
| # native AG lands, hybrid won't reach this fallback. | |
| if isinstance(quantizer, HybridQuantizer): | |
| prev_row, prev_col = quantizer.rowwise_usage, quantizer.columnwise_usage | |
| quantizer.set_usage(rowwise=True, columnwise=True) | |
| try: | |
| out = quantizer(out) | |
| finally: | |
| quantizer.set_usage(rowwise=prev_row, columnwise=prev_col) | |
| else: | |
| out = quantizer(out) | |
| out = quantizer(out) |
There was a problem hiding this comment.
I removed this HybridQuantizer special case, now high precision fallback calls the quantizer directly. This required updating the module wgrad AG calls to configure the required output explicitly (introduced set_quantizer_usage_for_wgrad_all_gather() in _common.py). Native hybrid AG is tracked in #3158.
| """Returns True if the quantizer supports only rowwise all-gather""" | ||
| return False | ||
|
|
||
| def allows_save_original_input_for_backward(self) -> bool: |
There was a problem hiding this comment.
I'm very suspicious of this. We're spilling out internal logic from the linear module into the quantized tensor. Why should the quantizer be aware of how it happens to be used within the linear autograd function?
Thinking through this logically, the real thing we are checking is whether the quantizer is deterministic on repeated calls. A function name like is_stateless would capture the desired behavior while keeping the logic properly scoped.
There was a problem hiding this comment.
replaced with is_requantization_safe(). Repeated quantization must reproduce every requested representation. save_original_input decision for hybrid: columnwise_source="original" repeats neither sub-quantizer, so reconstruction is always valid, "rowwise_dequantized" repeats only the rowwise quantizer, so only that quantizer must be safe to recall.
| if ( | ||
| save_original_input | ||
| and backward_needs_input | ||
| and input_quantizer is not None | ||
| and not input_quantizer.allows_save_original_input_for_backward() | ||
| ): | ||
| warnings.warn( | ||
| "Ignoring save_original_input=True because the input quantizer requires " | ||
| "the forward quantized activation for backward " | ||
| f"({input_quantizer}).", | ||
| stacklevel=2, | ||
| ) | ||
| save_original_input = False |
There was a problem hiding this comment.
Is this necessary? If the user has specified save_original_input=True, then presumably they know what they're getting themselves in to. We document that it doesn't work with FP8 DS, so expanding that comment with other stateful quantizers seems sufficient:
As a stylistic nit, I also don't like how we split up the save_original_input logic into multiple places. If we do add a run-time check, it's more readable to put it here:
TransformerEngine/transformer_engine/pytorch/module/linear.py
Lines 290 to 293 in a503009
There was a problem hiding this comment.
This runtime fallback is necessary because Megatron enables save_original_input automatically for proj, so this is not always an informed user opt-in. Moved Linear’s save_original_input policy below backward_override, and moved the reconstruction eligibility check into a shared module helper used by Linear and GroupedLinear.
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
|
/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
|
/te-ci L0 L1 |
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: