[Pytorch][Common] Hybrid quantization by negvet · Pull Request #2817 · NVIDIA/TransformerEngine

negvet · 2026-03-31T14:53:28Z

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:

# CustomRecipe calls quantization_factory(role) for each quantized tensor
# Factory chooses formats

def hybrid_factory(role):
    is_grouped_linear = role is not None and role.module_type == "grouped_linear"
    is_linear = role is not None and role.module_type == "linear"

    if is_grouped_linear and role.tensor_type == "input":
        return HybridQuantizer(
            rowwise_quantizer=NVFP4Quantizer(row_scaled_nvfp4=True, ...),
            columnwise_quantizer=MXFP8Quantizer(...),
        )

    if is_grouped_linear and role.tensor_type == "weight":
        return HybridQuantizer(
            rowwise_quantizer=NVFP4Quantizer(...),
            columnwise_quantizer=MXFP8Quantizer(...),
        )

    if is_grouped_linear and role.tensor_type == "grad_output":
        return MXFP8Quantizer(...)

    if is_linear:
        return MXFP8Quantizer(...)

    return MXFP8Quantizer(...)

recipe = CustomRecipe(qfactory=hybrid_factory)
with autocast(recipe=recipe):
    y = model(x)

By default, the above factory uses columnwise_source="original", so MXFP8 backward operands are quantized from the original high-precision tensor. Use columnwise_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:

Convergence of base (non-hybrid) recipes
HybridFloat8BlockScaling is xfailed under FSDP2 because dim-0 shards can split 128-row block-scale tiles, producing all-gathered scale buffers whose shape does not match the global tensor.
Delayed scaling
Mid-training recipe change

Follow-up issue tracker #3158.

Integration

Ecosystem integration (all functional, unit-tested):

[Done] quantized_model_init
[Done] FSDP2 (TODO: optimize communication buffers)
[Done] CPU offloading
[Done] Activation recomputation
[Done] TP/SP (TODO: enable quantized AG)

Megatron-LM integration status:

