[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 PyTorch TransformerEngine, allowing different quantization formats (e.g., MXFP8 rowwise + NVFP4 columnwise, or Float8 + Identity) to be used per-direction within the same tensor. It adds
Confidence Score: 3/5Large, well-structured feature PR with two previously-flagged P1 issues still present (FSDP current-scaling Hopper guard in utils.py, MXFP8 floor-division truncation in mxfp8_tensor_storage.py); safe to continue review but those issues should be resolved before merge. The core architecture is sound and the new FSDP2 protocol, GEMM dispatch, and distopt routing are well-designed. However, two previously-identified P1 issues remain unresolved in HEAD: the FSDP shard path in _cast_master_weights_to_fp8_current_scaling still lacks the Hopper columnwise-only guard, and mxfp8_tensor_storage.py still uses floor division (//) for _columnwise_scale_inv truncation. No new blocking issues were found in the hybrid/identity tensor implementations themselves. transformer_engine/pytorch/tensor/utils.py (FSDP current-scaling Hopper guard) and transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py (scale floor-division truncation) need attention before merge. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant User as User Code
participant HQ as HybridQuantizer
participant HT as HybridQuantizedTensor
participant RQ as RowwiseQuantizer
participant CQ as ColumnwiseQuantizer
participant GEMM as general_gemm / grouped_gemm
participant FSDP as FSDP2 AllGather
User->>HQ: quantize(input)
HQ->>RQ: "quantize(input, rowwise=True)"
RQ-->>HQ: rowwise_storage
alt "columnwise_source == rowwise_dequantized"
HQ->>HQ: dequantize rowwise_storage to hp_input
HQ->>CQ: "quantize(hp_input, columnwise=True)"
else "columnwise_source == input"
HQ->>CQ: "quantize(input, columnwise=True)"
end
CQ-->>HQ: columnwise_storage
HQ-->>HT: HybridQuantizedTensor(rowwise_storage, columnwise_storage)
Note over GEMM: GEMM operand unwrapping
User->>GEMM: "general_gemm(A=HybridTensor, B=HybridTensor, layout)"
GEMM->>GEMM: "_unwrap_hybrid_A: layout[0]==T → rowwise, N → columnwise"
GEMM->>GEMM: "_unwrap_hybrid_B: layout[1]==N → rowwise, T → columnwise"
GEMM-->>User: result
Note over FSDP: FSDP2 AllGather protocol
FSDP->>HT: fsdp_pre_all_gather()
HT->>HT: rowwise_storage.fsdp_extract_buffers()
HT->>HT: columnwise_storage.fsdp_extract_buffers()
HT-->>FSDP: param_buffers for all-gather
FSDP->>FSDP: all-gather across ranks
FSDP->>HT: fsdp_post_all_gather(gathered_buffers)
HT->>HT: rowwise_storage.fsdp_assign_gathered()
HT->>HT: columnwise_storage.fsdp_assign_gathered()
HT->>HT: _sync_usage() clears stale _transpose cache
HT-->>FSDP: done
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant User as User Code
participant HQ as HybridQuantizer
participant HT as HybridQuantizedTensor
participant RQ as RowwiseQuantizer
participant CQ as ColumnwiseQuantizer
participant GEMM as general_gemm / grouped_gemm
participant FSDP as FSDP2 AllGather
User->>HQ: quantize(input)
HQ->>RQ: "quantize(input, rowwise=True)"
RQ-->>HQ: rowwise_storage
alt "columnwise_source == rowwise_dequantized"
HQ->>HQ: dequantize rowwise_storage to hp_input
HQ->>CQ: "quantize(hp_input, columnwise=True)"
else "columnwise_source == input"
HQ->>CQ: "quantize(input, columnwise=True)"
end
CQ-->>HQ: columnwise_storage
HQ-->>HT: HybridQuantizedTensor(rowwise_storage, columnwise_storage)
Note over GEMM: GEMM operand unwrapping
User->>GEMM: "general_gemm(A=HybridTensor, B=HybridTensor, layout)"
GEMM->>GEMM: "_unwrap_hybrid_A: layout[0]==T → rowwise, N → columnwise"
GEMM->>GEMM: "_unwrap_hybrid_B: layout[1]==N → rowwise, T → columnwise"
GEMM-->>User: result
Note over FSDP: FSDP2 AllGather protocol
FSDP->>HT: fsdp_pre_all_gather()
HT->>HT: rowwise_storage.fsdp_extract_buffers()
HT->>HT: columnwise_storage.fsdp_extract_buffers()
HT-->>FSDP: param_buffers for all-gather
FSDP->>FSDP: all-gather across ranks
FSDP->>HT: fsdp_post_all_gather(gathered_buffers)
HT->>HT: rowwise_storage.fsdp_assign_gathered()
HT->>HT: columnwise_storage.fsdp_assign_gathered()
HT->>HT: _sync_usage() clears stale _transpose cache
HT-->>FSDP: done
Reviews (27): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
timmoon10
left a comment
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>
for more information, see https://pre-commit.ci
Signed-off-by: root <root@prenyx0286.a51.clusters.nvidia.com>
| # 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 |
There was a problem hiding this comment.
Do you intend to close those TODOs?
There was a problem hiding this comment.
Yes, but in the follow-up PRs, see #3158. for now it falls back to high-precision AG
| 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.
Could this check be moved to the quantizer contruction rather than being called on every forward pass?
save_original_input is dependent on module options and recipe.
There was a problem hiding this comment.
I would keep this check at runtime for now. save_original_input depends on thiings that might be changed at runtime right?
| qkv_quantizer = recipe.qfactory( | ||
| QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) | ||
| ) |
There was a problem hiding this comment.
This does not look safe (and also is wasteful) to call this every time. E.g. if the custom quantizer
is stateful or uses RNG etc. we would be redoing seeding of it every time potentially changing the
result depending on how many forwards of MHA we called.
There was a problem hiding this comment.
Fixed. The QKV quantizer is now materialized once through DPA recipe state and reused. MHA no longer calls qfactory to probe it.
| """ | ||
|
|
||
| _hp_data: Optional[torch.Tensor] | ||
| _quantizer: Optional[Quantizer] |
There was a problem hiding this comment.
_quantizer is already a field in the base class.
| """Return the held high-precision tensor (no-op dequantization).""" | ||
| if self._hp_data is None: | ||
| raise RuntimeError("IdentityTensorStorage has no data to dequantize") | ||
| if dtype is not None and self._hp_data.dtype != dtype: |
There was a problem hiding this comment.
Fixed. It now defaults to its nominal dtype
| # Return early if recipe state matches recipe | ||
| if self.fp8_meta_tensors_initialized: | ||
| recipe_state = self.fp8_meta[fp8_meta_tensor_key] | ||
| # Follow-up (#3157): Match built-in recipes by full config, not just RecipeState type, so |
There was a problem hiding this comment.
Can we change "Follow-up" to "TODO"? It is more discoverable that way.
| return False | ||
|
|
||
|
|
||
| def _has_identity_quantizer_list(quantizers): |
There was a problem hiding this comment.
This should not be called every time and instead we should only update this when the recipe changes.
| non_none = [q for q in quantizers if q is not None] | ||
| if not non_none: | ||
| return False | ||
| hybrid_count = sum(1 for q in non_none if isinstance(q, HybridQuantizer)) | ||
| if hybrid_count == 0: | ||
| return False |
There was a problem hiding this comment.
This is pretty convoluted and inefficient.
| """ | ||
| from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage | ||
|
|
||
| if not all(isinstance(q, HybridQuantizer) for q in quantizers): |
There was a problem hiding this comment.
But you already check that in the function above, why checking it again?
| from ...debug.pytorch.debug_quantization import DebugQuantizer | ||
| from ...debug.pytorch.debug_state import TEDebugState | ||
|
|
||
|
|
There was a problem hiding this comment.
A general comment to this entire file: this introduces quite a lot of CPU overhead to this module.
We should just limit this to only allow the same quantizers everywhere rather than checking every
time.
There was a problem hiding this comment.
General comment for the grouped_linear.py
This and the three related comments above are valid and resolved. See c7e11d2
| if out_data is not None: | ||
| out_shape = out_data.size() | ||
| else: | ||
| out_shape = torch.empty(tuple(self.size()), device="meta").view(shape).shape |
There was a problem hiding this comment.
Huh. That's pretty roundabout way of calculating that...
There was a problem hiding this comment.
Right, take a look at the fix at c7e11d2
| """The sub-storage providing columnwise quantized data.""" | ||
| return self._columnwise_storage | ||
|
|
||
| def update_usage( |
There was a problem hiding this comment.
Asking for a storage that does not exist should result in error.
There was a problem hiding this comment.
Added a check for missing directions
| def _canonical_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: | ||
| """Resolve PyTorch view shape syntax, including ``-1`` inference.""" | ||
| return torch.empty(tuple(input_shape), device="meta").view(*shape).shape |
There was a problem hiding this comment.
That's the same formula that was used in one of the other files. If useful, it should be somewhere where it could be used without rediscovering, but still I think it is pretty roundabout...
There was a problem hiding this comment.
Additionally using torch.empty, view etc would introduce quite a bit of CPU overheads.. We can define a function to do manual size resolution in python that resolves -1, 0 etc.
Here something Gemini gave me
def get_actual_view_shape_full(input_shape: tuple, view_shape: tuple) -> tuple:
# 1. First, resolve any '0' dimensions by copying from the input_shape
resolved_zeros = []
for i, dim in enumerate(view_shape):
if dim == 0:
if i >= len(input_shape):
raise ValueError("Dimension '0' exceeds input shape dimensions")
resolved_zeros.append(input_shape[i])
else:
resolved_zeros.append(dim)
# 2. Calculate element totals
total_elements = math.prod(input_shape)
known_elements = math.prod(d for d in resolved_zeros if d != -1)
# 3. Resolve the '-1' dimension if it exists
if -1 in resolved_zeros:
if resolved_zeros.count(-1) > 1:
raise ValueError("Only one dimension can be -1")
if total_elements % known_elements != 0:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
missing_dim = total_elements // known_elements
return tuple(missing_dim if d == -1 else d for d in resolved_zeros)
# 4. Strict validation if no -1 is used
if total_elements != known_elements:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
return tuple(resolved_zeros)
There was a problem hiding this comment.
Now float8 tensor and storage share the same better shaped utility, see c7e11d2
| ) | ||
|
|
||
| @classmethod | ||
| def __torch_dispatch__(cls, func, types, args, kwargs=None): |
There was a problem hiding this comment.
Do we need this function to be this big? We should be able to mostly just delegate to the held tensor, right?
There was a problem hiding this comment.
Good point, delegated
| "instance or a QuantizerRequest. For an intentional no-op " | ||
| "(high-precision / unquantized) slot, return an " | ||
| "IdentityQuantizer instead of None." |
There was a problem hiding this comment.
We could probably just allow None to mean the IdentityQuantizer here?
There was a problem hiding this comment.
I would not do that. Current behavior is explicit. Treating a factory return value of None as an Identity quantizer would silently broaden the custom recipe API and could hide a missing role. Also none quantizer has a semantic meaning in the modules, which is different from the IdentityQuantizer, right? I suggest Identity should be requested explicitly.
| def _canonical_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: | ||
| """Resolve PyTorch view shape syntax, including ``-1`` inference.""" | ||
| return torch.empty(tuple(input_shape), device="meta").view(*shape).shape |
There was a problem hiding this comment.
Additionally using torch.empty, view etc would introduce quite a bit of CPU overheads.. We can define a function to do manual size resolution in python that resolves -1, 0 etc.
Here something Gemini gave me
def get_actual_view_shape_full(input_shape: tuple, view_shape: tuple) -> tuple:
# 1. First, resolve any '0' dimensions by copying from the input_shape
resolved_zeros = []
for i, dim in enumerate(view_shape):
if dim == 0:
if i >= len(input_shape):
raise ValueError("Dimension '0' exceeds input shape dimensions")
resolved_zeros.append(input_shape[i])
else:
resolved_zeros.append(dim)
# 2. Calculate element totals
total_elements = math.prod(input_shape)
known_elements = math.prod(d for d in resolved_zeros if d != -1)
# 3. Resolve the '-1' dimension if it exists
if -1 in resolved_zeros:
if resolved_zeros.count(-1) > 1:
raise ValueError("Only one dimension can be -1")
if total_elements % known_elements != 0:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
missing_dim = total_elements // known_elements
return tuple(missing_dim if d == -1 else d for d in resolved_zeros)
# 4. Strict validation if no -1 is used
if total_elements != known_elements:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
return tuple(resolved_zeros)
| raise NotImplementedError( | ||
| "FSDP2 is only supported for MXFP8Tensors with compact scales" | ||
| ) | ||
| names = self.fsdp_buffer_fields() |
There was a problem hiding this comment.
I think it would be good to reuse fsdp extract buffers and fsdp_assign_gather for the current tensor class as well along with hybrid tensor(in this case mxfp8 and similarly for others)
| return dst | ||
|
|
||
| # ── FSDP2: new_zeros ───────────────────────────────────────── | ||
| if func == aten.new_zeros.default: |
There was a problem hiding this comment.
This is risky if we use it for anything apart from FSDP2. Might make sense to zero out the substorages.
| # Factories stay small; these tests target TP/SP plumbing. | ||
|
|
||
|
|
||
| def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): |
There was a problem hiding this comment.
We now recommend to use transformer_engine.pytorch.DType instead of tex.DType in TE. So we should change it here as well
Signed-off-by: Evgeny <etsykunov@nvidia.com>
…ation in GroupedLinear, etc 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>
|
/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 |
5 participants
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: