[Pytorch][Common] Hybrid quantization

qa/L0_pytorch_unittest/test.sh

@@ -48,6 +48,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py"
 NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py"
 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py"

qa/L1_pytorch_distributed_unittest/test.sh

@@ -27,6 +27,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PAT
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hybrid_tp_sp.xml $TE_PATH/tests/pytorch/distributed/test_hybrid_tp_sp.py || test_fail "test_hybrid_tp_sp.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py"

tests/pytorch/distributed/fsdp2_tests/conftest.py

@@ -45,6 +45,13 @@ def _check_nvfp4_support():
     ("NVFP4BlockScaling", _check_nvfp4_support),
 ]
 
+_HYBRID_RECIPE_CONFIGS = [
+    ("HybridFP8CurrentScaling", fp8.check_fp8_support),
+    ("HybridMXFP8", fp8.check_mxfp8_support),
+    ("HybridFloat8BlockScaling", fp8.check_fp8_block_scaling_support),
+    ("HybridMixed_MXFP8_FP8", fp8.check_mxfp8_support),
+]
+
 
 def _parametrize_recipes():
     params = []
@@ -56,6 +63,16 @@ def _parametrize_recipes():
     return params
 
 
+def _parametrize_hybrid_recipes():
+    params = []
+    for name, check_fn in _HYBRID_RECIPE_CONFIGS:
+        supported, reason = check_fn()
+        params.append(
+            pytest.param(name, id=name, marks=pytest.mark.skipif(not supported, reason=reason))
+        )
+    return params
+
+
 # ── Session / per-test fixtures ──────────────────────────────────────
 @pytest.fixture(scope="session", autouse=True)
 def dist_init():
@@ -83,3 +100,8 @@ def _cleanup():
 @pytest.fixture(params=_parametrize_recipes())
 def recipe_name(request):
     return request.param
+
+
+@pytest.fixture(params=_parametrize_hybrid_recipes())
+def hybrid_recipe_name(request):
+    return request.param

tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py

@@ -5,13 +5,99 @@
 """Shared utility functions for FSDP2 distributed tests."""
 
 import transformer_engine.common.recipe
-from transformer_engine.pytorch import QuantizedTensor
+from transformer_engine.pytorch import HybridQuantizer, QuantizedTensor
+from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import (
+    current_scaling_quantizer_factory,
+    float8_block_scaling_quantizer_factory,
+    mxfp8_quantizer_factory,
+)
 
 
 def get_recipe_from_string(recipe):
     return getattr(transformer_engine.common.recipe, recipe)()
 
 
+# ── Hybrid qfactories ─────────────────────────────────────────────────
+#
+# Module-level (picklable) qfactories used by ``get_hybrid_recipe_from_string``.
+# Each factory composes one or two role-aware base factories from
+# ``quantization_recipes_base`` per direction. Per-role behavior is delegated
+# to the base factory — the hybrid layer only decides direction pairing.
+#
+# 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``.
+
+
+def _hybrid_fp8_current_qfactory(role):
+    """FP8 current-scaling rowwise + FP8 current-scaling columnwise."""
+    is_linear = role is not None and role.module_type in ("linear", "grouped_linear")
+    if is_linear and role.tensor_type in ("input", "weight", "output"):
+        return HybridQuantizer(
+            rowwise_quantizer=current_scaling_quantizer_factory(role),
+            columnwise_quantizer=current_scaling_quantizer_factory(role),
+        )
+    return current_scaling_quantizer_factory(role)
+
+
+def _hybrid_mxfp8_qfactory(role):
+    """MXFP8 rowwise + MXFP8 columnwise."""
+    is_linear = role is not None and role.module_type in ("linear", "grouped_linear")
+    if is_linear and role.tensor_type in ("input", "weight", "output"):
+        return HybridQuantizer(
+            rowwise_quantizer=mxfp8_quantizer_factory(role),
+            columnwise_quantizer=mxfp8_quantizer_factory(role),
+        )
+    return mxfp8_quantizer_factory(role)
+
+
+def _hybrid_float8_block_qfactory(role):
+    """Float8 block-scaling rowwise + Float8 block-scaling columnwise."""
+    is_linear = role is not None and role.module_type in ("linear", "grouped_linear")
+    if is_linear and role.tensor_type in ("input", "weight", "output"):
+        return HybridQuantizer(
+            rowwise_quantizer=float8_block_scaling_quantizer_factory(role),
+            columnwise_quantizer=float8_block_scaling_quantizer_factory(role),
+        )
+    return float8_block_scaling_quantizer_factory(role)
+
+
+def _hybrid_mixed_mxfp8_fp8_qfactory(role):
+    """MXFP8 rowwise + FP8 current columnwise (cross-format hybrid)."""
+    is_linear = role is not None and role.module_type in ("linear", "grouped_linear")
+    if is_linear and role.tensor_type in ("input", "weight", "output"):
+        return HybridQuantizer(
+            rowwise_quantizer=mxfp8_quantizer_factory(role),
+            columnwise_quantizer=current_scaling_quantizer_factory(role),
+        )
+    return current_scaling_quantizer_factory(role)
+
+
+_HYBRID_QFACTORIES = {
+    "HybridFP8CurrentScaling": _hybrid_fp8_current_qfactory,
+    "HybridMXFP8": _hybrid_mxfp8_qfactory,
+    "HybridFloat8BlockScaling": _hybrid_float8_block_qfactory,
+    "HybridMixed_MXFP8_FP8": _hybrid_mixed_mxfp8_fp8_qfactory,
+}
+
+
+def get_hybrid_recipe_from_string(recipe):
+    """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory.
+
+    Supported values:
+        "HybridFP8CurrentScaling" — FP8 current for both directions
+        "HybridMXFP8"             — MXFP8 for both directions
+        "HybridFloat8BlockScaling" — Float8 block scaling for both directions
+        "HybridMixed_MXFP8_FP8"   — MXFP8 rowwise + FP8 current columnwise
+    """
+    if recipe not in _HYBRID_QFACTORIES:
+        raise ValueError(
+            f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_HYBRID_QFACTORIES.keys())}"
+        )
+    return transformer_engine.common.recipe.CustomRecipe(qfactory=_HYBRID_QFACTORIES[recipe])
+
+
 def save_custom_attrs(module):
     custom_attrs = {}
     for name, param in module.named_parameters():