[Done] 1 GPU baseline
[Done] DP + distributed optimizer
[TODO] quantized_model_init + --fp{4,8}-param-gather + dist opt (persistent low-precision params via quantized_model_init + sharded-master FP32 → quantized cast via quantize_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.
[TODO] Megatron-FSDP + --fp{4,8}-param-gather (fix private attribute access)
[TODO] Torch FSDP2 + --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
[Done] Activation recompute
[Done] CPU offload
[Done] TP/SP/PP
[Done] MoE + EP + grouped GEMM (qwen3 MoE; _hybrid_split_quantize under Megatron MoE)

Review

Total diff +14000
New hybrid source (hybrid_tensor.py, hybrid_tensor_storage.py, identity_tensor.py, identity_tensor_storage.py) ~1800
Adjacent modifications ~1500
Tests are the rest (~10K)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: Evgeny <etsykunov@nvidia.com>

for more information, see https://pre-commit.ci

greptile-apps · 2026-03-31T15:02:52Z

Greptile Summary

This PR introduces hybrid (per-direction) quantization for Transformer Engine, allowing rowwise and columnwise directions of a tensor to use different quantization formats (e.g. NVFP4 forward + MXFP8 backward) via HybridQuantizer / HybridQuantizedTensor and a companion IdentityQuantizer / IdentityTensor high-precision passthrough. The implementation spans ~14 000 diff lines covering new tensor types, FSDP2 sub-storage protocol, CPU offloading, activation recompute, TP/SP, and Megatron-LM distopt integration.

New tensor hierarchy: HybridQuantizedTensorStorage / HybridQuantizedTensor compose two independently-quantized sub-storages; IdentityTensorStorage / IdentityTensor provide BF16/FP32 passthrough within the same interface. The columnwise_source flag controls whether the columnwise direction is quantized from the original input or from the dequantized rowwise result.
FSDP2 sub-storage protocol: New fsdp_buffer_fields / fsdp_extract_buffers / fsdp_assign_gathered overrides on Float8TensorStorage, MXFP8TensorStorage, and Float8BlockwiseQTensorStorage enable direction-aware all-gather; Hopper/L40 columnwise-only tensors are rejected upfront for unsupported FSDP combinations.
GroupedLinear hybrid dispatch: _hybrid_split_quantize and _split_quantize_with_identity_fallback handle hybrid and identity operands; generation-guarded _validate_grouped_quantizer_list enforces per-direction backend uniformity across experts.

Confidence Score: 4/5

Safe to merge as an experimental feature; the core quantize/GEMM/FSDP2 paths are well-guarded and previously-raised crashes have been addressed in subsequent commits within the PR.

The two remaining findings are in the FSDP2 storage protocol (Float8BlockwiseQTensorStorage.fsdp_extract_buffers bidirectional inconsistency) and the aten.clone quantizer aliasing in HybridQuantizedTensor. Neither is exercised by any current code path, but both represent contract violations that could silently corrupt state in future callers or under concurrent quantizer mutation.

transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py (fsdp_buffer_fields/fsdp_extract_buffers bidirectional contract) and transformer_engine/pytorch/tensor/hybrid_tensor.py (aten.clone shared quantizer reference).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/hybrid_tensor.py	New HybridQuantizer + HybridQuantizedTensor core: columnwise_source handling, FSDP2 pre/post all-gather, TP/SP gather dispatch, and torch_dispatch for FSDP2 op set. aten.clone shares quantizer by reference, inconsistent with new_zeros which copies it.
transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py	New HybridQuantizedTensorStorage: prepare_for_saving/restore_from_saved, clear, view, and dequantize delegation to sub-storages. Logic is clean.
transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py	Adds FSDP2 sub-storage protocol. fsdp_buffer_fields can return 4 fields for bidirectional storage but fsdp_extract_buffers only returns 2 buffers; contract is violated if called on a bidirectional instance.
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Direction-aware fsdp_buffer_fields/fsdp_extract_buffers/fsdp_assign_gathered added for columnwise-only Hopper/L40 sub-storages with correct M-major transport layout.
transformer_engine/pytorch/tensor/float8_tensor.py	Handles _data=None in clone, update_usage, view, slice/select, and split ops; adds _columnwise_shape_for helper. Shape inference in columnwise-only paths is correct.
transformer_engine/pytorch/tensor/utils.py	Adds _route_hybrid_to_buckets, _validate_per_tensor_fp8_fsdp_hopper_policy, _update_transpose_only_float8_flat_fragment, and _cast_master_weights_to_identity with upfront validation guards.
transformer_engine/pytorch/module/grouped_linear.py	Adds _hybrid_split_quantize, _split_quantize_with_identity_fallback, _validate_grouped_quantizer_list with generation-guarded validation, and hybrid branches for forward/backward quantization.
transformer_engine/pytorch/cpp_extensions/gemm.py	Adds _unwrap_hybrid_A/B and _unwrap_identity_tensor to extract native sub-storage before GEMM dispatch. _reject_unsupported_output_quantizer correctly blocks Hybrid/Identity as output quantizers.
transformer_engine/pytorch/distributed.py	gather_along_first_dim hybrid branch temporarily sets usage to both directions via try/finally. Correctly handles the BF16 fallback path for hybrid tensors without native AG dispatch.
transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py	Adds fsdp_buffer_fields/fsdp_extract_buffers/fsdp_assign_gathered for MXFP8 sub-storages; uses math.ceil for columnwise scale truncation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[HybridQuantizer.quantize_impl] --> B{rowwise_usage?}
    B -- yes --> C[rowwise_quantizer.quantize]
    B -- no --> D[rowwise_result = None]
    C --> E{columnwise_source?}
    D --> E
    E -- original --> F[columnwise input = tensor]
    E -- rowwise_dequantized --> G[dequantize rowwise_result]
    G --> F2[columnwise input = dequantized]
    F --> H{columnwise_usage?}
    F2 --> H
    H -- yes --> I[columnwise_quantizer.quantize]
    H -- no --> J[columnwise_result = None]
    I --> K[HybridQuantizedTensor]
    J --> K
    K --> L{GEMM dispatch}
    L --> M[_unwrap_hybrid_A/B]
    M --> N[_unwrap_identity_tensor]
    N --> O[cuBLAS GEMM]
    K --> P{FSDP2 all-gather}
    P --> Q[fsdp_pre_all_gather]
    Q --> R[all_gather_into_tensor]
    R --> S[fsdp_post_all_gather + _sync_usage]
    S --> T[Reconstructed full-param HybridQuantizedTensor]

_{Reviews (30): Last reviewed commit: "Fix delayed-scaling hybrid FP8 FSDP guar..." | Re-trigger Greptile}

timmoon10

Overall I think this moves us in a good direction. I see some minor bugs, as well as bugs reported by @greptile-apps.

timmoon10 · 2026-04-02T01:19:37Z

+        rowwise_result = self.rowwise_quantizer.quantize(tensor)
+        columnwise_result = self.columnwise_quantizer.quantize(tensor)


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

Fixed in 4858491

timmoon10 · 2026-04-02T01:22:54Z

+        requires_grad: bool = False,
+        pin_memory: bool = False,
+    ) -> HybridQuantizedTensor:
+        self.rowwise_quantizer.internal = True


Could we just set internal=True in the constructor? I don't think we ever need PyTorch tensor functionality in the per-usage data.

This would not work under FSDP2.

timmoon10 · 2026-04-02T01:32:10Z

+        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


This is horrifying. Good test.

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>

for more information, see https://pre-commit.ci

Signed-off-by: Evgeny <etsykunov@nvidia.com>

for more information, see https://pre-commit.ci

ptrendx · 2026-05-29T16:17:48Z

+# 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``.


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 · 2026-05-29T16:34:14Z

+    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}"


That's not a very strict test, is there a way for us to do some numerical correctness comparisons?

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.

ptrendx · 2026-07-08T19:53:15Z

+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


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...

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)

Now float8 tensor and storage share the same better shaped utility, see c7e11d2

ptrendx · 2026-07-08T19:55:42Z

+        )
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args, kwargs=None):


Do we need this function to be this big? We should be able to mostly just delegate to the held tensor, right?

Good point, delegated

ptrendx · 2026-07-08T20:02:13Z

+                    "instance or a QuantizerRequest. For an intentional no-op "
+                    "(high-precision / unquantized) slot, return an "
+                    "IdentityQuantizer instead of None."


We could probably just allow None to mean the IdentityQuantizer here?

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.

vthumbe1503 · 2026-07-09T05:08:47Z

+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


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)

vthumbe1503 · 2026-07-09T06:00:35Z

+            raise NotImplementedError(
+                "FSDP2 is only supported for MXFP8Tensors with compact scales"
+            )
+        names = self.fsdp_buffer_fields()


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)

This makes sense. I would make a standalone PR under #3158. I updated #3158 with this.

vthumbe1503 · 2026-07-09T06:04:11Z

+                return dst
+
+        # ── FSDP2: new_zeros ─────────────────────────────────────────
+        if func == aten.new_zeros.default:


This is risky if we use it for anything apart from FSDP2. Might make sense to zero out the substorages.

Good catch, fixed!

vthumbe1503 · 2026-07-09T06:51:42Z

+# Factories stay small; these tests target TP/SP plumbing.
+
+
+def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3):


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>

negvet · 2026-07-17T18:00:51Z

/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

negvet · 2026-07-20T18:35:54Z

/te-ci pytorch L1

Signed-off-by: Evgeny <etsykunov@nvidia.com>

vthumbe1503

Have a few comments on the tests. Apart from that LGTM from FSDP2 side of things

vthumbe1503 · 2026-07-21T18:30:11Z

+            module.unshard()
+    # Hybrid sub-storages (e.g. MXFP8 scale unpad/repad through FSDP2) can introduce
+    # small numerical differences vs the manual dequantize-then-allgather path.
+    tols = dict(atol=5e-4, rtol=5e-3)


instead of replicating the function, could we please add tols as argument to the function?

Good point. I reuse the existing all-gather checker and add an optional tols argument.

vthumbe1503 · 2026-07-21T18:36:47Z

+
+
+@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}")
+@xfail_hopper_columnwise_per_tensor_fp8


Doesnt this increase the test timings by quite a bit, given that we are executing torchrun for each and every pytest?

Given that we use the same number of processes for all these torchrun tests, it might be better to use a single torchrun that shares multiple pytest, similar to how we are doing it with FSDP2 tests.

Not a blocker for this PR if running this test is say under 2 mins. But I am not sure if that is the case.

Addressed in 461daa1. Tests using the same quantization recipe are now passed together to one torchrun. The complete two-GPU suite takes about 87 seconds on the Blackwell node, so it remains under two minutes.

vthumbe1503 · 2026-07-21T18:44:05Z

    return getattr(transformer_engine.common.recipe, recipe)()


+def _hybrid_fp8_current_qfactory(role):


I think it makes sense to abstract away and put the creation of HybridQuantizerFactory in a single test utily file that both distributed and non-distrbuted tests can use. I see the similar kind of functions defined in tp_sp distributed tests, cpu offloading tests etc.

I am not sure if this makes sense, but we can also potentially expose a HybridQuantizerFactory from transformer_engine.common, something like this

@dataclass class HybridQuantizerFactory: """Compose two role-aware base factories into HybridQuantizer for selected roles.""" rowwise_factory: Callable[[Optional[QuantizerRole]], Quantizer] columnwise_factory: Callable[[Optional[QuantizerRole]], Quantizer] fallback_factory: Callable[[Optional[QuantizerRole]], Quantizer] | None = None # roles that become HybridQuantizer hybrid_module_types: frozenset[str] = frozenset({"linear", "grouped_linear"}) hybrid_tensor_types: frozenset[str] = frozenset({"input", "weight", "output"}) columnwise_source: Literal["original", "rowwise_dequantized"] = "original" # optional overrides: tensor_type -> factory (e.g. grads -> Identity) role_overrides: dict[str, Callable] = field(default_factory=dict) def __call__(self, role: Optional[QuantizerRole]): if role is not None and role.tensor_type in self.role_overrides: return self.role_overrides[role.tensor_type](role) if ( role is not None and role.module_type in self.hybrid_module_types and role.tensor_type in self.hybrid_tensor_types ): return HybridQuantizer( rowwise_quantizer=self.rowwise_factory(role), columnwise_quantizer=self.columnwise_factory(role), columnwise_source=self.columnwise_source, ) fallback = self.fallback_factory or self.columnwise_factory return fallback(role)

And the tests in this file for example can boil down to

_HYBRID_QFACTORIES = { "HybridFP8CurrentScaling": HybridQuantizerFactory( current_scaling_quantizer_factory, current_scaling_quantizer_factory, ), "HybridMXFP8": HybridQuantizerFactory( mxfp8_quantizer_factory, mxfp8_quantizer_factory, ), "HybridFloat8BlockScaling": HybridQuantizerFactory( float8_block_scaling_quantizer_factory, float8_block_scaling_quantizer_factory, ), "HybridMixed_MXFP8_FP8": HybridQuantizerFactory( mxfp8_quantizer_factory, current_scaling_quantizer_factory, fallback_factory=current_scaling_quantizer_factory, ), "HybridFP8CurrentScalingIdentity": HybridQuantizerFactory( current_scaling_quantizer_factory, high_precision_factory, # or a tiny identity_factory fallback_factory=current_scaling_quantizer_factory, role_overrides={ "grad_output": high_precision_factory, "grad_input": high_precision_factory, }, ), "Identity": high_precision_factory, }

This might not be general enough. I can imagine user might wanting to have specific quantizer for a specific (tensor_type, module_type) combination. So I would leave it upto you to make the judgement @negvet . But my main concern is the code duplication of the hybrid quantization factory in the tests.

I agree with the duplication concern. A fully general public factory abstraction raises broader questions around GEMM-centric configuration, role targeting, and non-Linear modules, etc, so I would prefer to think about it later. For this PR, I’ll move the common, explicit hybrid qfactory functions into a shared test utility and reuse them across FSDP2, TP/SP, CPU-offload, and unit tests. Specialized factories with different fallback semantics will remain local.

Signed-off-by: Evgeny <etsykunov@nvidia.com>

negvet · 2026-07-21T19:28:44Z

/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>

negvet and others added 2 commits March 31, 2026 14:51

Stage 1&2: Python containers + quantize/gemm dispatch/unwrap

7553e6a

Signed-off-by: Evgeny <etsykunov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

19acc5e

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Mar 31, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py Outdated

Comment thread transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py

Comment thread transformer_engine/pytorch/tensor/hybrid_tensor.py Outdated

timmoon10 reviewed Apr 2, 2026

View reviewed changes

negvet and others added 10 commits April 6, 2026 10:26

Merge branch 'main' into hybrid_quantization

3f33f40

Enable quantized_model_init

f80f5d0

Signed-off-by: Evgeny <etsykunov@nvidia.com>

FSDP support

2185b30

Signed-off-by: Evgeny <etsykunov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

f22a395

for more information, see https://pre-commit.ci

Merge branch 'main' into hybrid_quantization

c9754e9

Enable CPU offloading

103fffe

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Activation recomputation

16fb371

Signed-off-by: Evgeny <etsykunov@nvidia.com>

TP/SP

a50fd63

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Resolve comments: hybrid uniform list, make_empty try, __repr__, etc.

2214843

Signed-off-by: Evgeny <etsykunov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

88fe467

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Apr 24, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py Outdated

Merge branch 'main' into hybrid_quantization

f649cc9

greptile-apps Bot reviewed Apr 29, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/tensor/hybrid_tensor.py

negvet and others added 2 commits April 29, 2026 16:02

Respect usage

4858491

Signed-off-by: Evgeny <etsykunov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

ef31a9a

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Apr 29, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py

negvet added 3 commits May 13, 2026 12:34

Merge main

e7136fd

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Misc minor fixes: comments, tests, etc.

c7da5b2

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Towards MLM integration

a164cd3

Signed-off-by: Evgeny <etsykunov@nvidia.com>

negvet requested a review from ksivaman as a code owner May 21, 2026 13:53

[pre-commit.ci] auto fixes from pre-commit.com hooks

62e7668

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed May 21, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/tensor/float8_tensor.py

negvet mentioned this pull request May 21, 2026

[Common] Enable NVFP4 2D block scaling in columnwise only #3027

Merged

13 tasks

ptrendx reviewed May 29, 2026

View reviewed changes

ptrendx reviewed Jul 8, 2026

View reviewed changes

vthumbe1503 reviewed Jul 9, 2026

View reviewed changes

negvet and others added 5 commits July 13, 2026 15:15

Improve numerical correctness of the tests

4e54f02

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Resolve comments: fix rowwise_dequantized and improve quantizer valid…

c7e11d2

…ation in GroupedLinear, etc Signed-off-by: Evgeny <etsykunov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

05f4d1e

for more information, see https://pre-commit.ci

Merge branch 'main' into hybrid_quantization

023bbab

Finish comment resolving: materialize quantizer once, + minor fixes

0dd5a5f

Signed-off-by: Evgeny <etsykunov@nvidia.com>

negvet requested review from ptrendx and vthumbe1503 July 17, 2026 10:45

pre-commit-ci Bot and others added 3 commits July 17, 2026 10:46

[pre-commit.ci] auto fixes from pre-commit.com hooks

77ce9ff

for more information, see https://pre-commit.ci

Add automatic alignment contract for CustomRecipe

d0f604e

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Misc fixes: float8 fsdp, column-only float8, checks, etc.

cf92c3d

Signed-off-by: Evgeny <etsykunov@nvidia.com>

pre-commit-ci Bot and others added 4 commits July 17, 2026 18:02

[pre-commit.ci] auto fixes from pre-commit.com hooks

6fa849a

for more information, see https://pre-commit.ci

Hopper fixes

495abd5

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Merge branch 'main' into hybrid_quantization

65fc07f

[pre-commit.ci] auto fixes from pre-commit.com hooks

416288b

for more information, see https://pre-commit.ci

negvet added 3 commits July 21, 2026 13:08

More Hopper fixes

a3f5258

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Guard hybrid partial-master distopt updates

77e0e1c

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Blackwell CI fixes

461daa1

Signed-off-by: Evgeny <etsykunov@nvidia.com>

vthumbe1503 reviewed Jul 21, 2026

View reviewed changes

Docs minor update

67f0f48

Signed-off-by: Evgeny <etsykunov@nvidia.com>

greptile-apps Bot reviewed Jul 21, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/tensor/utils.py

negvet and others added 3 commits July 22, 2026 12:32

Improve fsdp2 and other tests

a6fe344

Signed-off-by: Evgeny <etsykunov@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

04caed2

for more information, see https://pre-commit.ci

Fix delayed-scaling hybrid FP8 FSDP guard

a503009

Signed-off-by: Evgeny <etsykunov@nvidia.com>

		rowwise_result = self.rowwise_quantizer.quantize(tensor)
		columnwise_result = self.columnwise_quantizer.quantize(tensor)

		# Factories stay small; these tests target TP/SP plumbing.


		def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3):



		@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}")
		@xfail_hopper_columnwise_per_tensor_fp8

		return getattr(transformer_engine.common.recipe, recipe)()


		def _hybrid_fp8_current_qfactory(role):

Uh oh!

Conversation

negvet commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Integration

Review

Suggested reading order

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet commented Mar 31, 2026 •

edited

Loading

greptile-apps Bot commented Mar 31, 2026 •

edited

Loading

negvet May 21, 2026 •

edited

Loading

negvet Jul 15, 2026 •

edited

Loading