diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 5d767ba4d1..ba8f7114ae 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -51,6 +51,8 @@ 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_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 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" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 50a51353d1..ea89ba71af 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -45,6 +45,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_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py" diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index bf9db094d2..0d1953e4c9 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -14,6 +14,7 @@ import torch import torch.distributed as dist from transformer_engine.pytorch import fp8 +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported # Ensure the correct CUDA device is active before _parametrize_recipes() # runs at collection time, since the session-scoped dist_init fixture @@ -45,6 +46,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 +64,26 @@ def _parametrize_recipes(): return params +def _parametrize_hybrid_recipes(): + params = [] + for name, check_fn in _HYBRID_RECIPE_CONFIGS: + supported, reason = check_fn() + marks = [pytest.mark.skipif(not supported, reason=reason)] + if name == "HybridFP8CurrentScaling" and not is_non_tn_fp8_gemm_supported(): + marks.append( + pytest.mark.xfail( + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 " + "quantization; tracked by NVIDIA/TransformerEngine#3158" + ), + ) + ) + params.append(pytest.param(name, id=name, marks=marks)) + return params + + # ── Session / per-test fixtures ────────────────────────────────────── @pytest.fixture(scope="session", autouse=True) def dist_init(): @@ -83,3 +111,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 diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 178ce62375..042333843b 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -7,11 +7,57 @@ import transformer_engine.common.recipe from transformer_engine.pytorch import QuantizedTensor +from hybrid_quantization_utils import ( + hybrid_float8_block_qfactory, + hybrid_fp8_current_identity_qfactory, + hybrid_fp8_current_qfactory, + hybrid_mixed_mxfp8_fp8_qfactory, + hybrid_mxfp8_qfactory, + identity_qfactory, +) + def get_recipe_from_string(recipe): return getattr(transformer_engine.common.recipe, recipe)() +# CustomRecipe has dynamic TE extra-state handling. Once FP8 state is +# initialized, TE's get_extra_state() pickles the recipe on save, so +# checkpoint-test qfactories must be module-level and picklable. On load, +# payloads without delayed-scaling state are identified and ignored without +# unpickling. See ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. +_HYBRID_QFACTORIES = { + "HybridFP8CurrentScaling": hybrid_fp8_current_qfactory, + "HybridMXFP8": hybrid_mxfp8_qfactory, + "HybridFloat8BlockScaling": hybrid_float8_block_qfactory, + "HybridMixed_MXFP8_FP8": hybrid_mixed_mxfp8_fp8_qfactory, + "HybridFP8CurrentScalingIdentity": hybrid_fp8_current_identity_qfactory, + "Identity": identity_qfactory, +} + + +def get_hybrid_recipe_from_string(recipe): + """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory. + + Each hybrid qfactory composes one or two role-aware base factories from + ``quantization_factory_base`` per direction; per-role behavior is delegated + to the base factory and the hybrid layer only decides the direction pairing. + + 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 + "HybridFP8CurrentScalingIdentity" — FP8 current forward + Identity backward + "Identity" — high-precision passthrough for every slot + """ + 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(): diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 1abb49e98c..08e762045a 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -44,12 +44,16 @@ from torch.distributed.tensor import DTensor import transformer_engine.pytorch as te -from transformer_engine.pytorch import QuantizedTensor +from transformer_engine.pytorch import HybridQuantizedTensor, QuantizedTensor +from transformer_engine.pytorch.tensor import ( + Float8BlockwiseQTensorStorage, + Float8TensorStorage, + MXFP8TensorStorage, +) import transformer_engine.common.recipe from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs - HIDDEN_SIZE = 256 FFN_HIDDEN_SIZE = 1024 NUM_ATTENTION_HEADS = 8 @@ -113,7 +117,7 @@ def _build_model( return model -def _shard_model(model, world_size): +def _shard_model(model, world_size, reshard_after_forward=None): """Apply FSDP2 sharding with save/restore custom attrs. If the model was created on the meta device (e.g. for FP8 init), @@ -122,12 +126,24 @@ def _shard_model(model, world_size): restore_custom_attrs is called last so it applies to the final parameter objects. For meta-device models, reset_parameters() replaces params via module_setattr (base.py:1336-1339), so attrs must be restored afterward. + + Parameters + ---------- + reshard_after_forward : bool, optional + Passed through to ``fully_shard``. ``None`` (default) keeps FSDP2's + own default: ``True`` for child modules, ``False`` for the root. + ``False`` on child modules keeps the full-precision gathered weight + alive through backward, exercising the iter-2+ buffer-reuse path + inside the same forward/backward rather than across training steps. """ has_meta_params = any(p.is_meta for p in model.parameters()) custom_attrs = save_custom_attrs(model) mesh = DeviceMesh("cuda", list(range(world_size))) + shard_kwargs = {"mesh": mesh} + if reshard_after_forward is not None: + shard_kwargs["reshard_after_forward"] = reshard_after_forward for child in model.children(): - fully_shard(child, mesh=mesh) + fully_shard(child, **shard_kwargs) fully_shard(model, mesh=mesh) if has_meta_params: for module in model.modules(): @@ -148,6 +164,185 @@ def _get_dist_info(): return world_size, device +def _collective_assert(condition, message): + """Raise on every rank only after all ranks report structural status.""" + failed = torch.tensor([not condition], dtype=torch.uint8, device="cuda") + if not condition: + print(f"[rank {dist.get_rank()}] {message}", flush=True) + dist.all_reduce(failed, dist.ReduceOp.MAX) + assert not bool(failed.item()), f"{message}: failed on at least one rank" + + +def _collective_assert_same_shape(tensor, message): + """Check tensor ranks/shapes before entering a shape-sensitive collective.""" + # Quantized parameter buffers in these tests are at most 2-D, but leave + # room for future formats without introducing object collectives. + layout = torch.full((9,), -1, dtype=torch.int64, device=tensor.device) + _collective_assert(tensor.ndim < layout.numel(), f"{message}: ndim={tensor.ndim}") + layout[0] = tensor.ndim + layout[1 : tensor.ndim + 1] = torch.tensor( + tensor.shape, dtype=torch.int64, device=tensor.device + ) + min_layout = layout.clone() + max_layout = layout.clone() + dist.all_reduce(min_layout, dist.ReduceOp.MIN) + dist.all_reduce(max_layout, dist.ReduceOp.MAX) + _collective_assert( + torch.equal(min_layout, max_layout), + f"{message}: tensor shape differs across ranks", + ) + + +def _record_exact(errors, actual, expected, label): + """Record an exact mismatch without interrupting later distributed work.""" + if actual is None or expected is None: + if actual is not expected: + errors.append(f"{label}: one tensor is None") + return + try: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + except (AssertionError, TypeError) as exc: + errors.append(f"{label}: {exc}") + + +def _raise_collective_errors(errors, context): + """Aggregate delayed exact failures after every rank finishes collectives.""" + if errors: + print(f"[rank {dist.get_rank()}] {context}:\n" + "\n".join(errors), flush=True) + failed = torch.tensor([bool(errors)], dtype=torch.uint8, device="cuda") + dist.all_reduce(failed, dist.ReduceOp.MAX) + assert not bool(failed.item()), f"{context}: exact check failed on at least one rank" + + +def _check_hybrid_direction_buffers( + local_sub, + full_sub, + expected_type, + *, + direction, + param_name, + world_size, + errors, +): + """Compare one direction's gathered raw data, scale buffers, and metadata.""" + _collective_assert( + isinstance(local_sub, expected_type), + f"{param_name}: {direction} local storage is {type(local_sub).__name__}, " + f"expected {expected_type.__name__}", + ) + _collective_assert( + isinstance(full_sub, expected_type), + f"{param_name}: {direction} full storage is {type(full_sub).__name__}, " + f"expected {expected_type.__name__}", + ) + + local_buffers, local_meta = local_sub.fsdp_extract_buffers() + full_buffers, full_meta = full_sub.fsdp_extract_buffers() + local_count = torch.tensor([len(local_buffers)], dtype=torch.int64, device="cuda") + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() == len(full_buffers), + f"{param_name}: {direction} buffer count mismatch: local={len(local_buffers)}, " + f"full={len(full_buffers)}, rank range=({min_count.item()}, {max_count.item()})", + ) + + expected_buffers = [] + for buffer in local_buffers: + _collective_assert( + buffer is not None, + f"{param_name}: {direction} unexpectedly exposed a None FSDP buffer", + ) + _collective_assert_same_shape(buffer, f"{param_name}: {direction} FSDP buffer") + gathered = [torch.zeros_like(buffer) for _ in range(world_size)] + dist.all_gather(gathered, buffer) + expected_buffers.append(torch.cat(gathered, dim=0)) + + for key in ("field_names", "direction"): + if key in local_meta or key in full_meta: + if local_meta.get(key) != full_meta.get(key): + errors.append( + f"{param_name}: {direction} metadata {key!r} differs: " + f"{full_meta.get(key)!r} != {local_meta.get(key)!r}" + ) + for index, (actual, expected) in enumerate(zip(full_buffers, expected_buffers)): + field_names = local_meta.get("field_names", ()) + field = field_names[index] if index < len(field_names) else f"buffer[{index}]" + _record_exact(errors, actual, expected, f"{param_name}: {direction} {field}") + + # Per-tensor FP8 scale is metadata rather than an FSDP buffer. It must be + # identical on every shard and preserved on the reconstructed full tensor. + if isinstance(local_sub, Float8TensorStorage): + local_scale_inv = getattr(local_sub, "_scale_inv", None) + _collective_assert( + isinstance(local_scale_inv, torch.Tensor), + f"{param_name}: {direction} Float8 storage has no tensor _scale_inv", + ) + local_scale = local_scale_inv.detach().clone() + _collective_assert_same_shape(local_scale, f"{param_name}: {direction} Float8 scale") + gathered_scales = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered_scales, local_scale) + for rank, scale in enumerate(gathered_scales[1:], start=1): + _record_exact( + errors, + scale, + gathered_scales[0], + f"{param_name}: {direction} scale rank {rank}", + ) + _record_exact( + errors, + full_sub._scale_inv, + gathered_scales[0], + f"{param_name}: {direction} reconstructed scale", + ) + + +def _manual_reconstruct_hybrid(local, *, param_name, world_size): + """Run the Hybrid FSDP pre/post protocol on manually gathered raw buffers.""" + _collective_assert( + isinstance(local, HybridQuantizedTensor), + f"{param_name}: local shard is {type(local).__name__}, expected HybridQuantizedTensor", + ) + sharded_tensors, metadata = local.fsdp_pre_all_gather( + mesh=None, + orig_size=local.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + local_count = torch.tensor([len(sharded_tensors)], dtype=torch.int64, device=local.device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{param_name}: Hybrid FSDP buffer count range is ({min_count.item()}, {max_count.item()})", + ) + + gathered_outputs = [] + for index, shard in enumerate(sharded_tensors): + _collective_assert( + shard is not None, + f"{param_name}: Hybrid FSDP buffer {index} is None", + ) + _collective_assert_same_shape(shard, f"{param_name}: Hybrid FSDP buffer {index}") + gathered = [torch.zeros_like(shard) for _ in range(world_size)] + dist.all_gather(gathered, shard) + gathered_outputs.append(torch.cat(gathered, dim=0)) + + reconstructed, _ = local.fsdp_post_all_gather( + tuple(gathered_outputs), metadata, local.dtype, out=None + ) + _collective_assert( + isinstance(reconstructed, HybridQuantizedTensor), + f"{param_name}: Hybrid FSDP reconstruction returned {type(reconstructed).__name__}", + ) + return reconstructed + + def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). @@ -1040,6 +1235,1047 @@ def test_dcp_resharding_load(recipe_name): os.remove(ref_output_path) +# --------------------------------------------------------------------------- +# Hybrid quantization + FSDP2 tests +# --------------------------------------------------------------------------- + + +def _build_hybrid_model(hybrid_recipe, use_meta_device=True): + """Build a model with quantized_model_init using a hybrid CustomRecipe.""" + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if use_meta_device: + kwargs["device"] = "meta" + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(NUM_LAYERS) + ] + ) + return model + + +def test_fused_adam_hybrid_master_weights(hybrid_recipe_name): + """FusedAdam + master_weights + FSDP2 + hybrid quantized_model_init. + + Verifies: + - Params are DTensors wrapping HybridQuantizedTensor local shards + - Training loop completes without error + - Optimizer states are FP32 + - Loss decreases over training steps + """ + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + hybrid_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + 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 + + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + +def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name): + """Hybrid FusedAdam training must be numerically invariant to FSDP2's + ``reshard_after_forward`` schedule. + + ``reshard_after_forward`` only changes *when* the gathered weight is + materialized/freed, not the math: ``True`` (FSDP2's child-module default) + drops the gathered weight after forward and re-gathers it in backward -- + invoking ``fsdp_post_all_gather(out=...)`` twice per step -- while ``False`` + keeps the gathered copy alive through backward (one gather per step). The + gathered quantized bytes are identical either way, so both schedules must + produce bitwise-identical outputs, losses, input/parameter gradients, + master parameters, moments, and step counters. + + Strictly stronger than "loss decreased": it locks in that the hybrid + all-gather hooks are schedule-invariant across both FSDP2 passes, and + regression-guards the future P1.1 buffer-split bandwidth optimization. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + # Shared, fixed input/target so the two schedules are compared on identical data. + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def assert_exact(actual, expected, path): + if torch.is_tensor(expected): + torch.testing.assert_close( + actual, + expected, + rtol=0.0, + atol=0.0, + msg=lambda m: f"reshard schedule changed {path}: {m}", + ) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys(), f"{path}: keys differ" + for key in expected: + assert_exact(actual[key], expected[key], f"{path}.{key}") + elif isinstance(expected, (list, tuple)): + assert len(actual) == len(expected), f"{path}: lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + assert_exact(actual_item, expected_item, f"{path}[{index}]") + else: + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + + def run(reshard_after_forward): + # Re-seed so both schedules get identical weight init from reset_parameters(). + torch.manual_seed(42) + torch.cuda.manual_seed(42) + model = _shard_model( + _build_hybrid_model(hybrid_recipe), + world_size, + reshard_after_forward=reshard_after_forward, + ) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + run_x = x.detach().clone().requires_grad_() + artifacts = { + "outputs": [], + "losses": [], + "input_grads": [], + "param_grads": [], + "optimizer": [], + } + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + run_x.grad = None + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(run_x) + artifacts["outputs"].append(output.detach().clone()) + loss = F.mse_loss(output, target) + artifacts["losses"].append(loss.detach().clone()) + loss.backward() + artifacts["input_grads"].append(run_x.grad.detach().clone()) + step_grads = [] + for param in model.parameters(): + grad = param.grad + if grad is not None: + grad = grad.to_local() if isinstance(grad, DTensor) else grad + grad = grad.detach().clone() + step_grads.append(grad) + artifacts["param_grads"].append(step_grads) + optimizer.step() + step_state = [] + for param in model.parameters(): + param_state = {} + for key, value in optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + param_state[key] = value + step_state.append(param_state) + artifacts["optimizer"].append(step_state) + return artifacts + + artifacts_resharded = run(reshard_after_forward=True) # re-gather in backward + artifacts_kept = run(reshard_after_forward=False) # keep gathered weight through backward + + losses_resharded = [loss.item() for loss in artifacts_resharded["losses"]] + assert all( + losses_resharded[i + 1] < losses_resharded[i] for i in range(NUM_STEPS - 1) + ), f"reshard_after_forward=True loss not strictly decreasing: {losses_resharded}" + assert_exact(artifacts_resharded, artifacts_kept, "training") + + +def test_fused_adam_hybrid_bf16_vs_hybrid_parity(hybrid_recipe_name): + """Compare hybrid+FSDP2 loss trajectory against BF16+FSDP2 within tolerance. + + This is a sanity check that hybrid quantized training converges similarly + to BF16 training, not a bitwise-exact comparison. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def run_training(model, recipe_for_autocast): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + losses = [] + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=(recipe_for_autocast is not None), recipe=recipe_for_autocast): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + return losses + + # BF16 baseline + torch.manual_seed(42) + torch.cuda.manual_seed(42) + bf16_model = _build_model(fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_losses = run_training(bf16_model, None) + + # Hybrid + torch.manual_seed(42) + torch.cuda.manual_seed(42) + hybrid_model = _build_hybrid_model(hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_losses = run_training(hybrid_model, hybrid_recipe) + + assert hybrid_losses[-1] < hybrid_losses[0], f"Hybrid loss did not decrease: {hybrid_losses}" + assert bf16_losses[-1] < bf16_losses[0], f"BF16 loss did not decrease: {bf16_losses}" + + # Hybrid stays within a few % of bf16 (seed-fixed). + rel_tol = 0.10 + for step, (h_loss, b_loss) in enumerate(zip(hybrid_losses, bf16_losses)): + rel_diff = abs(h_loss - b_loss) / max(abs(b_loss), 1e-10) + assert rel_diff < rel_tol, ( + f"Step {step}: hybrid loss ({h_loss:.4f}) vs bf16 ({b_loss:.4f}) " + f"differ by {rel_diff * 100:.2f}% (> {rel_tol * 100:.0f}%)" + ) + + +# Same-format hybrid -> the vanilla recipe it must match bitwise. Cross-format +# hybrids (e.g. HybridMixed_MXFP8_FP8) have no single-format vanilla equivalent. +_HYBRID_TO_BASE_RECIPE = { + "HybridFP8CurrentScaling": "Float8CurrentScaling", + "HybridMXFP8": "MXFP8BlockScaling", + "HybridFloat8BlockScaling": "Float8BlockScaling", +} + + +def _build_linear_parity_stack(recipe): + """Two bare ``te.Linear`` layers under ``quantized_model_init`` for + hybrid-vs-vanilla bitwise parity. + """ + with te.quantized_model_init(enabled=True, recipe=recipe): + return torch.nn.Sequential( + te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=torch.bfloat16, device="meta"), + te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=torch.bfloat16, device="meta"), + ) + + +def test_fused_adam_hybrid_vs_base_recipe_parity(hybrid_recipe_name): + """Same-format hybrid must match its vanilla recipe bitwise through the full + FSDP2 + FusedAdam loop. + + Every output, loss, input/parameter gradient, master parameter, optimizer + moment, and step counter is asserted bitwise-identical -- a regression guard + for both amax reduction and master-weight requantization. Uses a bare + ``te.Linear`` stack (see ``_build_linear_parity_stack``) to isolate + GEMM-operand quantization. + """ + if hybrid_recipe_name not in _HYBRID_TO_BASE_RECIPE: + pytest.skip( + f"{hybrid_recipe_name} is cross-format; no single-format vanilla " + "recipe to compare against." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + base_recipe_name = _HYBRID_TO_BASE_RECIPE[hybrid_recipe_name] + world_size, device = _get_dist_info() + + # Shared, fixed input/target; the comparison is per-rank (base vs hybrid on + # the same rank), so cross-rank input consistency does not matter. + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def run_training(build_fn, recipe_for_autocast): + def snapshot_optimizer(model, optimizer): + snapshots = [] + for param in model.parameters(): + state = {} + for key, value in optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + state[key] = value + snapshots.append(state) + return snapshots + + # Re-seed so both models get identical init from reset_parameters() (run + # after sharding); with same-format quantization and a dropout-free loop + # the full trajectory is then deterministic. + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + model = _shard_model(build_fn(), world_size) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + run_x = x.detach().clone().requires_grad_() + outputs = [] + losses = [] + input_grads = [] + grads_per_step = [] + optimizer_states = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + run_x.grad = None + with te.autocast(enabled=True, recipe=recipe_for_autocast): + output = model(run_x) + outputs.append(output.detach().clone()) + loss = F.mse_loss(output, target) + losses.append(loss.detach().clone()) + loss.backward() + input_grads.append(run_x.grad.detach().clone()) + # Snapshot grad local shards before the optimizer consumes them + # (p.grad is a DTensor under FSDP2) to assert backward parity directly. + step_grads = [] + for p in model.parameters(): + g = p.grad + if g is None: + step_grads.append(None) + else: + g = g.to_local() if isinstance(g, DTensor) else g + step_grads.append(g.detach().clone()) + grads_per_step.append(step_grads) + optimizer.step() + optimizer_states.append(snapshot_optimizer(model, optimizer)) + return outputs, losses, input_grads, grads_per_step, optimizer_states + + base_recipe = get_recipe_from_string(base_recipe_name) + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + + base_outputs, base_losses, base_input_grads, base_grads, base_opt_states = run_training( + lambda: _build_linear_parity_stack(base_recipe), base_recipe + ) + ( + hybrid_outputs, + hybrid_losses, + hybrid_input_grads, + hybrid_grads, + hybrid_opt_states, + ) = run_training(lambda: _build_linear_parity_stack(hybrid_recipe), hybrid_recipe) + + # (1) Every forward: bitwise-identical before and after optimizer updates. + for step, (base_output, hybrid_output) in enumerate(zip(base_outputs, hybrid_outputs)): + torch.testing.assert_close( + hybrid_output, + base_output, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} forward output not bitwise-identical: {m}", + ) + + # (2) Every per-step loss: bitwise-identical across the whole optimizer loop. + for step, (b_loss, h_loss) in enumerate(zip(base_losses, hybrid_losses)): + torch.testing.assert_close( + h_loss, + b_loss, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} loss not bitwise-identical to the vanilla recipe: {m}", + ) + + # (3) Backward: every weight-gradient shard at every step bitwise-identical + for step, (base_grad, hybrid_grad) in enumerate(zip(base_input_grads, hybrid_input_grads)): + torch.testing.assert_close( + hybrid_grad, + base_grad, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} input gradient not bitwise-identical: {m}", + ) + + # (implied by the loss trajectory, but asserted directly to be explicit). + for step, (b_step, h_step) in enumerate(zip(base_grads, hybrid_grads)): + for i, (b_grad, h_grad) in enumerate(zip(b_step, h_step)): + assert (b_grad is None) == (h_grad is None), ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} param {i}" + " gradient presence differs between hybrid and vanilla" + ) + if b_grad is None: + continue + torch.testing.assert_close( + h_grad, + b_grad, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step, i=i: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} gradient not bitwise-identical to the vanilla recipe: {m}", + ) + + # Optimizer continuation state: FP32 master params, moments, and counters. + assert len(base_opt_states) == len(hybrid_opt_states) + for step, (base_step, hybrid_step) in enumerate(zip(base_opt_states, hybrid_opt_states)): + assert len(base_step) == len(hybrid_step) + for param_idx, (base_state, hybrid_state) in enumerate(zip(base_step, hybrid_step)): + assert base_state.keys() == hybrid_state.keys(), ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} param {param_idx} " + f"optimizer state keys differ: {base_state.keys()} != {hybrid_state.keys()}" + ) + for key in base_state: + base_value = base_state[key] + hybrid_value = hybrid_state[key] + if torch.is_tensor(base_value): + torch.testing.assert_close( + hybrid_value, + base_value, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step, i=param_idx, k=key: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} optimizer state {k!r} not bitwise-identical: {m}", + ) + else: + assert hybrid_value == base_value, ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} " + f"param {param_idx} optimizer state {key!r} differs: " + f"{hybrid_value!r} != {base_value!r}" + ) + + +def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): + """Per-tensor hybrid weights must share ONE amax-reduced scale across FSDP2 + shards -- tolerance-free regression guard for the amax-reduction fix. + + Without cross-shard reduction each rank quantizes its shard with a local amax + and the scales differ; with the fix they match. Checked directly on the + sharded weight (no forward). Block-scaled formats (MXFP8) are skipped. + """ + if hybrid_recipe_name != "HybridFP8CurrentScaling": + pytest.skip("scale-uniformity check applies to per-tensor current scaling only") + + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + if world_size < 2: + pytest.skip("needs >=2 ranks to compare shard scales") + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + checked = {"rowwise": 0, "columnwise": 0} + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ): + continue + for direction in checked: + sub_storage = getattr(param._local_tensor, f"_{direction}_storage") + scale_inv = getattr(sub_storage, "_scale_inv", None) + if scale_inv is None: + continue + local_scale = scale_inv.detach().reshape(-1).clone() + gathered = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered[r], + gathered[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r, d=direction: f"{n}: rank {r} {d} _scale_inv differs from rank 0 -- cross-shard amax reduction was not applied to the hybrid current-scaling weight: {m}", + ) + checked[direction] += 1 + assert all( + count > 0 for count in checked.values() + ), f"missing hybrid current-scaling directions: {checked}" + + +def test_fused_adam_hybrid_identity_fp8_master_weights(): + """FSDP2 + FusedAdam with Hybrid(FP8 current rowwise, Identity columnwise). + + Covers the Identity sub-storage in hybrid FSDP2 all-gather while the FP8 + current rowwise direction validates cross-shard amax reduction via scale + uniformity. + """ + from transformer_engine.pytorch import HybridQuantizedTensor + from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( + IdentityTensorStorage, + ) + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + if world_size < 2: + pytest.skip("needs >=2 ranks to validate cross-shard amax reduction") + + hybrid_recipe = get_hybrid_recipe_from_string("HybridFP8CurrentScalingIdentity") + model = _build_linear_parity_stack(hybrid_recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + identity_steps = 2 + for step in range(identity_steps): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + if step < identity_steps - 1: + optimizer.step() + + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Hybrid Identity/FP8 loss not strictly decreasing: {losses}" + + checked_identity = 0 + checked_scale = 0 + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ): + continue + local = param._local_tensor + assert isinstance(local._columnwise_storage, IdentityTensorStorage) + + scale_inv = getattr(local._rowwise_storage, "_scale_inv", None) + if scale_inv is not None: + local_scale = scale_inv.detach().reshape(-1).clone() + gathered_scales = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered_scales, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered_scales[r], + gathered_scales[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r: f"{n}: rank {r} rowwise _scale_inv differs from rank 0 for Hybrid(FP8Current, Identity): {m}", + ) + checked_scale += 1 + + local_identity = local._columnwise_storage.dequantize().contiguous() + gathered_identity = [torch.zeros_like(local_identity) for _ in range(world_size)] + dist.all_gather(gathered_identity, local_identity) + manual_full = torch.cat(gathered_identity, dim=0) + + sharded_tensors, metadata = local.fsdp_pre_all_gather( + mesh=None, + orig_size=local.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + all_gather_outputs = [] + for shard in sharded_tensors: + gathered = [torch.zeros_like(shard) for _ in range(world_size)] + dist.all_gather(gathered, shard) + all_gather_outputs.append(torch.cat(gathered, dim=0)) + fsdp_full, _ = local.fsdp_post_all_gather( + tuple(all_gather_outputs), metadata, local.dtype, out=None + ) + assert isinstance(fsdp_full, HybridQuantizedTensor) + full_identity = fsdp_full._columnwise_storage.dequantize() + torch.testing.assert_close( + manual_full.float(), + full_identity[: manual_full.shape[0]].float(), + rtol=0.0, + atol=0.0, + msg=lambda m, n=name: f"{n}: Identity columnwise all-gather mismatch: {m}", + ) + checked_identity += 1 + + assert checked_identity > 0, "no Hybrid(FP8Current, Identity) params found" + assert checked_scale > 0, "no FP8 current rowwise scales found to check" + + +def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): + """Validate both Hybrid directions through FSDP2 at raw-buffer precision. + + The rowwise and columnwise sub-storages are checked independently. Every + gathered data buffer, scale buffer, field layout, and per-tensor scale + metadata value must exactly match a manual dim-0 all-gather. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + expected_types = { + "HybridFP8CurrentScaling": (Float8TensorStorage, Float8TensorStorage), + "HybridMXFP8": (MXFP8TensorStorage, MXFP8TensorStorage), + "HybridFloat8BlockScaling": ( + Float8BlockwiseQTensorStorage, + Float8BlockwiseQTensorStorage, + ), + "HybridMixed_MXFP8_FP8": (MXFP8TensorStorage, Float8TensorStorage), + } + row_type, col_type = expected_types[hybrid_recipe_name] + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + _ = model(x) + + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{hybrid_recipe_name}: Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) + + errors = [] + for name, param in params: + local = param._local_tensor + reconstructed = _manual_reconstruct_hybrid(local, param_name=name, world_size=world_size) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{name}: public full_tensor vs Hybrid reconstruction", + ) + _check_hybrid_direction_buffers( + local._rowwise_storage, + reconstructed._rowwise_storage, + row_type, + direction="rowwise", + param_name=name, + world_size=world_size, + errors=errors, + ) + _check_hybrid_direction_buffers( + local._columnwise_storage, + reconstructed._columnwise_storage, + col_type, + direction="columnwise", + param_name=name, + world_size=world_size, + errors=errors, + ) + + _raise_collective_errors(errors, f"{hybrid_recipe_name} raw Hybrid all-gather") + + +def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): + """Exercise MXFP8 block-scale unpad/pad on a sharded Linear whose shard + dim-0 is block-aligned (divisible by 32) but NOT divisible by 128. + + MXFP8 block scales are stored with ``[128, 4]`` / ``[4, 128]`` alignment + padding, which must be stripped before FSDP2's dim-0 all-gather and + re-applied after. With ``HIDDEN_SIZE`` and ``FFN_HIDDEN_SIZE`` both + divisible by 128, the default model never forces this code path, so this + test uses a hand-picked Linear size. + + Regression test for the "pre-fix" bug where + ``HybridQuantizedTensor.fsdp_pre_all_gather`` pulled raw tensor fields via + ``get_metadata()`` without unpadding the scale — the padded bytes would + have been interleaved at every rank boundary in the gather output. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + supported, reason = te.is_mxfp8_available(return_reason=True) + if not supported: + pytest.skip(f"MXFP8: {reason}") + + world_size, device = _get_dist_info() + + # FSDP2 shards a Linear weight of shape (out_features, in_features) along + # dim-0, so each rank holds `out_features / world_size` rows. Pick + # per-rank shard dim-0 = 96: divisible by MXFP8_BLOCK_SCALING_SIZE (32) + # so data alignment holds, but NOT divisible by 128 so the rowwise + # scale-inv needs alignment padding on the sharded copy. This is the + # shape that exercises the unpad-before-gather / pad-after-gather + # behaviour in MXFP8TensorStorage.fsdp_{extract,assign}_buffers. + per_rank_out = 96 + out_features = per_rank_out * world_size + in_features = 128 # arbitrary, divisible by 32; not sharded by FSDP2 here + assert per_rank_out % 32 == 0, ( + f"Test setup error: per_rank_out={per_rank_out} (= out_features / world_size, " + f"world_size={world_size}) must be a multiple of the MXFP8 block size (32) so the " + "sharded weight's data stays block-aligned. Pick a per_rank_out divisible by 32." + ) + assert per_rank_out % 128 != 0, ( + f"Test setup error: per_rank_out={per_rank_out} must NOT be a multiple of 128, or the " + "rowwise scale-inv needs no alignment padding and this test stops exercising the MXFP8 " + "unpad-before-gather / pad-after-gather path it exists to cover. Pick a per_rank_out " + "divisible by 32 but not 128 (e.g. 96)." + ) + + for recipe_name in ("HybridMXFP8", "HybridMixed_MXFP8_FP8"): + hybrid_recipe = get_hybrid_recipe_from_string(recipe_name) + + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + te.Linear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ), + ) + model = _shard_model(model, world_size) + + # Batch (leading) dim must be divisible by MXFP8_BLOCK_SCALING_SIZE (32). + x = torch.randn(32, in_features, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + out = model(x) + out.sum().backward() + + col_type = MXFP8TensorStorage if recipe_name == "HybridMXFP8" else Float8TensorStorage + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{recipe_name}: awkward-shape Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) + + errors = [] + for name, param in params: + local = param._local_tensor + label = f"{recipe_name}:{name}" + reconstructed = _manual_reconstruct_hybrid( + local, param_name=label, world_size=world_size + ) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{label}: public full_tensor vs Hybrid reconstruction", + ) + _check_hybrid_direction_buffers( + local._rowwise_storage, + reconstructed._rowwise_storage, + MXFP8TensorStorage, + direction="rowwise", + param_name=label, + world_size=world_size, + errors=errors, + ) + _check_hybrid_direction_buffers( + local._columnwise_storage, + reconstructed._columnwise_storage, + col_type, + direction="columnwise", + param_name=label, + world_size=world_size, + errors=errors, + ) + _raise_collective_errors(errors, f"{recipe_name} awkward raw Hybrid all-gather") + + +def test_fused_adam_hybrid_float8_block_unaligned_shard_shape(): + """Unaligned local shards are rejected before gathering incompatible scale tiles.""" + from transformer_engine.pytorch import fp8 + from fsdp2_utils import get_hybrid_recipe_from_string + + supported, reason = fp8.check_fp8_block_scaling_support() + if not supported: + pytest.skip(reason) + + world_size, device = _get_dist_info() + if world_size != 2: + pytest.skip("test shape is defined for exactly two FSDP ranks") + + in_features = 256 + per_rank_out = 192 + out_features = per_rank_out * world_size + assert out_features % 128 == 0 + assert per_rank_out % 128 != 0 + + hybrid_recipe = get_hybrid_recipe_from_string("HybridFloat8BlockScaling") + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + te.Linear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ) + ) + model = _shard_model(model, world_size) + + x = torch.randn(128, in_features, dtype=torch.bfloat16, device=device) + with pytest.raises( + RuntimeError, + match="local flattened M dimension.*not a multiple of 128", + ): + with te.autocast(enabled=True, recipe=hybrid_recipe): + model(x) + + +def test_hybrid_dcp_output_parity(hybrid_recipe_name): + """DCP roundtrip and exact forked optimizer continuation. + + Trains and checkpoints a hybrid model plus FusedAdam, loads both into fresh + objects, and compares the identical next step against uninterrupted training: + output, loss, input/parameter gradients, master params, moments, counters, + and post-step output must all match bitwise. + """ + import torch.distributed.checkpoint as dcp + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + # Deterministic, rank-agnostic checkpoint dir so all ranks read/write + # the same DCP path. ``os.getpid()`` differs per rank under torchrun. + checkpoint_dir = f"/tmp/te_test_fsdp2_hybrid_dcp_parity_{hybrid_recipe_name}" + + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + dist.barrier() + + try: + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + failures = [] + + def snapshot_optimizer(current_model, current_optimizer): + snapshots = [] + for param in current_model.parameters(): + state = {} + for key, value in current_optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + state[key] = value + snapshots.append(state) + return snapshots + + def check_optimizer_state(actual, expected, label): + if len(actual) != len(expected): + failures.append( + f"{label}: parameter count differs: {len(actual)} != {len(expected)}" + ) + for param_idx, (actual_state, expected_state) in enumerate(zip(actual, expected)): + if actual_state.keys() != expected_state.keys(): + failures.append( + f"{label}: param {param_idx} state keys differ: " + f"{actual_state.keys()} != {expected_state.keys()}" + ) + for key in actual_state.keys() & expected_state.keys(): + actual_value = actual_state[key] + expected_value = expected_state[key] + if torch.is_tensor(expected_value): + _record_exact( + failures, + actual_value, + expected_value, + f"{label}: param {param_idx} optimizer state {key!r}", + ) + elif actual_value != expected_value: + failures.append( + f"{label}: param {param_idx} optimizer state {key!r} differs: " + f"{actual_value!r} != {expected_value!r}" + ) + + def run_continuation_step(current_model, current_optimizer): + step_x = x.detach().clone().requires_grad_() + current_optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + step_output = current_model(step_x) + step_loss = F.mse_loss(step_output, target) + step_loss.backward() + step_grads = [] + for param in current_model.parameters(): + grad = param.grad + if grad is not None: + grad = grad.to_local() if isinstance(grad, DTensor) else grad + grad = grad.detach().clone() + step_grads.append(grad) + current_optimizer.step() + with torch.no_grad(), te.autocast(enabled=True, recipe=hybrid_recipe): + post_step_output = current_model(x).detach().clone() + return { + "output": step_output.detach().clone(), + "loss": step_loss.detach().clone(), + "input_grad": None if step_x.grad is None else step_x.grad.detach().clone(), + "param_grads": step_grads, + "optimizer": snapshot_optimizer(current_model, current_optimizer), + "post_step_output": post_step_output, + } + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + F.mse_loss(output, target).backward() + optimizer.step() + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=hybrid_recipe): + ref_output = model(x).clone() + + save_state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + } + dcp.save(save_state, checkpoint_id=checkpoint_dir) + saved_optimizer_state = snapshot_optimizer(model, optimizer) + + model2 = _build_hybrid_model(hybrid_recipe) + model2 = _shard_model(model2, world_size) + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + state_to_load = { + "model": model2.state_dict(), + "optimizer": optimizer2.state_dict(), + } + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict(state_to_load["model"]) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=hybrid_recipe): + loaded_output = model2(x) + check_optimizer_state( + snapshot_optimizer(model2, optimizer2), + saved_optimizer_state, + "DCP optimizer state immediately after load", + ) + _record_exact(failures, loaded_output, ref_output, "DCP roundtrip output") + _raise_collective_errors(failures, "DCP state immediately after load") + failures.clear() + + # Fork from the checkpoint and execute the identical next step. This + # catches an ignored optimizer checkpoint even when model-only output + # parity succeeds immediately after load. + torch.manual_seed(9876) + torch.cuda.manual_seed(9876) + reference_step = run_continuation_step(model, optimizer) + torch.manual_seed(9876) + torch.cuda.manual_seed(9876) + resumed_step = run_continuation_step(model2, optimizer2) + + for key in ("output", "loss", "input_grad", "post_step_output"): + _record_exact( + failures, + resumed_step[key], + reference_step[key], + f"DCP continuation {key}", + ) + + resumed_grads = resumed_step["param_grads"] + reference_grads = reference_step["param_grads"] + if len(resumed_grads) != len(reference_grads): + failures.append( + "DCP continuation parameter-gradient count differs: " + f"{len(resumed_grads)} != {len(reference_grads)}" + ) + for param_idx, (resumed_grad, reference_grad) in enumerate( + zip(resumed_grads, reference_grads) + ): + _record_exact( + failures, + resumed_grad, + reference_grad, + f"DCP continuation param {param_idx} gradient", + ) + + check_optimizer_state( + resumed_step["optimizer"], + reference_step["optimizer"], + "DCP optimizer state after continuation", + ) + _raise_collective_errors(failures, "DCP forked continuation") + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + + TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, @@ -1053,6 +2289,24 @@ def test_dcp_resharding_load(recipe_name): "dcp_resharding_save": test_dcp_resharding_save, "dcp_resharding_load": test_dcp_resharding_load, "safetensors_fp32_export": test_safetensors_fp32_export, + "fused_adam_hybrid_master_weights": test_fused_adam_hybrid_master_weights, + "fused_adam_hybrid_bf16_vs_hybrid_parity": test_fused_adam_hybrid_bf16_vs_hybrid_parity, + "fused_adam_hybrid_vs_base_recipe_parity": test_fused_adam_hybrid_vs_base_recipe_parity, + "fused_adam_hybrid_scale_uniform_across_shards": ( + test_fused_adam_hybrid_scale_uniform_across_shards + ), + "fused_adam_hybrid_identity_fp8_master_weights": ( + test_fused_adam_hybrid_identity_fp8_master_weights + ), + "fused_adam_hybrid_allgather_correctness": test_fused_adam_hybrid_allgather_correctness, + "fused_adam_hybrid_mxfp8_awkward_shard_shape": test_fused_adam_hybrid_mxfp8_awkward_shard_shape, + "hybrid_dcp_output_parity": test_hybrid_dcp_output_parity, +} + +# Hybrid tests that are NOT parametrized by recipe (they sweep internally). +_HYBRID_NON_PARAMETRIZED_TESTS = { + "fused_adam_hybrid_identity_fp8_master_weights", + "fused_adam_hybrid_mxfp8_awkward_shard_shape", } @@ -1070,6 +2324,11 @@ def test_dcp_resharding_load(recipe_name): "Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling", + "HybridFP8CurrentScaling", + "HybridMXFP8", + "HybridFloat8BlockScaling", + "HybridMixed_MXFP8_FP8", + "HybridFP8CurrentScalingIdentity", ], ) args = parser.parse_args() @@ -1079,7 +2338,10 @@ def test_dcp_resharding_load(recipe_name): torch.manual_seed(42) torch.cuda.manual_seed(42) try: - TESTS[args.test](args.recipe) + if args.test in _HYBRID_NON_PARAMETRIZED_TESTS: + TESTS[args.test]() + else: + TESTS[args.test](args.recipe) finally: if dist.is_initialized(): dist.destroy_process_group() diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index b5436e2709..e943d48d4a 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -452,12 +452,169 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in ) +# ── Hybrid quantization memory tests ───────────────────────────────── + + +def _build_hybrid_model(num_layers, hybrid_recipe, use_meta_device=True): + """Build a model with quantized_model_init using a hybrid CustomRecipe.""" + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if use_meta_device: + kwargs["device"] = "meta" + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(num_layers) + ] + ) + return model + + +def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): + """Hybrid quantized weights should not accumulate across layers during forward. + + Same methodology as test_fp8_temp_accumulation_across_layers but for + hybrid quantized tensors. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # bf16 baseline + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_increments = _measure_forward_increments(bf16_model, bf16_optimizer, None, x, target) + bf16_avg = sum(bf16_increments) / len(bf16_increments) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # Hybrid model + hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_optimizer = te.optimizers.FusedAdam( + hybrid_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) + hybrid_increments = _measure_forward_increments( + hybrid_model, + hybrid_optimizer, + hybrid_recipe, + x, + target, + ) + hybrid_avg = sum(hybrid_increments) / len(hybrid_increments) + + excess_per_layer = hybrid_avg - bf16_avg + # Basis: forward growth is constant per layer (no accumulation) for both bf16 and + # hybrid; the excess is just hybrid's extra per-layer quantized buffers. Measured + # excess: ~3 KiB (FP8 current) / ~7 KiB (mixed MXFP8+FP8) / ~12 KiB (MXFP8). A + # leaked layer's quantized weights would be hundreds of KiB, so 50 KiB sits above + # the real per-layer overhead and well below a leak. + tolerance_per_layer = 50 * 1024 # 50 KiB + + assert excess_per_layer <= tolerance_per_layer, ( + "Hybrid per-layer forward memory increment exceeds bf16 baseline by " + f"{excess_per_layer/1024:.1f} KiB/layer (tolerance: {tolerance_per_layer/1024:.1f} KiB). " + f"bf16 avg: {bf16_avg/1024:.1f} KiB/layer, hybrid avg: {hybrid_avg/1024:.1f} KiB/layer." + ) + + +def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): + """Detect transpose caches from hybrid sub-storages persisting after backward.""" + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # bf16 baseline + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_bwd_delta = _measure_backward_memory_delta(bf16_model, bf16_optimizer, None, x, target) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # Hybrid model + hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_optimizer = te.optimizers.FusedAdam( + hybrid_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) + hybrid_bwd_delta = _measure_backward_memory_delta( + hybrid_model, + hybrid_optimizer, + hybrid_recipe, + x, + target, + ) + + excess = hybrid_bwd_delta - bf16_bwd_delta + # Basis: hybrid retains no more than bf16 after backward+step — measured excess is + # slightly negative (~-0.02..-0.09 MiB vs a ~2 MiB bf16 delta). The tolerance only + # absorbs allocator/measurement noise; a genuinely retained gathered weight or + # transpose cache would be MiB-scale (>> 256 KiB). + tolerance = 256 * 1024 # 256 KiB + + assert excess <= tolerance, ( + f"Hybrid backward retains {excess/1024**2:.2f} MiB more than bf16 baseline. " + f"bf16 backward delta: {bf16_bwd_delta/1024**2:.2f} MiB, " + f"hybrid backward delta: {hybrid_bwd_delta/1024**2:.2f} MiB." + ) + + # ── Standalone runner ──────────────────────────────────────────────── TESTS = { "bf16_no_excess_forward_memory": test_bf16_no_excess_forward_memory, "bf16_no_excess_backward_memory": test_bf16_no_excess_backward_memory, "fp8_temp_accumulation_across_layers": test_fp8_temp_accumulation_across_layers, "transpose_cache_retained_after_backward": test_transpose_cache_retained_after_backward, + "hybrid_no_excess_forward_memory": test_hybrid_no_excess_forward_memory, + "hybrid_transpose_cache_after_backward": test_hybrid_transpose_cache_after_backward, } if __name__ == "__main__": @@ -473,6 +630,10 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in "Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling", + "HybridFP8CurrentScaling", + "HybridMXFP8", + "HybridFloat8BlockScaling", + "HybridMixed_MXFP8_FP8", ], ) parser.add_argument("--quantized-model-init", action="store_true", default=False) @@ -488,11 +649,17 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in "fp8_temp_accumulation_across_layers", "transpose_cache_retained_after_backward", } + _HYBRID_PARAMETRIZED_TESTS = { + "hybrid_no_excess_forward_memory", + "hybrid_transpose_cache_after_backward", + } try: test_fn = TESTS[args.test] if args.test in _PARAMETRIZED_TESTS: test_fn(args.recipe, args.quantized_model_init) + elif args.test in _HYBRID_PARAMETRIZED_TESTS: + test_fn(args.recipe) else: test_fn() finally: diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 36ac307b90..39d825e701 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -27,6 +27,7 @@ """ import gc +import math import os import sys import argparse @@ -211,10 +212,14 @@ def shard_model_with_fsdp2(model, mesh): @torch.no_grad() -def _check_fp8_fsdp2_allgather(model): - # Do manual allgather in fp32 and match against fp8 allgather done - # with fsdp2 - # FP32 manual weight allgather +def _check_fp8_fsdp2_allgather(model, *, tols=None): + """Compare FSDP2's quantized all-gather against a manual HP all-gather. + + ``tols=None`` preserves the format-specific defaults used by the original + FP8 checker. Callers for other quantized tensor compositions can provide a + single tolerance dictionary that applies to every parameter. + """ + # Manual high-precision weight all-gather. fp32_allgathered_params = {} for name, param in model.named_parameters(): assert isinstance(param, DTensor) @@ -225,29 +230,32 @@ def _check_fp8_fsdp2_allgather(model): if device_mesh.ndim > 1 else device_mesh.get_group() ) - # Perform manual allgather on local_tensor. zeros_like will create hp tensor since - # torch_dispatch for local_tensor will go down the dequantization route. + # Materialize high precision explicitly so this also works for composite + # quantized tensors such as HybridQuantizedTensor. + local_hp = local_tensor.dequantize() gathered_tensor = [ - torch.zeros_like(local_tensor) for _ in range(dist.get_world_size(group=dist_group)) + torch.zeros_like(local_hp) for _ in range(dist.get_world_size(group=dist_group)) ] - dist.all_gather(gathered_tensor, local_tensor.dequantize(), group=dist_group) + dist.all_gather(gathered_tensor, local_hp, group=dist_group) full_tensor = torch.cat(gathered_tensor, dim=0) fp32_allgathered_params[name] = full_tensor - # FP8 allgather using FSDP2 + # Quantized all-gather using FSDP2. for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. if hasattr(module, "unshard"): module.unshard() - # Make sure allgathered parameters match exactly + # Make sure all-gathered parameters match. for name, param in model.named_parameters(): - # NVFP4 scale unpad/repad through FSDP2 introduces small numerical - # differences vs the manual dequantize-then-allgather path. - if isinstance(param, NVFP4Tensor): - tols = dict(atol=5e-4, rtol=5e-3) + if tols is not None: + param_tols = tols + elif isinstance(param, NVFP4Tensor): + # NVFP4 scale unpad/repad through FSDP2 introduces small numerical + # differences vs the manual dequantize-then-allgather path. + param_tols = dict(atol=5e-4, rtol=5e-3) else: - tols = {} - torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **tols) - # Revert model to original sharded state + param_tols = {} + torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **param_tols) + # Revert model to original sharded state. for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. if hasattr(module, "reshard"): @@ -405,5 +413,254 @@ def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): _run_training(args) +def test_distributed_hybrid(hybrid_recipe_name): + """FSDP2 training with hybrid quantized_model_init. + + Uses quantized_model_init with a hybrid CustomRecipe on a TransformerLayer + model: + - params are DTensors wrapping HybridQuantizedTensor local shards, + - the loss is finite and strictly decreasing on fixed data (grads flow), + - params keep their hybrid quantized type across optimizer.step(), + - FSDP2's quantized all-gather matches a manual fp32 dequant-then-allgather. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # Float8 block scaling requires every local FSDP shard's flattened M + # dimension to contain whole 128-row scale tiles. Keep the historical + # 512-wide model on up to four ranks, and scale it for larger worlds so + # the projection, fused-QKV, and MLP output dimensions all stay aligned. + shard_alignment = 128 * world_size + hidden_size = ((512 + shard_alignment - 1) // shard_alignment) * shard_alignment + ffn_hidden_size = 4 * hidden_size + assert hidden_size % shard_alignment == 0 + + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[te.TransformerLayer(hidden_size, ffn_hidden_size, 8, **kwargs) for _ in range(2)] + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + model = shard_model_with_fsdp2(model, mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + from transformer_engine.pytorch import HybridQuantizedTensor + + def _hybrid_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + + # quantized_model_init must produce HybridQuantizedTensor local shards. + hybrid_count = _hybrid_param_count() + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer_lr = 1e-3 if hidden_size <= 512 else 1e-4 + optimizer = optim.Adam(model.parameters(), lr=optimizer_lr) + + input_data = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(3): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Hybrid iteration {iteration} completed with loss {loss_val}") + + # Training must actually progress on fixed data: strictly monotonic decrease. + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + # Params must stay HybridQuantizedTensor after the optimizer step -- guards a + # silent dequantize-to-bf16 through the FSDP2 / optimizer path. + assert ( + _hybrid_param_count() == hybrid_count + ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" + + # FSDP2 quantized all-gather must match a manual fp32 dequant-then-allgather. + # Hybrid sub-storages (e.g. MXFP8 scale unpad/repad through FSDP2) can + # introduce small differences vs the manual dequantize-then-allgather path. + _check_fp8_fsdp2_allgather(model, tols=dict(atol=5e-4, rtol=5e-3)) + + +def test_distributed_hybrid_identity_all(): + """FSDP2 training/all-gather with an all-Identity CustomRecipe. + + This is the high-precision passthrough baseline for the hybrid/identity + tensor plumbing: quantized_model_init should produce IdentityTensor local + shards, optimizer steps should preserve that type, and FSDP2 all-gather + should reconstruct the same high-precision values as a manual gather. + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + identity_recipe = get_hybrid_recipe_from_string("Identity") + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + with te.quantized_model_init(enabled=True, recipe=identity_recipe): + model = torch.nn.Sequential( + *[te.TransformerLayer(512, 2048, 8, **kwargs) for _ in range(2)] + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + model = shard_model_with_fsdp2(model, mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + def _identity_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, IdentityTensor) + ) + + identity_count = _identity_param_count() + assert identity_count > 0, "No IdentityTensor local tensors after FSDP2 sharding" + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + input_data = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(3): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=identity_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite Identity loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Identity iteration {iteration} completed with loss {loss_val}") + + assert losses[-1] < losses[0], f"Identity loss did not decrease: {losses}" + assert ( + _identity_param_count() == identity_count + ), "IdentityTensor params lost their quantized type after optimizer.step()" + + _check_fp8_fsdp2_allgather(model, tols={}) + + +def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): + """FSDP2 training with hybrid params and reshard_after_forward=True. + + A single LayerNormLinear as the root module gets reshard_after_forward=True + from FSDP2. This exercises the forward-reshard-backward-reshard cycle where + split/as_strided/slice dispatch ops fire every iteration. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # Keep dim-0 FSDP shards aligned to Float8 block scaling's 128-row tiles. + shard_alignment = 128 * world_size + in_features = ((512 + shard_alignment - 1) // shard_alignment) * shard_alignment + out_features = in_features * 3 + assert in_features % shard_alignment == 0 + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = te.LayerNormLinear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + fully_shard(model, mesh=mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + from transformer_engine.pytorch import HybridQuantizedTensor + + def _hybrid_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + + hybrid_count = _hybrid_param_count() + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + + x = torch.randn(128, 16, in_features, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, out_features, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(5): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Hybrid reshard_after_fwd iter {iteration}, loss {loss_val:.4f}") + + # The forward-reshard-backward-reshard cycle must still train: strict decrease. + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + # Params must survive the split/as_strided/slice reshard dispatch ops with + # their hybrid quantized type intact. + assert ( + _hybrid_param_count() == hybrid_count + ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" + + if __name__ == "__main__": sys.exit(_train(_parse_args())) diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py new file mode 100644 index 0000000000..6094b6cc51 --- /dev/null +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -0,0 +1,826 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Distributed TP/SP coverage for hybrid quantization.""" + +import argparse +import datetime +import os +import sys + +import torch +import torch.distributed as dist +from torch import nn + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe as te_recipe + +from hybrid_quantization_utils import ( + hybrid_fp8_current_e5m2_grads_qfactory, + hybrid_fp8_current_identity_qfactory, + hybrid_mxfp8_identity_qfactory, + hybrid_mxfp8_qfactory, + hybrid_nvfp4_qfactory, + hybrid_tp_mxfp8_nvfp4_qfactory, + identity_qfactory, +) +from distributed.run_layer_with_overlap import _compare_tensors + +# ── Global state ───────────────────────────────────────────────────── + +SEQ_LEN = 32 +BATCH_SIZE = 32 +HIDDEN_SIZE = 128 +FFN_HIDDEN_SIZE = 128 +NR_HEADS = 4 + +WORLD_RANK = None +WORLD_SIZE = None +NCCL_WORLD = None +QUANTIZATION = None + +LOSS_FN = nn.MSELoss() + + +def hybrid_recipe(): + """Return a fresh CustomRecipe for the selected test recipe.""" + if QUANTIZATION == "hybrid_fp8": + return te_recipe.CustomRecipe(qfactory=hybrid_fp8_current_e5m2_grads_qfactory) + if QUANTIZATION == "hybrid_mxfp8": + return te_recipe.CustomRecipe(qfactory=hybrid_mxfp8_qfactory) + if QUANTIZATION == "hybrid_fp8_identity": + return te_recipe.CustomRecipe(qfactory=hybrid_fp8_current_identity_qfactory) + if QUANTIZATION == "hybrid_mxfp8_identity": + return te_recipe.CustomRecipe(qfactory=hybrid_mxfp8_identity_qfactory) + if QUANTIZATION == "identity": + return te_recipe.CustomRecipe(qfactory=identity_qfactory) + if QUANTIZATION == "hybrid_nvfp4": + return te_recipe.CustomRecipe(qfactory=hybrid_nvfp4_qfactory) + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + return te_recipe.CustomRecipe(qfactory=hybrid_tp_mxfp8_nvfp4_qfactory) + raise ValueError(f"Unknown hybrid QUANTIZATION={QUANTIZATION!r}") + + +# ── Tolerances ─────────────────────────────────────────────────────── +# +# Mostly upstream ``run_numerics.py`` tolerances. NVFP4 needs a slightly +# larger atol for the measured column+SP Linear output. + + +def _get_tolerances(): + if QUANTIZATION == "identity": + # BF16 TP reductions accumulate in a different order. + return {"rtol": 1.6e-2, "atol": 1.0e-5} + if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): + # Loose because of sequence parallel & amax reduction (fp8_cs). + return {"rtol": 0.4, "atol": 0.25} + if QUANTIZATION in ("hybrid_mxfp8", "hybrid_mxfp8_identity"): + return {"rtol": 0.125, "atol": 0.0625} + if QUANTIZATION == "hybrid_nvfp4": + # Measured column+SP Linear output max abs is ~0.144. + return {"rtol": 0.125, "atol": 0.15} + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + # Backward GEMMs run in NVFP4 -> inherit the (looser) NVFP4 bounds. + return {"rtol": 0.125, "atol": 0.12} + raise ValueError(f"No tolerances for QUANTIZATION={QUANTIZATION!r}") + + +# ── Distributed helpers ────────────────────────────────────────────── + + +def dist_print(msg, src=None, error=False): + stream = sys.stderr if error else sys.stdout + if WORLD_RANK == (0 if src is None else src): + stream.write(f"[rank{WORLD_RANK}] {msg}\n") + stream.flush() + + +def _gather(tensor, dim=0): + """All-gather with ``run_numerics.py`` gradient scaling.""" + + class HalfGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad_output): + return grad_output / WORLD_SIZE + + tensor = HalfGradient.apply(tensor) + gathered = torch.distributed.nn.functional.all_gather(tensor, group=NCCL_WORLD) + return torch.cat(gathered, dim=dim) + + +def _copy_params(model_distributed, model_single): + """Copy single-node parameters into the local TP shard.""" + for dp, sp in zip(model_distributed.parameters(), model_single.parameters()): + with torch.no_grad(): + to_copy = sp + for dim, _ in enumerate(dp.shape): + if dp.shape[dim] != sp.shape[dim]: + start = WORLD_RANK * dp.shape[dim] + end = (WORLD_RANK + 1) * dp.shape[dim] + indices = [slice(None)] * max(min(dim, len(dp.shape) - 1), 0) + indices.append(slice(start, end)) + if dim < len(dp.shape) - 1: + indices.append(slice(None)) + to_copy = sp[tuple(indices)] + dp.copy_(to_copy) + + +def _match_param_sizes(dist_param, single_param): + indices = [slice(None)] * len(single_param.shape) + for i in range(len(dist_param.shape)): + if dist_param.shape[i] != single_param.shape[i]: + start = WORLD_RANK * dist_param.shape[i] + end = (WORLD_RANK + 1) * dist_param.shape[i] + indices[i] = slice(start, end) + return single_param[tuple(indices)] + + +def _check_outputs(output_single, output_dist, label="outputs", *, tolerances=None): + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + f, info = _compare_tensors( + label, output_dist, output_single, **(tolerances or _get_tolerances()) + ) + if f: + dist_print(info, src=WORLD_RANK, error=True) + failed[0] = int(f) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: numerical check failed on at least one rank" + + +def _collective_assert(condition, label): + """Raise on every rank only after all ranks report structural status.""" + failed = torch.tensor([not condition], dtype=torch.uint8, device="cuda") + if not condition: + dist_print(label, src=WORLD_RANK, error=True) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: failed on at least one rank" + + +def _check_gradients( + model_dist, model_single, *, bitwise=False, reduce_replicated=False, tolerances=None +): + """Compare every parameter gradient, including presence and parameter count. + + Sequence-parallel replicated parameters accumulate a partial gradient on + each rank. ``reduce_replicated`` reconstructs their full reference gradient + without mutating the gradient held by the model. + """ + dist_params = list(model_dist.named_parameters()) + single_params = list(model_single.named_parameters()) + local_counts = torch.tensor( + [len(dist_params), len(single_params)], dtype=torch.int64, device="cuda" + ) + min_counts = local_counts.clone() + max_counts = local_counts.clone() + dist.all_reduce(min_counts, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_counts, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_counts, max_counts) and len(dist_params) == len(single_params), + "parameter count mismatch: " + f"local distributed={len(dist_params)}, local single={len(single_params)}, " + f"global min={min_counts.tolist()}, global max={max_counts.tolist()}", + ) + for i, ((name, pd), (single_name, ps)) in enumerate(zip(dist_params, single_params)): + _collective_assert( + name == single_name, + f"parameter {i} name mismatch: {name!r} != {single_name!r}", + ) + local_presence = torch.tensor( + [pd.grad is not None, ps.grad is not None], dtype=torch.uint8, device="cuda" + ) + min_presence = local_presence.clone() + max_presence = local_presence.clone() + dist.all_reduce(min_presence, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_presence, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_presence, max_presence) + and local_presence[0].item() == local_presence[1].item(), + f"grad[{i}].{name}: gradient presence differs locally or across ranks: " + f"local={local_presence.tolist()}, global min={min_presence.tolist()}, " + f"global max={max_presence.tolist()}", + ) + if pd.grad is None: + continue + pd_grad = pd.grad + local_reduce = torch.tensor( + [reduce_replicated and pd_grad.shape == ps.grad.shape], + dtype=torch.uint8, + device="cuda", + ) + min_reduce = local_reduce.clone() + max_reduce = local_reduce.clone() + dist.all_reduce(min_reduce, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_reduce, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_reduce, max_reduce), + f"grad[{i}].{name}: replicated-gradient reduction branch differs across ranks", + ) + if local_reduce.item(): + pd_grad = pd_grad.detach().clone() + dist.all_reduce(pd_grad, dist.ReduceOp.SUM, NCCL_WORLD) + ps_grad = _match_param_sizes(pd_grad, ps.grad) + label = f"grad[{i}].{name}" + if bitwise: + _check_bitwise(pd_grad, ps_grad, label) + else: + _check_outputs(ps_grad, pd_grad, label=label, tolerances=tolerances) + + +def _check_input_gradient(inp_dist, inp_single, label, *, bitwise=False): + """Compare the local TP/SP input-gradient shard with its full reference.""" + local_presence = torch.tensor( + [inp_dist.grad is not None, inp_single.grad is not None], + dtype=torch.uint8, + device="cuda", + ) + min_presence = local_presence.clone() + max_presence = local_presence.clone() + dist.all_reduce(min_presence, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_presence, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_presence, max_presence) + and local_presence[0].item() == local_presence[1].item(), + f"{label}: input-gradient presence differs locally or across ranks: " + f"local={local_presence.tolist()}, global min={min_presence.tolist()}, " + f"global max={max_presence.tolist()}", + ) + _collective_assert( + inp_dist.grad is not None, + f"{label}: both input gradients are unexpectedly absent", + ) + expected = _match_param_sizes(inp_dist.grad, inp_single.grad) + if bitwise: + _check_bitwise(inp_dist.grad, expected, label) + else: + _check_outputs(expected, inp_dist.grad, label=label) + + +def _apply_models(model_single, model_dist, inp_single, inp_dist, **kwargs): + """Run both models with fresh CustomRecipe instances.""" + inp_single.requires_grad_() + inp_dist.requires_grad_() + with te.autocast(enabled=True, recipe=hybrid_recipe()): + out_single = model_single(inp_single, **kwargs) + with te.autocast(enabled=True, recipe=hybrid_recipe()): + out_dist = model_dist(inp_dist, **kwargs) + return out_single, out_dist + + +def _loss_backward(out_single, out_dist): + target = torch.randn_like(out_single) + LOSS_FN(out_single, target).backward() + LOSS_FN(out_dist, target).backward() + + +# ── Test 1: te.Linear TP (column + row) × SP (on/off) ──────────────── + + +def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, amax_stress=False): + dist_print( + f"linear: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" + f" dtype={params_dtype} amax_stress={amax_stress}" + ) + + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + + model_single = te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.Linear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode=parallel_mode, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + # Match run_numerics._test_linear input layouts. + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if parallel_mode == "row": + split = HIDDEN_SIZE // WORLD_SIZE + inp_dist = inp_single[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() + elif parallel_mode == "column": + if sequence_parallel: + # SP column: input is sharded along batch/sequence dim 0. + inp_single = torch.empty((WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if amax_stress and WORLD_RANK == WORLD_SIZE - 1: + # One-rank outlier before SP gather. + inp_dist[-1, -1] = 1.0e3 + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_dist = inp_single.clone() + else: + raise ValueError(parallel_mode) + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # For column-parallel: output is split along feature dim 1; gather. + # For row-parallel + SP: output is split along seq dim 0; gather. + if parallel_mode == "column" or (sequence_parallel and parallel_mode == "row"): + gather_dim = 1 if parallel_mode == "column" else 0 + out_dist = _gather(out_dist, dim=gather_dim) + + _loss_backward(out_single, out_dist) + _check_outputs( + out_single, + out_dist, + label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}]", + ) + + _check_input_gradient( + inp_dist, + inp_single, + label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}] dgrad", + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_linear(): + for parallel_mode in ["column", "row"]: + for sequence_parallel in [False, True]: + _test_linear(parallel_mode, sequence_parallel) + # Current-scaling amax stress: one rank owns an outlier before SP gather. + if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): + _test_linear("column", True, amax_stress=True) + + +# ── Test 1b: te.Linear hybrid-vs-vanilla bitwise operand equivalence ─ + + +def vanilla_recipe(): + """Built-in recipe for same-format hybrid-vs-vanilla checks.""" + if QUANTIZATION == "hybrid_fp8": + return te_recipe.Float8CurrentScaling() + if QUANTIZATION == "hybrid_mxfp8": + return te_recipe.MXFP8BlockScaling() + if QUANTIZATION == "hybrid_nvfp4": + return te_recipe.NVFP4BlockScaling() + if QUANTIZATION == "identity": + return None + raise ValueError(f"No vanilla recipe for QUANTIZATION={QUANTIZATION!r}") + + +def _backward_not_bitwise_comparable(): + """NVFP4 backward consumes SR RNG differently in hybrid vs vanilla.""" + return QUANTIZATION == "hybrid_nvfp4" + + +def _check_bitwise(actual, expected, label): + """Assert bitwise equality (rtol=0, atol=0), all-reduced across ranks.""" + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + try: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + except AssertionError as exc: + dist_print(f"{label}: {exc}", src=WORLD_RANK, error=True) + failed[0] = 1 + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: not bitwise-identical on at least one rank" + + +def _test_linear_vs_vanilla(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): + """Same-topology Linear check: forward bitwise, backward where comparable.""" + dist_print( + f"linear_vs_vanilla: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" + ) + + def run(recipe): + # Fresh model per recipe (re-seeded for identical weights): TE caches a + # quantized weight workspace on the module, so reusing one model would + # let the first recipe's cached weight contaminate the second. + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + model = te.Linear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode=parallel_mode, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + torch.manual_seed(34567) + torch.cuda.manual_seed(34567) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if parallel_mode == "row": + split = HIDDEN_SIZE // WORLD_SIZE + inp = inp[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() + inp.requires_grad_() + + with te.autocast(enabled=recipe is not None, recipe=recipe): + out = model(inp) + # Fixed, recipe-independent target so both backward graphs match. + torch.manual_seed(54321) + torch.cuda.manual_seed(54321) + target = torch.randn_like(out) + LOSS_FN(out, target).backward() + weight_grads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return out.detach().clone(), inp.grad.detach().clone(), weight_grads + + out_h, dinp_h, wgrads_h = run(hybrid_recipe()) + out_v, dinp_v, wgrads_v = run(vanilla_recipe()) + + tag = f"linear_vs_vanilla[{parallel_mode},sp={sequence_parallel}]" + + # Same-topology fprop operands should match bitwise. + _check_bitwise(out_h, out_v, f"{tag} forward") + + # NVFP4 backward consumes columnwise stochastic-rounding RNG differently. + if not _backward_not_bitwise_comparable(): + _check_bitwise(dinp_h, dinp_v, f"{tag} dgrad") + assert len(wgrads_h) == len(wgrads_v), f"{tag}: weight-grad count mismatch" + for i, (gh, gv) in enumerate(zip(wgrads_h, wgrads_v)): + _check_bitwise(gh, gv, f"{tag} wgrad[{i}]") + + +def test_linear_vs_vanilla(): + # These recipes have no same-format vanilla bitwise target. + if QUANTIZATION in ( + "hybrid_mxfp8_nvfp4", + "hybrid_fp8_identity", + "hybrid_mxfp8_identity", + ): + dist_print("linear_vs_vanilla: skipped for hybrid without a vanilla equivalent") + return + for parallel_mode in ["column", "row"]: + for sequence_parallel in [False, True]: + _test_linear_vs_vanilla(parallel_mode, sequence_parallel) + + +def _same_format_parity_supported(): + return QUANTIZATION in ("hybrid_fp8", "hybrid_mxfp8", "identity") + + +def _check_same_topology_parity( + out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, tolerances=None +): + check = ( + _check_bitwise + if tolerances is None + else lambda a, e, label: _check_outputs(e, a, label=label, tolerances=tolerances) + ) + check(out_h, out_v, f"{tag} forward") + check(dinp_h, dinp_v, f"{tag} dgrad") + _check_gradients(model_h, model_v, bitwise=tolerances is None, tolerances=tolerances) + + +def _test_layernorm_linear_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): + if not _same_format_parity_supported(): + dist_print("layernorm_linear_vs_vanilla: skipped for recipe without vanilla equivalent") + return + dist_print(f"layernorm_linear_vs_vanilla: sequence_parallel={sequence_parallel}") + + def run(recipe_obj): + torch.manual_seed(23456) + torch.cuda.manual_seed(23456) + model = te.LayerNormLinear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode="column", + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + torch.manual_seed(45670) + torch.cuda.manual_seed(45670) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp.requires_grad_() + with te.autocast(enabled=recipe_obj is not None, recipe=recipe_obj): + out = model(inp) + torch.manual_seed(45671) + torch.cuda.manual_seed(45671) + LOSS_FN(out, torch.randn_like(out)).backward() + return model, out.detach().clone(), inp.grad.detach().clone() + + model_h, out_h, dinp_h = run(hybrid_recipe()) + model_v, out_v, dinp_v = run(vanilla_recipe()) + _check_same_topology_parity( + out_h, + dinp_h, + model_h, + out_v, + dinp_v, + model_v, + f"layernorm_linear_vs_vanilla[sp={sequence_parallel}]", + ) + + +def test_layernorm_linear_vs_vanilla(): + for sequence_parallel in [False, True]: + _test_layernorm_linear_vs_vanilla(sequence_parallel) + + +def _test_layernorm_mlp_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): + if not _same_format_parity_supported(): + dist_print("layernorm_mlp_vs_vanilla: skipped for recipe without vanilla equivalent") + return + dist_print(f"layernorm_mlp_vs_vanilla: sequence_parallel={sequence_parallel}") + + def run(recipe_obj): + torch.manual_seed(45678) + torch.cuda.manual_seed(45678) + model = te.LayerNormMLP( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + torch.manual_seed(56780) + torch.cuda.manual_seed(56780) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp.requires_grad_() + with te.autocast(enabled=recipe_obj is not None, recipe=recipe_obj): + out = model(inp) + torch.manual_seed(56781) + torch.cuda.manual_seed(56781) + LOSS_FN(out, torch.randn_like(out)).backward() + return model, out.detach().clone(), inp.grad.detach().clone() + + model_h, out_h, dinp_h = run(hybrid_recipe()) + model_v, out_v, dinp_v = run(vanilla_recipe()) + _check_same_topology_parity( + out_h, + dinp_h, + model_h, + out_v, + dinp_v, + model_v, + f"layernorm_mlp_vs_vanilla[sp={sequence_parallel}]", + # Hybrid CustomRecipe disables quantized-norm fusion while the native + # recipe uses it. Both paths consume identical quantized GEMM operands, + # but the BF16 norm boundary can differ by one rounding step. + tolerances=(None if QUANTIZATION == "identity" else {"rtol": 2**-7, "atol": 2**-10}), + ) + + +def test_layernorm_mlp_vs_vanilla(): + for sequence_parallel in [False, True]: + _test_layernorm_mlp_vs_vanilla(sequence_parallel) + + +# ── Test 2: te.LayerNormLinear column + SP ────────────────────────── + + +def _test_layernorm_linear(sequence_parallel, params_dtype=torch.bfloat16): + """Column-parallel LayerNormLinear with optional SP.""" + dist_print(f"layernorm_linear: parallel_mode=column sequence_parallel={sequence_parallel}") + + torch.manual_seed(23456) + torch.cuda.manual_seed(23456) + + model_single = te.LayerNormLinear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.LayerNormLinear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode="column", + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + if sequence_parallel: + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # Column-parallel output: gather along dim 1. + out_dist = _gather(out_dist, dim=1) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"layernorm_linear[sp={sequence_parallel}]") + + _check_input_gradient( + inp_dist, inp_single, label=f"layernorm_linear[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_layernorm_linear(): + for sequence_parallel in [False, True]: + _test_layernorm_linear(sequence_parallel) + + +# ── Test 3: te.LayerNormMLP + TP + SP ─────────────────────────────── + + +def _test_layernorm_mlp(sequence_parallel, params_dtype=torch.bfloat16): + """LayerNormMLP with set_parallel_mode=True and optional SP.""" + dist_print(f"layernorm_mlp: parallel_mode=set sequence_parallel={sequence_parallel}") + + torch.manual_seed(45678) + torch.cuda.manual_seed(45678) + + model_single = te.LayerNormMLP(HIDDEN_SIZE, FFN_HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.LayerNormMLP( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + if sequence_parallel: + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # Row-parallel FC2 output is in the full hidden space; with SP it is + # reduce-scattered along the token dim 0, so gather it back. + if sequence_parallel: + out_dist = _gather(out_dist, dim=0) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"layernorm_mlp[sp={sequence_parallel}]") + + _check_input_gradient( + inp_dist, inp_single, label=f"layernorm_mlp[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_layernorm_mlp(): + for sequence_parallel in [False, True]: + _test_layernorm_mlp(sequence_parallel) + + +# ── Test 4: te.TransformerLayer + TP + SP ─────────────────────────── + + +def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): + """TransformerLayer integration with TP and optional SP.""" + dist_print(f"transformer_layer: parallel_mode=set sequence_parallel={sequence_parallel}") + + torch.manual_seed(34567) + torch.cuda.manual_seed(34567) + + model_single = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NR_HEADS, + attention_dropout=0.0, + hidden_dropout=0.0, + fuse_qkv_params=True, + params_dtype=params_dtype, + ).cuda() + model_dist = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NR_HEADS, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + seq_length=WORLD_SIZE * SEQ_LEN if sequence_parallel else None, + attention_dropout=0.0, + hidden_dropout=0.0, + fuse_qkv_params=True, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + inp_single = ( + torch.randn((WORLD_SIZE * SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + ) + if sequence_parallel: + inp_dist = inp_single[WORLD_RANK * SEQ_LEN : (WORLD_RANK + 1) * SEQ_LEN, :, :].contiguous() + else: + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + if sequence_parallel: + out_dist = _gather(out_dist, dim=0) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"transformer_layer[sp={sequence_parallel}]") + + _check_input_gradient( + inp_dist, inp_single, label=f"transformer_layer[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_transformer_layer(): + for sequence_parallel in [False, True]: + _test_transformer_layer(sequence_parallel) + + +# ── Driver ─────────────────────────────────────────────────────────── + + +def main(argv=None): + global WORLD_RANK, WORLD_SIZE, NCCL_WORLD, QUANTIZATION + + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + + assert WORLD_SIZE == LOCAL_SIZE, "This test is single-node only" + assert LOCAL_SIZE <= torch.cuda.device_count() + + torch.cuda.set_device(LOCAL_RANK) + dist.init_process_group( + backend="nccl", + rank=WORLD_RANK, + world_size=WORLD_SIZE, + timeout=datetime.timedelta(seconds=60), + init_method="env://", + device_id=torch.device(f"cuda:{LOCAL_RANK}"), + ) + NCCL_WORLD = dist.new_group(backend="nccl") + + parser = argparse.ArgumentParser() + parser.add_argument( + "--quantization", + type=str, + required=True, + choices=[ + "hybrid_fp8", + "hybrid_mxfp8", + "hybrid_fp8_identity", + "hybrid_mxfp8_identity", + "identity", + "hybrid_nvfp4", + "hybrid_mxfp8_nvfp4", + ], + ) + parser.add_argument( + "--test", + type=str, + nargs="+", + default=["all"], + choices=[ + "all", + "linear", + "linear_vs_vanilla", + "layernorm_linear_vs_vanilla", + "layernorm_mlp_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", + ], + help="Run one or more named tests in the same distributed process group", + ) + args = parser.parse_args(argv) + QUANTIZATION = args.quantization + + test_map = { + "linear": test_linear, + "linear_vs_vanilla": test_linear_vs_vanilla, + "layernorm_linear_vs_vanilla": test_layernorm_linear_vs_vanilla, + "layernorm_mlp_vs_vanilla": test_layernorm_mlp_vs_vanilla, + "layernorm_linear": test_layernorm_linear, + "layernorm_mlp": test_layernorm_mlp, + "transformer_layer": test_transformer_layer, + } + if "all" in args.test: + if len(args.test) != 1: + parser.error("--test all cannot be combined with named tests") + tests_to_run = list(test_map.values()) + else: + tests_to_run = [test_map[name] for name in args.test] + + for test_fn in tests_to_run: + dist_print(f"=== Starting {test_fn.__name__} ===") + test_fn() + dist.barrier() + dist_print(f"=== Passed {test_fn.__name__} ===") + + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index 15ae2dae63..6af992dac2 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -58,7 +58,7 @@ def get_nvfp4_quantizer_factory(): Mirrors the canonical "branch on what we care about, default fall-through" pattern from - ``transformer_engine.pytorch.custom_recipes.quantization_recipes_base``; + ``transformer_engine.pytorch.custom_recipes.quantization_factory_base``; every slot gets a real :class:`NVFP4QuantizerRef` (``CustomRecipeState`` rejects ``None`` returns). diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py new file mode 100644 index 0000000000..6030e9d32f --- /dev/null +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -0,0 +1,146 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest launcher for hybrid TP/SP distributed tests.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch +import transformer_engine.pytorch as te +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +if torch.cuda.device_count() < 2: + pytest.skip( + "Distributed TP/SP tests need at least 2 GPUs.", + allow_module_level=True, + ) + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +TEST_ROOT = Path(__file__).parent.resolve() +NUM_PROCS = min(2, torch.cuda.device_count()) +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] + +xfail_hopper_columnwise_per_tensor_fp8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + + +def _run_tests(quantization: str, tests: tuple[str, ...] = ("all",)): + """Run related cases under one torchrun/process-group startup.""" + script = TEST_ROOT / "run_hybrid_tp_sp.py" + cmd = LAUNCH_CMD + [ + str(script), + "--quantization", + quantization, + "--test", + *tests, + ] + result = subprocess.run(cmd, env=os.environ, check=False) + assert result.returncode == 0, ( + f"run_hybrid_tp_sp.py (quantization={quantization}, tests={list(tests)})" + f" exited with code {result.returncode}" + ) + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid FP8 current scaling +# ────────────────────────────────────────────────────────────────────── +# Exercises TP amax reduction and SP gather paths. + + +_SAME_FORMAT_TESTS = ( + "linear", + "linear_vs_vanilla", + "layernorm_linear_vs_vanilla", + "layernorm_mlp_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", +) + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 +def test_hybrid_fp8(): + """Hybrid FP8 TP/SP coverage and same-topology vanilla parity.""" + _run_tests("hybrid_fp8", _SAME_FORMAT_TESTS) + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_identity_linear(): + """Linear TP/SP coverage for FP8 forward plus Identity backward.""" + _run_tests("hybrid_fp8_identity", ("linear",)) + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_identity_all_modules(): + """All-Identity TP/SP end-to-end coverage for every supported TE module.""" + _run_tests("identity") + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid MXFP8 +# ────────────────────────────────────────────────────────────────────── +# Covers per-block scale layout through TP shards. + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8(): + _run_tests("hybrid_mxfp8", _SAME_FORMAT_TESTS) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_identity_linear(): + """Linear TP/SP coverage for MXFP8 forward plus Identity backward.""" + _run_tests("hybrid_mxfp8_identity", ("linear",)) + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid NVFP4 +# ────────────────────────────────────────────────────────────────────── +# Same-format NVFP4 coverage with base role-wise settings. + + +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4(): + _run_tests( + "hybrid_nvfp4", + ( + "linear", + "linear_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", + ), + ) + + +# ────────────────────────────────────────────────────────────────────── +# Cross-format hybrid: MXFP8 rowwise + NVFP4 columnwise +# ────────────────────────────────────────────────────────────────────── +# No single vanilla recipe exists for bitwise comparison. + +_cross_format_available = mxfp8_available and nvfp4_available +_reason_for_no_cross_format = reason_for_no_mxfp8 if not mxfp8_available else reason_for_no_nvfp4 + + +@pytest.mark.skipif( + not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" +) +def test_hybrid_mxfp8_nvfp4(): + _run_tests( + "hybrid_mxfp8_nvfp4", + ("linear", "layernorm_linear", "layernorm_mlp", "transformer_layer"), + ) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index f386659b6c..02b1cb037e 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -42,6 +42,8 @@ def test_fsdp2_model_tests(): "-v", "-s", "--tb=short", + "-k", + "not hybrid", ], valid_returncodes=(0, 5), env=os.environ, @@ -165,6 +167,57 @@ def test_fsdp2_fused_adam_dcp_resharding(recipe): assert result.returncode == 0, f"DCP resharding load phase failed: {result.returncode}" +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_hybrid_fused_adam_tests(): + """FSDP2 FusedAdam tests with hybrid quantized params (parametrized by hybrid recipe).""" + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) + run_distributed( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + "-k", + "hybrid", + ], + valid_returncodes=(0, 5), + env=os.environ, + timeout=600, + ) + + +@pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_hybrid_model_tests(): + """FSDP2 model tests with hybrid quantized params (parametrized by hybrid recipe).""" + test_path = _FSDP2_DIR / "run_fsdp2_model.py" + run_distributed( + [ + "torchrun", + f"--nproc_per_node={NUM_PROCS}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + "-k", + "hybrid", + ], + valid_returncodes=(0, 5), + env=os.environ, + timeout=600, + ) + + def test_dummy() -> None: """Dummy test diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py new file mode 100644 index 0000000000..dfbe8d4c34 --- /dev/null +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -0,0 +1,358 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared explicit quantizer factories for hybrid quantization tests. + +These are intentionally ordinary module-level functions rather than a public +factory abstraction. Keeping them module-level makes ``CustomRecipe`` objects +that reference them picklable in distributed checkpoint tests. +""" + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + current_scaling_quantizer_factory, + float8_block_scaling_quantizer_factory, + mxfp8_quantizer_factory, + nvfp4_quantizer_factory, +) + +_LINEAR_MODULE_TYPES = ("linear", "grouped_linear") +_FORWARD_TENSOR_TYPES = ("input", "weight", "output") +_GRAD_TENSOR_TYPES = ("grad_output", "grad_input") + + +def _is_linear_role(role): + return role is not None and role.module_type in _LINEAR_MODULE_TYPES + + +def _make_fp8_current(*, fp8_dtype=te.DType.kFloat8E4M3): + return te.Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _make_mxfp8(*, fp8_dtype=te.DType.kFloat8E4M3): + return te.MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def fp8_e4m3_factory(): + """Construct the default E4M3 current-scaling test quantizer.""" + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") + + +def fp8_e5m2_factory(): + """Construct the default E5M2 current-scaling test quantizer.""" + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + + +def mxfp8_e4m3_factory(): + """Construct the default E4M3 MXFP8 test quantizer.""" + return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + +def make_fp8_quantizer(*, rowwise=True, columnwise=True): + """Construct the standard current-scaling FP8 test quantizer.""" + return te.Float8CurrentScalingQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + device="cuda", + rowwise=rowwise, + columnwise=columnwise, + ) + + +def make_nvfp4_quantizer(*, rowwise=True, columnwise=True): + """Construct the standard NVFP4 test quantizer.""" + return te.NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def make_hybrid_quantizer_fp8_row_fp4_col(): + """Construct a hybrid quantizer with FP8 rowwise and NVFP4 columnwise.""" + return te.HybridQuantizer( + rowwise_quantizer=make_fp8_quantizer(), + columnwise_quantizer=make_nvfp4_quantizer(), + ) + + +def as_data_tensor_tuple(storage): + """Return a storage's raw buffers as a tuple without copying them.""" + if storage is None: + return () + tensors = storage.get_data_tensors() + return tensors if isinstance(tensors, tuple) else (tensors,) + + +def snapshot_storage_tensor_metadata(storage, *, clone=False): + """Capture every tensor-valued concrete-storage metadata field.""" + if storage is None: + return None + tensor_metadata = {} + for name, value in storage.get_metadata().items(): + if isinstance(value, torch.Tensor) or value is None: + snapshot_value = value.detach().clone() if clone and value is not None else value + if ( + value is not None + and getattr(storage, "_is_2D_scaled", False) + and name in ("rowwise_scale_inv", "columnwise_scale_inv") + ): + # Float8Block 2D scales pad one tile dimension. Kernels do not + # initialize or consume that padding, so canonicalize it. + snapshot_value = value.detach().clone() + m, n = storage._fsdp_logical_mn() + block_len = storage._FSDP_BLOCK_LEN + m_tiles = (m + block_len - 1) // block_len + n_tiles = (n + block_len - 1) // block_len + if name == "rowwise_scale_inv": + snapshot_value[:, n_tiles:] = 0 + else: + snapshot_value[:, m_tiles:] = 0 + tensor_metadata[name] = snapshot_value + return { + "storage_type": type(storage).__name__, + "tensor_metadata": tensor_metadata, + } + + +def assert_nested_state_exact(actual, expected, *, path="state"): + """Recursively compare nested state, using zero tolerance for tensors.""" + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor), f"{path}: expected Tensor, got {type(actual)}" + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, msg=path) + return + if isinstance(expected, dict): + assert isinstance(actual, dict), f"{path}: expected dict, got {type(actual)}" + assert actual.keys() == expected.keys(), f"{path}: dictionary keys differ" + for key in expected: + assert_nested_state_exact(actual[key], expected[key], path=f"{path}.{key}") + return + if isinstance(expected, (list, tuple)): + assert isinstance( + actual, type(expected) + ), f"{path}: expected {type(expected).__name__}, got {type(actual).__name__}" + assert len(actual) == len(expected), f"{path}: sequence lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + assert_nested_state_exact( + actual_item, + expected_item, + path=f"{path}[{index}]", + ) + return + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + + +def assert_storage_data_exact(actual, expected, *, context): + """Assert every tensor-valued data, scale, and amax metadata field exactly.""" + assert_nested_state_exact( + snapshot_storage_tensor_metadata(actual), + snapshot_storage_tensor_metadata(expected), + path=context, + ) + + +def assert_hybrid_tensor_exact(actual, expected, *, context): + """Compare both hybrid directions, including metadata and dequantization.""" + for direction in ("rowwise", "columnwise"): + actual_storage = getattr(actual, f"{direction}_sub_storage") + expected_storage = getattr(expected, f"{direction}_sub_storage") + assert_storage_data_exact( + actual_storage, + expected_storage, + context=f"{context} {direction}", + ) + if expected_storage is None: + continue + try: + expected_dequantized = expected_storage.dequantize() + except NotImplementedError: + continue + torch.testing.assert_close( + actual_storage.dequantize(), + expected_dequantized, + rtol=0.0, + atol=0.0, + msg=f"{context} {direction} dequantized value differs", + ) + + +def make_role_aware_quantizer(factory, role): + """Construct a quantizer with the standard Block-FP8 GEMM geometry.""" + quantizer = factory() + if isinstance(quantizer, te.Float8BlockQuantizer): + is_weight = ( + role is not None + and role.module_type in _LINEAR_MODULE_TYPES + and role.tensor_type == "weight" + ) + quantizer.block_scaling_dim = 2 if is_weight else 1 + return quantizer + + +def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): + """Build a CustomRecipe with hybrid forward and configurable grad quantizers.""" + if grad_factory is None: + grad_factory = col_factory + + def qfactory(role): + is_linear = _is_linear_role(role) + if is_linear and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=make_role_aware_quantizer(row_factory, role), + columnwise_quantizer=make_role_aware_quantizer(col_factory, role), + ) + if is_linear and role.tensor_type in _GRAD_TENSOR_TYPES: + return make_role_aware_quantizer(grad_factory, role) + return make_role_aware_quantizer(row_factory, role) + + return recipe.CustomRecipe(qfactory=qfactory) + + +def hybrid_fp8_current_qfactory(role): + """FP8 current scaling in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=current_scaling_quantizer_factory(role), + columnwise_quantizer=current_scaling_quantizer_factory(role), + ) + return current_scaling_quantizer_factory(role) + + +def hybrid_fp8_current_e5m2_grads_qfactory(role): + """FP8 current-scaling hybrid with E5M2 for both grad boundary roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_fp8_current(), + columnwise_quantizer=_make_fp8_current(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return _make_fp8_current(fp8_dtype=te.DType.kFloat8E5M2) + return _make_fp8_current() + + +def hybrid_mxfp8_qfactory(role): + """MXFP8 in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.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 in both hybrid directions for forward roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.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_block_fp8_e4m3_qfactory(role): + """Hybrid E4M3 Block-FP8 with role-aware 1D/2D block geometry.""" + is_linear = _is_linear_role(role) + is_weight = is_linear and role.tensor_type == "weight" + block_scaling_dim = 2 if is_weight else 1 + + def make_quantizer(): + return te.Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=block_scaling_dim, + ) + + if is_linear and role.tensor_type in _GRAD_TENSOR_TYPES: + return make_quantizer() + return te.HybridQuantizer( + rowwise_quantizer=make_quantizer(), + columnwise_quantizer=make_quantizer(), + ) + + +def hybrid_mixed_mxfp8_fp8_qfactory(role): + """MXFP8 rowwise plus FP8 current-scaling columnwise.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=current_scaling_quantizer_factory(role), + ) + return current_scaling_quantizer_factory(role) + + +def hybrid_fp8_current_identity_qfactory(role): + """FP8 current-scaling forward plus Identity backward.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=current_scaling_quantizer_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return te.IdentityQuantizer() + return current_scaling_quantizer_factory(role) + + +def hybrid_mxfp8_identity_qfactory(role): + """MXFP8 forward plus Identity backward.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return te.IdentityQuantizer() + return mxfp8_quantizer_factory(role) + + +def identity_qfactory(role): # pylint: disable=unused-argument + """High-precision passthrough for every quantizer slot.""" + return te.IdentityQuantizer() + + +def hybrid_nvfp4_qfactory(role): + """NVFP4 in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + return nvfp4_quantizer_factory(role) + + +def hybrid_tp_mxfp8_nvfp4_qfactory(role): + """TP/SP MXFP8 rowwise plus NVFP4 columnwise, including boundary roles.""" + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return nvfp4_quantizer_factory(role) + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + + +def hybrid_fp8_mxfp8_qfactory(role): + """FP8 current-scaling rowwise plus MXFP8 columnwise for CPU-offload tests.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_fp8_current(), + columnwise_quantizer=_make_mxfp8(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return _make_mxfp8(fp8_dtype=te.DType.kFloat8E5M2) + return _make_fp8_current() + + +def hybrid_mxfp8_nvfp4_qfactory(role): + """MXFP8 rowwise plus NVFP4 columnwise for CPU-offload tests.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_mxfp8(), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return nvfp4_quantizer_factory(role) + return _make_mxfp8() diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index e4e36da38a..3de06f3847 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -19,6 +19,10 @@ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe +from hybrid_quantization_utils import ( + hybrid_fp8_mxfp8_qfactory, + hybrid_mxfp8_nvfp4_qfactory, +) from utils import ModelConfig, recipe_id, skip_unsupported_backward_override # Check supported quantization schemes @@ -65,6 +69,10 @@ def nvfp4_4over6(): quantization_recipes.append(recipe.NVFP4BlockScaling()) quantization_recipes.append(nvfp4_4over6()) quantization_recipes.append(nvfp4_row_scaled()) +if fp8_available and mxfp8_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_fp8_mxfp8_qfactory)) +if mxfp8_available and nvfp4_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_mxfp8_nvfp4_qfactory)) model_config = { @@ -219,6 +227,17 @@ def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) nvfp4_use_4over6=use_4over6, ) return quantizer(tensor) + elif recipe.custom(): + # CustomRecipe: invoke the qfactory for the linear weight role + # as a representative quantizer (returns a HybridQuantizer for the + # hybrid factories registered at module scope). + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = recipe.qfactory(QuantizerRole(module_type="linear", tensor_type="weight")) + if quantizer is None: + # Fallback: factory did not supply a weight quantizer. + return tensor.requires_grad_() if requires_grad else tensor + return quantizer(tensor) @staticmethod def create_recipe_ctx(recipe: Optional[recipe.Recipe]): @@ -492,6 +511,16 @@ def test_sanity(self, layer_type, recipe, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Skip hybrid (CustomRecipe) on ops-based LayerNormMLP: the ops-based + # LayerNorm passes the quantizer directly to the fused C++ kernel which + # does not recognize HybridQuantizer (cf. design-doc TODO; the regular + # layernorm_mlp.py has an unfused fallback but the ops path does not + # yet). Unrelated to CPU offload. + # grouped_linear is NOT skipped here — it passes test_sanity with + # hybrid; only memory-accounting assertions trip it in test_memory / + # test_manual_synchronization. + if layer_type in ("layernorm_mlp_ops",) and recipe is not None and recipe.custom(): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") recipe_ctx = Utils.create_recipe_ctx(recipe) init_cuda_memory = Utils.get_cuda_memory_mb() @@ -540,6 +569,17 @@ def test_memory(self, layer_type, recipe, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Memory-accounting checks fail for grouped_linear with hybrid because + # `_hybrid_split_quantize` produces per-group HybridQuantizedTensorStorage + # whose individual sub-buffers don't all cross the 256K-element offload + # threshold — the net GPU memory drop after offload is smaller than the + # analytical estimate. Correctness (test_sanity, test_numerics) passes. + if ( + layer_type in ("layernorm_mlp_ops", "grouped_linear") + and recipe is not None + and recipe.custom() + ): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") offload_ctx, sync_function = get_cpu_offload_context( enabled=True, @@ -633,6 +673,13 @@ def test_manual_synchronization(self, recipe, layer_type, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Same memory-accounting caveat as test_memory (see comment there). + if ( + layer_type in ("layernorm_mlp_ops", "grouped_linear") + and recipe is not None + and recipe.custom() + ): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") offload_ctx, sync_function, manual_controller = get_cpu_offload_context( enabled=True, @@ -712,6 +759,8 @@ def test_numerics( and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + if layer_type in ("layernorm_mlp_ops",) and recipe is not None and recipe.custom(): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") recipe_ctx = Utils.create_recipe_ctx(recipe) @@ -781,44 +830,79 @@ def forward(self, x): ): param_offload.data.copy_(param_no_offload.data) - x = Utils.create_tensor(None) + x = Utils.create_tensor(None, requires_grad=True) if use_cuda_graphs: callable_offload = te.make_graphed_callables( callable_offload, (x,), enabled=recipe is not None, - recipe=(Utils.create_recipe_ctx(recipe) if recipe is not None else None), + recipe=recipe, ) # warm up (for example to compute sf for delayed scaling) for _ in range(4): + callable_offload.zero_grad(set_to_none=True) + callable_no_offload.zero_grad(set_to_none=True) + x.grad = None out = callable_offload(x) out.sum().backward() + x.grad = None out = callable_no_offload(x) out.sum().backward() callable_offload.zero_grad(set_to_none=True) + callable_no_offload.zero_grad(set_to_none=True) + x.grad = None + rng_state = torch.cuda.get_rng_state() out_offload = callable_offload(x) out_offload.sum().backward() - # save out and gradients - offload_outs = [out_offload] + # Clone before the no-offload pass. CUDA graphs may reuse static output + # buffers, and the comparison must cover computed values rather than + # aliases or unchanged parameters. + offload_out = out_offload.detach().clone() + assert x.grad is not None + offload_input_grad = x.grad.detach().clone() + offload_param_grads = [] for param in callable_offload.parameters(): - offload_outs.append(param.detach().clone()) + assert param.grad is not None + offload_param_grads.append(param.grad.detach().clone()) torch.cuda.reset_peak_memory_stats() + torch.cuda.set_rng_state(rng_state) + x.grad = None out_no_offload = callable_no_offload(x) out_no_offload.sum().backward() - # collect gradients - no_offload_outs = [out_no_offload] + no_offload_out = out_no_offload.detach().clone() + assert x.grad is not None + no_offload_input_grad = x.grad.detach().clone() + no_offload_param_grads = [] for param in callable_no_offload.parameters(): - no_offload_outs.append(param.detach().clone()) - - # check if tensors are the same - for i in range(len(offload_outs)): - assert torch.allclose(offload_outs[i], no_offload_outs[i]), f"Error in tensor {i}." + assert param.grad is not None + no_offload_param_grads.append(param.grad.detach().clone()) + + # CPU offload is byte-preserving transport. It must not alter forward + # results, input gradients, or any parameter gradient. + torch.testing.assert_close(offload_out, no_offload_out, rtol=0.0, atol=0.0) + torch.testing.assert_close( + offload_input_grad, + no_offload_input_grad, + rtol=0.0, + atol=0.0, + ) + assert len(offload_param_grads) == len(no_offload_param_grads) + for index, (offload_grad, no_offload_grad) in enumerate( + zip(offload_param_grads, no_offload_param_grads) + ): + torch.testing.assert_close( + offload_grad, + no_offload_grad, + rtol=0.0, + atol=0.0, + msg=f"Parameter gradient {index} differs with CPU offload", + ) torch.cuda.synchronize() diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index 153bceca7d..19271a5f42 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -12,6 +12,10 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe +from hybrid_quantization_utils import ( + hybrid_fp8_mxfp8_qfactory, + hybrid_mxfp8_nvfp4_qfactory, +) from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported from utils import ModelConfig, get_available_attention_backends @@ -19,10 +23,20 @@ # Check supported quantization schemes fp8_available = te.is_fp8_available() mxfp8_available = te.is_mxfp8_available() +nvfp4_available = te.is_nvfp4_available() + quantization_recipes: Optional[recipe.Recipe] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) +if fp8_available and mxfp8_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_fp8_mxfp8_qfactory)) +if mxfp8_available and nvfp4_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_mxfp8_nvfp4_qfactory)) + +hybrid_quantization_recipes = [ + item for item in quantization_recipes if item is not None and item.custom() +] model_config = { "small": ModelConfig(8, 512, 8, 64, num_layers=5, eps=0.1), @@ -100,6 +114,15 @@ def _estimate_cached_weight_size( if quantization_recipe is None: return 0 + # Hybrid (CustomRecipe) caches two sub-storages per weight with + # potentially different formats. Returning ``None`` signals the caller + # to skip the exact memory-accounting assertion — the ``memory_with_offload + # < memory_without_offload`` check still applies. Deriving an analytical + # estimate here is blocked on the FSDP2-style packing optimization still + # being a TODO in hybrid_quantization_design.md. + if quantization_recipe.custom(): + return None + # Count number of weight param elements param_elements = 0 for module in modules: @@ -184,6 +207,19 @@ def _measure_cached_memory( def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: str) -> None: """Check that CPU offloading runs and has expected memory usage.""" + # Skip hybrid (CustomRecipe) on module types whose integration with hybrid + # is not yet complete (preexisting, independent of CPU offload): + # - layernorm_mlp_ops: the ops-based LayerNorm passes the quantizer + # directly to the fused C++ kernel which does not recognize + # HybridQuantizer (cf. design doc; the regular layernorm_mlp.py has + # an unfused fallback but the ops path does not yet). + if ( + model_name in ("layernorm_mlp_ops",) + and quantization_recipe is not None + and quantization_recipe.custom() + ): + pytest.skip(f"Hybrid CustomRecipe + {model_name} integration is not yet complete") + # Construct model modules_list = [model_types[model_name]() for _ in range(NUM_LAYERS)] if model_name in ["multihead_attention", "transformer_layer"]: @@ -212,4 +248,98 @@ def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: s modules_list, quantization_recipe, ) - assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + # ``_estimate_cached_weight_size`` returns ``None`` for recipes whose + # analytical cached-weight size is not worked out (CustomRecipe / hybrid); + # in that case the memory-savings assertion above is the only check. + if memory_from_cached_weights is not None: + assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + + +@pytest.mark.parametrize( + "quantization_recipe", + hybrid_quantization_recipes, + ids=[item.qfactory.__name__ for item in hybrid_quantization_recipes], +) +def test_hybrid_cpu_offload_numerics_exact(quantization_recipe: recipe.Recipe) -> None: + """V1 offload must preserve hybrid forward and backward values bitwise.""" + + def run(modules, inp, grad_output, *, cpu_offload): + if cpu_offload: + offload_context, sync_function = te.get_cpu_offload_context( + enabled=True, + num_layers=len(modules), + model_layers=len(modules) + 1, + offload_activations=True, + offload_weights=False, + ) + else: + offload_context = contextlib.nullcontext() + sync_function = lambda tensor: tensor + + out = inp + for module in modules: + with te.autocast(enabled=True, recipe=quantization_recipe), offload_context: + out = module(out) + out = sync_function(out) + # Commit the final offload group, mirroring the memory test above. + with offload_context: + out = out.clone() + out = sync_function(out) + out.backward(grad_output) + + grads = {} + for module_index, module in enumerate(modules): + for name, param in module.named_parameters(): + assert param.grad is not None + grads[f"{module_index}.{name}"] = param.grad.detach().clone() + assert inp.grad is not None + return out.detach().clone(), inp.grad.detach().clone(), grads + + torch.manual_seed(9000) + reference_modules = [model_types["linear"]() for _ in range(2)] + torch.manual_seed(9001) + offload_modules = [model_types["linear"]() for _ in range(2)] + with torch.no_grad(): + for offload_module, reference_module in zip(offload_modules, reference_modules): + for offload_param, reference_param in zip( + offload_module.parameters(), reference_module.parameters() + ): + offload_param.copy_(reference_param) + + # V1 requires one pass to initialize cached quantized weight workspaces. + warmup_rng_state = torch.cuda.get_rng_state() + _warmup_model(offload_modules, quantization_recipe) + for module in offload_modules: + module.zero_grad(set_to_none=True) + torch.cuda.set_rng_state(warmup_rng_state) + _warmup_model(reference_modules, quantization_recipe) + for module in reference_modules: + module.zero_grad(set_to_none=True) + + torch.manual_seed(9002) + x = torch.randn((8, SIZE, SIZE), device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn_like(x) + x_offload = x.detach().clone().requires_grad_(True) + x_reference = x.detach().clone().requires_grad_(True) + + # Replay identical stochastic-rounding/RHT randomness in both paths. + rng_state = torch.cuda.get_rng_state() + offload_out, offload_input_grad, offload_param_grads = run( + offload_modules, x_offload, grad_output, cpu_offload=True + ) + torch.cuda.set_rng_state(rng_state) + reference_out, reference_input_grad, reference_param_grads = run( + reference_modules, x_reference, grad_output, cpu_offload=False + ) + + torch.testing.assert_close(offload_out, reference_out, rtol=0.0, atol=0.0) + torch.testing.assert_close(offload_input_grad, reference_input_grad, rtol=0.0, atol=0.0) + assert offload_param_grads.keys() == reference_param_grads.keys() + for name, reference_grad in reference_param_grads.items(): + torch.testing.assert_close( + offload_param_grads[name], + reference_grad, + rtol=0.0, + atol=0.0, + msg=f"V1 CPU-offload parameter gradient mismatch for {name}", + ) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 3e6fdb816b..e88ef93c95 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -11,20 +11,29 @@ from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx from transformer_engine.pytorch import ( autocast, + Fp8Padding, + Fp8Unpadding, Linear, LayerNormLinear, LayerNormMLP, GroupedLinear, Float8CurrentScalingQuantizer, ) -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.quantization import ( + QuantizerRole, + get_align_size_for_quantization, +) import transformer_engine.pytorch.ops as te_ops -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( current_scaling_quantizer_factory, mxfp8_quantizer_factory, float8_block_scaling_quantizer_factory, nvfp4_quantizer_factory, delayed_scaling_quantizer_factory, + high_precision_factory, +) +from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, ) from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import ( nvfp4_ref_rht_2d_quantizer_factory, @@ -477,6 +486,40 @@ def test_factory_matches_nvfp4(): _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) +def test_factory_matches_high_precision(): + """high_precision_factory (all IdentityQuantizer) should produce bit-identical + results to running with no quantization (plain BF16, no autocast). + + Unlike the other ``test_factory_matches_*`` cases, the reference here is not a + built-in recipe but the unquantized path: IdentityQuantizer leaves tensors + untouched, so the GEMMs run in high precision exactly as they would with + autocast disabled. + + Note: the custom side still runs inside ``autocast(enabled=True, ...)``, which + asserts FP8 availability on entry, so this test skips on non-FP8 hardware even + though no tensor is actually quantized. + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + # Reference: no autocast -> high-precision GEMMs, no quantization. + with autocast(enabled=False): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + out_ref = out_ref.clone() + grad_ref = inp_ref.grad.clone() + pgrads_ref = {n: p.grad.clone() for n, p in model_ref.named_parameters() if p.grad is not None} + + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=high_precision_factory) + ) + + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + def test_custom_recipe_quantization_targets(): """Validate fine-grained per-module quantization targeting via QuantizerRole. @@ -1063,7 +1106,7 @@ def test_role_change_does_not_invalidate_when_role_unchanged(): def test_custom_recipe_dpa_fp8(): - """DotProductAttention forward+backward with CustomRecipe and role-based mixed quantizers. + """DotProductAttention forward+backward with CustomRecipe and role-based DPA quantizers. Uses the nvfp4_linear_fp8_dpa_factory which dispatches: * DPA S/dP slots -> DelayedScalingRequest (stateful) @@ -1090,7 +1133,7 @@ def test_custom_recipe_dpa_fp8(): Float8Quantizer, Float8CurrentScalingQuantizer, ) - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -1217,7 +1260,7 @@ def test_custom_recipe_dpa_mxfp8(): from transformer_engine.pytorch.quantization import CustomRecipeState from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_mxfp8_dpa_factory, ) @@ -1849,3 +1892,94 @@ def test_slot_role_supports_module_type_only_role(): assert resolved.name == "" # Tensor-type-only recipes fall back to positional for this slot. assert state._slot_tensor_type(0) == "input" + + +_alignment_fp8_available, _alignment_fp8_reason = te.is_fp8_available(return_reason=True) +_alignment_mxfp8_available, _alignment_mxfp8_reason = te.is_mxfp8_available(return_reason=True) +_alignment_nvfp4_available, _alignment_nvfp4_reason = te.is_nvfp4_available(return_reason=True) + + +def test_custom_recipe_quantization_alignment_contract(): + """The alignment contract is safe, configurable, and does not invoke qfactory.""" + calls = [] + + def qfactory(role): + calls.append(role) + return current_scaling_quantizer_factory(role) + + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + assert get_align_size_for_quantization(custom_recipe) == 128 + assert calls == [] + + custom_recipe = recipe.CustomRecipe(qfactory=qfactory, quantization_alignment=32) + assert get_align_size_for_quantization(custom_recipe) == 32 + assert calls == [] + + with pytest.raises(ValueError, match="quantization_alignment must be positive"): + recipe.CustomRecipe(qfactory=qfactory, quantization_alignment=0) + + +@pytest.mark.parametrize( + "qfactory", + [ + pytest.param( + current_scaling_quantizer_factory, + marks=pytest.mark.skipif( + not _alignment_fp8_available, + reason=_alignment_fp8_reason, + ), + id="fp8_current_scaling", + ), + pytest.param( + mxfp8_quantizer_factory, + marks=pytest.mark.skipif( + not _alignment_mxfp8_available, + reason=_alignment_mxfp8_reason, + ), + id="mxfp8", + ), + pytest.param( + nvfp4_quantizer_factory, + marks=pytest.mark.skipif( + not _alignment_nvfp4_available, + reason=_alignment_nvfp4_reason, + ), + id="nvfp4", + ), + pytest.param( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, + marks=pytest.mark.skipif( + not (_alignment_mxfp8_available and _alignment_nvfp4_available), + reason=f"MXFP8: {_alignment_mxfp8_reason}; NVFP4: {_alignment_nvfp4_reason}", + ), + id="hybrid_mxfp8_nvfp4", + ), + ], +) +def test_custom_recipe_automatic_padding_bad_splits(qfactory): + """Automatic alignment makes misaligned CustomRecipe grouped GEMMs valid.""" + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + padding = Fp8Padding(2) + unpadding = Fp8Unpadding(2) + grouped_linear = GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ) + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + original_splits = [17, 15] + + with autocast(enabled=True, recipe=custom_recipe): + padded_inp, padded_splits = padding(inp, original_splits) + assert padding.align_size == get_align_size_for_quantization(custom_recipe) == 128 + assert padded_splits == [128, 128] + padded_out = grouped_linear(padded_inp, padded_splits) + out = unpadding(padded_out, original_splits) + + out.sum().backward() + assert out.shape == (sum(original_splits), 128) + assert torch.isfinite(out).all() + assert torch.isfinite(inp.grad).all() diff --git a/tests/pytorch/test_float8blockwisetensor.py b/tests/pytorch/test_float8blockwisetensor.py index 1b5e17f463..ab00786e4f 100644 --- a/tests/pytorch/test_float8blockwisetensor.py +++ b/tests/pytorch/test_float8blockwisetensor.py @@ -96,6 +96,50 @@ def test_constructor( assert tensor.dtype == dtype, "Incorrect nominal dtype" assert tensor.is_cuda, "Incorrect device" + def test_fsdp_assign_gathered_rejects_invalid_scale_geometry(self): + """Reject gathered scales that cannot describe the gathered data tiles.""" + shape = (128, 256) + quantizer = Float8BlockQuantizer( + fp8_dtype=DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ) + tensor = Float8BlockwiseQTensor( + shape=shape, + dtype=torch.bfloat16, + rowwise_data=torch.zeros(shape, device="cuda", dtype=torch.uint8), + rowwise_scale_inv=torch.ones( + quantizer.get_scale_shape(shape, columnwise=False), + device="cuda", + dtype=torch.float32, + ), + columnwise_data=torch.zeros( + quantizer.get_columnwise_shape(shape), + device="cuda", + dtype=torch.uint8, + ), + columnwise_scale_inv=torch.ones( + quantizer.get_scale_shape(shape, columnwise=True), + device="cuda", + dtype=torch.float32, + ), + fp8_dtype=DType.kFloat8E4M3, + is_2D_scaled=True, + quantizer=quantizer, + ) + original_data = tensor._rowwise_data + gathered_data = torch.zeros((256, 256), device="cuda", dtype=torch.uint8) + invalid_scale = torch.ones((3, 2), device="cuda", dtype=torch.float32) + + with pytest.raises(RuntimeError, match="gathered scale geometry"): + tensor.fsdp_assign_gathered( + (gathered_data, invalid_scale), + {"direction": "rowwise"}, + ) + + assert tensor._rowwise_data is original_data + def _test_quantize_dequantize( self, quantizer: Float8BlockQuantizer, diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py new file mode 100644 index 0000000000..84036ae84f --- /dev/null +++ b/tests/pytorch/test_hybrid_quantization.py @@ -0,0 +1,7333 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for hybrid quantization (mixed rowwise/columnwise formats).""" + +import io +import warnings +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex + +from hybrid_quantization_utils import ( + as_data_tensor_tuple as _as_data_tensor_tuple, + assert_hybrid_tensor_exact as _assert_hybrid_tensor_exact, + assert_nested_state_exact as _assert_nested_state_exact, + assert_storage_data_exact as _assert_storage_data_exact, + fp8_e4m3_factory as _fp8_row_factory, + fp8_e5m2_factory as _fp8_grad_factory, + hybrid_block_fp8_e4m3_qfactory as _hybrid_block_fp8_qfactory, + hybrid_custom_recipe as _hybrid_custom_recipe, + hybrid_fp8_current_e5m2_grads_qfactory as _hybrid_fp8_current_qfactory, + make_fp8_quantizer as _make_fp8_quantizer, + make_hybrid_quantizer_fp8_row_fp4_col as _make_hybrid_quantizer_fp8_row_fp4_col, + make_nvfp4_quantizer as _make_nvfp4_quantizer, + make_role_aware_quantizer as _make_role_aware_quantizer, + mxfp8_e4m3_factory as _mxfp8_factory, + snapshot_storage_tensor_metadata as _snapshot_storage_tensor_metadata, +) +from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + nvfp4_quantizer_factory, +) +from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, +) +from transformer_engine.pytorch import ( + autocast, + quantized_model_init, + Linear, + LayerNormLinear, + LayerNormMLP, + TransformerLayer, + GroupedLinear, + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, + HybridQuantizer, + HybridQuantizedTensor, + IdentityQuantizer, + HybridQuantizedTensorStorage, + Float8Tensor, + Float8TensorStorage, + NVFP4Tensor, + NVFP4TensorStorage, + QuantizedTensor, +) +from transformer_engine.pytorch.cpp_extensions.gemm import ( + _unwrap_tensor, + _validate_native_gemm_output_quantizer, +) +from transformer_engine.pytorch.quantized_tensor import Quantizer +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +_fp8_col_factory = _fp8_row_factory + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR = ( + "Columnwise-only per-tensor FP8 quantization is not implemented" +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + +requires_fp8 = pytest.mark.skipif( + not fp8_available, + reason=f"FP8: {reason_for_no_fp8}", +) + +requires_nvfp4 = pytest.mark.skipif( + not nvfp4_available, + reason=f"NVFP4: {reason_for_no_nvfp4}", +) + +requires_fp8_and_nvfp4 = pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) + +requires_mxfp8_and_nvfp4 = pytest.mark.skipif( + not (mxfp8_available and nvfp4_available), + reason=f"MXFP8: {reason_for_no_mxfp8}; NVFP4: {reason_for_no_nvfp4}", +) + + +def test_native_gemm_output_quantizer_support_is_opt_in(): + """Unknown quantizers must be rejected before entering the native C++ path.""" + _validate_native_gemm_output_quantizer(None) + + unknown_quantizer = Quantizer(rowwise=True, columnwise=False) + with pytest.raises( + NotImplementedError, + match="Quantizer is not supported as a native GEMM output quantizer", + ): + _validate_native_gemm_output_quantizer(unknown_quantizer) + + +def test_hybrid_storage_snapshots_parent_quantizer(): + """Caller mutations must not change an existing Hybrid tensor's behavior.""" + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + tensor = quantizer(torch.ones((2, 2))) + + assert tensor._quantizer is not quantizer + assert tensor._quantizer.rowwise_quantizer is not quantizer.rowwise_quantizer + assert tensor._quantizer.columnwise_quantizer is not quantizer.columnwise_quantizer + + # Make the two Identity directions independent so a skipped update is observable. + tensor._rowwise_storage._hp_data = tensor._rowwise_storage._hp_data.clone() + tensor._columnwise_storage._hp_data = tensor._columnwise_storage._hp_data.clone() + + quantizer.set_usage(rowwise=False, columnwise=True) + tensor.copy_(torch.full(tensor.shape, 2.0)) + + assert tensor._quantizer.get_usages() == {"rowwise": True, "columnwise": True} + torch.testing.assert_close( + tensor._rowwise_storage.dequantize(), + torch.full(tensor.shape, 2.0), + ) + torch.testing.assert_close( + tensor._columnwise_storage.dequantize(), + torch.full(tensor.shape, 2.0), + ) + + +def _make_hybrid_quantizer_fp4_row_fp8_col(): + """NVFP4 rowwise + FP8 columnwise (reversed direction).""" + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + + +def _clone_nested_state(value): + """Clone tensors in nested checkpoint state so later steps cannot mutate snapshots.""" + if isinstance(value, torch.Tensor): + return value.detach().clone() + if isinstance(value, dict): + return {key: _clone_nested_state(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_nested_state(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_nested_state(item) for item in value) + return value + + +def _snapshot_model_parameters(model): + """Capture normal values and all tensor metadata in both hybrid directions.""" + snapshot = {} + for name, param in model.named_parameters(): + if isinstance(param, HybridQuantizedTensor): + snapshot[name] = { + "rowwise": _snapshot_storage_tensor_metadata(param.rowwise_sub_storage, clone=True), + "columnwise": _snapshot_storage_tensor_metadata( + param.columnwise_sub_storage, clone=True + ), + } + else: + snapshot[name] = param.detach().clone() + return snapshot + + +class _CountingIdentityQuantizer(IdentityQuantizer): + """Replay-safe identity quantizer that counts quantize calls.""" + + def __init__(self, counter, *, dtype=None, rowwise=True, columnwise=True): + super().__init__(dtype=dtype, rowwise=rowwise, columnwise=columnwise) + self.counter = counter + + def copy(self): + quantizer = type(self)( + self.counter, + dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def quantize_impl(self, tensor): + self.counter["calls"] += 1 + return super().quantize_impl(tensor) + + +class _CountingUnsafeIdentityQuantizer(_CountingIdentityQuantizer): + """Non-replay-safe identity quantizer that counts quantize calls.""" + + def is_requantization_safe(self): + self.counter["safety_calls"] = self.counter.get("safety_calls", 0) + 1 + return False + + +class _CountingPythonQuantizer(Quantizer): + """Custom Python quantizer used to verify fallback dispatch.""" + + def __init__(self, calls, *, rowwise=True, columnwise=True): + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.calls = calls + + def copy(self): + quantizer = type(self)( + self.calls, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def quantize_impl(self, tensor): + self.calls.append(tensor) + fallback = IdentityQuantizer( + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + fallback.internal = self.internal + return fallback.quantize_impl(tensor) + + +@requires_mxfp8_and_nvfp4 +class TestComposerStyleFactory: + """Composer 2-style row-scaled NVFP4 forward + MXFP8 backward dispatch.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory, + ) + + return nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory(role) + + @pytest.mark.parametrize( + "tensor_type,row_scaled", + [("input", True), ("weight", False)], + ) + def test_grouped_linear_forward_roles_use_nvfp4_rowwise_mxfp8_columnwise( + self, tensor_type, row_scaled + ): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory( + QuantizerRole(module_type="grouped_linear", tensor_type=tensor_type) + ) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) + assert isinstance(quantizer.columnwise_quantizer, MXFP8Quantizer) + assert quantizer.rowwise_quantizer.row_scaled_nvfp4 is row_scaled + assert quantizer.rowwise_quantizer.with_rht is False + assert quantizer.rowwise_quantizer.with_post_rht_amax is False + assert quantizer.rowwise_quantizer.with_2d_quantization is False + assert quantizer.rowwise_quantizer.stochastic_rounding is False + assert quantizer.rowwise_quantizer.rowwise_usage is True + assert quantizer.rowwise_quantizer.columnwise_usage is False + assert quantizer.columnwise_quantizer.rowwise_usage is False + assert quantizer.columnwise_quantizer.columnwise_usage is True + + @pytest.mark.parametrize("tensor_type", ["input", "output", "weight"]) + def test_regular_linear_forward_roles_fall_back_to_mxfp8(self, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type="linear", tensor_type=tensor_type)) + + assert isinstance(quantizer, MXFP8Quantizer) + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["grad_output", "grad_input"]) + def test_backward_roles_use_mxfp8(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) + + assert isinstance(quantizer, MXFP8Quantizer) + + def test_non_linear_roles_fall_back_to_mxfp8(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type="dpa", tensor_type="qkv")) + + assert isinstance(quantizer, MXFP8Quantizer) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestMXFP8FwdDequantizedBwdFactory: + """MXFP8 forward + dequantized backward qfactory dispatch.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + ) + + return mxfp8_fwd_dequantized_bwd_quantizer_factory(role) + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["input", "weight"]) + def test_forward_roles_are_requantization_safe(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) + assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) + assert quantizer.is_requantization_safe() is True + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + def test_grad_output_role_is_requantization_safe(self, module_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="grad_output")) + + assert isinstance(quantizer, IdentityQuantizer) + assert quantizer.is_requantization_safe() is True + + +@requires_nvfp4 +class TestNVFP4WeightDoubleQuantFactory: + """Base NVFP4 recipe with W.T sourced from dequantized 1D W.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_1d_double_quantized_weight_quantizer_factory, + ) + + return nvfp4_1d_double_quantized_weight_quantizer_factory(role) + + @staticmethod + def _assert_plain_1d_nvfp4(quantizer): + assert isinstance(quantizer, NVFP4Quantizer) + assert quantizer.row_scaled_nvfp4 is False + assert quantizer.with_rht is False + assert quantizer.with_post_rht_amax is False + assert quantizer.with_2d_quantization is False + assert quantizer.stochastic_rounding is False + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + def test_weight_roles_use_rowwise_dequantized_source(self, module_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="weight")) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + self._assert_plain_1d_nvfp4(quantizer.rowwise_quantizer) + self._assert_plain_1d_nvfp4(quantizer.columnwise_quantizer) + assert quantizer.rowwise_quantizer.rowwise_usage is True + assert quantizer.rowwise_quantizer.columnwise_usage is False + assert quantizer.columnwise_quantizer.rowwise_usage is False + assert quantizer.columnwise_quantizer.columnwise_usage is True + + @staticmethod + def _assert_matches_base_nvfp4_recipe(quantizer, expected): + assert isinstance(quantizer, NVFP4Quantizer) + assert isinstance(expected, NVFP4Quantizer) + assert quantizer.dtype == expected.dtype + assert quantizer.rowwise_usage == expected.rowwise_usage + assert quantizer.columnwise_usage == expected.columnwise_usage + assert quantizer.with_rht == expected.with_rht + assert quantizer.with_post_rht_amax == expected.with_post_rht_amax + assert quantizer.with_2d_quantization == expected.with_2d_quantization + assert quantizer.stochastic_rounding == expected.stochastic_rounding + assert quantizer.row_scaled_nvfp4 == expected.row_scaled_nvfp4 + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["input", "output", "grad_output", "grad_input"]) + def test_non_weight_linear_roles_match_base_nvfp4_recipe(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + nvfp4_quantizer_factory, + ) + + role = QuantizerRole(module_type=module_type, tensor_type=tensor_type) + quantizer = self._factory(role) + expected = nvfp4_quantizer_factory(role) + + self._assert_matches_base_nvfp4_recipe(quantizer, expected) + + def test_non_linear_roles_match_base_nvfp4_recipe(self): + from transformer_engine.pytorch.quantization import QuantizerRole + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + nvfp4_quantizer_factory, + ) + + role = QuantizerRole(module_type="dpa", tensor_type="qkv") + quantizer = self._factory(role) + expected = nvfp4_quantizer_factory(role) + + self._assert_matches_base_nvfp4_recipe(quantizer, expected) + + def test_weight_quantizer_produces_both_nvfp4_storages(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + torch.manual_seed(2026) + src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + quantizer = self._factory(QuantizerRole(module_type="linear", tensor_type="weight")) + + out = quantizer.quantize(src) + + assert isinstance(out.rowwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + assert isinstance(out.columnwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + + +@requires_fp8_and_nvfp4 +class TestHybridQuantizerConstruction: + """Test construction and basic properties of hybrid quantizer.""" + + def test_creation(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + assert isinstance(hq, HybridQuantizer) + assert hq.rowwise_usage is True + assert hq.columnwise_usage is True + assert isinstance(hq.rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(hq.columnwise_quantizer, NVFP4Quantizer) + + def test_hybrid_storage_and_tensor_require_parent_quantizer(self): + with pytest.raises(TypeError, match="requires a parent HybridQuantizer"): + HybridQuantizedTensorStorage( + rowwise_storage=None, + columnwise_storage=None, + quantizer=None, + ) + with pytest.raises(TypeError, match="requires a parent HybridQuantizer"): + HybridQuantizedTensor( + shape=(1, 1), + dtype=torch.bfloat16, + rowwise_storage=None, + columnwise_storage=None, + quantizer=None, + ) + + def test_rejects_same_sub_quantizer_instance_for_both_directions(self): + quantizer = _make_fp8_quantizer() + + with pytest.raises(ValueError, match="requires distinct rowwise and columnwise"): + HybridQuantizer(rowwise_quantizer=quantizer, columnwise_quantizer=quantizer) + + assert quantizer.rowwise_usage is True + assert quantizer.columnwise_usage is True + + def test_compatible_recipe_is_custom_recipe(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + assert hq._get_compatible_recipe() is recipe.CustomRecipe + + def test_supports_only_rowwise_all_gather_nvfp4_columnwise(self): + """NVFP4 columnwise sub-quantizer forces rowwise-only AG. + + ``NVFP4Tensor.dequantize()`` raises ``NotImplementedError`` for + columnwise-only data, so the BF16 fallback in + ``gather_along_first_dim`` cannot operate on a columnwise-only + NVFP4 hybrid sub-storage. ``HybridQuantizer.supports_only_rowwise_all_gather`` + must return True in this case so ``_linear_forward_impl`` / + ``_linear_backward`` preserve rowwise data (which NVFP4 can + dequantize) instead. + """ + hq = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + def test_supports_only_rowwise_all_gather_mxfp8_both(self): + """MXFP8 in both directions → columnwise dequant works → default + False so the save-columnwise (for wgrad) path stays active.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + hq = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + assert hq.supports_only_rowwise_all_gather() is False + + def test_supports_only_rowwise_all_gather_fp8_current_propagates(self): + """Float8CurrentScalingQuantizer returns True for its own flag; + hybrid must propagate (not swallow) that semantics.""" + hq = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + def test_supports_only_rowwise_all_gather_nvfp4_both(self): + """NVFP4 in both directions → columnwise sub-quantizer is NVFP4 + → forces rowwise-only AG regardless of rowwise flag.""" + hq = HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + +@requires_fp8 +class TestHybridColumnwiseSource: + """Test columnwise source provenance in HybridQuantizer.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(90210) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + @staticmethod + def _make_quantizer(columnwise_source="original"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source=columnwise_source, + ) + + def test_default_columnwise_source_original(self, input_tensor): + hq = self._make_quantizer() + out = hq.quantize(input_tensor) + + assert hq.columnwise_source == "original" + torch.testing.assert_close( + out.columnwise_sub_storage.dequantize(), input_tensor, rtol=0.0, atol=0.0 + ) + + def test_invalid_columnwise_source_raises(self): + with pytest.raises(ValueError, match="columnwise_source"): + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="dequantized", + ) + + def test_default_quantizer_is_requantization_safe(self): + assert IdentityQuantizer().is_requantization_safe() is True + + @pytest.mark.parametrize("columnwise_source", ("original", "rowwise_dequantized")) + def test_hybrid_requantization_safe_with_safe_sub_quantizers(self, columnwise_source): + hq = self._make_quantizer(columnwise_source=columnwise_source) + + assert hq.is_requantization_safe() is True + + def test_hybrid_requantization_checks_requested_sub_quantizers(self): + from transformer_engine.pytorch.module._common import ( + can_reconstruct_wgrad_input_from_original, + ) + + hq = HybridQuantizer( + rowwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_source="original", + ) + + assert hq.is_requantization_safe() is False + assert can_reconstruct_wgrad_input_from_original(hq) is True + + @pytest.mark.parametrize("columnwise_source", ("original", "rowwise_dequantized")) + @pytest.mark.parametrize("rowwise_safe", (True, False)) + @pytest.mark.parametrize("columnwise_safe", (True, False)) + def test_requantization_and_wgrad_reconstruction_policy_matrix( + self, + columnwise_source, + rowwise_safe, + columnwise_safe, + ): + from transformer_engine.pytorch.module._common import ( + can_reconstruct_wgrad_input_from_original, + ) + + rowwise_cls = ( + _CountingIdentityQuantizer if rowwise_safe else _CountingUnsafeIdentityQuantizer + ) + columnwise_cls = ( + _CountingIdentityQuantizer if columnwise_safe else _CountingUnsafeIdentityQuantizer + ) + hq = HybridQuantizer( + rowwise_quantizer=rowwise_cls({"calls": 0}), + columnwise_quantizer=columnwise_cls({"calls": 0}), + columnwise_source=columnwise_source, + ) + + # General requantization reproduces both requested representations. + assert hq.is_requantization_safe() is (rowwise_safe and columnwise_safe) + # Wgrad reconstruction only repeats the rowwise stage for a + # rowwise-dequantized columnwise source. + expected_reconstruction = columnwise_source == "original" or rowwise_safe + assert can_reconstruct_wgrad_input_from_original(hq) is expected_reconstruction + + def test_hybrid_rowwise_source_requires_safe_rowwise_quantizer(self): + hq = HybridQuantizer( + rowwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + assert hq.is_requantization_safe() is False + + def test_copy_preserves_columnwise_source(self): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.set_usage(rowwise=False, columnwise=True) + + copied = hq.copy() + + assert copied.columnwise_source == "rowwise_dequantized" + assert copied.rowwise_usage is False + assert copied.columnwise_usage is True + + def test_rowwise_dequantized_identity_columnwise_matches_rowwise(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + out = hq.quantize(input_tensor) + + rowwise_dq = out.rowwise_sub_storage.dequantize() + columnwise_dq = out.columnwise_sub_storage.dequantize() + torch.testing.assert_close(columnwise_dq, rowwise_dq, rtol=0.0, atol=0.0) + + def test_internal_rowwise_storage_preserves_input_dtype(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.rowwise_quantizer.internal = True + + columnwise_src = hq._columnwise_src_from_rowwise(input_tensor, None) + + assert columnwise_src.dtype == input_tensor.dtype + + def test_columnwise_only_uses_transient_rowwise_source(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.set_usage(rowwise=False, columnwise=True) + expected = _make_fp8_quantizer().quantize(input_tensor).dequantize() + + out = hq.quantize(input_tensor) + + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is not None + torch.testing.assert_close( + out.columnwise_sub_storage.dequantize(), expected, rtol=0.0, atol=0.0 + ) + + def test_update_quantized_columnwise_only_uses_transient_rowwise_source(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + dst = hq.quantize(input_tensor) + rowwise_before = dst.rowwise_sub_storage.dequantize().clone() + new_src = torch.randn_like(input_tensor) * 8 + expected = _make_fp8_quantizer().quantize(new_src).dequantize() + + hq.set_usage(rowwise=False, columnwise=True) + hq.update_quantized(new_src, dst) + + torch.testing.assert_close( + dst.rowwise_sub_storage.dequantize(), rowwise_before, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + dst.columnwise_sub_storage.dequantize(), expected, rtol=0.0, atol=0.0 + ) + + def test_update_quantized_uses_updated_rowwise_storage(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + dst = hq.quantize(input_tensor) + new_src = torch.randn_like(input_tensor) * 8 + + hq.set_usage(rowwise=True, columnwise=True) + hq.update_quantized(new_src, dst) + + torch.testing.assert_close( + dst.columnwise_sub_storage.dequantize(), + dst.rowwise_sub_storage.dequantize(), + rtol=0.0, + atol=0.0, + ) + + +@requires_fp8 +class TestHybridSaveOriginalInputPolicy: + """Module-level save_original_input policy for hybrid qfactory inputs.""" + + def test_linear_high_precision_override_skips_requantization_safety_veto(self): + counters = {"rowwise": {"calls": 0}, "columnwise": {"calls": 0}} + custom_recipe = self._counting_identity_hybrid_recipe(counters) + custom_recipe.backward_override = "high_precision" + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + out.float().sum().backward() + + assert counters["rowwise"].get("safety_calls", 0) == 0 + + def test_linear_custom_delayed_scaling_disables_save_original_input(self): + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + delayed_scaling_quantizer_factory, + ) + + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + custom_recipe = recipe.CustomRecipe(qfactory=delayed_scaling_quantizer_factory) + + with pytest.warns( + UserWarning, + match="save_original_input is incompatible with delayed-scaling quantizers", + ): + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + out.float().sum().backward() + + def test_linear_builtin_delayed_scaling_rejects_save_original_input(self): + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with pytest.raises( + ValueError, + match="DelayedScaling recipe is not supported with save_original_input", + ): + with autocast(enabled=True, recipe=recipe.DelayedScaling()): + model(inp) + + def test_grouped_linear_classifies_requantization_safety_once_per_generation(self): + counters = [{"calls": 0}, {"calls": 0}] + input_quantizers = [_CountingUnsafeIdentityQuantizer(counter) for counter in counters] + generation = [] + for input_quantizer in input_quantizers: + generation.extend((input_quantizer, IdentityQuantizer(), IdentityQuantizer())) + + module = GroupedLinear( + 2, + 16, + 16, + bias=False, + device="meta", + ) + module.quantizers["scaling_fwd"] = generation + + module._validate_quantizer_generation(fwd=True) + assert module._unsafe_requantization_input_quantizer is input_quantizers[0] + assert [counter.get("safety_calls", 0) for counter in counters] == [1, 0] + + # The generation list is stable between forwards, so the O(1) identity + # guard must avoid re-running capability checks. + module._validate_quantizer_generation(fwd=True) + assert [counter.get("safety_calls", 0) for counter in counters] == [1, 0] + + @staticmethod + def _counting_identity_hybrid_recipe( + counters, + *, + columnwise_source="rowwise_dequantized", + rowwise_safe=False, + columnwise_safe=True, + ): + def factory(role): + if role is not None and role.module_type == "linear" and role.tensor_type == "input": + rowwise_cls = ( + _CountingIdentityQuantizer if rowwise_safe else _CountingUnsafeIdentityQuantizer + ) + columnwise_cls = ( + _CountingIdentityQuantizer + if columnwise_safe + else _CountingUnsafeIdentityQuantizer + ) + return HybridQuantizer( + rowwise_quantizer=rowwise_cls(counters["rowwise"]), + columnwise_quantizer=columnwise_cls(counters["columnwise"]), + columnwise_source=columnwise_source, + ) + return IdentityQuantizer() + + return recipe.CustomRecipe(qfactory=factory) + + def test_linear_save_original_input_veto_uses_saved_forward_quantized_input(self): + torch.manual_seed(205) + counters = {"rowwise": {"calls": 0}, "columnwise": {"calls": 0}} + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast( + enabled=True, + recipe=self._counting_identity_hybrid_recipe(counters), + ): + out = model(inp) + + calls_after_forward = counters["rowwise"]["calls"] + assert calls_after_forward > 0 + + out.float().sum().backward() + + assert counters["rowwise"]["calls"] == calls_after_forward + + @pytest.mark.parametrize( + "columnwise_source,rowwise_safe,columnwise_safe,optimization_enabled," + "expected_rowwise_calls", + ( + ("original", True, True, True, 1), + ("original", False, False, True, 1), + ("rowwise_dequantized", True, True, True, 2), + ("rowwise_dequantized", True, False, True, 2), + ("rowwise_dequantized", False, True, False, 1), + ), + ) + def test_linear_save_original_input_policy_is_bitwise_exact( + self, + columnwise_source, + rowwise_safe, + columnwise_safe, + optimization_enabled, + expected_rowwise_calls, + ): + torch.manual_seed(206) + in_features, out_features, batch = 128, 128, 32 + reference = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + candidate = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + candidate.load_state_dict(reference.state_dict()) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + reference_input = base_input.clone().detach().requires_grad_(True) + candidate_input = base_input.clone().detach().requires_grad_(True) + reference_counters = { + "rowwise": {"calls": 0}, + "columnwise": {"calls": 0}, + } + candidate_counters = { + "rowwise": {"calls": 0}, + "columnwise": {"calls": 0}, + } + recipe_kwargs = { + "columnwise_source": columnwise_source, + "rowwise_safe": rowwise_safe, + "columnwise_safe": columnwise_safe, + } + + with autocast( + enabled=True, + recipe=self._counting_identity_hybrid_recipe(reference_counters, **recipe_kwargs), + ): + reference_output = reference(reference_input) + + candidate_recipe = self._counting_identity_hybrid_recipe( + candidate_counters, + **recipe_kwargs, + ) + if optimization_enabled: + with autocast(enabled=True, recipe=candidate_recipe): + candidate_output = candidate(candidate_input) + else: + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast(enabled=True, recipe=candidate_recipe): + candidate_output = candidate(candidate_input) + + assert torch.equal(reference_output, candidate_output) + + reference_output.float().sum().backward() + candidate_output.float().sum().backward() + + assert reference_input.grad is not None and candidate_input.grad is not None + assert torch.equal(reference_input.grad, candidate_input.grad) + reference_weight_grad = dict(reference.named_parameters())["weight"].grad + candidate_weight_grad = dict(candidate.named_parameters())["weight"].grad + assert reference_weight_grad is not None and candidate_weight_grad is not None + assert torch.equal(reference_weight_grad, candidate_weight_grad) + + # The columnwise quantizer is always invoked exactly once. Only a + # rowwise-dequantized reconstruction safely repeats the rowwise stage. + assert candidate_counters["rowwise"]["calls"] == expected_rowwise_calls + assert candidate_counters["columnwise"]["calls"] == 1 + + +@requires_fp8_and_nvfp4 +class TestHybridQuantize: + """Test quantization via HybridQuantizer.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_quantize_returns_hybrid_tensor(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert isinstance(result, HybridQuantizedTensor) + + def test_quantize_shape_preserved(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert result.shape == input_tensor.shape + + def test_quantize_dtype_preserved(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert result.dtype == input_tensor.dtype + + def test_sub_storage_types_fp8_row_fp4_col(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + row_storage = result.rowwise_sub_storage + col_storage = result.columnwise_sub_storage + assert isinstance(row_storage, (Float8TensorStorage, Float8Tensor)) + assert isinstance(col_storage, (NVFP4TensorStorage, NVFP4Tensor)) + + def test_sub_storage_types_reversed(self, input_tensor): + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + result = hq.quantize(input_tensor) + row_storage = result.rowwise_sub_storage + col_storage = result.columnwise_sub_storage + assert isinstance(row_storage, (NVFP4TensorStorage, NVFP4Tensor)) + assert isinstance(col_storage, (Float8TensorStorage, Float8Tensor)) + + def test_quantize_internal_returns_storage(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(input_tensor) + assert isinstance(result, HybridQuantizedTensorStorage) + assert not isinstance(result, HybridQuantizedTensor) + hq.internal = False + + +@requires_fp8_and_nvfp4 +class TestHybridDequantize: + """Test dequantization round-trip.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_dequantize_shape(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + assert dequantized.shape == input_tensor.shape + + def test_dequantize_dtype(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + assert dequantized.dtype == input_tensor.dtype + + def test_dequantize_explicit_dtype(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize(dtype=torch.float32) + assert dequantized.dtype == torch.float32 + assert dequantized.shape == input_tensor.shape + + def test_dequantize_close_to_original(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + torch.testing.assert_close( + dequantized.float(), input_tensor.float(), rtol=0.125, atol=0.0675 + ) + + def test_dequantize_reversed_close_to_original(self, input_tensor): + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + torch.testing.assert_close(dequantized.float(), input_tensor.float(), rtol=0.5, atol=1.0) + + def test_storage_dequantize(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(input_tensor) + dequantized = result.dequantize(dtype=torch.bfloat16) + assert dequantized.shape == input_tensor.shape + hq.internal = False + + +class TestHybridUpdateUsage: + """Test update_usage semantics and sub-storage cleanup.""" + + @pytest.fixture + def hybrid_tensor(self): + inp = torch.randn(4, 8) + hq = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + return hq.quantize(inp) + + def test_initial_usages(self, hybrid_tensor): + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is True + assert usages["columnwise"] is True + + def test_drop_rowwise(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + assert hybrid_tensor.rowwise_sub_storage is None + assert hybrid_tensor.columnwise_sub_storage is not None + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is False + assert usages["columnwise"] is True + + def test_drop_columnwise(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + assert hybrid_tensor.columnwise_sub_storage is None + assert hybrid_tensor.rowwise_sub_storage is not None + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is True + assert usages["columnwise"] is False + + def test_drop_both(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False, columnwise_usage=False) + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is False + assert usages["columnwise"] is False + + def test_request_true_is_noop(self, hybrid_tensor): + row_before = hybrid_tensor.rowwise_sub_storage + col_before = hybrid_tensor.columnwise_sub_storage + hybrid_tensor.update_usage(rowwise_usage=True, columnwise_usage=True) + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is col_before + + def test_request_missing_columnwise_raises(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + + with pytest.raises(RuntimeError, match="no columnwise sub-storage"): + hybrid_tensor.update_usage(columnwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is not None + assert hybrid_tensor.columnwise_sub_storage is None + + def test_request_missing_rowwise_raises(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + + with pytest.raises(RuntimeError, match="no rowwise sub-storage"): + hybrid_tensor.update_usage(rowwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is None + assert hybrid_tensor.columnwise_sub_storage is not None + + def test_missing_direction_request_is_atomic(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + row_before = hybrid_tensor.rowwise_sub_storage + + with pytest.raises(RuntimeError, match="no columnwise sub-storage"): + hybrid_tensor.update_usage(rowwise_usage=False, columnwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is None + + def test_none_preserves_missing_direction(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + row_before = hybrid_tensor.rowwise_sub_storage + + hybrid_tensor.update_usage(rowwise_usage=None, columnwise_usage=None) + + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is None + + def test_repr_after_drop(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + r = repr(hybrid_tensor) + assert "HybridQuantizedTensor" in r + assert "rowwise=None" in r + + hybrid_tensor.update_usage(columnwise_usage=False) + r = repr(hybrid_tensor) + assert "rowwise=None" in r + assert "columnwise=None" in r + + +requires_mxfp8 = pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", +) + + +@requires_fp8_and_nvfp4 +class TestHybridClear: + """Test HybridQuantizedTensorStorage.clear() — buffer deallocation. + + ``clear()`` is invoked by cpu_offload_v1 after the offloader has taken + its own reference to the extracted buffers, to free the GPU-resident + originals. HybridQuantizedTensorStorage delegates to each sub-storage's + own clear(), which replaces primary data buffers with empty tensors. + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + @staticmethod + def _primary_data_numels(sub_storage): + """Collect numel() of primary data buffers on a sub-storage. + + After ``clear()`` every entry should be 0 (native sub-storages set + ``t.data = _empty_tensor()`` on the primary buffers). + """ + if sub_storage is None: + return [] + data = sub_storage.get_data_tensors() + if not isinstance(data, tuple): + data = (data,) + return [t.numel() for t in data if t is not None] + + def test_clear_delegates_to_both_sub_storages(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + + row_before = self._primary_data_numels(ht.rowwise_sub_storage) + col_before = self._primary_data_numels(ht.columnwise_sub_storage) + assert row_before and all(n > 0 for n in row_before) + assert col_before and all(n > 0 for n in col_before) + + ht.clear() + + row_after = self._primary_data_numels(ht.rowwise_sub_storage) + col_after = self._primary_data_numels(ht.columnwise_sub_storage) + assert all(n == 0 for n in row_after) + assert all(n == 0 for n in col_after) + + @requires_mxfp8 + def test_clear_delegates_mxfp8_nvfp4(self, input_tensor): + """Per-block sub-storage path (MXFP8 rowwise + NVFP4 columnwise).""" + hq = HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + ht = hq.quantize(input_tensor) + ht.clear() + for sub in (ht.rowwise_sub_storage, ht.columnwise_sub_storage): + for n in self._primary_data_numels(sub): + assert n == 0 + + def test_clear_with_rowwise_only(self, input_tensor): + """columnwise sub-storage is None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(columnwise_usage=False) + assert ht.columnwise_sub_storage is None + + ht.clear() + + assert all(n == 0 for n in self._primary_data_numels(ht.rowwise_sub_storage)) + + def test_clear_with_columnwise_only(self, input_tensor): + """rowwise sub-storage is None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(rowwise_usage=False) + assert ht.rowwise_sub_storage is None + + ht.clear() + + assert all(n == 0 for n in self._primary_data_numels(ht.columnwise_sub_storage)) + + def test_clear_with_both_sub_storages_dropped(self, input_tensor): + """Both sub-storages are None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(rowwise_usage=False, columnwise_usage=False) + assert ht.rowwise_sub_storage is None + assert ht.columnwise_sub_storage is None + + ht.clear() # must not raise + + def test_clear_is_idempotent(self, input_tensor): + """Calling clear() twice must not raise and leaves buffers empty.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.clear() + ht.clear() + for sub in (ht.rowwise_sub_storage, ht.columnwise_sub_storage): + for n in self._primary_data_numels(sub): + assert n == 0 + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridTensorShapeOps: + """Shape ops that preserve supported Hybrid sub-storages.""" + + def test_fp8_current_non_noop_slice_and_narrow_preserve_hybrid(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 0, 32, 1) + narrowed = tensor.narrow(0, 16, 32) + + assert isinstance(sliced, HybridQuantizedTensor) + assert isinstance(narrowed, HybridQuantizedTensor) + torch.testing.assert_close(sliced.dequantize(), dequantized[:32], rtol=0, atol=0) + torch.testing.assert_close(narrowed.dequantize(), dequantized[16:48], rtol=0, atol=0) + + def test_full_span_step_slice_is_not_treated_as_noop(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 0, tensor.size(0), 2) + + assert isinstance(sliced, HybridQuantizedTensor) + torch.testing.assert_close(sliced.dequantize(), dequantized[::2], rtol=0, atol=0) + + def test_same_shape_as_strided_with_offset_is_not_treated_as_noop(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + base_view = torch.ops.aten.slice.Tensor(tensor, 0, 0, 63, 1) + shifted = torch.ops.aten.as_strided.default( + base_view, base_view.shape, base_view.stride(), x.stride(0) + ) + + assert isinstance(shifted, HybridQuantizedTensor) + torch.testing.assert_close(shifted.dequantize(), dequantized[1:], rtol=0, atol=0) + + +@requires_fp8_and_nvfp4 +class TestHybridDetachIsolation: + """``HybridQuantizedTensor.detach()`` must produce a hybrid whose + sub-storage wrappers are NOT shared with ``self`` — so that + ``detached.prepare_for_saving()`` does not null out fields on the + original. + + This is the property cpu_offload_v2 relies on at + ``cpu_offload.py:378-382``: + + tensor_copy = tensor.detach() + saved_tensors, _ = tensor_copy.prepare_for_saving() + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_detach_produces_distinct_sub_storage_wrappers(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + detached = ht.detach() + + assert detached is not ht + assert detached._rowwise_storage is not ht._rowwise_storage + assert detached._columnwise_storage is not ht._columnwise_storage + + def test_detach_prepare_for_saving_does_not_affect_original(self, input_tensor): + """prepare_for_saving on the detach() copy must not null the original. + + This is the specific invariant the cpu_offload_v2 push/reload cycle + depends on. Without it, a second push on the same tensor — or even + a bare ``.device`` read during offload eligibility checks — hits + `` has no data!``. + """ + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + + detached = ht.detach() + _ = detached.prepare_for_saving() + + # Original must still be usable: dequantize, .device, repeated clone + _ = ht.dequantize() + _ = ht.device + + def test_detach_shares_underlying_buffers(self, input_tensor): + """Buffer tensors are shared (no GPU allocation) — only wrappers differ.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + detached = ht.detach() + + orig_row_buffers = ht._rowwise_storage.get_data_tensors() + new_row_buffers = detached._rowwise_storage.get_data_tensors() + if not isinstance(orig_row_buffers, tuple): + orig_row_buffers = (orig_row_buffers,) + new_row_buffers = (new_row_buffers,) + for a, b in zip(orig_row_buffers, new_row_buffers): + if a is None and b is None: + continue + assert a is b, "detach() must share buffer tensors, not copy them" + + @pytest.mark.parametrize("direction", ("rowwise", "columnwise")) + def test_detach_rejects_storage_only_sub_storage(self, input_tensor, direction): + """Storage-only children have no supported wrapper-copy protocol.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + getattr(hq, f"{direction}_quantizer").internal = True + ht = hq.quantize(input_tensor) + + with pytest.raises( + NotImplementedError, + match=rf"storage-only {direction} sub-storage", + ): + ht.detach() + + +@requires_fp8_and_nvfp4 +class TestHybridSaveRestore: + """Test prepare_for_saving / restore_from_saved round-trip.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_save_restore_roundtrip(self, hybrid_tensor): + dq_before = hybrid_tensor.dequantize() + expected_metadata = { + "rowwise": _snapshot_storage_tensor_metadata( + hybrid_tensor.rowwise_sub_storage, clone=True + ), + "columnwise": _snapshot_storage_tensor_metadata( + hybrid_tensor.columnwise_sub_storage, clone=True + ), + } + buffers_before = _as_data_tensor_tuple(hybrid_tensor) + tensors, obj = hybrid_tensor.prepare_for_saving() + assert isinstance(tensors, list) + assert all(t is None or isinstance(t, torch.Tensor) for t in tensors) + + remainder = obj.restore_from_saved(tensors) + assert isinstance(remainder, list) + assert len(remainder) == 0 + + dq_after = hybrid_tensor.dequantize() + buffers_after = _as_data_tensor_tuple(hybrid_tensor) + assert len(buffers_after) == len(buffers_before) + for before, after in zip(buffers_before, buffers_after): + assert after is before, "Direct save/restore must reattach the exact saved buffer" + torch.testing.assert_close(dq_before, dq_after, rtol=0.0, atol=0.0) + for direction in ("rowwise", "columnwise"): + actual_metadata = _snapshot_storage_tensor_metadata( + getattr(hybrid_tensor, f"{direction}_sub_storage"), clone=False + ) + _assert_nested_state_exact( + actual_metadata, + expected_metadata[direction], + path=f"save/restore {direction}", + ) + + def test_save_clears_data(self, hybrid_tensor): + tensors, obj = hybrid_tensor.prepare_for_saving() + row_storage = hybrid_tensor.rowwise_sub_storage + row_data_tensors = row_storage.get_data_tensors() + if isinstance(row_data_tensors, tuple): + assert all(t is None for t in row_data_tensors) + else: + assert row_data_tensors is None + # Restore to clean up + obj.restore_from_saved(tensors) + + +@requires_fp8_and_nvfp4 +class TestHybridMakeEmpty: + """Test HybridQuantizer.make_empty().""" + + def test_make_empty_shape(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert isinstance(empty, HybridQuantizedTensor) + assert empty.shape == torch.Size(shape) + + def test_make_empty_dtype(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert empty.dtype == torch.bfloat16 + + def test_make_empty_has_sub_storages(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is not None + assert empty.columnwise_sub_storage is not None + + def test_make_empty_internal_returns_storage(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + + assert isinstance(empty, HybridQuantizedTensorStorage) + assert not isinstance(empty, HybridQuantizedTensor) + assert empty.size() == torch.Size((128, 256)) + assert empty.dequantize().dtype == torch.bfloat16 + + +@requires_fp8_and_nvfp4 +class TestHybridUsageFlagsRespected: + """``HybridQuantizer`` must skip directions whose parent usage flag is + False. Native quantizers honor ``rowwise_usage`` / ``columnwise_usage`` + inside the C++ kernel; hybrid sub-quantizers are pinned to one direction + in ``__init__``, so the parent's flags never reach C++ — the equivalent + skip lives in the Python composition layer. Modules call ``set_usage`` + extensively before each ``quantize`` (inference, output / grad_input + quantizers, AG paths), so honoring the flags avoids 2x quantization waste. + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + # ── quantize_impl ──────────────────────────────────────────── + + def test_quantize_rowwise_only(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is None + + def test_quantize_columnwise_only(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=True) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is not None + + def test_quantize_both_false(self, input_tensor): + """``set_usage(False, False)`` mirrors ``update_usage(False, False)`` — + both produce an empty hybrid. No defensive assert (matches native).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=False) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is None + + def test_quantize_both_true_default(self, input_tensor): + """Default state (both flags True) keeps both directions populated.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is not None + + def test_quantize_internal_storage_rowwise_only(self, input_tensor): + """Internal storage path (used by FSDP2 / make_like flows) also + honors the gate.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + hq.internal = True + try: + out = hq.quantize(input_tensor) + assert isinstance(out, HybridQuantizedTensorStorage) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is None + finally: + hq.internal = False + + def test_quantize_flag_change_between_calls(self, input_tensor): + """A single quantizer can be re-used with different flags across + calls (which is exactly how modules use one ``input_quantizer`` / + ``weight_quantizer`` across forward / backward phases).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + + hq.set_usage(rowwise=True, columnwise=False) + out_row = hq.quantize(input_tensor) + assert out_row.rowwise_sub_storage is not None + assert out_row.columnwise_sub_storage is None + + hq.set_usage(rowwise=False, columnwise=True) + out_col = hq.quantize(input_tensor) + assert out_col.rowwise_sub_storage is None + assert out_col.columnwise_sub_storage is not None + + hq.set_usage(rowwise=True, columnwise=True) + out_both = hq.quantize(input_tensor) + assert out_both.rowwise_sub_storage is not None + assert out_both.columnwise_sub_storage is not None + + # ── make_empty ─────────────────────────────────────────────── + + def test_make_empty_rowwise_only(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is not None + assert empty.columnwise_sub_storage is None + + def test_make_empty_columnwise_only(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=True) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is None + assert empty.columnwise_sub_storage is not None + + def test_make_empty_both_false(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=False) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is None + assert empty.columnwise_sub_storage is None + + # ── update_quantized ───────────────────────────────────────── + # + # Comparison strategy: snapshot raw data buffers via ``get_data_tensors()`` + # and compare bytes pre/post-update (same pattern as ``TestHybridClear``). + # Avoids per-format ``dequantize()`` limitations (NVFP4 columnwise raises + # NotImplementedError) and is a strictly stronger check — if the kernel + # writes, raw bytes differ regardless of whether dequant is reversible. + + @staticmethod + def _clone_data_tensors(sub_storage): + """Deep-clone the primary data buffers of a sub-storage.""" + if sub_storage is None: + return () + data = sub_storage.get_data_tensors() + if not isinstance(data, tuple): + data = (data,) + return tuple(t.clone() if t is not None else None for t in data) + + @staticmethod + def _assert_data_tensors_equal(snapshot, sub_storage): + """Assert sub-storage's current data buffers byte-match a prior snapshot.""" + assert sub_storage is not None + current = sub_storage.get_data_tensors() + if not isinstance(current, tuple): + current = (current,) + assert len(snapshot) == len( + current + ), f"Buffer count changed: {len(snapshot)} → {len(current)}" + for before, after in zip(snapshot, current): + if before is None: + assert after is None + continue + assert after is not None + torch.testing.assert_close(before, after, rtol=0, atol=0) + + @staticmethod + def _assert_data_tensors_differ(snapshot, sub_storage): + """Assert at least one buffer changed bytes vs the prior snapshot.""" + assert sub_storage is not None + current = sub_storage.get_data_tensors() + if not isinstance(current, tuple): + current = (current,) + any_changed = False + for before, after in zip(snapshot, current): + if before is None or after is None: + continue + if not torch.equal(before, after): + any_changed = True + break + assert any_changed, "Expected at least one data buffer to change but none did" + + def test_update_quantized_rowwise_only_preserves_columnwise_data(self, input_tensor): + """``update_quantized`` must not refresh a direction whose parent flag + is False, even if the dst storage has that direction allocated. + + Mirrors how native ``tex.quantize(src, quantizer, dst, noop_flag)`` + skips a direction when ``quantizer.rowwise_usage=False`` even if the + dst storage has that direction allocated. + """ + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + # Fully populate both directions + dst = hq.quantize(input_tensor) + # Snapshot the columnwise raw buffers before the targeted rowwise-only update + col_before = self._clone_data_tensors(dst._columnwise_storage) + + # Switch to rowwise-only refresh and feed a substantially different src + hq.set_usage(rowwise=True, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + # Both sub-storage objects survive in-place; columnwise bytes untouched + self._assert_data_tensors_equal(col_before, dst._columnwise_storage) + + def test_update_quantized_columnwise_only_preserves_rowwise_data(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + + hq.set_usage(rowwise=False, columnwise=True) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + self._assert_data_tensors_equal(row_before, dst._rowwise_storage) + + def test_update_quantized_both_false_is_noop(self, input_tensor): + """``set_usage(False, False)`` then ``update_quantized`` must leave + both sub-storages' bytes untouched.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + col_before = self._clone_data_tensors(dst._columnwise_storage) + + hq.set_usage(rowwise=False, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + self._assert_data_tensors_equal(row_before, dst._rowwise_storage) + self._assert_data_tensors_equal(col_before, dst._columnwise_storage) + + def test_update_quantized_actually_refreshes_requested(self, input_tensor): + """Sanity check: when the parent flag is True, the corresponding + sub-storage IS refreshed (otherwise the previous tests would pass + vacuously by not refreshing anything).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + + hq.set_usage(rowwise=True, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + # Rowwise bytes must differ — confirms update_quantized actually ran + self._assert_data_tensors_differ(row_before, dst._rowwise_storage) + + # ── te.Linear integration: inference path takes rowwise-only ─ + + def test_te_linear_inference_workspace_rowwise_only(self): + """``te.Linear`` forward under ``torch.no_grad()`` with a hybrid + ``CustomRecipe`` must produce a rowwise-only weight workspace. + ``linear.py:266-274`` sets ``weight_quantizer.set_usage(columnwise=False)`` + in inference; without the parent-flag gate, hybrid would still allocate + both directions. + """ + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + ) + torch.manual_seed(2026) + model = Linear(128, 256, bias=False, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda") + + # is_first_microbatch=True forces the cache_name="weight" path + # (see linear.py:1631) so the hybrid workspace persists in + # model._fp8_workspaces and we can inspect its sub-storages. + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + _ = model(x, is_first_microbatch=True) + + ws = model._fp8_workspaces.get("weight") + assert isinstance( + ws, HybridQuantizedTensorStorage + ), f"Expected hybrid weight workspace, got {type(ws).__name__}" + assert ws.rowwise_sub_storage is not None, "Rowwise sub-storage must be populated for fprop" + assert ws.columnwise_sub_storage is None, ( + "Inference forward must produce rowwise-only hybrid weight workspace; " + "columnwise quantization should have been skipped per " + "weight_quantizer.set_usage(rowwise=True, columnwise=False)." + ) + + +@requires_fp8_and_nvfp4 +class TestHybridTorchDispatch: + """Test torch dispatch operations.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_detach(self, hybrid_tensor): + detached = hybrid_tensor.detach() + assert isinstance(detached, HybridQuantizedTensor) + assert not detached.requires_grad + + def test_repr(self, hybrid_tensor): + r = repr(hybrid_tensor) + assert "HybridQuantizedTensor" in r + + +@requires_fp8_and_nvfp4 +class TestHybridGetDataTensors: + """Test get_data_tensors returns data from both sub-storages.""" + + def test_get_data_tensors(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(inp) + data_tensors = result.get_data_tensors() + row_tensors = _as_data_tensor_tuple(result.rowwise_sub_storage) + col_tensors = _as_data_tensor_tuple(result.columnwise_sub_storage) + + assert isinstance(data_tensors, tuple) + assert row_tensors and col_tensors + assert len(data_tensors) == len(row_tensors) + len(col_tensors) + assert all( + actual is expected for actual, expected in zip(data_tensors, row_tensors + col_tensors) + ), "Hybrid get_data_tensors must concatenate both sub-storages in direction order" + + +@requires_fp8_and_nvfp4 +class TestHybridDeviceAndSize: + """Test device and size properties.""" + + def test_device(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(inp) + assert result.device.type == "cuda" + + def test_size_from_storage(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(inp) + size = result.size() + assert size == torch.Size([128, 256]) + hq.internal = False + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridGemmBitwiseIdentical: + """Hybrid quantizer with same FP8 format in both directions must produce + bitwise-identical results to the vanilla Float8CurrentScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_fp8(self): + torch.manual_seed(123) + + in_features = 64 + out_features = 64 + batch = 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.Float8CurrentScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + loss_ref = out_ref.float().sum() + loss_ref.backward() + + def hybrid_fp8_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight", "output") + ): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_fp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + loss_hybrid = out_hybrid.float().sum() + loss_hybrid.backward() + + # Forward outputs must be bitwise identical + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + + # Input gradients must be bitwise identical + assert inp_ref.grad is not None and inp_hybrid.grad is not None + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + + # Parameter gradients must be bitwise identical + ref_params = dict(model_ref.named_parameters()) + hybrid_params = dict(model_hybrid.named_parameters()) + for name, p_ref in ref_params.items(): + p_hyb = hybrid_params[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestHybridGemmBitwiseIdenticalMXFP8: + """Hybrid quantizer with MXFP8 in both directions must produce + bitwise-identical results to the vanilla MXFP8BlockScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_mxfp8(self): + torch.manual_seed(200) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.MXFP8BlockScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_mxfp8_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_mxfp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwise(self): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + ) + + torch.manual_seed(204) + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + model_qfactory = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + model_qfactory.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_qfactory = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.MXFP8BlockScaling() + ref_recipe.backward_override = "dequantized" + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + + qfactory_recipe = recipe.CustomRecipe(qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory) + with autocast(enabled=True, recipe=qfactory_recipe): + out_qfactory = model_qfactory(inp_qfactory) + + assert torch.equal( + out_ref, out_qfactory + ), f"Forward mismatch: max diff = {(out_ref - out_qfactory).abs().max().item()}" + + out_ref.float().sum().backward() + out_qfactory.float().sum().backward() + + assert inp_ref.grad is not None and inp_qfactory.grad is not None + assert torch.equal(inp_ref.grad, inp_qfactory.grad), ( + "Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_qfactory.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_qfactory = dict(model_qfactory.named_parameters())[name] + assert ( + p_ref.grad is not None and p_qfactory.grad is not None + ), f"Missing gradient for param {name!r}" + assert torch.equal(p_ref.grad, p_qfactory.grad), ( + f"Param {name!r} grad mismatch: max diff = " + f"{(p_ref.grad - p_qfactory.grad).abs().max().item()}" + ) + + +@requires_fp8 +class TestCustomDPALocalRecipeCache: + """Custom-DPA native recipe labels track the quantizer rebuild.""" + + def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, monkeypatch): + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + from transformer_engine.pytorch.module.base import TransformerEngineBaseModule + from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + custom_recipe = recipe.CustomRecipe(qfactory=lambda _role: IdentityQuantizer()) + monkeypatch.setattr( + FP8GlobalStateManager, + "get_fp8_recipe", + classmethod(lambda _cls: custom_recipe), + ) + + state = [object()] + quantizer = [object()] + + def fake_base_init(module, num_gemms=1): # pylint: disable=unused-argument + module.fp8_meta["scaling_fwd"] = state[0] + module.quantizers["scaling_fwd"] = [quantizer[0]] + + monkeypatch.setattr( + TransformerEngineBaseModule, + "init_fp8_metadata", + fake_base_init, + ) + + inferred_labels = [recipe.MXFP8BlockScaling()] + mutated_recipe_labels = [recipe.MXFP8BlockScaling(fp8_mha=True)] + inference_results = iter((inferred_labels, mutated_recipe_labels, None)) + inference_calls = 0 + + def fake_infer(*_args, **_kwargs): + nonlocal inference_calls + inference_calls += 1 + return next(inference_results) + + monkeypatch.setattr(dpa_module, "_infer_custom_dpa_local_recipes", fake_infer) + + dpa = te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + ) + dpa.init_fp8_metadata() + assert dpa.fp8_meta["local_recipes"] is inferred_labels + dpa.init_fp8_metadata() + assert inference_calls == 1 + assert dpa.fp8_meta["local_recipes"] is inferred_labels + + # Native labels also copy these mutable fields from CustomRecipe. They + # must refresh even when the quantizer generation itself is unchanged. + custom_recipe.fp8_mha = True + dpa.init_fp8_metadata() + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + dpa.init_fp8_metadata() + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + + # A rebuilt recipe state/quantizer list invalidates the cache. If the + # new family has no native label, the old label must not survive. + state[0] = object() + quantizer[0] = object() + dpa.init_fp8_metadata() + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + dpa.init_fp8_metadata() + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + + @pytest.mark.parametrize( + "factory_name,expected", + [ + ("current_scaling_quantizer_factory", (True, False)), + ("delayed_scaling_quantizer_factory", (False, False)), + ], + ) + def test_qkv_capabilities_reuse_canonical_quantizer(self, factory_name, expected): + """Capability queries must not call qfactory outside recipe-state setup.""" + from transformer_engine.pytorch.custom_recipes import quantization_factory_base + + base_qfactory = getattr(quantization_factory_base, factory_name) + calls = [] + + def counting_qfactory(role): + calls.append(role) + # Model a factory with observable RNG state. An extra classification + # probe would consume another value and change subsequent results. + _ = torch.rand((), device="cuda") + return base_qfactory(role) + + custom_recipe = recipe.CustomRecipe( + qfactory=counting_qfactory, + fp8_dpa=True, + fp8_mha=True, + ) + dpa = te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="counted_dpa", + ).cuda() + + with autocast(enabled=True, recipe=custom_recipe): + first = dpa.get_qkv_quantization_capabilities() + canonical_qkv = dpa._qkv_capabilities_quantizer + calls_after_first = len(calls) + second = dpa.get_qkv_quantization_capabilities() + + assert first == expected + assert second == first + assert dpa._qkv_capabilities_quantizer is canonical_qkv + assert len(calls) == calls_after_first + + # A recipe-state rebuild creates a new canonical slot and must + # invalidate the capability cache automatically. + def rebuilt_qfactory(role): + return counting_qfactory(role) + + rebuilt_recipe = recipe.CustomRecipe( + qfactory=rebuilt_qfactory, + fp8_dpa=True, + fp8_mha=True, + ) + with autocast(enabled=True, recipe=rebuilt_recipe): + rebuilt = dpa.get_qkv_quantization_capabilities() + rebuilt_qkv = dpa._qkv_capabilities_quantizer + + assert rebuilt == first + assert rebuilt_qkv is not canonical_qkv + assert len(calls) > calls_after_first + + @pytest.mark.parametrize("fp8_mha", [False, True]) + def test_mha_forward_does_not_probe_qfactory(self, monkeypatch, fp8_mha): + """Repeated MHA forwards must not create discarded DPA quantizers.""" + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + current_scaling_quantizer_factory, + delayed_scaling_quantizer_factory, + ) + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") + calls = [] + + def valid_fp8_dpa_qfactory(role): + # Hopper FP8 attention supports DelayedScaling. Use it for the + # complete MHA recipe so that all requests share one coherent + # delayed-scaling state, including the DPA boundary tensors. + if cc < (10, 0): + return delayed_scaling_quantizer_factory(role) + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + if is_dpa or is_dpa_boundary: + return nvfp4_linear_fp8_dpa_factory(role) + return current_scaling_quantizer_factory(role) + + def counting_qfactory(role): + calls.append(role) + _ = torch.rand((), device="cuda") + return valid_fp8_dpa_qfactory(role) + + custom_recipe = recipe.CustomRecipe( + qfactory=counting_qfactory, + fp8_dpa=True, + fp8_mha=fp8_mha, + ) + model = te.MultiheadAttention( + hidden_size=128, + num_attention_heads=2, + kv_channels=64, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="counted_mha", + ).cuda() + inp = torch.randn(128, 2, 128, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(), autocast(enabled=True, recipe=custom_recipe): + model(inp) + calls_after_first = len(calls) + model(inp) + + assert len(calls) == calls_after_first + + +@requires_fp8_and_nvfp4 +class TestAttentionFactoryNativeRecipeParity: + """Linear + DPA qfactories should match native DPA recipe switches bitwise.""" + + batch = 2 + seq_len = 128 + hidden_size = 128 + num_heads = 4 + kv_channels = hidden_size // num_heads + + class _LinearDPALinear(torch.nn.Module): + def __init__(self, hidden_size, num_heads, kv_channels): + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + self.kv_channels = kv_channels + self.qkv_proj = Linear( + hidden_size, + 3 * hidden_size, + params_dtype=torch.bfloat16, + bias=False, + name="qkv", + ).cuda() + self.dpa = te.DotProductAttention( + num_heads, + kv_channels, + attention_dropout=0.0, + qkv_format="bshd", + name="core_attention", + ).cuda() + self.out_proj = Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + bias=False, + name="proj", + ).cuda() + + def forward(self, inp): + batch, seq_len, _ = inp.shape + qkv = self.qkv_proj(inp).view( + batch, + seq_len, + 3, + self.num_heads, + self.kv_channels, + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + attn_out = self.dpa(q, k, v, qkv_format="bshd").reshape( + batch, + seq_len, + self.hidden_size, + ) + return self.out_proj(attn_out) + + @staticmethod + def _set_native_dpa_recipe(monkeypatch, recipe_name): + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + monkeypatch.setenv("NVTE_DPA_FP8_RECIPE", recipe_name) + monkeypatch.setenv("NVTE_DPA_FP8_FORMAT", "HYBRID") + monkeypatch.setenv("NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent") + monkeypatch.setenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1") + monkeypatch.setenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", recipe_name) + monkeypatch.setattr(dpa_module, "_dpa_fp8_format", recipe.Format.HYBRID) + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_amax_algo", "most_recent") + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_amax_histlen", 1) + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_reduce_amax", True) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", recipe_name) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_dpa", False) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_mha", False) + + @staticmethod + def _clear_native_dpa_recipe(monkeypatch): + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + + monkeypatch.delenv("NVTE_DPA_FP8_RECIPE", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8_FORMAT", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_AMAX_ALGO", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_REDUCE_AMAX", raising=False) + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", "") + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_dpa", False) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_mha", False) + + @staticmethod + def _assert_equal(actual, expected, label): + assert torch.equal( + actual, expected + ), f"{label} mismatch: max diff = {(actual.float() - expected.float()).abs().max().item()}" + + def _run_model(self, model, inp, grad, fp8_recipe, seed): + run_inp = inp.clone().detach().requires_grad_(True) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + with autocast(enabled=True, recipe=fp8_recipe): + out = model(run_inp) + out.backward(grad) + local_recipes = [type(r).__name__ for r in model.dpa.fp8_meta.get("local_recipes", [])] + return ( + out.detach().clone(), + run_inp.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + local_recipes, + ) + + @pytest.mark.parametrize( + "case_name,native_dpa_recipe,qfactory_name", + [ + ( + "fp8_dpa", + "Float8CurrentScaling", + "nvfp4_linear_fp8_dpa_factory", + ), + ( + "mxfp8_dpa", + "MXFP8BlockScaling", + "nvfp4_linear_mxfp8_dpa_factory", + ), + ], + ) + def test_linear_dpa_linear_matches_native_env_recipe_bitwise( + self, + monkeypatch, + case_name, + native_dpa_recipe, + qfactory_name, + ): + if case_name == "mxfp8_dpa" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.custom_recipes import quantization_factory_zoo + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) + + torch.manual_seed(2201) + model_native = self._LinearDPALinear( + self.hidden_size, + self.num_heads, + self.kv_channels, + ) + model_qfactory = self._LinearDPALinear( + self.hidden_size, + self.num_heads, + self.kv_channels, + ) + model_qfactory.load_state_dict(model_native.state_dict()) + + torch.manual_seed(2202) + base_inp = torch.randn( + self.batch, + self.seq_len, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(base_inp) + + native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) + qfactory = getattr(quantization_factory_zoo, qfactory_name) + qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) + + native_out, native_dx, native_grads, native_local_recipes = self._run_model( + model_native, + base_inp, + grad, + native_recipe, + seed=2203, + ) + self._clear_native_dpa_recipe(monkeypatch) + qfactory_out, qfactory_dx, qfactory_grads, qfactory_local_recipes = self._run_model( + model_qfactory, + base_inp, + grad, + qfactory_recipe, + seed=2203, + ) + + expected_local_recipes = ( + ["Float8CurrentScaling", "DelayedScaling"] + if native_dpa_recipe == "Float8CurrentScaling" + else ["MXFP8BlockScaling"] + ) + assert native_local_recipes == expected_local_recipes + assert qfactory_local_recipes == expected_local_recipes + + self._assert_equal(qfactory_out, native_out, f"{case_name} output") + self._assert_equal(qfactory_dx, native_dx, f"{case_name} input grad") + assert qfactory_grads.keys() == native_grads.keys() + for name, native_grad in native_grads.items(): + self._assert_equal( + qfactory_grads[name], + native_grad, + f"{case_name} param grad {name}", + ) + + def _run_mha_model(self, model, inp, grad, fp8_recipe, seed): + run_inp = inp.clone().detach().requires_grad_(True) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + with autocast(enabled=True, recipe=fp8_recipe): + out = model(run_inp, attn_mask_type="no_mask") + if isinstance(out, tuple): + out = out[0] + out.backward(grad) + return ( + out.detach().clone(), + run_inp.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + ) + + @pytest.mark.parametrize( + "case_name,native_dpa_recipe,qfactory_name,expected_flags", + [ + ( + "fp8_dpa", + "Float8CurrentScaling", + "nvfp4_linear_fp8_dpa_factory", + (False, False, False), + ), + ( + "mxfp8_dpa", + "MXFP8BlockScaling", + "nvfp4_linear_mxfp8_dpa_factory", + (False, False, False), + ), + ], + ) + def test_multihead_attention_matches_native_env_recipe_bitwise( + self, + monkeypatch, + case_name, + native_dpa_recipe, + qfactory_name, + expected_flags, + ): + if case_name == "mxfp8_dpa" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.custom_recipes import quantization_factory_zoo + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + recorded_flags = [] + orig_update_roles = mha_module.MultiheadAttention._update_output_quantizer_roles + + def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output): + recorded_flags.append((qkv_fp8_output, dpa_fp8_output, proj_fp8_grad)) + return orig_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + + monkeypatch.setattr( + mha_module.MultiheadAttention, + "_update_output_quantizer_roles", + _recording_update_roles, + ) + + self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) + + torch.manual_seed(2301) + model_native = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + model_qfactory = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + model_qfactory.load_state_dict(model_native.state_dict()) + + torch.manual_seed(2302) + base_inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(base_inp) + + native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) + qfactory = getattr(quantization_factory_zoo, qfactory_name) + qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) + + native_out, native_dx, native_grads = self._run_mha_model( + model_native, + base_inp, + grad, + native_recipe, + seed=2303, + ) + native_flags = recorded_flags[-1] + self._clear_native_dpa_recipe(monkeypatch) + qfactory_out, qfactory_dx, qfactory_grads = self._run_mha_model( + model_qfactory, + base_inp, + grad, + qfactory_recipe, + seed=2303, + ) + qfactory_flags = recorded_flags[-1] + + assert native_flags == expected_flags + assert qfactory_flags == expected_flags + self._assert_equal(qfactory_out, native_out, f"{case_name} MHA output") + self._assert_equal(qfactory_dx, native_dx, f"{case_name} MHA input grad") + assert qfactory_grads.keys() == native_grads.keys() + for name, native_grad in native_grads.items(): + self._assert_equal( + qfactory_grads[name], + native_grad, + f"{case_name} MHA param grad {name}", + ) + + def test_update_output_quantizer_roles_wires_independent_boundaries(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + model = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + qkv = model.layernorm_qkv if model.input_layernorm else model.qkv + + expected_qkv = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=model.core_attention.name or "", + ) + expected_do = QuantizerRole( + module_type="dpa", + tensor_type="do", + name=model.core_attention.name or "", + ) + expected_o = QuantizerRole( + module_type="linear", + tensor_type="input", + name=model.proj.name or "", + ) + expected_dqkv = QuantizerRole( + module_type="linear", + tensor_type="grad_output", + name=qkv.name or "", + ) + + def boundary_roles(): + return ( + qkv.output_quantizer_role, + model.proj.grad_input_quantizer_role, + model.core_attention.output_quantizer_role, + model.core_attention.grad_input_quantizer_role, + ) + + model._update_output_quantizer_roles(True, False, False) + assert boundary_roles() == (expected_qkv, None, None, None) + + model._update_output_quantizer_roles(False, True, False) + assert boundary_roles() == (None, expected_do, None, None) + + model._update_output_quantizer_roles(False, False, True) + assert boundary_roles() == (None, None, expected_o, expected_dqkv) + + model._update_output_quantizer_roles(False, False, False) + assert boundary_roles() == (None, None, None, None) + + def test_mxfp8_qfactory_uses_plain_bf16_mha_boundaries(self, monkeypatch): + """MXFP8 DPA stays internal; MHA boundary tensors remain plain BF16.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.attention.dot_product_attention import ( + backends as dpa_backends, + ) + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_linear_mxfp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + self._clear_native_dpa_recipe(monkeypatch) + torch.manual_seed(2401) + model = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + + boundary_tensors = {} + saved_dpa_tensors = {} + orig_prepare_for_saving = dpa_backends.prepare_for_saving + + def _record_prepare_for_saving(*tensors, **kwargs): + saved_dpa_tensors["fp8_o"] = tensors[3] + saved_dpa_tensors["f16_o"] = tensors[7] + return orig_prepare_for_saving(*tensors, **kwargs) + + monkeypatch.setattr(dpa_backends, "prepare_for_saving", _record_prepare_for_saving) + + def _record_grad(name): + def _hook(grad): + boundary_tensors[name] = grad + return grad + + return _hook + + def _record_dpa_inputs(_module, inputs): + for name, tensor in zip(("q", "k", "v"), inputs[:3]): + boundary_tensors[name] = tensor + tensor.register_hook(_record_grad(f"d{name}")) + + def _record_dpa_output(_module, _inputs, output): + boundary_tensors["o"] = output + output.register_hook(_record_grad("do")) + + def _record_projection_input(_module, inputs): + boundary_tensors["proj_input"] = inputs[0] + boundary_tensors["o_is_proj_input"] = boundary_tensors["o"] is inputs[0] + + handles = ( + model.core_attention.register_forward_pre_hook(_record_dpa_inputs), + model.core_attention.register_forward_hook(_record_dpa_output), + model.proj.register_forward_pre_hook(_record_projection_input), + ) + + torch.manual_seed(2402) + inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(inp) + qfactory_recipe = recipe.CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + try: + self._run_mha_model(model, inp, grad, qfactory_recipe, seed=2403) + finally: + for handle in handles: + handle.remove() + + assert boundary_tensors["o_is_proj_input"] + assert saved_dpa_tensors["fp8_o"] is None + saved_o = saved_dpa_tensors["f16_o"] + assert type(saved_o) is torch.Tensor + assert saved_o.dtype is torch.bfloat16 + assert ( + saved_o.untyped_storage().data_ptr() + == boundary_tensors["o"].untyped_storage().data_ptr() + ) + + for name in ("q", "k", "v", "o", "proj_input", "dq", "dk", "dv", "do"): + tensor = boundary_tensors[name] + assert type(tensor) is torch.Tensor, f"{name} is {type(tensor)}" + assert tensor.dtype is torch.bfloat16, f"{name} has dtype {tensor.dtype}" + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +class TestHybridGemmBitwiseIdenticalBlockFP8: + """Hybrid quantizer with Block FP8 in both directions must produce + bitwise-identical results to the vanilla Float8BlockScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_block_fp8(self): + torch.manual_seed(201) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.Float8BlockScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_block_fp8_factory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + is_weight = is_linear and role.tensor_type == "weight" + dim = 2 if is_weight else 1 + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, + ) + return HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, + ), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_block_fp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestHybridGemmBitwiseIdenticalNVFP4: + """Same-format hybrid NVFP4 must match vanilla with seeded SR.""" + + def test_linear_fwd_bwd_matches_vanilla_nvfp4(self): + torch.manual_seed(202) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.NVFP4BlockScaling() + torch.manual_seed(1202) + torch.cuda.manual_seed_all(1202) + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_nvfp4_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "grad_output" + ): + return nvfp4_quantizer_factory(role) + return HybridQuantizer( + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_factory) + torch.manual_seed(1202) + torch.cuda.manual_seed_all(1202) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + def test_linear_fwd_bwd_all_roles_hybrid(self): + """Exercise grad_output as a HybridQuantizer too.""" + torch.manual_seed(203) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.NVFP4BlockScaling() + torch.manual_seed(1203) + torch.cuda.manual_seed_all(1203) + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_nvfp4_all_roles_factory(role): + return HybridQuantizer( + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_all_roles_factory) + torch.manual_seed(1203) + torch.cuda.manual_seed_all(1203) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@requires_fp8_and_nvfp4 +class TestHybridGemmMixedFormat: + """FP8 rowwise + NVFP4 columnwise through te.Linear forward+backward.""" + + def test_linear_fwd_bwd_fp8_row_nvfp4_col(self): + torch.manual_seed(42) + + in_features = 128 + out_features = 128 + batch = 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, + in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + def mixed_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return _make_nvfp4_quantizer() + return _make_fp8_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert out.dtype == torch.bfloat16 + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert inp.grad.shape == inp.shape + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + assert not torch.isinf(inp.grad).any(), "Input gradient contains Inf" + + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + assert not torch.isinf(p.grad).any(), f"Gradient for '{name}' contains Inf" + + def test_numerical_sanity_against_bf16(self): + """Mixed-format output should be within reasonable tolerance of BF16 baseline.""" + torch.manual_seed(42) + + in_features = 128 + out_features = 128 + batch = 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + # BF16 baseline (no quantization) + with torch.no_grad(): + out_bf16 = model(inp) + + def mixed_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + return _make_fp8_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + with torch.no_grad(): + with autocast(enabled=True, recipe=mixed_recipe): + out_mixed = model(inp) + + # FP8/FP4 quantization introduces error, but the result should be + # in the same ballpark as BF16 + torch.testing.assert_close( + out_mixed.float(), + out_bf16.float(), + rtol=0.25, + atol=0.5, + ) + + +@requires_fp8_and_nvfp4 +class TestUnwrapTensor: + """Test GEMM input dispatch by rowwise or columnwise usage.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_hybrid_rowwise_returns_rowwise_sub_storage(self, hybrid_tensor): + assert _unwrap_tensor(hybrid_tensor, "rowwise") is hybrid_tensor.rowwise_sub_storage + + def test_hybrid_columnwise_returns_columnwise_sub_storage(self, hybrid_tensor): + assert _unwrap_tensor(hybrid_tensor, "columnwise") is hybrid_tensor.columnwise_sub_storage + + def test_rowwise_sub_storage_type(self, hybrid_tensor): + assert isinstance( + _unwrap_tensor(hybrid_tensor, "rowwise"), + (Float8TensorStorage, Float8Tensor), + ) + + def test_columnwise_sub_storage_type(self, hybrid_tensor): + assert isinstance( + _unwrap_tensor(hybrid_tensor, "columnwise"), + (NVFP4TensorStorage, NVFP4Tensor), + ) + + def test_non_hybrid_passthrough(self): + plain = torch.randn(4, 4, device="cuda") + for usage in ("rowwise", "columnwise"): + assert _unwrap_tensor(plain, usage) is plain + + def test_fp8_tensor_passthrough(self): + quantizer = _make_fp8_quantizer() + inp = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + fp8 = quantizer.quantize(inp) + for usage in ("rowwise", "columnwise"): + assert _unwrap_tensor(fp8, usage) is fp8 + + def test_identity_falls_back_to_high_precision(self): + inp = torch.randn(4, 4, device="cuda") + identity = IdentityQuantizer()(inp) + + result = _unwrap_tensor(identity, "rowwise") + + assert isinstance(result, torch.Tensor) + assert not isinstance(result, QuantizedTensor) + assert torch.equal(result, inp) + + def test_custom_tensor_passthrough(self): + from transformer_engine.pytorch.custom_recipes.quantization_ref_current_scaling import ( + CurrentScalingTensorRef, + ) + + custom = CurrentScalingTensorRef() + assert _unwrap_tensor(custom, "rowwise") is custom + + def test_invalid_usage_raises(self): + with pytest.raises(ValueError, match="Unsupported GEMM tensor usage"): + _unwrap_tensor(torch.empty(0), "diagonal") + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridBiasGradient: + """Verify bias gradients are computed correctly with HybridQuantizer. + + tex.bgrad_quantize doesn't recognize HybridQuantizer, so the unfused + bgrad path is used instead. + """ + + def _make_uniform_hybrid_factory(self): + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + + return factory + + def test_bias_grad_matches_vanilla_fp8(self): + torch.manual_seed(456) + in_features, out_features, batch = 64, 64, 16 + + model_ref = Linear(in_features, out_features, bias=True, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear( + in_features, out_features, bias=True, params_dtype=torch.bfloat16 + ).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + # Reference + inp_ref = base_inp.clone().detach().requires_grad_(True) + with autocast(enabled=True, recipe=recipe.Float8CurrentScaling()): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + # Hybrid + inp_hyb = base_inp.clone().detach().requires_grad_(True) + with autocast( + enabled=True, recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory()) + ): + out_hyb = model_hybrid(inp_hyb) + out_hyb.float().sum().backward() + + ref_bias_grad = dict(model_ref.named_parameters())["bias"].grad + hyb_bias_grad = dict(model_hybrid.named_parameters())["bias"].grad + assert ref_bias_grad is not None and hyb_bias_grad is not None + assert torch.equal( + ref_bias_grad, hyb_bias_grad + ), f"Bias grad mismatch: max diff = {(ref_bias_grad - hyb_bias_grad).abs().max().item()}" + + def test_no_bias_fwd_bwd(self): + """Linear with bias=False skips bgrad_quantize entirely.""" + torch.manual_seed(42) + in_features, out_features, batch = 64, 64, 16 + + model = Linear(in_features, out_features, bias=False, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + + with autocast( + enabled=True, recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory()) + ): + out = model(inp) + out.float().sum().backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + + +@requires_fp8_and_nvfp4 +class TestHybridScalingModeCompatibility: + """cuBLAS requires matching scaling modes within a single GEMM. + + For hybrid quantization, this means the columnwise format for + linear_input/linear_weight must match the columnwise format for + linear_grad_output — otherwise the wgrad GEMM (NT layout) fails. + """ + + def test_matching_columnwise_formats_succeed(self): + """Both operands use NVFP4 columnwise → wgrad GEMM succeeds.""" + torch.manual_seed(42) + # NVFP4 GEMM requires dimensions ≥ 128 for cuBLAS support. + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return _make_nvfp4_quantizer() + return _make_fp8_quantizer() + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + out.float().sum().backward() + assert inp.grad is not None + + def test_mismatched_columnwise_formats_raise(self): + """NVFP4 input × FP8 grad_output columnwise → cuBLAS rejects.""" + torch.manual_seed(42) + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return _make_fp8_quantizer() + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + with pytest.raises(RuntimeError, match="scaling_mode"): + out.float().sum().backward() + + +@requires_fp8_and_nvfp4 +class TestHybridReversedDirection: + """Reversed hybrid: NVFP4 rowwise (fprop) + FP8 columnwise (backward). + + Exercises NVFP4×NVFP4 in the fprop (TN) GEMM and FP8×FP8 in the + dgrad (NN) and wgrad (NT) GEMMs — the opposite of the primary + FP8-row/NVFP4-col configuration. + """ + + def test_nvfp4_row_fp8_col_forward_only(self): + """Forward (TN) with NVFP4×NVFP4 rowwise succeeds.""" + torch.manual_seed(99) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + return _make_nvfp4_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with torch.no_grad(): + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + + def test_nvfp4_row_fp8_col_full_fwd_bwd(self): + """Full fwd+bwd with NVFP4 rowwise (fprop) + FP8 columnwise (backward).""" + torch.manual_seed(99) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return _make_fp8_quantizer() + return _make_nvfp4_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridMixedWithNonHybrid: + """Only one operand is hybrid; the other uses a plain TE quantizer. + + Exercises _unwrap_hybrid passthrough for the non-hybrid operand. + All roles must use compatible scaling modes for each GEMM: + fprop (TN): all rowwise formats must match + dgrad (NN): weight rowwise must match grad_output rowwise + wgrad (NT): input columnwise must match grad_output columnwise + """ + + @staticmethod + def _assert_native_fp8_parity(hybrid_role, *, seed): + torch.manual_seed(seed) + in_features, out_features, batch = 128, 128, 32 + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_native = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_native.load_state_dict(model_hybrid.state_dict()) + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + def mixed_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type == hybrid_role: + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=mixed_factory), + seed=seed + 100, + ) + native_result = _run_linear_forward_backward( + model_native, + base_input, + grad_output, + recipe.Float8CurrentScaling(), + seed=seed + 100, + ) + _assert_linear_results_exact( + hybrid_result, + native_result, + output=True, + input_grad=True, + param_grads=True, + ) + + def test_hybrid_input_plain_weight_fwd_bwd(self): + """Input is hybrid (FP8 row / FP8 col), weight + grad_output plain FP8. + + Wgrad columnwise: FP8 (input.col) × FP8 (grad_output.col) → compatible. + """ + self._assert_native_fp8_parity("input", seed=77) + + def test_plain_input_hybrid_weight_fwd_bwd(self): + """Input is plain FP8, weight is hybrid (FP8 row / FP8 col).""" + self._assert_native_fp8_parity("weight", seed=88) + + +# --------------------------------------------------------------------------- +# Parametrized cross-format tests (stateless quantizers) +# --------------------------------------------------------------------------- + + +def _make_mxfp8_quantizer(*, rowwise=True, columnwise=True): + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_mxfp8_quantizer_e5m2(*, rowwise=True, columnwise=True): + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_block_quantizer(*, rowwise=True, columnwise=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_block_quantizer_e5m2(*, rowwise=True, columnwise=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=rowwise, + columnwise=columnwise, + ) + + +# (fwd_e4m3_factory, bwd_e5m2_factory, skip_condition, skip_reason) +_QUANTIZER_CONFIGS = { + "fp8_current": ( + _make_fp8_quantizer, + lambda **kw: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda", **kw), + not fp8_available, + f"FP8: {reason_for_no_fp8}", + ), + "mxfp8": ( + _make_mxfp8_quantizer, + _make_mxfp8_quantizer_e5m2, + not mxfp8_available, + f"MXFP8: {reason_for_no_mxfp8}", + ), + "block_fp8": ( + _make_block_quantizer, + _make_block_quantizer_e5m2, + not fp8_block_scaling_available, + reason_for_no_fp8_block_scaling, + ), + "nvfp4": ( + _make_nvfp4_quantizer, + None, # NVFP4 has no E5M2 variant + not (fp8_available and nvfp4_available), + f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), +} + + +def _build_cross_format_params(): + """Build parametrize list for all stateless cross-format hybrid combos.""" + combos = [ + ("fp8_current", "mxfp8"), + ("fp8_current", "nvfp4"), + ("fp8_current", "block_fp8"), + ("mxfp8", "fp8_current"), + ("mxfp8", "mxfp8"), + ("mxfp8", "nvfp4"), + ("mxfp8", "block_fp8"), + ("block_fp8", "fp8_current"), + ("block_fp8", "mxfp8"), + ("block_fp8", "nvfp4"), + ("block_fp8", "block_fp8"), + ("nvfp4", "fp8_current"), + ("nvfp4", "mxfp8"), + ("nvfp4", "block_fp8"), + ] + params = [] + for row, col in combos: + row_cfg = _QUANTIZER_CONFIGS[row] + col_cfg = _QUANTIZER_CONFIGS[col] + hw_skip = row_cfg[2] or col_cfg[2] + hw_reason = "; ".join( + filter(None, [row_cfg[3] if row_cfg[2] else "", col_cfg[3] if col_cfg[2] else ""]) + ) + marks = [] + if hw_skip: + marks.append(pytest.mark.skipif(True, reason=hw_reason or "N/A")) + if col == "fp8_current": + marks.append(_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8) + params.append(pytest.param(row, col, id=f"{row}_row_x_{col}_col", marks=marks)) + return params + + +def _set_quantization_test_seed(seed): + """Reset every RNG that a quantizer or GEMM helper may consume.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _run_linear_forward(model, base_input, fp8_recipe, *, seed): + """Run a deterministic forward pass and return a detached result.""" + _set_quantization_test_seed(seed) + with torch.no_grad(): + with autocast(enabled=True, recipe=fp8_recipe): + output = model(base_input) + return output.detach().clone() + + +def _run_linear_forward_backward(model, base_input, grad_output, fp8_recipe, *, seed): + """Run Linear with an external gradient and capture every numerical result.""" + model.zero_grad(set_to_none=True) + run_input = base_input.clone().detach().requires_grad_(True) + _set_quantization_test_seed(seed) + with autocast(enabled=True, recipe=fp8_recipe): + output = model(run_input) + output.backward(grad_output) + return ( + output.detach().clone(), + run_input.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + ) + + +def _plain_linear_qfactory(operand_factory, grad_factory): + """Build a non-hybrid factory with role-correct operand and grad dtypes.""" + + def factory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _make_role_aware_quantizer(grad_factory, role) + return _make_role_aware_quantizer(operand_factory, role) + + return factory + + +def _assert_linear_results_exact(actual, expected, *, output, input_grad, param_grads): + """Compare selected Linear results with zero tolerance.""" + if output: + torch.testing.assert_close(actual[0], expected[0], rtol=0.0, atol=0.0) + if input_grad: + torch.testing.assert_close(actual[1], expected[1], rtol=0.0, atol=0.0) + if param_grads: + assert actual[2].keys() == expected[2].keys() + for name in actual[2]: + torch.testing.assert_close( + actual[2][name], + expected[2][name], + rtol=0.0, + atol=0.0, + msg=f"Parameter gradient {name!r} differs", + ) + + +class TestHybridCrossFormatParametrized: + """Parametrized fwd+bwd over all stateless quantizer cross-format pairs.""" + + @pytest.mark.parametrize("row_name,col_name", _build_cross_format_params()) + def test_fwd_bwd(self, row_name, col_name): + torch.manual_seed(42) + in_features, out_features, batch = 128, 128, 32 + + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_bwd_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref.load_state_dict(model_hybrid.state_dict()) + model_bwd_ref.load_state_dict(model_hybrid.state_dict()) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + row_cfg = _QUANTIZER_CONFIGS[row_name] + col_cfg = _QUANTIZER_CONFIGS[col_name] + make_row_operand = row_cfg[0] + make_row_grad = row_cfg[1] if row_cfg[1] is not None else row_cfg[0] + make_col_operand = col_cfg[0] + make_col_grad = col_cfg[1] if col_cfg[1] is not None else col_cfg[0] + + def hybrid_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type in ("input", "weight"): + return HybridQuantizer( + rowwise_quantizer=_make_role_aware_quantizer(make_row_operand, role), + columnwise_quantizer=_make_role_aware_quantizer(make_col_operand, role), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _make_role_aware_quantizer(make_col_grad, role) + return _make_role_aware_quantizer(make_row_operand, role) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_factory) + fprop_ref_recipe = recipe.CustomRecipe( + qfactory=_plain_linear_qfactory(make_row_operand, make_row_grad) + ) + bwd_ref_recipe = recipe.CustomRecipe( + qfactory=_plain_linear_qfactory(make_col_operand, make_col_grad) + ) + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + hybrid_recipe, + seed=1234, + ) + fprop_ref = _run_linear_forward( + model_fprop_ref, + base_input, + fprop_ref_recipe, + seed=1234, + ) + bwd_ref = _run_linear_forward_backward( + model_bwd_ref, + base_input, + grad_output, + bwd_ref_recipe, + seed=1234, + ) + + torch.testing.assert_close(hybrid_result[0], fprop_ref, rtol=0.0, atol=0.0) + _assert_linear_results_exact( + hybrid_result, + bwd_ref, + output=False, + input_grad=True, + param_grads=True, + ) + + +# --------------------------------------------------------------------------- +# CPU offload push/pop protocol (v2 OffloadableLayerState path) +# --------------------------------------------------------------------------- + + +class TestHybridCpuOffloadPushPop: + """Exercise the cpu_offload_v2 push/pop protocol on HybridQuantizedTensor. + + Uses :class:`OffloadableLayerState` directly — same pattern as + ``test_cpu_offloading.py::TestsOffloadableLayerState::test_general``. + Each test runs the full cycle: + + push → start_offload → release_activation_forward_gpu_memory + → start_reload → pop → release_all_memory + + The push path decomposes the hybrid via ``prepare_for_saving`` + (HybridQuantizedTensorStorage), recursively pushes each sub-storage + buffer, then reconstructs on pop via ``restore_from_saved``. Sub-buffers + below the 256K-element offload threshold (e.g. small block scales) are + returned unchanged; large data buffers round-trip through CPU. + """ + + # Hybrid tensor shape — each sub-storage primary buffer must exceed the + # cpu_offload _check_if_offload threshold (256K elements) so the path is + # actually exercised end-to-end. + _SHAPE = (1024, 1024) + + def _run_roundtrip(self, hybrid_tensor): + """Push → offload → release → reload → pop one hybrid tensor. + + Returns the reloaded tensor (a new HybridQuantizedTensor instance + reconstructed from the gathered-back buffers). + """ + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + + tid = state.push_tensor(hybrid_tensor) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + reloaded = state.pop_tensor(tid) + torch.cuda.synchronize() + + try: + return reloaded + finally: + state.release_all_memory() + + @pytest.mark.parametrize("row_name,col_name", _build_cross_format_params()) + def test_push_pop_roundtrip(self, row_name, col_name): + """Dequantize-equivalence round-trip across the full 14-pair matrix.""" + torch.manual_seed(42) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + row_cfg = _QUANTIZER_CONFIGS[row_name] + col_cfg = _QUANTIZER_CONFIGS[col_name] + hq = HybridQuantizer( + rowwise_quantizer=row_cfg[0](), + columnwise_quantizer=col_cfg[0](), + ) + hybrid = hq.quantize(inp) + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload roundtrip") + assert isinstance(reloaded, HybridQuantizedTensor) + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) + + @requires_fp8_and_nvfp4 + def test_push_pop_preserves_sub_storage_types(self): + """Reconstructed hybrid preserves each sub-storage's concrete type.""" + torch.manual_seed(7) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + row_type = type(hybrid.rowwise_sub_storage) + col_type = type(hybrid.columnwise_sub_storage) + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded.rowwise_sub_storage, row_type) + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload storage types") + assert isinstance(reloaded.columnwise_sub_storage, col_type) + + @requires_fp8_and_nvfp4 + def test_push_pop_with_rowwise_only(self): + """Columnwise sub-storage dropped pre-push — roundtrip still works.""" + torch.manual_seed(11) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + hybrid.update_usage(columnwise_usage=False) + assert hybrid.columnwise_sub_storage is None + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert reloaded.columnwise_sub_storage is None + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload rowwise-only") + assert reloaded.rowwise_sub_storage is not None + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) + + @requires_fp8_and_nvfp4 + def test_push_pop_with_columnwise_only(self): + """Rowwise sub-storage dropped pre-push — roundtrip still works. + + Uses the reversed hybrid (NVFP4 rowwise + FP8 columnwise) so that + ``hybrid.dequantize()`` can fall through to the columnwise sub-storage. + ``HybridQuantizedTensorStorage.dequantize`` prefers rowwise and only + falls back to columnwise when rowwise is ``None``; NVFP4 does not yet + support columnwise-only dequantize, but Float8 does. + """ + torch.manual_seed(13) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + hybrid = hq.quantize(inp) + hybrid.update_usage(rowwise_usage=False) + assert hybrid.rowwise_sub_storage is None + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert reloaded.rowwise_sub_storage is None + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload columnwise-only") + assert reloaded.columnwise_sub_storage is not None + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) + + @requires_fp8_and_nvfp4 + def test_push_pop_roundtrip_does_not_leak_intermediate_buffers(self): + """After release_all_memory the offloader holds no hybrid buffers. + + Sanity check that the v2 cycle completes cleanly — no dangling CPU + pinned buffers left behind on a one-shot push/pop. + """ + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + torch.manual_seed(17) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + + tid = state.push_tensor(hybrid) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + _ = state.pop_tensor(tid) + torch.cuda.synchronize() + state.release_all_memory() + + assert len(state.fwd_gpu_tensor_group.tensor_list) == 0 + assert len(state.cpu_tensor_group.tensor_list) == 0 + assert len(state.bwd_gpu_tensor_group.tensor_list) == 0 + assert state.state == "not_offloaded" + + +# --------------------------------------------------------------------------- +# 3-format hybrid: different quantization for fprop, dgrad, wgrad +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (fp8_available and mxfp8_available and nvfp4_available), + reason="Requires FP8 + MXFP8 + NVFP4", +) +class TestHybridThreeFormats: + """Three distinct formats: FormatA (fprop), FormatB (dgrad), FormatC (wgrad). + + Per-operand unwrap selects the correct sub-storage per GEMM: + fprop TN: weight.row(A) × input.row(A) → FormatA × FormatA + dgrad NN: weight.col(B) × grad_output.row(B) → FormatB × FormatB + wgrad NT: input.col(C) × grad_output.col(C) → FormatC × FormatC + + grad_output is itself hybrid (FormatB row + FormatC col) when B ≠ C. + """ + + @staticmethod + def _assert_three_format_routing( + make_fprop, make_dgrad, make_wgrad, *, seed, plain_grad_output=False + ): + in_features, out_features, batch = 128, 128, 32 + torch.manual_seed(seed) + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_dgrad_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_wgrad_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + state_dict = model_hybrid.state_dict() + model_fprop_ref.load_state_dict(state_dict) + model_dgrad_ref.load_state_dict(state_dict) + model_wgrad_ref.load_state_dict(state_dict) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + def hybrid_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=make_fprop(), + columnwise_quantizer=make_dgrad(), + ) + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=make_fprop(), + columnwise_quantizer=make_wgrad(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + if plain_grad_output: + return make_dgrad() + return HybridQuantizer( + rowwise_quantizer=make_dgrad(), + columnwise_quantizer=make_wgrad(), + ) + return make_fprop() + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=hybrid_factory), + seed=seed + 100, + ) + fprop_ref = _run_linear_forward( + model_fprop_ref, + base_input, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_fprop, make_fprop)), + seed=seed + 100, + ) + dgrad_ref = _run_linear_forward_backward( + model_dgrad_ref, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_dgrad, make_dgrad)), + seed=seed + 100, + ) + wgrad_ref = _run_linear_forward_backward( + model_wgrad_ref, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_wgrad, make_wgrad)), + seed=seed + 100, + ) + + torch.testing.assert_close(hybrid_result[0], fprop_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(hybrid_result[1], dgrad_ref[1], rtol=0.0, atol=0.0) + torch.testing.assert_close( + hybrid_result[2]["weight"], + wgrad_ref[2]["weight"], + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + hybrid_result[2]["bias"], + dgrad_ref[2]["bias"], + rtol=0.0, + atol=0.0, + ) + + def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): + """FP8 current (fprop) + MXFP8 (dgrad) + NVFP4 (wgrad).""" + self._assert_three_format_routing( + _make_fp8_quantizer, + _make_mxfp8_quantizer, + _make_nvfp4_quantizer, + seed=300, + ) + + def test_nvfp4_fprop_fp8_dgrad_mxfp8_wgrad(self): + """NVFP4 (fprop) + FP8 current (dgrad) + MXFP8 (wgrad).""" + self._assert_three_format_routing( + _make_nvfp4_quantizer, + _make_fp8_quantizer, + _make_mxfp8_quantizer, + seed=301, + ) + + def test_same_dgrad_wgrad_reduces_to_plain_grad(self): + """When dgrad format == wgrad format, grad_output can be a plain quantizer.""" + self._assert_three_format_routing( + _make_nvfp4_quantizer, + _make_mxfp8_quantizer, + _make_mxfp8_quantizer, + plain_grad_output=True, + seed=302, + ) + + +# --------------------------------------------------------------------------- +# All-modules test: hybrid quantization through every TE module type +# --------------------------------------------------------------------------- + + +def _make_hybrid_fp8_factory(): + """Factory returning HybridQuantizer(FP8 row + FP8 col) for fwd roles, + plain FP8 E5M2 for bwd roles.""" + + def factory(role): + 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=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ) + + return factory + + +@requires_fp8 +@pytest.mark.parametrize("norm_cls", (te.ops.LayerNorm, te.ops.RMSNorm)) +def test_fusible_norm_hybrid_output_falls_back_to_explicit_quantize(norm_cls): + """Unsupported fused Hybrid output is quantized by the following operation.""" + + def qfactory(_role): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=IdentityQuantizer(), + ) + + hidden_size = 128 + norm = norm_cls( + hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + forward = te.ops.Sequential(norm, te.ops.Quantize()) + inp = torch.randn( + 16, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=qfactory)): + out = forward(inp) + + assert isinstance(out, HybridQuantizedTensor) + out.backward(torch.randn_like(inp)) + assert inp.grad is not None + assert norm.weight.grad is not None + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridAllModules: + """Hybrid quantization through all TE module types (not just Linear). + + Uses FP8 in both hybrid directions so the test validates module integration + without introducing cross-format scaling-mode concerns. + """ + + hidden_size = 128 + ffn_hidden_size = 128 + num_heads = 4 + batch = 16 + seq_len = 8 + + def _run_fwd_bwd(self, model, inp, *model_args, output_atol=0.0, param_atols=None): + """Compare same-format hybrid numerics against the native FP8 recipe.""" + import copy + + model_native = copy.deepcopy(model) + param_atols = {} if param_atols is None else param_atols + grad_output = torch.randn_like(inp) + + def run(run_model, run_recipe): + run_model.zero_grad(set_to_none=True) + run_input = inp.detach().clone().requires_grad_(True) + _set_quantization_test_seed(3370) + with autocast(enabled=True, recipe=run_recipe): + output = run_model(run_input, *model_args) + output.backward(grad_output) + return ( + output.detach().clone(), + run_input.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in run_model.named_parameters() + if param.grad is not None + }, + ) + + native_result = run(model_native, recipe.Float8CurrentScaling()) + hybrid_result = run( + model, + recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()), + ) + + torch.testing.assert_close(hybrid_result[0], native_result[0], rtol=0.0, atol=output_atol) + torch.testing.assert_close(hybrid_result[1], native_result[1], rtol=0.0, atol=0.0) + assert hybrid_result[2].keys() == native_result[2].keys() + for name in hybrid_result[2]: + torch.testing.assert_close( + hybrid_result[2][name], + native_result[2][name], + rtol=0.0, + atol=param_atols.get(name, 0.0), + msg=f"Parameter gradient {name!r} differs", + ) + + def test_linear(self): + torch.manual_seed(500) + model = Linear( + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_layernorm_linear(self): + torch.manual_seed(501) + model = LayerNormLinear( + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_layernorm_mlp(self): + torch.manual_seed(502) + model = LayerNormMLP( + hidden_size=self.hidden_size, + ffn_hidden_size=self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + # Native fuses LN+quantize while Hybrid takes the explicit quantize path. + # The only non-bitwise results are output (1 BF16 quantum here) and + # fc2_weight grad; all other gradients remain zero-tolerance checks. + self._run_fwd_bwd( + model, + inp, + output_atol=0.0009765625, + param_atols={"fc2_weight": 0.0234375}, + ) + + def test_grouped_linear(self): + torch.manual_seed(504) + num_gemms = 3 + model = GroupedLinear( + num_gemms, + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + # Hopper FP8 grouped wgrad uses each expert's token count as a + # leading dimension, which cuBLAS requires to be divisible by 16. + m_splits = [16] * num_gemms + inp = torch.randn( + sum(m_splits), + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + self._run_fwd_bwd(model, inp, m_splits) + + def test_transformer_layer(self): + torch.manual_seed(503) + model = TransformerLayer( + self.hidden_size, + self.ffn_hidden_size, + self.num_heads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + # The LayerNormMLP submodule has the same fused-vs-explicit ordering. + # Keep dgrad and every other parameter exact; bound only the observed + # BF16 output and final projection-weight accumulation roundoff. + self._run_fwd_bwd( + model, + inp, + output_atol=0.015625, + param_atols={"layernorm_mlp.fc2_weight": 0.140625}, + ) + + +@requires_fp8 +class TestHybridGroupedLinearValidation: + """GroupedLinear generation-validation and split-dispatch coverage. + + Structural compatibility is validated once per real quantizer generation. + Steady-state dispatch reads the first expert after that uniformity check.""" + + @pytest.mark.parametrize( + "quantizers", + [ + pytest.param( + [_make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3)], + id="hybrid", + ), + pytest.param([_make_fp8_quantizer() for _ in range(3)], id="plain"), + pytest.param([None, None, None], id="none"), + ], + ) + def test_uniform_lists_validate(self, quantizers): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_plain_custom_quantizer_uses_python_split_fallback(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + monkeypatch.setattr( + grouped_linear.tex, + "split_quantize", + lambda *args, **kwargs: pytest.fail("entered native split_quantize"), + ) + calls = [] + tensor = torch.randn((4, 8)) + quantizers = [_CountingPythonQuantizer(calls) for _ in range(2)] + + out = grouped_linear._split_quantize_non_hybrid( + tensor, + [2, 2], + quantizers, + tensor.dtype, + ) + + assert len(calls) == 2 + for actual, expected in zip(out, torch.split(tensor, [2, 2])): + torch.testing.assert_close(actual.dequantize(), expected, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("direction", ("rowwise", "columnwise")) + def test_hybrid_custom_child_uses_python_split_fallback( + self, + monkeypatch, + direction, + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + monkeypatch.setattr( + grouped_linear.tex, + "split_quantize", + lambda *args, **kwargs: pytest.fail("entered native split_quantize"), + ) + calls = [] + quantizers = [] + for _ in range(2): + rowwise = IdentityQuantizer() + columnwise = IdentityQuantizer() + if direction == "rowwise": + rowwise = _CountingPythonQuantizer(calls) + else: + columnwise = _CountingPythonQuantizer(calls) + quantizer = HybridQuantizer( + rowwise_quantizer=rowwise, + columnwise_quantizer=columnwise, + ) + quantizers.append(quantizer) + + out = grouped_linear._split_quantize_hybrid( + torch.randn((4, 8)), + [2, 2], + quantizers, + ) + + assert len(calls) == 2 + assert all(result.rowwise_sub_storage is not None for result in out) + assert all(result.columnwise_sub_storage is not None for result in out) + + def test_mixed_hybrid_and_plain_raises(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + _make_fp8_quantizer(), + _make_hybrid_quantizer_fp8_row_fp4_col(), + ] + with pytest.raises(ValueError, match="mix HybridQuantizer and non-hybrid"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_none_plus_hybrid_raises(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + None, + _make_hybrid_quantizer_fp8_row_fp4_col(), + ] + with pytest.raises(ValueError, match="mix None and concrete quantizers"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_mixed_identity_dtype_raises(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + IdentityQuantizer(dtype=torch.bfloat16), + IdentityQuantizer(dtype=torch.float16), + ] + with pytest.raises(ValueError, match="incompatible plain backend configurations"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_distinct_delayed_scaling_state_is_allowed(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [_make_delayed_quantizer(), _make_delayed_quantizer()] + quantizers[1].scale.fill_(2.0) + quantizers[1].amax.fill_(3.0) + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.parametrize( + ("usage", "expected"), + [ + pytest.param((True, False), {"rowwise": True, "columnwise": False}, id="rowwise"), + pytest.param( + (False, True), + {"rowwise": False, "columnwise": True}, + id="columnwise", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + (True, True), + {"rowwise": True, "columnwise": True}, + id="both", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + ], + ) + def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected): + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_hybrid, + ) + + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=usage[0], columnwise=usage[1]) + + out = _split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert [storage.get_usages() for storage in out] == [expected, expected] + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_columnwise_only_rowwise_dequantized_uses_transient_grouped_row(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + result = real_split_quantize(tensor, m_splits, quantizers, **kwargs) + calls.append((tensor, result)) + return result + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=False, columnwise=True) + + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert len(calls) == 2 + expected_columnwise_source = torch.cat( + [result.dequantize(dtype=tensor.dtype) for result in calls[0][1]], dim=0 + ) + torch.testing.assert_close(calls[1][0], expected_columnwise_source, rtol=0.0, atol=0.0) + assert calls[1][0] is not tensor + assert all(storage.rowwise_sub_storage is None for storage in out) + assert all(storage.columnwise_sub_storage is not None for storage in out) + + @requires_nvfp4 + @pytest.mark.parametrize("m_splits", ([128, 128], [0, 128])) + def test_nvfp4_rowwise_dequantized_preserves_source_dtype(self, monkeypatch, m_splits): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + input_dtypes = [] + + def tracked_split_quantize(tensor, splits, quantizers, **kwargs): + input_dtypes.append(tensor.dtype) + return real_split_quantize(tensor, splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn( + sum(m_splits), + 128, + dtype=torch.bfloat16, + device="cuda", + ) + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in m_splits + ] + + out = grouped_linear._split_quantize_hybrid(tensor, m_splits, quantizers) + + assert input_dtypes == [tensor.dtype, tensor.dtype] + assert len(out) == len(m_splits) + + def test_rowwise_only_skips_columnwise_quantization(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + calls.append(tensor) + return real_split_quantize(tensor, m_splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=True, columnwise=False) + + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert calls == [tensor] + assert all(storage.rowwise_sub_storage is not None for storage in out) + assert all(storage.columnwise_sub_storage is None for storage in out) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_original_source_preserves_two_bulk_call_fast_path(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + calls.append(tensor) + return real_split_quantize(tensor, m_splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="original", + ) + for _ in range(2) + ] + + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert calls == [tensor, tensor] + assert all(storage.rowwise_sub_storage is not None for storage in out) + assert all(storage.columnwise_sub_storage is not None for storage in out) + + def test_validation_rejects_mixed_columnwise_source_policies(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=source, + ) + for source in ("original", "rowwise_dequantized") + ] + + with pytest.raises(ValueError, match="mixed columnwise source policies"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_validation_rejects_same_family_config_mismatch(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [_make_fp8_quantizer(), _make_fp8_quantizer()] + quantizers[1].force_pow_2_scales = True + + with pytest.raises( + ValueError, + match="incompatible plain backend configurations", + ): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_validation_runs_only_with_quantizer_generation(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + def make_qfactory(columnwise_source): + def qfactory(_role): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=columnwise_source, + ) + + return qfactory + + model = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16).cuda() + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([64, 64], dtype=torch.int64) + original_recipe = recipe.CustomRecipe(qfactory=make_qfactory("original")) + + real_validate = grouped_linear._validate_grouped_quantizer_list + validation_calls = [] + + def tracked_validate(quantizers, *, operand_name="operand"): + validation_calls.append((operand_name, id(quantizers[0]))) + return real_validate(quantizers, operand_name=operand_name) + + monkeypatch.setattr( + grouped_linear, + "_validate_grouped_quantizer_list", + tracked_validate, + ) + + with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): + model(tensor, m_splits) + first_call_count = len(validation_calls) + first_generation = model._validated_quantizer_generations["scaling_fwd"] + assert first_call_count > 0 + + with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): + model(tensor, m_splits) + assert len(validation_calls) == first_call_count + assert model._validated_quantizer_generations["scaling_fwd"] is first_generation + + rebuilt_recipe = recipe.CustomRecipe(qfactory=make_qfactory("rowwise_dequantized")) + with torch.no_grad(), autocast(enabled=True, recipe=rebuilt_recipe): + model(tensor, m_splits) + rebuilt_generation = model._validated_quantizer_generations["scaling_fwd"] + assert len(validation_calls) > first_call_count + assert rebuilt_generation is not first_generation + assert rebuilt_generation[0].columnwise_source == "rowwise_dequantized" + + input_count = 0 + + def mixed_source_qfactory(role): + nonlocal input_count + source = "original" + if role is not None and role.tensor_type == "input": + source = "original" if input_count == 0 else "rowwise_dequantized" + input_count += 1 + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=source, + ) + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_source_qfactory) + # A failed generation is never marked validated. Base metadata can then + # early-return on retry, so the O(1) guard must validate it again. + for _ in range(2): + with pytest.raises(ValueError, match="mixed columnwise source policies"): + with torch.no_grad(), autocast(enabled=True, recipe=mixed_recipe): + model(tensor, m_splits) + assert model._validated_quantizer_generations["scaling_fwd"] is rebuilt_generation + + # Stale invalid recipe metadata must not affect the non-quantized path. + with torch.no_grad(): + model(tensor, m_splits) + + @requires_fp8_and_nvfp4 + def test_hybrid_split_quantize_honors_rowwise_dequantized_source(self): + """NVFP4 column data must derive from the actual grouped row result.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_hybrid, + ) + + torch.manual_seed(3598) + # NVFP4 grouped split-quantize requires each M split to be a multiple + # of 64. + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + + def make_quantizer(): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + + quantizers = [make_quantizer(), make_quantizer()] + actual = _split_quantize_hybrid( + tensor, + [64, 64], + quantizers, + ) + + for index, (actual_part, quantizer) in enumerate(zip(actual, quantizers)): + expected_columnwise = quantizer.columnwise_quantizer.quantize( + actual_part.rowwise_sub_storage.dequantize() + ) + _assert_storage_data_exact( + actual_part.columnwise_sub_storage, + expected_columnwise, + context=f"GroupedLinear split {index} columnwise provenance", + ) + + +# =========================================================================== +# Quantized Parameters (quantized_model_init) tests for hybrid quantization +# =========================================================================== + +# --------------------------------------------------------------------------- +# 1. quantized_model_init: model creation and parameter type verification +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridQuantizedModelInit: + """Verify that quantized_model_init with a hybrid CustomRecipe produces + HybridQuantizedTensor parameters.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_linear_weight_is_hybrid_quantized_tensor(self): + """model.weight should be a HybridQuantizedTensor after quantized_model_init.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance( + weight, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + assert isinstance( + weight, QuantizedTensor + ), "HybridQuantizedTensor should be a QuantizedTensor subclass" + + def test_linear_weight_has_both_sub_storages(self): + """Quantized param should have rowwise and columnwise sub-storages.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert weight.rowwise_sub_storage is not None, "Missing rowwise sub-storage" + assert weight.columnwise_sub_storage is not None, "Missing columnwise sub-storage" + + def test_linear_weight_shape_preserved(self): + """Quantized param should retain its logical shape.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 256, params_dtype=torch.bfloat16).cuda() + + assert model.weight.shape == torch.Size([256, 128]) + + def test_linear_bias_stays_bf16(self): + """Bias should remain BF16 (not quantized).""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, bias=True, params_dtype=torch.bfloat16).cuda() + + assert not isinstance(model.bias, QuantizedTensor), "Bias should not be a QuantizedTensor" + assert model.bias.dtype == torch.bfloat16 + + def test_layernorm_linear_weight_is_hybrid(self): + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = LayerNormLinear(128, 128, params_dtype=torch.bfloat16).cuda() + + assert isinstance(model.weight, HybridQuantizedTensor) + + def test_dequantize_close_to_original(self): + """Dequantized hybrid param should be close to the BF16 init values.""" + hybrid_recipe = self._hybrid_fp8_recipe() + + # Create a non-quantized reference + torch.manual_seed(42) + ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + ref_weight = ref_model.weight.detach().clone() + + # Create quantized model with the same seed + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + dq_weight = model.weight.dequantize() + torch.testing.assert_close(dq_weight.float(), ref_weight.float(), rtol=0.125, atol=0.1) + + def test_preserve_high_precision_init_val(self): + """preserve_high_precision_init_val should store original BF16 on CPU.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init( + enabled=True, + recipe=hybrid_recipe, + preserve_high_precision_init_val=True, + ): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor) + assert hasattr(weight, "get_high_precision_init_val") + hp_val = weight.get_high_precision_init_val() + assert hp_val is not None, "High-precision init val should be stored" + assert hp_val.device.type == "cpu" + assert hp_val.shape == weight.shape + + +# --------------------------------------------------------------------------- +# 2. get_weight_workspace cache invalidation for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridWeightWorkspaceCache: + """Test that get_weight_workspace handles HybridQuantizedTensorStorage + correctly for the quantized-params early-return path and the BF16 cache path.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_quantized_param_skips_workspace(self): + """When weight is already a HybridQuantizedTensor (quantized params), + get_weight_workspace should return it directly without creating a workspace.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp, is_first_microbatch=True) + + assert out.shape == (32, 128) + assert "weight" not in model._fp8_workspaces + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_bf16_weight_creates_hybrid_workspace(self): + """When weight is BF16 and recipe produces HybridQuantizer, the workspace + should be a HybridQuantizedTensor.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp, is_first_microbatch=True) + + assert out.shape == (32, 128) + workspace = model._fp8_workspaces.get("weight") + assert isinstance(workspace, HybridQuantizedTensorStorage) + assert workspace.rowwise_sub_storage is not None + assert workspace.columnwise_sub_storage is not None + + def test_workspace_cache_reuse_across_microbatches(self): + """Cached hybrid workspace should be reused on 2nd+ microbatches.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with autocast(enabled=True, recipe=hybrid_recipe): + with torch.no_grad(): + out1 = model(inp, is_first_microbatch=True) + workspace = model._fp8_workspaces["weight"] + buffers = _as_data_tensor_tuple(workspace) + out2 = model(inp, is_first_microbatch=False) + + assert isinstance(workspace, HybridQuantizedTensorStorage) + assert model._fp8_workspaces["weight"] is workspace + current_buffers = _as_data_tensor_tuple(workspace) + assert len(current_buffers) == len(buffers) + assert all(current is cached for current, cached in zip(current_buffers, buffers)) + torch.testing.assert_close(out1, out2, rtol=0.0, atol=0.0) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_workspace_cache_invalidation_on_usage_change(self): + """If usage requirements change (e.g. inference→training), the cache + should be invalidated and a fresh workspace created.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + # First pass: inference (no columnwise needed) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + model(inp, is_first_microbatch=True) + inference_workspace = model._fp8_workspaces["weight"] + assert isinstance(inference_workspace, HybridQuantizedTensorStorage) + assert inference_workspace.rowwise_sub_storage is not None + assert inference_workspace.columnwise_sub_storage is None + + # Second pass: training (columnwise now needed for backward) + with autocast(enabled=True, recipe=hybrid_recipe): + out_train = model(inp, is_first_microbatch=True) + training_workspace = model._fp8_workspaces["weight"] + + assert isinstance(training_workspace, HybridQuantizedTensorStorage) + assert training_workspace is not inference_workspace + assert training_workspace.rowwise_sub_storage is not None + assert training_workspace.columnwise_sub_storage is not None + + loss = out_train.float().sum() + loss.backward() + + assert inp.grad is not None + + +# --------------------------------------------------------------------------- +# 3. _update_weight_quantizers for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridUpdateWeightQuantizers: + """Test that quantizer refresh propagates correctly to hybrid sub-quantizers.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_quantized_param_survives_multiple_forward_passes(self): + """Weight should remain a HybridQuantizedTensor across multiple forward passes, + each of which triggers init_fp8_metadata → potential quantizer updates.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + for i in range(3): + inp_i = inp.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp_i) + out.float().sum().backward() + assert not torch.isnan(out).any(), f"NaN at iteration {i}" + assert inp_i.grad is not None, f"No input grad at iteration {i}" + + assert isinstance( + model.weight, HybridQuantizedTensor + ), "Weight lost HybridQuantizedTensor type after multiple passes" + + +# --------------------------------------------------------------------------- +# 3b. quantize_master_weights + post_all_gather_processing for hybrid params +# +# Covers the supported (same-format) cases and the rejected (cross-format, +# missing sub-storage, unsupported sub-quantizer) cases. The supported subset +# is the first incremental hybrid integration with the distributed-optimizer +# quantized-param all-gather flow. Cross-format support is deferred to +# follow-up #3158; the tests below pin the NotImplementedError contract so the +# rejection messaging stays clear as the feature evolves. +# --------------------------------------------------------------------------- + + +def _ensure_single_rank_dp_group(): + """Return a single-rank NCCL process group for hybrid quantize_master_weights + tests. Mirrors the local-pytest setup in + `tests/pytorch/distributed/test_cast_master_weights_to_fp8.py` so we can call + `torch.distributed.all_reduce` against a trivial group from inside the + per-format helpers. The group is created lazily on first call and reused + across tests within the same pytest process. + """ + # pylint: disable=import-outside-toplevel + import tempfile + import pathlib + + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _hybrid_recipe_fp8_current(): + """Same-format Float8CurrentScaling on both directions (supported).""" + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + ) + + +def _make_delayed_quantizer(fp8_dtype=None, *, rowwise=True, columnwise=True): + """Construct a ``Float8Quantizer`` (delayed scaling) with locally-allocated + scale/amax buffers for single-shot unit tests. + + The full delayed-scaling lifecycle (``FP8GlobalStateManager`` updating + ``amax_history`` -> ``scale`` across iterations) is out of scope here; for + ``quantize_master_weights`` we only need the helper to read/write + ``quantizer.amax`` / ``quantizer.scale`` / ``model_weight._scale_inv``, + which works with any pair of 1-element float32 tensors. Initial scale=1.0 + and amax=0.0 mirror the cold-start state ``FP8GlobalStateManager`` would + initialize for the first iteration. + """ + if fp8_dtype is None: + fp8_dtype = tex.DType.kFloat8E4M3 + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _hybrid_recipe_fp8_delayed(): + """Same-format Float8 delayed scaling on both directions (supported).""" + return _hybrid_custom_recipe( + row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_fp8_delayed_row_current_col(): + """Cross-format per-tensor Float8: delayed rowwise + current columnwise. + + Routed per-direction: row sub-storage -> delayed bucket, col sub-storage + -> current bucket. The two helpers run independently (no shared state), + so each direction's scale is computed via its own scaling lifecycle. + """ + return _hybrid_custom_recipe( + row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + # grad_factory matches the columnwise direction so the wgrad GEMM's + # grad_output sub-quantizer pairs with the input/weight col format. + grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + ) + + +def _hybrid_recipe_fp8_current_row_delayed_col(): + """Cross-format per-tensor Float8: current rowwise + delayed columnwise. + + Reversed variant of ``_hybrid_recipe_fp8_delayed_row_current_col``: row + sub-storage -> current bucket, col sub-storage -> delayed bucket. + """ + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_mxfp8(): + """Same-format MXFP8 on both directions (rejected today; TODO #3158).""" + return _hybrid_custom_recipe( + row_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_blockwise(): + """Same-format Float8Blockwise on both directions (rejected today; TODO #3158).""" + return _hybrid_custom_recipe( + row_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + col_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + grad_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True + ), + ) + + +def _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe): + """Build a `HybridQuantizedTensor` weight via `quantized_model_init`. + + Returns (weight, fp32_high_precision_init_val) where the high-precision + init val is on GPU so we can use it as the "master" weight in + quantize_master_weights tests. + """ + torch.manual_seed(42) + with quantized_model_init( + enabled=True, + recipe=hybrid_recipe, + preserve_high_precision_init_val=True, + ): + model = Linear(in_features, out_features, bias=False, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance( + weight, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + hp_init_cpu = weight.get_high_precision_init_val() + assert hp_init_cpu is not None, "preserve_high_precision_init_val should populate the cpu val" + hp_init = hp_init_cpu.to(weight.device).float() + return weight, hp_init + + +def _hybrid_param_for(out_features, in_features, hybrid_recipe): + """Same as `_build_hybrid_linear_weight` but discards the init val.""" + weight, _ = _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe) + return weight + + +@requires_fp8 +class TestHybridQuantizeMasterWeights: + """`quantize_master_weights` + `post_all_gather_processing` for hybrid params. + + Dispatch is per-direction: each sub-storage is routed independently into the + per-format bucket matching its own sub-quantizer type. Currently-supported + sub-quantizer types can mix freely across directions (e.g. Float8 delayed + row + Float8 current col), single-direction hybrid (one sub-storage dropped + via ``update_usage``) routes the live direction(s) only; per-block sub- + quantizers (MXFP8, NVFP4, Float8Blockwise) raise NotImplementedError + regardless of which direction they appear in. + + Supported subset (per-tensor Float8) -- positive tests verify the present + sub-storage(s) dequantize close to the master weight after the cast: + + * Float8CurrentScaling on both directions (same-format, full master) + * Float8 delayed scaling on both directions (same-format) + * Float8 delayed row + Float8 current col (cross-format; row -> delayed + bucket, col -> current bucket) + * Float8 current row + Float8 delayed col (cross-format, reversed) + * Single-direction (rowwise-only) hybrid via ``update_usage`` + * Single-direction (columnwise-only) hybrid via ``update_usage`` + + Rejected subset (NotImplementedError / ValueError) -- negative tests pin + the per-direction rejection contract and the both-None guardrail: + * MXFP8 as a hybrid sub-quantizer (rowwise OR columnwise) + * NVFP4 as a hybrid sub-quantizer (rowwise OR columnwise) + * Float8Blockwise as a hybrid sub-quantizer + * A live ``rowwise_dequantized`` column (deferred to #3158) + * A partial ``original``-source master without a two-direction FSDP shard + * Both sub-storages dropped (caller bug: nothing left to cast) + """ + + @staticmethod + def _make_transpose_only_float8_weight(shape, quantizer, *, fill_value=173): + """Build a Hopper/L40-style columnwise-only Float8Tensor. + + Blackwell keeps ``_data`` populated for columnwise-only Float8, so these + tests synthesize the older architecture layout directly: logical + ``[M, K]`` shape with the only live FP8 bytes in ``_transpose[K, M]``. + """ + rows, cols = shape + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=torch.full((cols, rows), fill_value, dtype=torch.uint8, device="cuda"), + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + + @staticmethod + def _scatter_expected_logical_bytes(initial_transpose, fp8_bytes, logical_shape, start_offset): + rows, cols = logical_shape + expected = initial_transpose.clone() + expected_2d = expected.reshape(cols, rows) + remaining = fp8_bytes.numel() + logical_offset = start_offset + src_offset = 0 + while remaining > 0: + row = logical_offset // cols + col = logical_offset % cols + n = min(remaining, cols - col) + expected_2d[col : col + n, row].copy_(fp8_bytes[src_offset : src_offset + n]) + logical_offset += n + src_offset += n + remaining -= n + return expected + + @staticmethod + def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): + quantizer = Float8Quantizer( + scale=scale, + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + raw = torch.empty((1, master.numel()), dtype=torch.uint8, device="cuda") + temp = quantizer.create_tensor_from_data(raw, dtype) + quantizer.update_quantized(master.reshape(1, -1), temp) + return temp._data.reshape(-1) + + @pytest.mark.parametrize("partial_master", (False, True)) + def test_rowwise_dequantized_master_update_raises_before_mutation( + self, + partial_master, + ): + """Defer row-to-column sequencing to the #3158 follow-up.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + quantizer = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + weight = quantizer(source) + state_before = { + "rowwise": _snapshot_storage_tensor_metadata( + weight._rowwise_storage, + clone=True, + ), + "columnwise": _snapshot_storage_tensor_metadata( + weight._columnwise_storage, + clone=True, + ), + } + start_offset = shape[-1] if partial_master else 0 + master = torch.randn(weight.numel(), dtype=torch.float32, device="cuda") + master = master[start_offset:].contiguous() + + with pytest.raises( + NotImplementedError, + match="rowwise update/all-gather.*#3158", + ): + quantize_master_weights( + [weight], + [master], + [start_offset], + group=group, + ) + + state_after = { + "rowwise": _snapshot_storage_tensor_metadata( + weight._rowwise_storage, + clone=True, + ), + "columnwise": _snapshot_storage_tensor_metadata( + weight._columnwise_storage, + clone=True, + ), + } + _assert_nested_state_exact(state_after, state_before) + + def test_fp8_current_original_partial_master_raises(self): + """Reject an independently sourced Hybrid column from a partial master shard. + + A one-payload distributed optimizer only communicates the rowwise payload. + Re-creating an independently quantized column from a partial master shard can + therefore produce a different value after checkpoint resume. Construct the + Hopper transpose-only column explicitly so this safety regression does not + depend on columnwise-only kernel support. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + shape = (4, 8) + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + quantizer = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + columnwise_source="original", + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise = row_quantizer(source) + columnwise = self._make_transpose_only_float8_weight(shape, col_quantizer) + weight = HybridQuantizedTensor( + shape=shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + master_shard = source.reshape(-1)[shape[-1] :].float().contiguous() + + with pytest.raises( + ValueError, + match="partial master shard.*columnwise_source='original'.*full-master data", + ): + quantize_master_weights( + [weight], + [master_shard], + [shape[-1]], + group=None, + ) + + @pytest.mark.skipif( + is_non_tn_fp8_gemm_supported(), + reason="Hopper-only: Blackwell supports columnwise per-tensor FP8 FSDP updates", + ) + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_per_tensor_fp8_fsdp_transpose_only_raises_before_mutation(self, scaling): + """Reject Hopper FSDP updates that cannot flatten columnwise-only storage.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + shape = (4, 8) + shard_shape = (2, 8) + if scaling == "current": + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + else: + row_quantizer = _make_delayed_quantizer(rowwise=True, columnwise=False) + col_quantizer = _make_delayed_quantizer(rowwise=False, columnwise=True) + quantizer = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + shard_source = source[: shard_shape[0]].contiguous() + weight = HybridQuantizedTensor( + shape=shape, + dtype=source.dtype, + rowwise_storage=row_quantizer(source), + columnwise_storage=self._make_transpose_only_float8_weight(shape, col_quantizer), + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + shard = HybridQuantizedTensor( + shape=shard_shape, + dtype=source.dtype, + rowwise_storage=row_quantizer(shard_source), + columnwise_storage=self._make_transpose_only_float8_weight(shard_shape, col_quantizer), + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + weight_col = weight._columnwise_storage + shard_col = shard._columnwise_storage + weight_scale_before = weight_col._scale_inv.clone() + weight_transpose_before = weight_col._transpose.clone() + shard_scale_before = shard_col._scale_inv.clone() + shard_transpose_before = shard_col._transpose.clone() + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantize_master_weights( + [weight], + [shard_source.float()], + [0], + group=None, + fsdp_shard_model_weights=[shard], + ) + + assert weight_col._data is None + assert shard_col._data is None + torch.testing.assert_close(weight_col._scale_inv, weight_scale_before, rtol=0, atol=0) + torch.testing.assert_close(shard_col._scale_inv, shard_scale_before, rtol=0, atol=0) + assert torch.equal(weight_col._transpose, weight_transpose_before) + assert torch.equal(shard_col._transpose, shard_transpose_before) + + @staticmethod + def _logical_float8_bytes(storage): + """Return FP8 payload in the tensor's logical row-major order.""" + if storage._data is not None: + return storage._data.reshape(-1) + assert storage._transpose is not None + return storage._transpose.transpose(-2, -1).contiguous().reshape(-1) + + def test_fp8_current_transpose_only_nonzero_offset(self): + """Current-scaling distopt update handles Hopper-style columnwise storage. + + Regression for ``model_weight.reshape(-1)`` reaching + ``Float8Tensor._ReshapeFunc.forward`` and dereferencing ``_data=None``. + The nonzero offset spans multiple logical rows, so the test also checks + row-major shard bytes are scattered into transposed storage correctly. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + start_offset = 5 + master = torch.linspace(-2.0, 2.0, steps=17, dtype=torch.float32, device="cuda") + quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", rowwise=False, columnwise=True + ) + weight = self._make_transpose_only_float8_weight(shape, quantizer) + initial = weight._transpose.clone() + + quantize_master_weights([weight], [master], [start_offset], group=group) + + scale = torch.reciprocal(weight._scale_inv.detach().clone()) + fp8_bytes = self._reference_fp8_bytes(master.to(weight.dtype), scale, weight.dtype) + expected = self._scatter_expected_logical_bytes(initial, fp8_bytes, shape, start_offset) + assert weight._data is None + assert weight._transpose_invalid is False + assert torch.equal(weight._transpose, expected) + + def test_fp8_delayed_transpose_only_nonzero_offset(self): + """Delayed-scaling distopt update handles Hopper-style columnwise storage. + + Regression for the direct ``model_weight._data.view(-1)`` path in the + delayed-scaling helper. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + start_offset = 6 + master = torch.linspace(-3.0, 1.0, steps=15, dtype=torch.float32, device="cuda") + quantizer = _make_delayed_quantizer(tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=False, columnwise=True) + weight = self._make_transpose_only_float8_weight(shape, quantizer) + initial = weight._transpose.clone() + + quantize_master_weights([weight], [master], [start_offset], group=group) + + fp8_bytes = self._reference_fp8_bytes( + master.to(weight.dtype), weight._get_quantizer().scale, weight.dtype + ) + expected = self._scatter_expected_logical_bytes(initial, fp8_bytes, shape, start_offset) + assert weight._data is None + assert weight._transpose_invalid is False + assert torch.equal(weight._transpose, expected) + + # ---------- Positive tests (same-format) ---------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_current_same_format_full_master(self): + """Full master (start_offset=0) routes both sub-storages through the + existing per-format current-scaling helper. Verifies both directions + dequantize close to the master weight after the cast. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + # Distributed-optimizer convention: master weight is the flat FP32 shard + # owned by the current rank (or the full param for non-distributed cases). + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + master_bf16 = master_flat.to(weight.dtype).reshape(weight.shape) + expected_row = weight._quantizer.rowwise_quantizer.copy().quantize(master_bf16) + expected_column = weight._quantizer.columnwise_quantizer.copy().quantize(master_bf16) + _assert_storage_data_exact( + weight._rowwise_storage, + expected_row, + context="full-master rowwise independent quantization", + ) + _assert_storage_data_exact( + weight._columnwise_storage, + expected_column, + context="full-master columnwise independent quantization", + ) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + # FP8 E4M3 round-trip; matches the loose tolerance the equivalent + # native-FP8-current test uses (e.g. test_dequantize_close_to_original). + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_delayed_same_format_full_master(self): + """Same-format delayed scaling on both directions. Both sub-storages + route into the delayed-scaling bucket as independent entries; the + helper processes them with a single bucket-wide amax all-reduce. + Verifies each direction dequantizes close to the master weight. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_delayed() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_delayed_row_current_col_full_master(self): + """Cross-format per-tensor Float8: delayed row + current col. + + Pins the new per-direction routing: row sub-storage goes to the + delayed bucket, col sub-storage goes to the current bucket. Each + helper runs independently on its single-entry bucket, with no + cross-pollination between the two scaling lifecycles. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_delayed_row_current_col() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + assert isinstance(weight._quantizer.rowwise_quantizer, Float8Quantizer) + assert isinstance(weight._quantizer.columnwise_quantizer, Float8CurrentScalingQuantizer) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_current_row_delayed_col_full_master(self): + """Cross-format per-tensor Float8: current row + delayed col. + + Reversed variant of the test above — pins that the per-direction + loop's second iteration (col) reaches the delayed dispatch arm + independently of what the rowwise iteration did. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current_row_delayed_col() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + assert isinstance(weight._quantizer.rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(weight._quantizer.columnwise_quantizer, Float8Quantizer) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + # NOTE: Per-block sub-quantizers (MXFP8, NVFP4, Float8Blockwise) are not + # supported as hybrid sub-quantizers by this initial integration, regardless + # of which direction they appear in. See the per-direction rejection tests + # below (``test_mxfp8_*_raises`` covers both rowwise and columnwise rejection + # of MXFP8; ``test_nvfp4_*_raises`` and ``test_blockwise_*_raises`` similarly). + # The TODO #3158 block above ``_route_hybrid_to_buckets`` in tensor/utils.py + # documents the upstream constraints (single-direction cast helper / kernel + # support) whose unblocker drops per-block format support in for free. + + # ---------- Negative tests (per-direction rejection contract) ---------- + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_mxfp8_rowwise_raises(self): + """MXFP8 in the rowwise sub-quantizer is rejected per-direction. + + ``_cast_master_weights_to_fp8_mxfp8_scaling`` assumes each entry's + ``model_weight`` has BOTH ``_rowwise_*`` and ``_columnwise_*`` populated + (the underlying partial-cast kernel is bidirectional), while a hybrid + sub-storage is single-direction by construction. See + ``TODO(#3158, hybrid-mxfp8-distopt)`` in tensor/utils.py for the unblocker shape. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_mxfp8() + # Shape must be a multiple of MXFP8 block size (32) on both axes. + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="MXFP8Quantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_mxfp8_columnwise_raises(self): + """MXFP8 in the columnwise sub-quantizer is rejected per-direction. + + Pairs FP8 current scaling in the rowwise slot (supported) with MXFP8 + in the columnwise slot (rejected). The rowwise iteration of + ``_route_hybrid_to_buckets`` routes the FP8 sub-storage into the + current-scaling bucket cleanly; the columnwise iteration then hits + MXFP8 and raises. Pins that per-direction dispatch visits and rejects + the columnwise sub-quantizer too — not just the rowwise one. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + # Shape must be a multiple of MXFP8 block size (32) on both axes. + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="MXFP8Quantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_nvfp4_rowwise_raises(self): + """NVFP4 in the rowwise sub-quantizer is rejected per-direction. + + The NVFP4 cast path is blocked on a pair of upstream constraints + documented in the TODO #3158 block above ``_route_hybrid_to_buckets`` in + tensor/utils.py. + + NOTE: after PR #3027, single-direction 2D NVFP4 construction works, + so this test now reaches the intended ``quantize_master_weights`` + rejection while using the base weight scaling mode. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + col_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + grad_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + ), + ) + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="NVFP4Quantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_nvfp4_columnwise_raises(self): + """NVFP4 in only the columnwise slot is rejected per-direction.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="NVFP4Quantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", + ) + def test_blockwise_rowwise_raises(self): + """Float8BlockQuantizer in the rowwise sub-quantizer is rejected + per-direction (no e2e factory uses it; TODO #3158 marker in tensor/utils.py). + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_blockwise() + weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="Float8BlockQuantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", + ) + def test_blockwise_columnwise_raises(self): + """Float8BlockQuantizer in only the columnwise slot is rejected.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="Float8BlockQuantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_rowwise_only_fp8_current_full_master(self): + """Single-direction hybrid: columnwise dropped via update_usage. + + Pins that the per-direction loop in `_route_hybrid_to_buckets` skips + the dropped direction silently and routes only the present (rowwise) + sub-storage. Useful for inference / memory-saving paths that + deliberately keep only the fprop-side direction. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=True, columnwise_usage=False) + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is None + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + # Columnwise stays dropped (the cast must not silently revive it). + assert weight._columnwise_storage is None + expected_row = weight._quantizer.rowwise_quantizer.copy().quantize( + master_flat.to(weight.dtype).reshape(weight.shape) + ) + _assert_storage_data_exact( + weight._rowwise_storage, + expected_row, + context="rowwise-only full-master independent quantization", + ) + # Rowwise is populated and dequantizes close to the master. + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_columnwise_only_fp8_current_full_master(self): + """Single-direction hybrid: rowwise dropped via update_usage. + + Reversed variant — verifies the column-only iteration of the per- + direction loop reaches the dispatch and routes correctly. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=False, columnwise_usage=True) + assert weight._rowwise_storage is None + assert weight._columnwise_storage is not None + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + # Rowwise stays dropped (the cast must not silently revive it). + assert weight._rowwise_storage is None + expected_column = weight._quantizer.columnwise_quantizer.copy().quantize( + master_flat.to(weight.dtype).reshape(weight.shape) + ) + _assert_storage_data_exact( + weight._columnwise_storage, + expected_column, + context="columnwise-only full-master independent quantization", + ) + # Columnwise is populated and dequantizes close to the master. + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_both_sub_storages_none_raises(self): + """Both sub-storages dropped via update_usage — nothing left to cast. + + This is the only remaining sub-storage-presence guardrail after the + single-direction enablement: a fully-dropped hybrid weight reaching + `quantize_master_weights` is a caller bug, not a deferred feature, + so we surface it as a ValueError. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=False, columnwise_usage=False) + assert weight._rowwise_storage is None + assert weight._columnwise_storage is None + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(ValueError, match="both rowwise and columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridPostAllGatherProcessing: + """Hybrid branch of `post_all_gather_processing` is exercised indirectly by + the positive `TestHybridQuantizeMasterWeights` tests; the case below pins + an additional invariant that the routing logic must preserve. + """ + + def test_post_ag_idempotent_for_fp8_current_hybrid(self): + """Calling post_all_gather_processing twice on a same-format Float8 + hybrid must not corrupt the sub-storages. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + dq_row_first = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col_first = weight._columnwise_storage.dequantize(dtype=torch.float32) + + post_all_gather_processing([weight]) + dq_row_second = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col_second = weight._columnwise_storage.dequantize(dtype=torch.float32) + + torch.testing.assert_close(dq_row_first, dq_row_second, rtol=0.0, atol=0.0) + torch.testing.assert_close(dq_col_first, dq_col_second, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# 4. Recipe correspondence validation +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridRecipeCorrespondence: + """Test _check_weight_tensor_recipe_correspondence with hybrid params.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_hybrid_param_with_matching_recipe_does_not_raise(self): + """Forward pass with matching recipe should not raise.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + assert not torch.isnan(out).any() + + def test_hybrid_param_with_mismatched_recipe_raises(self): + """Forward pass with a non-CustomRecipe on a hybrid param should raise.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + mismatch_recipe = recipe.Float8CurrentScaling() + with pytest.raises(RuntimeError, match="Recipe mismatch"): + with torch.no_grad(): + with autocast(enabled=True, recipe=mismatch_recipe): + model(inp) + + +# --------------------------------------------------------------------------- +# 5. quantize_ in-place update for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridQuantizeInPlace: + """Test in-place re-quantization (quantize_) for HybridQuantizedTensor. + + This is needed for the optimizer writeback path (param.quantize_(master_weight)) + and the workspace cache update path (out.quantize_(new_bf16_weight)). + """ + + def test_quantize_inplace_updates_data(self): + """quantize_() should re-quantize both sub-storages from new BF16 data.""" + torch.manual_seed(42) + hq = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor = hq.quantize(original) + + # Update with different data + new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + expected = hq.quantize(new_data) + result = tensor.quantize_(new_data) + + assert result is tensor + _assert_hybrid_tensor_exact(tensor, expected, context="quantize_") + + def test_quantize_inplace_preserves_tensor_identity(self): + """quantize_() should update in-place, not create a new tensor.""" + hq = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor = hq.quantize(original) + + new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + result = tensor.quantize_(new_data) + + assert result is tensor, "quantize_() must return the object it updated" + + # noop_flag is a delayed-scaling feature; not tested here since + # delayed scaling is out of scope for hybrid quantization. + + +# --------------------------------------------------------------------------- +# 6. FusedAdam with hybrid quantized params +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridFusedAdam: + """Test FusedAdam optimizer with HybridQuantizedTensor parameters.""" + + def _build_hybrid_model(self): + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_fused_adam_accepts_hybrid_params(self): + """FusedAdam should not crash when given HybridQuantizedTensor params.""" + model, _ = self._build_hybrid_model() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + assert optimizer is not None + + def test_fused_adam_master_weights_track_reference(self): + """FP32 master weights should closely track a reference Adam optimizer. + + Small divergence is expected because HybridQuantizedTensor.float() + may take a slightly different dequantization path than + detach().clone().float() through __torch_dispatch__. + """ + model, _ = self._build_hybrid_model() + + ref_params = [p.detach().clone().float() for p in model.parameters()] + + options = {"lr": 5e-4, "betas": (0.9, 0.999), "eps": 1e-8, "weight_decay": 0} + ref_optim = torch.optim.Adam(ref_params, **options) + tst_optim = te.optimizers.FusedAdam( + list(model.parameters()), + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + **options, + ) + + for _ in range(5): + for p_ref, p in zip(ref_params, model.parameters()): + p_ref.grad = torch.rand_like(p_ref) + p.decoupled_grad = p_ref.grad.clone() + ref_optim.step() + tst_optim.step() + + master_params = [ + tst_optim.get_unscaled_state(p, "master_param") for p in model.parameters() + ] + torch.testing.assert_close(ref_params, master_params, rtol=1e-3, atol=1e-3) + + def test_fused_adam_param_remains_hybrid_after_step(self): + """Weight params should still be HybridQuantizedTensors after optimizer step.""" + model, _ = self._build_hybrid_model() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + ) + + for _ in range(3): + for p in model.parameters(): + p.decoupled_grad = torch.rand_like(p.float()) + optimizer.step() + + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, HybridQuantizedTensor + ), f"{name} lost HybridQuantizedTensor type: {type(p).__name__}" + + def test_fused_adam_requires_master_weights(self): + """FusedAdam without master_weights should raise for hybrid quantized params.""" + model, _ = self._build_hybrid_model() + + with pytest.raises(RuntimeError, match="master_weights"): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=False, + ) + for p in model.parameters(): + p.grad = torch.rand_like(p.float()).to(p.dtype) + optimizer.step() + + +# --------------------------------------------------------------------------- +# 7. End-to-end training loop: fwd + bwd + optimizer step +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridQuantizedParamsEndToEnd: + """Full training loop: quantized_model_init + autocast fwd + bwd + FusedAdam.step().""" + + def _build_model_and_recipe(self): + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_training_loop_loss_decreases(self): + """Loss should decrease over a few training steps.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(7): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + + for name, p in model.named_parameters(): + assert p.grad is not None, f"Step {i}: {name} has no gradient" + assert torch.isfinite(p.grad).all(), f"Step {i}: {name} has non-finite grad" + + optimizer.step() + + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + def test_training_loop_params_remain_quantized(self): + """Params should remain HybridQuantizedTensors after training.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + for _ in range(3): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + optimizer.step() + + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, HybridQuantizedTensor + ), f"{name} is {type(p).__name__}, not HybridQuantizedTensor" + + def test_training_loop_optimizer_states_are_fp32(self): + """Optimizer states should be FP32.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + for _ in range(2): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + output.float().sum().backward() + optimizer.step() + + for name, p in model.named_parameters(): + state = optimizer.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "bias" not in name: + assert state["master_param"].dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# 8. Mixed-format quantized params (e.g. MXFP8 row + NVFP4 col) +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (mxfp8_available and nvfp4_available), + reason=f"MXFP8: {reason_for_no_mxfp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestHybridMixedFormatQuantizedParams: + """Quantized params with genuinely different formats per direction.""" + + def _build_mixed_model(self, in_features=256, out_features=256): + """MXFP8 rowwise (fprop) + role-aware NVFP4 columnwise (bwd).""" + + def qfactory(role): + 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=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + if is_linear and role.tensor_type == "grad_output": + return nvfp4_quantizer_factory(role) + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + + hybrid_recipe = recipe.CustomRecipe(qfactory=qfactory) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_mixed_format_param_creation(self): + """Model init with mixed MXFP8/NVFP4 hybrid should produce a + HybridQuantizedTensor parameter.""" + model, _ = self._build_mixed_model() + assert isinstance(model.weight, HybridQuantizedTensor) + + def test_mixed_format_forward_only(self): + """Forward pass with mixed-format quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + inp = torch.randn(32, 256, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + + assert out.shape == (32, 256) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_mixed_format_forward_backward(self): + """Full fwd+bwd with mixed-format quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + inp = torch.randn(32, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"No gradient for {name}" + assert not torch.isnan(p.grad).any(), f"NaN gradient for {name}" + + def test_mixed_format_training_loop(self): + """End-to-end training loop with mixed-format hybrid quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(5): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance(p, HybridQuantizedTensor), f"{name} is {type(p).__name__}" + + def test_mixed_format_sub_storage_types(self): + """Verify that sub-storages have the correct types (MXFP8 vs NVFP4).""" + model, _ = self._build_mixed_model() + weight = model.weight + from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, + ) + + row = weight.rowwise_sub_storage + col = weight.columnwise_sub_storage + assert isinstance(row, MXFP8TensorStorage) or hasattr( + row, "_rowwise_data" + ), f"Expected MXFP8 rowwise sub-storage, got {type(row).__name__}" + assert isinstance(col, NVFP4TensorStorage) or hasattr( + col, "_rowwise_data" + ), f"Expected NVFP4 columnwise sub-storage, got {type(col).__name__}" + + +# --------------------------------------------------------------------------- +# 9. Quantized params equivalence: vanilla vs hybrid (same format both dirs) +# --------------------------------------------------------------------------- + + +def _hybrid_mxfp8_qfactory(role): + """Hybrid MXFP8 (E4M3 both dirs).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + + +def _hybrid_nvfp4_qfactory(role): + """Hybrid NVFP4 (E2M1 both dirs, base role behavior).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "grad_output": + return nvfp4_quantizer_factory(role) + return HybridQuantizer( + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + + +class _QuantizedParamsEquivalenceBase: + """Base for comparing vanilla vs hybrid quantized params training. + + When the hybrid quantizer uses the same format in both directions, + the full quantized_model_init + training loop should produce + equivalent results to the vanilla (non-hybrid) quantized params path. + """ + + hidden_size = 256 + num_steps = 5 + + def _vanilla_recipe(self): + raise NotImplementedError + + def _hybrid_recipe(self): + raise NotImplementedError + + def _build_models(self): + """Create two models with identical init: one vanilla, one hybrid.""" + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=self._vanilla_recipe()): + model_ref = Linear( + self.hidden_size, + self.hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=self._hybrid_recipe()): + model_hyb = Linear( + self.hidden_size, + self.hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + + return model_ref, model_hyb + + def _run_training_loop(self, model, train_recipe, x, target, num_steps): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + trajectory = [] + hybrid_metadata_trajectory = [] + for step in range(num_steps): + optimizer.zero_grad(set_to_none=True) + step_input = x.detach().clone().requires_grad_(True) + _set_quantization_test_seed(199 + step) + with autocast(enabled=True, recipe=train_recipe): + output = model(step_input) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + gradients = { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + } + pre_step_cuda_rng_state = torch.cuda.get_rng_state() + optimizer.step() + post_step_cuda_rng_state = torch.cuda.get_rng_state() + + hybrid_parameters = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, HybridQuantizedTensor) + ] + expected_storages = {} + torch.cuda.set_rng_state(pre_step_cuda_rng_state) + try: + for name, param in hybrid_parameters: + master = optimizer.get_unscaled_state(param, "master_param") + expected_storages[name] = ( + param._quantizer.rowwise_quantizer.copy().quantize(master), + param._quantizer.columnwise_quantizer.copy().quantize(master), + ) + finally: + # The oracle must be observational only. Restore the state left + # by the real optimizer step after replaying NVFP4 stochastic + # rounding from its pre-step RNG state. + torch.cuda.set_rng_state(post_step_cuda_rng_state) + + hybrid_metadata = {} + for name, param in hybrid_parameters: + expected_row, expected_column = expected_storages[name] + _assert_storage_data_exact( + param.rowwise_sub_storage, + expected_row, + context=f"step {step} {name} rowwise writeback", + ) + _assert_storage_data_exact( + param.columnwise_sub_storage, + expected_column, + context=f"step {step} {name} columnwise writeback", + ) + hybrid_metadata[name] = { + "rowwise": _snapshot_storage_tensor_metadata( + param.rowwise_sub_storage, clone=True + ), + "columnwise": _snapshot_storage_tensor_metadata( + param.columnwise_sub_storage, clone=True + ), + } + hybrid_metadata_trajectory.append(hybrid_metadata) + logical_parameters = { + name: ( + param.dequantize(dtype=torch.float32).detach().clone() + if isinstance(param, QuantizedTensor) + else param.detach().float().clone() + ) + for name, param in model.named_parameters() + } + trajectory.append( + { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "input_gradient": step_input.grad.detach().clone(), + "parameter_gradients": gradients, + "logical_parameters": logical_parameters, + "optimizer": _clone_nested_state(optimizer.state_dict()), + } + ) + return trajectory, hybrid_metadata_trajectory + + def _test_equivalence(self): + model_ref, model_hyb = self._build_models() + + torch.manual_seed(99) + x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + ref_trajectory, ref_hybrid_metadata = self._run_training_loop( + model_ref, + self._vanilla_recipe(), + x, + target, + self.num_steps, + ) + hybrid_trajectory, hybrid_metadata_trajectory = self._run_training_loop( + model_hyb, + self._hybrid_recipe(), + x, + target, + self.num_steps, + ) + _assert_nested_state_exact( + hybrid_trajectory, + ref_trajectory, + path="same-format quantized-parameter trajectory", + ) + assert all(not step_metadata for step_metadata in ref_hybrid_metadata) + assert len(hybrid_metadata_trajectory) == self.num_steps + for step_metadata in hybrid_metadata_trajectory: + assert step_metadata.keys() == {"weight"} + assert step_metadata["weight"]["rowwise"] is not None + assert step_metadata["weight"]["columnwise"] is not None + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestQuantizedParamsEquivalenceFP8CurrentScaling(_QuantizedParamsEquivalenceBase): + """Vanilla versus same-format hybrid FP8-current training parity.""" + + def _vanilla_recipe(self): + return recipe.Float8CurrentScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_fp8_current_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestQuantizedParamsEquivalenceMXFP8(_QuantizedParamsEquivalenceBase): + """Vanilla MXFP8BlockScaling vs hybrid MXFP8 (same format both dirs).""" + + def _vanilla_recipe(self): + return recipe.MXFP8BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +class TestQuantizedParamsEquivalenceBlockFP8(_QuantizedParamsEquivalenceBase): + """Vanilla Float8BlockScaling vs hybrid block FP8 (same format both dirs).""" + + def _vanilla_recipe(self): + return recipe.Float8BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestQuantizedParamsEquivalenceNVFP4(_QuantizedParamsEquivalenceBase): + """Vanilla NVFP4BlockScaling vs same-format hybrid NVFP4.""" + + def _vanilla_recipe(self): + return recipe.NVFP4BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +# --------------------------------------------------------------------------- +# 10. State dict save/load (checkpointing) for hybrid quantized params +# --------------------------------------------------------------------------- + + +# Module-level qfactories give TE-to-TE quantized-param checkpoints a stable +# importable reference for any pickled quantizer/recipe metadata. Portable BF16 +# checkpoint loading should not depend on importing these factories. + + +@requires_mxfp8_and_nvfp4 +class TestHybridCheckpoint: + """Test state_dict save/load round-trips for models with hybrid quantized params.""" + + def _hybrid_checkpoint_recipe(self): + return recipe.CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) + + def test_state_dict_save_load_roundtrip(self): + """state_dict → save → load → same model should produce identical outputs.""" + torch.manual_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_before = model(inp) + + state_dict = model.state_dict() + + # Create a fresh model and load + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + model2.load_state_dict(state_dict) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_after = model2(inp) + + torch.testing.assert_close(out_before, out_after, rtol=0.0, atol=0.0) + + def test_state_dict_contains_weight(self): + """state_dict should contain the weight key.""" + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + sd = model.state_dict() + assert "weight" in sd, f"state_dict keys: {list(sd.keys())}" + + def test_load_bf16_state_dict_into_hybrid_model(self): + """Loading a BF16 state_dict into a hybrid quantized model should work. + + This is the common scenario: pretrained BF16 weights loaded into a + model initialized with quantized_model_init. + """ + torch.manual_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + + # Create BF16 model and get its state_dict + ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + bf16_state_dict = ref_model.state_dict() + + # Create hybrid quantized model + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + # Load BF16 weights into hybrid model + model.load_state_dict(bf16_state_dict) + + # Verify model produces valid output + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_state_dict_torch_save_load(self): + """Full round-trip through torch.save/torch.load (file-based).""" + import tempfile + import os + + torch.manual_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_before = model(inp) + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f: + torch.save(model.state_dict(), f.name) + tmp_path = f.name + + try: + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + state_dict = torch.load(tmp_path, weights_only=False) + model2.load_state_dict(state_dict) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_after = model2(inp) + + torch.testing.assert_close(out_before, out_after, rtol=0.0, atol=0.0) + finally: + os.unlink(tmp_path) + + @staticmethod + def _checkpoint_training_step(model, optimizer, x, target, hybrid_recipe): + optimizer.zero_grad(set_to_none=True) + step_input = x.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(step_input) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + gradients = { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + } + optimizer.step() + return { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "input_gradient": step_input.grad.detach().clone(), + "gradients": gradients, + "parameters": _snapshot_model_parameters(model), + "optimizer": _clone_nested_state(optimizer.state_dict()), + } + + def test_checkpoint_resume_training(self): + """Save mid-training, load into new model+optimizer, verify training continues.""" + import os + import tempfile + + _set_quantization_test_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + for _ in range(3): + self._checkpoint_training_step(model, optimizer, x, target, hybrid_recipe) + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as checkpoint_file: + torch.save( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "cpu_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.cuda.get_rng_state(), + }, + checkpoint_file.name, + ) + checkpoint_path = checkpoint_file.name + + try: + uninterrupted = self._checkpoint_training_step( + model, + optimizer, + x, + target, + hybrid_recipe, + ) + + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + resumed_model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + resumed_optimizer = te.optimizers.FusedAdam( + resumed_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + checkpoint = torch.load(checkpoint_path, weights_only=False) + resumed_model.load_state_dict(checkpoint["model"]) + resumed_optimizer.load_state_dict(checkpoint["optimizer"]) + torch.set_rng_state(checkpoint["cpu_rng_state"]) + torch.cuda.set_rng_state(checkpoint["cuda_rng_state"]) + + resumed = self._checkpoint_training_step( + resumed_model, + resumed_optimizer, + x, + target, + hybrid_recipe, + ) + + _assert_nested_state_exact( + resumed, + uninterrupted, + path="checkpoint continuation", + ) + finally: + os.unlink(checkpoint_path) + + +# --------------------------------------------------------------------------- +# 11. Activation recomputation (torch.utils.checkpoint / te.checkpoint) +# --------------------------------------------------------------------------- + + +def _reset_rng(seed: int = 1234): + """Reset deterministic RNG for reproducible forward/backward comparisons. + + Activation recompute relies on RNG equality between the first forward + and the recomputed forward. These tests use dropout-free modules, so + RNG advancement doesn't affect numerics, but we still reset between + runs so the reference (no-recompute) and checkpointed paths see + identical weight init, input, and grad_output seeds. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _collect_outputs(out, inp, model): + """Gather forward output, input grad, and parameter grads into a flat list. + + Mirrors ``test_numerics.py::_test_e2e_*_recompute`` conventions so the + comparison against a non-recomputed baseline is a simple zip. + """ + results = [out.detach().clone()] + if inp.grad is not None: + results.append(inp.grad.detach().clone()) + for _, p in model.named_parameters(): + if p.requires_grad and p.grad is not None: + results.append(p.grad.detach().clone()) + return results + + +def _assert_outputs_bitwise_equal(ref, test, label): + """All stateless same-format hybrid recipes should be bitwise-identical + under activation recompute: same input bytes → same quantized bytes → + same GEMM result. Any drift means the recompute path silently diverged + (e.g. fell back to a different quantization path).""" + assert len(ref) == len(test), f"{label}: output count mismatch" + for i, (r, t) in enumerate(zip(ref, test)): + torch.testing.assert_close( + t, r, rtol=0, atol=0, msg=f"{label}: bitwise mismatch at output {i}" + ) + + +@requires_fp8 +class TestHybridActivationRecompute: + """Activation recomputation around TE modules with a hybrid CustomRecipe. + + Probes the interaction between ``HybridQuantizedTensor`` / + ``HybridQuantizedTensorStorage`` and the three activation-checkpoint + paths in use today: + + * ``te.checkpoint(fn, ..., use_reentrant=True)`` — reentrant path; wraps + ``torch.autograd.Function`` that re-runs the forward under + ``activation_recompute_forward(recompute_phase=True)``. This is the + Megatron-style path. + * ``te.checkpoint(fn, ..., use_reentrant=False)`` — non-reentrant path; + uses ``_checkpoint_hook`` (torch saved-tensors hooks) to discard + saved tensors on the first forward and recompute them on unpack. + * ``torch.utils.checkpoint.checkpoint(fn, ..., use_reentrant=False)`` + — vanilla PyTorch path without TE wrapper. Exercised because users + (and some Megatron configs) invoke it directly around TE modules. + + Failure modes it catches: + + * Silent BF16 fallback during recompute (would break bitwise parity + but pass loose tolerance — hence the bitwise assertion for + same-format stateless recipes). + * ``HybridQuantizedTensorStorage.prepare_for_saving`` / + ``restore_from_saved`` chain losing a sub-storage across the + save-for-backward boundary. + * ``HybridQuantizedTensor`` subclass being stripped by the autograd + engine (would manifest as ``AttributeError`` on the recomputed + tensor). + """ + + in_features = 128 + out_features = 128 + batch = 32 + + # ----- helpers --------------------------------------------------- + + def _same_format_fp8_recipe(self): + """Same-format FP8 current scaling both directions → bitwise-safe + baseline. Matches + :class:`TestHybridGemmBitwiseIdentical` construction so + recompute parity can be asserted bitwise-equal.""" + return _hybrid_custom_recipe( + row_factory=_fp8_row_factory, + col_factory=_fp8_col_factory, + grad_factory=_fp8_grad_factory, + ) + + def _same_format_mxfp8_recipe(self): + """Same-format MXFP8 both directions — stateless, per-block scales + computed from the tensor content; bitwise-stable under recompute.""" + return _hybrid_custom_recipe( + row_factory=_mxfp8_factory, + col_factory=_mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + + def _cross_format_fp8_mxfp8_recipe(self): + """Cross-format FP8 row + MXFP8 col — the canonical hybrid + scenario. Numerical parity is not bitwise because the wgrad GEMM + uses MXFP8 scaling modes on both operands (so grad_output must be + MXFP8 columnwise), pairing differently from the fprop path.""" + return _hybrid_custom_recipe( + row_factory=_fp8_row_factory, + col_factory=_mxfp8_factory, + grad_factory=_mxfp8_factory, + ) + + def _run_linear(self, recipe_obj, *, checkpoint_fn=None): + """Build a fresh Linear, run forward+backward, return collected + outputs. ``checkpoint_fn`` is an optional callable of the form + ``fn(model, inp) -> output`` that wraps the forward in an + activation-checkpoint implementation; ``None`` is the reference + (non-recompute) baseline. + """ + _reset_rng(seed=4242) + model = Linear(self.in_features, self.out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + self.batch, + self.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + inp.retain_grad() + + with autocast(enabled=True, recipe=recipe_obj): + out = checkpoint_fn(model, inp) if checkpoint_fn is not None else model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + def _run_transformer_layer(self, recipe_obj, *, checkpoint_fn=None): + """Small TransformerLayer (no dropout, fuse_qkv) with optional + activation checkpointing around the whole block.""" + _reset_rng(seed=5151) + hidden = 128 + ffn = 128 + nheads = 4 + seq = 8 + bs = 4 + + model = TransformerLayer( + hidden, + ffn, + nheads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + + inp = torch.randn(seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + inp.retain_grad() + + with autocast(enabled=True, recipe=recipe_obj): + out = checkpoint_fn(model, inp) if checkpoint_fn is not None else model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + # ----- te.checkpoint, reentrant --------------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_linear_fp8_bitwise(self): + """te.checkpoint(use_reentrant=True) around te.Linear with + same-format FP8 hybrid → bitwise parity with non-recompute. + + This is the Megatron-style activation-recompute path. Bitwise + parity catches silent BF16 fallback (would pass loose tolerance). + """ + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_reentrant_linear_mxfp8_bitwise(self): + """Same as FP8 but MXFP8 hybrid — per-block scales must recompute + identically. Asserts that the MXFP8 path does not get disabled + during recompute.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) MXFP8") + + # ----- te.checkpoint, non-reentrant ----------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_non_reentrant_linear_fp8_bitwise(self): + """te.checkpoint(use_reentrant=False) — the saved-tensors-hooks + path. Different recompute infra (``_checkpoint_hook``) than the + reentrant path; validates the hybrid activation survives the + pack/unpack transport.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_non_reentrant_linear_mxfp8_bitwise(self): + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) MXFP8") + + # ----- torch.utils.checkpoint (vanilla, non-reentrant) ---------- + # + # These tests document a *known* TE-level incompatibility between + # vanilla ``torch.utils.checkpoint.checkpoint(..., use_reentrant=False)`` + # and TE's weight-workspace cache (``_linear_forward_impl`` in + # ``module/linear.py``). The mechanism: + # + # * First forward: ``quantize_weight`` takes the cache-miss path, + # creating a fresh hybrid workspace and threading it into + # ``prepare_for_saving`` → ``ctx.save_for_backward``. + # * Recompute forward: the workspace is already populated on the + # module, so ``quantize_weight`` takes the cache-hit path and + # saves a different tensor-count. + # + # Vanilla ``torch.utils.checkpoint`` (``use_reentrant=False``) + # enforces a strict count match between original-forward and + # recompute-forward ``save_for_backward`` calls, and rejects the + # discrepancy with ``CheckpointError: A different number of tensors + # was saved``. The 2:1 count ratio (``8`` forward vs ``4`` recompute) + # is a hybrid signature — both sub-storages are saved on cache-miss + # and only the remaining one on cache-hit. + # + # ``te.checkpoint`` avoids this by threading ``is_first_microbatch`` + # / ``skip_fp8_weight_update`` correctly across the recompute phase, + # which is why the ``te.checkpoint`` tests above pass bitwise. + # + # Keeping the xfail'd tests here (tracked by #3158): + # 1. pins the boundary — users hitting this failure get a clear + # diagnosis and pointer to ``te.checkpoint``; + # 2. becomes a regression signal if the underlying cache-vs- + # checkpoint interaction is ever resolved (the xfail flips to + # an unexpected pass). + # + # Not hybrid-specific *in nature* (any quantized TE module with + # weight-workspace caching hits it under vanilla torch checkpoint), + # but hybrid amplifies and surfaces it via the 2x sub-storage count. + + _TORCH_CHECKPOINT_FP8_XFAIL = pytest.mark.xfail( + raises=( + NotImplementedError + if not is_non_tn_fp8_gemm_supported() + else torch.utils.checkpoint.CheckpointError + ), + strict=True, + reason=( + "On Hopper, same-format FP8 hybrid reaches the unsupported columnwise-only " + "per-tensor quantization guard before checkpointing. On architectures that " + "support non-TN FP8 GEMMs, vanilla torch.utils.checkpoint(use_reentrant=False) " + "is incompatible with TE's weight-workspace cache. Tracked by #3158." + ), + ) + + _TORCH_CHECKPOINT_CACHE_XFAIL = pytest.mark.xfail( + raises=torch.utils.checkpoint.CheckpointError, + strict=True, + reason=( + "Vanilla torch.utils.checkpoint(use_reentrant=False) is incompatible " + "with TE's weight-workspace cache. Tracked by #3158." + ), + ) + + @_TORCH_CHECKPOINT_FP8_XFAIL + def test_torch_checkpoint_non_reentrant_linear_fp8_bitwise(self): + """Vanilla ``torch.utils.checkpoint.checkpoint`` without TE wrapper + around a hybrid-quantized te.Linear. + + Users invoke ``torch.utils.checkpoint`` directly in many Megatron + branches and custom recomputation schemes. Currently fails due to + the weight-workspace cache interaction documented above; pins the + boundary so a future fix would flip this to an unexpected pass. + """ + + def fn(model, inp): + return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "torch.utils.checkpoint FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + @_TORCH_CHECKPOINT_CACHE_XFAIL + def test_torch_checkpoint_non_reentrant_linear_mxfp8_bitwise(self): + def fn(model, inp): + return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "torch.utils.checkpoint MXFP8") + + # ----- cross-format + recompute (functional, loose tolerance) --- + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_reentrant_linear_cross_format(self): + """Cross-format hybrid (FP8 row + MXFP8 col) under activation + recompute. Numerics are allowed to drift from non-recompute only + through paths recompute is allowed to affect; in practice they + should still match tightly because the recipe is stateless. Loose + tolerance catches only catastrophic silent fallbacks.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._cross_format_fp8_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._cross_format_fp8_mxfp8_recipe(), checkpoint_fn=fn) + # Expected to match bitwise since both quantizers are stateless + # and the input bytes are identical between the two runs. Use a + # strict tolerance; if this ever drifts it's a real bug. + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) FP8xMXFP8 cross-format") + + # ----- TransformerLayer ----------------------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_transformer_layer_fp8(self): + """te.checkpoint(reentrant) around a full TransformerLayer under + hybrid FP8. Exercises LayerNormLinear + DPA + LayerNormMLP in one + shot — the ``with_quantized_norm=False`` unfused path for hybrid + in ``layernorm_linear.py`` / ``layernorm_mlp.py`` must produce + the same result when recomputed. + + Asserted bitwise: the module uses ``hidden_dropout=0.0``, + ``attention_dropout=0.0``, and ``te.checkpoint`` restores RNG + state before recompute, so every kernel sees identical inputs and + there are no stochastic ops. Non-determinism at this level would + indicate a real regression (e.g. a kernel quietly taking a + non-deterministic code path) — not measurement noise.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) TransformerLayer FP8") + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_non_reentrant_transformer_layer_fp8(self): + """Same TransformerLayer setup but through the non-reentrant + saved-tensors-hooks recompute path. Same bitwise-equality + rationale as the reentrant variant above.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(non-reentrant) TransformerLayer FP8" + ) + + # ----- quantized_model_init + recompute ------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_quantized_model_init_fp8_bitwise(self): + """Combine ``quantized_model_init`` (persistent + HybridQuantizedTensor weights) with activation recompute — + verifies the recompute path doesn't try to re-quantize an already- + quantized weight incorrectly, and the HybridQuantizer workspace + caching stays consistent across first-forward + recomputed-forward.""" + import transformer_engine.pytorch as te_pytorch + + hybrid_recipe = self._same_format_fp8_recipe() + + def _build_and_run(use_checkpoint): + _reset_rng(seed=7777) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear( + self.in_features, self.out_features, params_dtype=torch.bfloat16 + ).cuda() + inp = torch.randn( + self.batch, + self.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + inp.retain_grad() + with autocast(enabled=True, recipe=hybrid_recipe): + if use_checkpoint: + out = te_pytorch.checkpoint(model, inp, use_reentrant=True) + else: + out = model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + ref = _build_and_run(use_checkpoint=False) + test = _build_and_run(use_checkpoint=True) + _assert_outputs_bitwise_equal( + ref, test, "quantized_model_init + te.checkpoint(reentrant) FP8" + ) + + # ----- GroupedLinear + recompute -------------------------------- + + def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): + """Build a GroupedLinear, run forward+backward with optional + activation checkpointing around the module. Exercises the + ``_split_quantize_hybrid`` code path under recompute. + + GroupedLinear is the MoE token-dispatch kernel: a single batch + is split along dim-0 into ``num_gemms`` chunks and each chunk + goes through its own weight matrix. Under hybrid quantization, + ``_split_quantize_hybrid`` (``module/grouped_linear.py``) runs + ``tex.split_quantize`` twice (once per sub-quantizer direction) + and zips the results into a list of ``HybridQuantizedTensor`` + chunks — save-for-backward then receives a *list* of hybrid + tensors, not a single one, so the ``prepare_for_saving`` chain + has to handle an extended tensor-object list. + """ + _reset_rng(seed=9090) + num_gemms = 3 + hidden = 128 + ffn = 128 + bs = 24 + + model = GroupedLinear(num_gemms, hidden, ffn, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + inp.retain_grad() + base = bs // num_gemms + rem = bs % num_gemms + m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] + + with autocast(enabled=True, recipe=recipe_obj): + if checkpoint_fn is not None: + out = checkpoint_fn(model, inp, m_splits) + else: + out = model(inp, m_splits) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_grouped_linear_fp8_bitwise(self): + """GroupedLinear + te.checkpoint(reentrant) under same-format FP8 + hybrid. Exercises the MoE ``_split_quantize_hybrid`` + list-of- + hybrid-tensors save-for-backward path under recompute.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp, m_splits): + return te_pytorch.checkpoint(model, inp, m_splits, use_reentrant=True) + + ref = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) GroupedLinear FP8") + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_non_reentrant_grouped_linear_fp8_bitwise(self): + """Same GroupedLinear recompute setup but through the non- + reentrant saved-tensors-hooks path — verifies that the list of + hybrid activations survives the pack/unpack transport (one hook + invocation per split × per sub-storage buffer, not just one).""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp, m_splits): + return te_pytorch.checkpoint(model, inp, m_splits, use_reentrant=False) + + ref = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) GroupedLinear FP8") + + # ----- Selective attention recompute ---------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_selective_attention_recompute_transformer_layer_fp8_bitwise(self): + """``TransformerLayer(..., checkpoint_core_attention=True)`` — + the Megatron default memory-savings pattern. + + Unlike full-layer recompute (``te.checkpoint(layer, inp)``), + selective attention recompute is a TransformerLayer-internal + option: only the DPA (dot-product attention) block is wrapped + in a checkpoint, everything else runs normally. This is a + *different* code path in ``transformer.py`` from the + ``te.checkpoint(...)`` tests above — DPA internally invokes its + own checkpoint context around the attention kernel. + + For hybrid, the question is whether a hybrid activation produced + by LayerNormLinear (QKV projection) survives the DPA-internal + recompute boundary (which saves it for backward) and is + consumable by the backward GEMM unchanged. + + Bitwise because the model uses ``hidden_dropout=0.0``, + ``attention_dropout=0.0``, and the DPA checkpoint restores RNG + state — so reference and recomputed paths should be identical + to the last bit.""" + _reset_rng(seed=5151) + hidden = 128 + ffn = 128 + nheads = 4 + seq = 8 + bs = 4 + + def _run(checkpoint_core_attention): + _reset_rng(seed=5151) + model = TransformerLayer( + hidden, + ffn, + nheads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + inp.retain_grad() + with autocast(enabled=True, recipe=self._same_format_fp8_recipe()): + out = model(inp, checkpoint_core_attention=checkpoint_core_attention) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + ref = _run(checkpoint_core_attention=False) + test = _run(checkpoint_core_attention=True) + _assert_outputs_bitwise_equal(ref, test, "checkpoint_core_attention TransformerLayer FP8") + + # ----- Linear bitwise parametrized across all 4 stateless formats ----- + + @pytest.mark.parametrize( + "format_name,reentrant", + [ + pytest.param( + "fp8_current", + True, + id="fp8_current-reentrant", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + "fp8_current", + False, + id="fp8_current-nonreentrant", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + "mxfp8", + True, + id="mxfp8-reentrant", + marks=pytest.mark.skipif( + not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}" + ), + ), + pytest.param( + "mxfp8", + False, + id="mxfp8-nonreentrant", + marks=pytest.mark.skipif( + not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}" + ), + ), + pytest.param( + "block_fp8", + True, + id="block_fp8-reentrant", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"BlockFP8: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "block_fp8", + False, + id="block_fp8-nonreentrant", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"BlockFP8: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "nvfp4", + True, + id="nvfp4-reentrant", + marks=pytest.mark.skipif( + not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}" + ), + ), + pytest.param( + "nvfp4", + False, + id="nvfp4-nonreentrant", + marks=pytest.mark.skipif( + not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}" + ), + ), + ], + ) + def test_te_checkpoint_linear_all_stateless_formats_bitwise(self, format_name, reentrant): + """Bitwise parity of Linear + te.checkpoint across all four + stateless hybrid formats (FP8 current, MXFP8, BlockFP8, NVFP4), + both reentrant and non-reentrant. + + Each format has a distinct history of columnwise-only kernel + support — BlockFP8 required C++ null-check patches before + columnwise-only mode worked, NVFP4 has packed FP4 layout plus + optional RHT cache, MXFP8 has [128,4]/[4,128] scale padding. + The recompute path exercises columnwise-only sub-quantizers + (rowwise is freed after fprop and only recreated on backward), + so format-specific columnwise-only handling is on the critical + path. + + A regression in any of these would silently fall back to BF16 + during recompute; bitwise equality catches that immediately.""" + import transformer_engine.pytorch as te_pytorch + + row_factory, col_factory_for_grad, hw_skip, hw_reason = _QUANTIZER_CONFIGS[format_name] + # Most formats have a distinct E5M2 variant for grad; NVFP4 has + # only one format (col_factory_for_grad is None → reuse + # row_factory, which is what the existing hybrid NVFP4 tests do). + grad_factory = col_factory_for_grad if col_factory_for_grad is not None else row_factory + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=row_factory, + col_factory=row_factory, + grad_factory=grad_factory, + ) + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=reentrant) + + ref = self._run_linear(hybrid_recipe, checkpoint_fn=None) + test = self._run_linear(hybrid_recipe, checkpoint_fn=fn) + label = ( + f"te.checkpoint({'reentrant' if reentrant else 'non-reentrant'}) Linear {format_name}" + ) + _assert_outputs_bitwise_equal(ref, test, label) + + # ----- save_for_backward round-trip (unit-level) ---------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_prepare_restore_roundtrip_is_identity(self): + """Unit-level guarantee: the + ``prepare_for_saving`` / ``restore_from_saved`` chain used by + activation-recompute ``ctx.save_for_backward`` preserves both + sub-storages bitwise. + + This is the primitive the recompute path is built on; pinning it + here gives a focused failure signal independent of the module- + level recompute tests above.""" + torch.manual_seed(0) + inp = torch.randn(256, 256, dtype=torch.bfloat16, device="cuda") + hq = HybridQuantizer( + rowwise_quantizer=_fp8_row_factory(), + columnwise_quantizer=_fp8_col_factory(), + ) + hybrid = hq.quantize(inp) + expected = hybrid.dequantize() + + saved_tensors, saved_obj = hybrid.prepare_for_saving() + # Mimic the autograd ctx round-trip: all saved tensors pass + # through ``ctx.save_for_backward`` (a no-op for semantics). + leftover = saved_obj.restore_from_saved(list(saved_tensors)) + assert leftover == [], "restore_from_saved should consume every element" + torch.testing.assert_close(saved_obj.dequantize(), expected, rtol=0, atol=0) diff --git a/tests/pytorch/test_hybrid_quantization_fsdp2.py b/tests/pytorch/test_hybrid_quantization_fsdp2.py new file mode 100644 index 0000000000..cda3e7184a --- /dev/null +++ b/tests/pytorch/test_hybrid_quantization_fsdp2.py @@ -0,0 +1,1515 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for HybridQuantizedTensor behavior required by PyTorch FSDP2.""" + +import io + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex + +from hybrid_quantization_utils import ( + as_data_tensor_tuple as _as_data_tensor_tuple, + assert_hybrid_tensor_exact as _assert_hybrid_tensor_exact, + assert_storage_data_exact as _assert_storage_data_exact, + fp8_e4m3_factory as _fp8_row_factory, + fp8_e5m2_factory as _fp8_grad_factory, + hybrid_block_fp8_e4m3_qfactory as _hybrid_block_fp8_qfactory, + hybrid_custom_recipe as _hybrid_custom_recipe, + make_fp8_quantizer as _make_fp8_quantizer, + make_hybrid_quantizer_fp8_row_fp4_col as _make_hybrid_quantizer_fp8_row_fp4_col, + mxfp8_e4m3_factory as _mxfp8_factory, +) +from transformer_engine.common import recipe +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + Float8Tensor, + Float8TensorStorage, + HybridQuantizedTensor, + HybridQuantizer, + IdentityQuantizer, + Linear, + MXFP8Quantizer, + NVFP4Quantizer, + QuantizedTensor, + quantized_model_init, +) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +_fp8_col_factory = _fp8_row_factory + + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + +requires_fp8 = pytest.mark.skipif( + not fp8_available, + reason=f"FP8: {reason_for_no_fp8}", +) + +requires_fp8_and_nvfp4 = pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) + +requires_mxfp8 = pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", +) + + +# --------------------------------------------------------------------------- +# 1. __torch_dispatch__ operations that FSDP2 relies on +# --------------------------------------------------------------------------- + +aten = torch.ops.aten + + +def _make_hybrid_param_for_dispatch( + row_factory, col_factory, grad_factory=None, in_features=256, out_features=256 +): + """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" + hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_dispatch_configs = [ + pytest.param( + "fp8_fp8", + id="same-format-fp8", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), +] +if mxfp8_available: + _dispatch_configs.append(pytest.param("mxfp8_mxfp8", id="same-format-mxfp8")) + + +def _get_dispatch_hybrid_param(config_name): + """Return a HybridQuantizedTensor weight for the given config.""" + if config_name == "fp8_fp8": + return _make_hybrid_param_for_dispatch( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + elif config_name == "mxfp8_mxfp8": + return _make_hybrid_param_for_dispatch( + _mxfp8_factory, + _mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + else: + raise ValueError(f"Unknown config: {config_name}") + + +@requires_fp8 +class TestFloat8TransposeOnlySplit: + """Regression coverage for columnwise-only Float8 split metadata. + + A columnwise-only per-tensor Float8 sub-storage may have ``_data=None`` and + store its bytes in ``_transpose`` with physical shape ``[K, M]``. Splitting + that tensor must still produce pieces whose wrapper shape is the logical + row-major shape ``[M_i, K]``; otherwise HybridQuantizedTensor uses the + transposed shape when rowwise storage is absent. + """ + + @staticmethod + def _make_transpose_only_float8_tensor(shape=(12, 16)): + m, k = shape + data_transpose = torch.empty((k, m), dtype=torch.uint8, device="cuda") + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + requires_grad=False, + device="cuda", + ) + + @pytest.mark.parametrize( + "split_size,dim,expected_shapes,expected_transpose_shapes", + [ + (5, 0, [(5, 16), (5, 16), (2, 16)], [(16, 5), (16, 5), (16, 2)]), + (6, 1, [(12, 6), (12, 6), (12, 4)], [(6, 12), (6, 12), (4, 12)]), + ], + ) + def test_float8_split_uses_logical_shape_for_transpose_only_storage( + self, split_size, dim, expected_shapes, expected_transpose_shapes + ): + tensor = self._make_transpose_only_float8_tensor() + + pieces = torch.split(tensor, split_size, dim=dim) + + assert [tuple(piece.shape) for piece in pieces] == expected_shapes + assert [tuple(piece._transpose.shape) for piece in pieces] == expected_transpose_shapes + assert all(piece._data is None for piece in pieces) + assert all(piece._transpose_invalid is False for piece in pieces) + + def test_float8_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = tensor.view(12, 16) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_aten_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = torch.ops.aten.view.default(tensor, [12, 16]) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_storage_view_preserves_transpose_only_storage(self): + data_transpose = torch.empty((16, 12), dtype=torch.uint8, device="cuda") + storage = Float8TensorStorage( + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + fake_dtype=torch.bfloat16, + ) + + viewed = storage.view(torch.Size((12, 16))) + + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_shape_changing_view_raises_for_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + with pytest.raises(NotImplementedError, match="columnwise-only data"): + tensor.view(6, 32) + + def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): + columnwise = self._make_transpose_only_float8_tensor() + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + quantizer.set_usage(rowwise=False, columnwise=True) + hybrid = HybridQuantizedTensor( + shape=columnwise.shape, + dtype=columnwise.dtype, + rowwise_storage=None, + columnwise_storage=columnwise, + quantizer=quantizer, + device="cuda", + ) + + pieces = torch.split(hybrid, 5, dim=0) + + assert [tuple(piece.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert all(piece.rowwise_sub_storage is None for piece in pieces) + assert [tuple(piece.columnwise_sub_storage.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert [tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces] == [ + (16, 5), + (16, 5), + (16, 2), + ] + + @staticmethod + def _make_valid_transpose_only_float8_tensor(shape=(12, 16)): + """Build transpose-only storage with known, numerically valid FP8 bytes.""" + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + rowwise = rowwise_quantizer(source) + columnwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + tensor = Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.detach().clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=columnwise_quantizer, + requires_grad=False, + device="cuda", + ) + return tensor, rowwise.dequantize() + + @pytest.mark.parametrize("weights_only", (False, True)) + def test_serialization_preserves_transpose_only_payload(self, weights_only): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + buffer = io.BytesIO() + + torch.save(tensor, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=weights_only) + + assert loaded._data is None + assert loaded._transpose_invalid is False + assert torch.equal(loaded._transpose, tensor._transpose) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_dequantize_from_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + torch.testing.assert_close(tensor.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_slice_and_select_preserve_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + sliced = tensor[2:8] + selected = torch.select(tensor, 1, 3) + + assert isinstance(sliced, Float8Tensor) + assert sliced._data is None + assert sliced._transpose.shape == torch.Size((16, 6)) + torch.testing.assert_close(sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0) + assert isinstance(selected, Float8Tensor) + assert selected._data is None + assert selected._transpose.shape == torch.Size((12,)) + torch.testing.assert_close(selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0) + + def test_as_strided_row_shard_preserves_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + shard = torch.as_strided(tensor, (5, 16), (16, 1), 16) + + assert isinstance(shard, Float8Tensor) + assert shard._data is None + assert shard._transpose.shape == torch.Size((16, 5)) + torch.testing.assert_close(shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0) + + def test_as_strided_falls_back_for_nonrepresentable_layout(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + output = torch.as_strided(tensor, (16, 12), (1, 16), 0) + + assert type(output) is torch.Tensor + reference = torch.as_strided(expected, (16, 12), (1, 16), 0) + torch.testing.assert_close(output, reference, rtol=0.0, atol=0.0) + + def test_new_zeros_preserves_transpose_only_layout(self): + tensor, _ = self._make_valid_transpose_only_float8_tensor() + + output = tensor.new_zeros((5, 16)) + + assert isinstance(output, Float8Tensor) + assert output.shape == torch.Size((5, 16)) + assert output._data is None + assert output._transpose.shape == torch.Size((16, 5)) + assert output._transpose.dtype == torch.uint8 + assert torch.count_nonzero(output._transpose).item() == 0 + torch.testing.assert_close( + output.dequantize(), + torch.zeros((5, 16), dtype=output.dtype, device=output.device), + rtol=0.0, + atol=0.0, + ) + + +class TestHybridNewZeros: + """Public new_zeros semantics independent of FSDP-specific allocation.""" + + @staticmethod + def _make_identity_hybrid(): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + return quantizer, source + + @pytest.mark.parametrize( + "keep_rowwise,keep_columnwise", + [(True, True), (True, False), (False, True)], + ) + def test_initializes_every_present_direction_and_preserves_kwargs( + self, + keep_rowwise, + keep_columnwise, + ): + quantizer, source = self._make_identity_hybrid() + if not keep_rowwise: + source.update_usage(rowwise_usage=False) + if not keep_columnwise: + source.update_usage(columnwise_usage=False) + + # Direction selection must come from source storage, not mutable parent + # usage flags, and new_zeros must not change those flags. + quantizer.set_usage(rowwise=False, columnwise=False) + result = torch.ops.aten.new_zeros.default( + source, + [3, 5], + dtype=torch.float32, + device=source.device, + ) + + assert isinstance(result, HybridQuantizedTensor) + assert result.shape == torch.Size((3, 5)) + assert result.dtype == torch.float32 + assert result.device == source.device + assert (result.rowwise_sub_storage is not None) is keep_rowwise + assert (result.columnwise_sub_storage is not None) is keep_columnwise + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + assert result._quantizer is not quantizer + assert result._quantizer.get_usages() == { + "rowwise": keep_rowwise, + "columnwise": keep_columnwise, + } + + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), + torch.zeros((3, 5), dtype=torch.float32, device=source.device), + rtol=0.0, + atol=0.0, + ) + + # A plain-source copy routes through the result's parent quantizer. It + # must update every allocated direction even though the source parent + # had both of its mutable usage flags disabled before new_zeros. + plain_source = torch.full((3, 5), 7.0, dtype=torch.float32, device=source.device) + result.copy_(plain_source) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), plain_source, rtol=0.0, atol=0.0 + ) + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + + def test_identity_substorages_allow_integer_dtype(self): + _, source = self._make_identity_hybrid() + result = source.new_zeros((2, 3), dtype=torch.int32) + + assert result.dtype == torch.int32 + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + assert sub_storage.dequantize().dtype == torch.int32 + assert torch.count_nonzero(sub_storage.dequantize()).item() == 0 + + @requires_fp8 + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + @pytest.mark.parametrize("unsupported_dtype", (torch.int32, torch.float64, torch.bool)) + def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(TypeError, match="new_zeros only supports"): + source.new_zeros((2, 3), dtype=unsupported_dtype) + + def test_does_not_invoke_live_quantizers_or_consume_rng(self, monkeypatch): + quantizer, source = self._make_identity_hybrid() + cpu_rng_before = torch.get_rng_state().clone() + cuda_rng_before = torch.cuda.get_rng_state(source.device).clone() + + def fail_live_make_empty(*args, **kwargs): + raise AssertionError("new_zeros invoked a live quantizer") + + monkeypatch.setattr(quantizer, "make_empty", fail_live_make_empty) + monkeypatch.setattr( + source.rowwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + monkeypatch.setattr( + source.columnwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + + result = source.new_zeros((2, 6)) + + assert isinstance(result, HybridQuantizedTensor) + assert torch.equal(torch.get_rng_state(), cpu_rng_before) + assert torch.equal(torch.cuda.get_rng_state(source.device), cuda_rng_before) + + def test_rejects_empty_hybrid(self): + _, source = self._make_identity_hybrid() + source.update_usage(rowwise_usage=False, columnwise_usage=False) + + with pytest.raises(RuntimeError, match="at least one present sub-storage"): + source.new_zeros((2, 6)) + + @requires_fp8_and_nvfp4 + def test_initializes_nvfp4_data_scale_and_amax_buffers(self): + quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ) + def test_initializes_float8_block_data_and_scale_buffers(self): + quantizer = HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + ) + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + +@requires_fp8 +class TestHybridTorchDispatchFSDP2Ops: + """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. + + Each op is called directly via torch.ops.aten and the result is verified to + still be HybridQuantizedTensor with valid sub-storages. + """ + + @pytest.fixture(params=_dispatch_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _get_dispatch_hybrid_param(request.param) + + def test_split_preserves_hybrid_type(self, hybrid_param): + """torch.split must return a list of HybridQuantizedTensor pieces.""" + dim0 = hybrid_param.shape[0] + chunk_size = dim0 // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + expected_rowwise = torch.split(hybrid_param.rowwise_sub_storage, chunk_size, dim=0) + expected_columnwise = torch.split(hybrid_param.columnwise_sub_storage, chunk_size, dim=0) + assert len(pieces) == len(expected_rowwise) == len(expected_columnwise) + + assert len(pieces) >= 2 + for piece in pieces: + assert isinstance( + piece, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(piece).__name__}" + assert piece.rowwise_sub_storage is not None + assert piece.columnwise_sub_storage is not None + for index, (piece, expected_row, expected_column) in enumerate( + zip(pieces, expected_rowwise, expected_columnwise) + ): + _assert_storage_data_exact( + piece.rowwise_sub_storage, + expected_row, + context=f"split {index} rowwise", + ) + _assert_storage_data_exact( + piece.columnwise_sub_storage, + expected_column, + context=f"split {index} columnwise", + ) + + total_rows = sum(p.shape[0] for p in pieces) + assert total_rows == dim0 + + orig_deq = hybrid_param.dequantize() + reassembled = torch.cat([p.dequantize() for p in pieces], dim=0) + torch.testing.assert_close(orig_deq, reassembled, rtol=0.0, atol=0.0) + + def test_split_sub_storage_types_preserved(self, hybrid_param): + """After split, sub-storage types must match the original.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + chunk_size = hybrid_param.shape[0] // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + for piece in pieces: + assert type(piece.rowwise_sub_storage) is orig_row_type + assert type(piece.columnwise_sub_storage) is orig_col_type + + @requires_mxfp8 + def test_split_rejects_mxfp8_high_precision_fallback(self): + """Fail before wrapping unquantizable MXFP8 shards for Hybrid FSDP2.""" + quantizer = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + tensor = quantizer(torch.randn(64, 64, dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(NotImplementedError, match="local shape.*divisible by 32"): + torch.split(tensor, 16, dim=0) + + def test_view_preserves_hybrid_type(self, hybrid_param): + """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" + shape_2d = hybrid_param.shape + result = aten.view.default(hybrid_param, list(shape_2d)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_view_same_shape_preserves_hybrid(self, hybrid_param): + """view with same shape must return HybridQuantizedTensor.""" + shape_2d = list(hybrid_param.shape) + result = aten.view.default(hybrid_param, shape_2d) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + + def test_as_strided_noop_preserves_hybrid(self, hybrid_param): + """as_strided with matching shape/strides is a no-op that preserves type.""" + shape = tuple(hybrid_param.size()) + strides = (shape[-1], 1) + result = aten.as_strided.default(hybrid_param, list(shape), list(strides)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_slice_noop_preserves_hybrid(self, hybrid_param): + """slice with full range is a no-op that preserves type.""" + result = aten.slice.Tensor(hybrid_param, 0, 0, hybrid_param.size(0)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + + def test_copy_between_hybrid_tensors(self, hybrid_param): + """copy_ between compatible HybridQuantizedTensors copies quantized data directly.""" + src_deq = hybrid_param.dequantize().clone() + dst = hybrid_param._quantizer.make_empty( + shape=hybrid_param.shape, + dtype=hybrid_param.dtype, + device=hybrid_param.device, + ) + assert isinstance(dst, HybridQuantizedTensor) + + aten.copy_.default(dst, hybrid_param) + dst_deq = dst.dequantize() + torch.testing.assert_close(src_deq, dst_deq, rtol=0.0, atol=0.0) + _assert_hybrid_tensor_exact(dst, hybrid_param, context="hybrid copy_") + + def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): + quantizer = hybrid_param._quantizer.copy() + src = quantizer.quantize(torch.ones_like(hybrid_param.dequantize())) + dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) + src.update_usage(columnwise_usage=False) + dst_before = tuple( + None if tensor is None else tensor.clone() for tensor in _as_data_tensor_tuple(dst) + ) + + with pytest.raises( + NotImplementedError, + match="requires matching rowwise/columnwise usages", + ): + aten.copy_.default(dst, src) + + dst_after = _as_data_tensor_tuple(dst) + assert len(dst_after) == len(dst_before) + for before, after in zip(dst_before, dst_after): + if before is None: + assert after is None + else: + torch.testing.assert_close(after, before, rtol=0.0, atol=0.0) + + def test_copy_from_bf16_to_hybrid(self, hybrid_param): + """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" + param = hybrid_param.detach() + bf16_data = torch.randn_like(param.dequantize()) + expected = param._quantizer.quantize(bf16_data) + aten.copy_.default(param, bf16_data) + assert isinstance(param, HybridQuantizedTensor) + _assert_hybrid_tensor_exact(param, expected, context="BF16 copy_") + torch.testing.assert_close(param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0) + + def test_new_zeros_returns_hybrid(self, hybrid_param): + """new_zeros returns initialized storage that remains FSDP-copyable.""" + new_shape = list(hybrid_param.shape) + result = aten.new_zeros.default(hybrid_param, new_shape) + + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + assert type(result.rowwise_sub_storage) is type(hybrid_param.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(hybrid_param.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + + # FSDP2 overwrites the initialized destination via copy_ after gather. + aten.copy_.default(result, hybrid_param) + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="new_zeros then copy_") + + def test_empty_like_returns_hybrid(self, hybrid_param): + """empty_like must return a HybridQuantizedTensor.""" + result = aten.empty_like.default(hybrid_param) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + + def test_clone_returns_hybrid(self, hybrid_param): + """clone must return an independent HybridQuantizedTensor with same data.""" + result = aten.clone.default(hybrid_param) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result is not hybrid_param + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="clone") + + +# --------------------------------------------------------------------------- +# 2. fsdp_pre_all_gather protocol +# --------------------------------------------------------------------------- + + +def _make_fsdp_protocol_param(config_name): + """Create a HybridQuantizedTensor weight for FSDP protocol tests.""" + if config_name == "fp8_fp8": + r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "mxfp8_fp8": + r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "block_fp8": + r = recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + else: + raise ValueError(f"Unknown config: {config_name}") + with quantized_model_init(enabled=True, recipe=r): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_fsdp_protocol_configs = [ + pytest.param( + "fp8_fp8", + id="same-format", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ) +] +if mxfp8_available: + _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) +if fp8_block_scaling_available: + _fsdp_protocol_configs.append(pytest.param("block_fp8", id="same-format-block-fp8")) + + +@requires_fp8 +class TestHybridFsdpPreAllGatherProtocol: + """Test the fsdp_pre_all_gather method on HybridQuantizedTensor. + + These tests call the method directly (no actual all-gather communication) + to verify the protocol contract: returns (sharded_tensors, metadata) where + sharded_tensors is a tuple of plain torch.Tensor. + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_all_gather_returns_tuple_pair(self, hybrid_param): + """fsdp_pre_all_gather returns (sharded_tensors, metadata).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert isinstance( + sharded_tensors, tuple + ), f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" + assert len(sharded_tensors) > 0, "sharded_tensors should not be empty" + assert isinstance( + metadata, tuple + ), f"metadata should be tuple, got {type(metadata).__name__}" + + def test_pre_all_gather_buffers_are_plain_tensors(self, hybrid_param): + """Every element in sharded_tensors must be a plain torch.Tensor.""" + sharded_tensors, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for i, t in enumerate(sharded_tensors): + assert isinstance( + t, torch.Tensor + ), f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" + assert not isinstance( + t, QuantizedTensor + ), f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" + + def test_pre_all_gather_buffer_count_consistent(self, hybrid_param): + """Buffer count must be the same across repeated calls (FSDP2 buffer reuse).""" + sharded_1, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + sharded_2, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert len(sharded_1) == len( + sharded_2 + ), f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" + + def test_pre_all_gather_metadata_sufficient_for_reconstruction(self, hybrid_param): + """Metadata must contain enough info to reconstruct the tensor.""" + _, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert metadata is not None + assert len(metadata) > 0, "metadata should not be empty" + + +# --------------------------------------------------------------------------- +# 3. fsdp_post_all_gather protocol +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpPostAllGatherProtocol: + """Test the fsdp_post_all_gather method on HybridQuantizedTensor. + + Simulates the post-all-gather phase by passing the sharded_tensors + from pre_all_gather directly (mimicking a single-rank all-gather). + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): + """With out=None, post_all_gather returns (HybridQuantizedTensor, outputs).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, ag_outputs = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_post_all_gather_buffer_reuse(self, hybrid_param): + """On second call with out=previous, the same object is returned (buffer reuse).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=first_result, + ) + assert ( + second_result is first_result + ), "Buffer reuse: post_all_gather(out=prev) should return the same object" + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="post-all-gather reuse") + + def test_post_all_gather_dequantize_matches_original(self, hybrid_param): + """Reconstructed tensor should dequantize close to the original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + result_deq = result.dequantize() + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather") + torch.testing.assert_close(orig_deq, result_deq, rtol=0.0, atol=0.0) + + def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): + """Reconstructed tensor's sub-storages match the original types.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + assert type(result.rowwise_sub_storage) is orig_row_type + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather storage types") + assert type(result.columnwise_sub_storage) is orig_col_type + + +# --------------------------------------------------------------------------- +# 4. Pre/post all-gather roundtrip +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpRoundtrip: + """End-to-end single-process roundtrip (pre -> post) without communication.""" + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_post_roundtrip_preserves_data(self, hybrid_param): + """pre_all_gather -> post_all_gather(out=None) -> dequantize matches original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="FSDP roundtrip") + torch.testing.assert_close(orig_deq, result.dequantize(), rtol=0.0, atol=0.0) + + def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): + """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=first_result, + ) + assert second_result is first_result + torch.testing.assert_close( + hybrid_param.dequantize(), second_result.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="FSDP roundtrip reuse") + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_scale_refresh_across_iterations(self): + """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. + + Per-tensor FP8 does NOT include ``_scale_inv`` in ``fsdp_buffer_fields`` + (only ``_data`` is gathered; the scalar scale travels via iter-1 + metadata). This relies on the invariant that the sharded and gathered + ``Float8Tensor`` s share the same ``_scale_inv`` tensor object, and + that ``Float8CurrentScalingQuantizer.update_quantized`` writes the new + scale in place rather than replacing the tensor reference. If either + invariant broke, the gathered copy would carry a stale scale on + iter-2+ and silently apply the wrong dequantization. + + This test locks the invariant down by forcing a radically different + scale between iterations and asserting the gathered tensor's + dequantization tracks the sharded one. + """ + torch.manual_seed(42) + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + hybrid_param = model.weight + + # Iter-1 gather with the initial (small-magnitude) weights + sharded_tensors_1, metadata_1 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + gathered, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_1, + metadata_1, + hybrid_param.dtype, + out=None, + ) + + # Simulate an optimizer writeback that produces a much larger weight; + # Float8CurrentScalingQuantizer.update_quantized must recompute + # _scale_inv for this range. If the gathered copy didn't see the new + # scale, the dequantize below would disagree with the sharded copy. + huge_master = torch.randn_like(hybrid_param.dequantize()) * 100.0 + hybrid_param._quantizer.update_quantized(huge_master, hybrid_param) + + # Iter-2+ path: reuse the gathered buffer + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + gathered_refreshed, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=gathered, + ) + assert gathered_refreshed is gathered + + # The gathered copy must now reflect the new sharded scale, not the + # tiny original scale. + torch.testing.assert_close( + hybrid_param.dequantize(), + gathered_refreshed.dequantize(), + rtol=0.0, + atol=0.0, + ) + _assert_hybrid_tensor_exact(gathered_refreshed, hybrid_param, context="FSDP scale refresh") + # And the magnitude really did change (sanity: this test would pass + # vacuously if update_quantized didn't actually change anything). + assert gathered_refreshed.dequantize().abs().max() > 10.0, ( + "update_quantized did not produce a sufficiently different " + "weight; the scale-refresh invariant is not being exercised" + ) + + def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): + """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. + + NVIDIA/TransformerEngine#3158 tracks the missing NVFP4 FSDP2 support: + packed FP4 dim-0 alignment, columnwise dequantization, and RHT-cache + handling are not implemented. Until that lands, hybrid pre-all-gather must + refuse an NVFP4 sub-storage cleanly via the ``fsdp_buffer_fields`` + protocol rather than silently producing wrong data. + + This test pins that contract: any hybrid whose sub-storage does not + implement ``fsdp_buffer_fields`` raises ``NotImplementedError`` at + ``fsdp_pre_all_gather`` time. The prior version of this test + inadvertently asserted the opposite when buffer extraction used + implicit ``get_metadata()``-based tensor scanning. + """ + if not (fp8_available and nvfp4_available): + pytest.skip("Requires FP8 + NVFP4 support") + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer(), + col_factory=lambda: NVFP4Quantizer(), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + param = model.weight + + # Clean refusal: hybrid's pre_all_gather raises an NVFP4-specific + # message identifying the unsupported sub-storage protocol, not a generic + # "NVFP4Tensor does not implement fsdp_buffer_fields" from deep inside + # the base class. + with pytest.raises(NotImplementedError) as exc_info: + param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + msg = str(exc_info.value) + assert "NVFP4Tensor" in msg + assert "rowwise sub-storage" in msg + assert "Use a supported sub-quantizer" in msg + assert "fsdp_buffer_fields" in msg + + +# --------------------------------------------------------------------------- +# 5. make_like correctness +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridMakeLike: + """Test that make_like produces correct copies for __torch_dispatch__ usage.""" + + def _make_hybrid_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_make_like_preserves_sub_storages(self): + """make_like result has the same sub-storage types, quantizers, and dtype.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + + assert isinstance(copy, HybridQuantizedTensor) + assert copy.dtype == param.dtype + assert copy.shape == param.shape + assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) + assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) + _assert_hybrid_tensor_exact(copy, param, context="make_like") + torch.testing.assert_close(copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0) + + def test_make_like_is_independent(self): + """make_like result should not share the same tensor identity.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + assert copy is not param + + +# --------------------------------------------------------------------------- +# 5b. Hopper-only paths: columnwise-only Float8 sub-storage +# --------------------------------------------------------------------------- +# +# On architectures where ``is_non_tn_fp8_gemm_supported()`` returns False +# (Hopper sm_90, L40 sm_89), per-tensor FP8 GEMM only supports the TN +# layout — non-TN layouts are simulated by feeding pre-transposed data. +# So a columnwise-only ``Float8TensorStorage`` (used as a hybrid sub- +# storage) holds its quantized data in ``_transpose`` instead of +# ``_data``, with ``_data = None``. +# +# This is the exact layout the FSDP2 buffer protocol must recognize +# when the sub-storage is part of a ``HybridQuantizedTensor`` parameter. +# These tests pin the contracts that would break if the buffer +# protocol regressed to the unconditional ``("_data",)`` field name +# (which would all-gather a ``None`` tensor on Hopper). +# +# Skip on Blackwell where the C++ kernel always populates ``_data`` and +# the columnwise-only Float8 path doesn't exercise ``_transpose``. + +requires_hopper_fp8 = pytest.mark.skipif( + is_non_tn_fp8_gemm_supported() or not fp8_available, + reason=( + "Hopper-only: requires per-tensor FP8 with non-TN GEMM unsupported " + "(Hopper sm_90 / L40 sm_89). On Blackwell the C++ kernel populates " + "_data even for columnwise-only mode, so the _transpose-only path " + "is not exercised." + ), +) + + +@requires_hopper_fp8 +class TestHybridFloat8ColumnwiseOnlyHopperPath: + """Hopper transpose-only Float8 uses an M-major FSDP transport layout.""" + + @staticmethod + def _make_columnwise_only_float8_storage(shape=(32, 64), value=1.0, quantizer=None): + source = torch.full(shape, value, device="cuda", dtype=torch.bfloat16) + byte_quantizer = Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + rowwise = byte_quantizer(source) + if quantizer is None: + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + storage = Float8Tensor( + shape=shape, + dtype=source.dtype, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return storage, rowwise.dequantize() + + @classmethod + def _make_hybrid_shard(cls, value): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + source = torch.full((32, 64), value, device="cuda", dtype=torch.bfloat16) + rowwise = quantizer.rowwise_quantizer(source) + columnwise, expected = cls._make_columnwise_only_float8_storage( + value=value, quantizer=quantizer.columnwise_quantizer + ) + hybrid = HybridQuantizedTensor( + shape=source.shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return hybrid, expected + + def test_columnwise_only_float8_fsdp_buffer_fields_returns_transpose(self): + out, _ = self._make_columnwise_only_float8_storage() + + assert out._data is None + assert out._transpose is not None + assert not out._transpose_invalid + assert out.fsdp_buffer_fields() == ("_transpose",) + + def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): + out, _ = self._make_columnwise_only_float8_storage() + + buffers, metadata = out.fsdp_extract_buffers() + + assert len(buffers) == 1 + assert buffers[0].shape == out.shape + assert buffers[0].is_contiguous() + torch.testing.assert_close(buffers[0], out._transpose.movedim(0, -1).contiguous()) + assert metadata == { + "field_names": ("_transpose",), + "transport_layout": "columnwise_m_major", + } + + def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): + out, expected = self._make_columnwise_only_float8_storage() + buffers, metadata = out.fsdp_extract_buffers() + gathered = torch.cat((buffers[0], buffers[0]), dim=0) + rebuilt = Float8Tensor.make_like(out, shape=gathered.shape) + + rebuilt.fsdp_assign_gathered((gathered,), metadata) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._data is None + assert rebuilt._transpose.shape == torch.Size((64, 64)) + assert not rebuilt._transpose_invalid + torch.testing.assert_close(rebuilt._transpose, gathered.movedim(-1, 0).contiguous()) + torch.testing.assert_close(rebuilt.dequantize(), torch.cat((expected, expected), dim=0)) + + def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): + shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] + extracted = [ + shard.fsdp_pre_all_gather( + mesh=None, + orig_size=shard.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for shard, _ in shards + ] + gathered = tuple( + torch.cat((extracted[0][0][i], extracted[1][0][i]), dim=0) + for i in range(len(extracted[0][0])) + ) + expected = torch.cat((shards[0][1], shards[1][1]), dim=0) + + rebuilt, _ = shards[0][0].fsdp_post_all_gather(gathered, extracted[0][1], torch.bfloat16) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage._data is None + assert rebuilt._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(rebuilt._columnwise_storage.dequantize(), expected) + + reused, _ = shards[0][0].fsdp_post_all_gather( + gathered, extracted[0][1], torch.bfloat16, out=rebuilt + ) + assert reused is rebuilt + assert reused._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(reused._columnwise_storage.dequantize(), expected) + + def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + out = quantizer(torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)) + + assert out._data is not None + assert out._transpose is not None + assert out.fsdp_buffer_fields() == ("_data",) + buffers, metadata = out.fsdp_extract_buffers() + assert len(buffers) == 1 + assert buffers[0] is out._data + assert metadata == { + "field_names": ("_data",), + "transport_layout": "native", + } + + +@requires_hopper_fp8 +class TestPerTensorFloat8ColumnwiseOnlyGuard: + """Per-tensor FP8 cannot yet emit only the physical transpose.""" + + @staticmethod + def _make_quantizer(scaling): + if scaling == "current": + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + return Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=False, + columnwise=True, + ) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_initial_quantize_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer(source) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_update_quantized_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + output = quantizer.make_empty( + source.shape, + dtype=source.dtype, + device=source.device, + ) + assert output._data is None + assert output._transpose is not None + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer.update_quantized(source, output) + + +@requires_hopper_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridFsdpPostAllGatherUpdateUsage: + """``HybridQuantizedTensor.fsdp_post_all_gather`` must call + ``update_usage`` on each sub-storage after writing gathered data + (mirroring vanilla ``Float8Tensor.fsdp_post_all_gather:888``). + Without it, on Hopper a previously-cached ``_transpose`` from the + prior iteration is silently reused with the new ``_data``, producing + incorrect dgrad / wgrad GEMMs. + """ + + def _make_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(64, 64, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): + """Simulates iter-2+ buffer reuse: pre-existing ``out`` with a + possibly-stale ``_transpose`` cache; after ``fsdp_post_all_gather`` + the rowwise sub-storage's ``_transpose`` must be invalidated / + regenerated to match the freshly gathered ``_data``. + """ + param = self._make_param() + # Build a plausible iter-2+ "out" with stale state. + out = HybridQuantizedTensor.make_like(param) + # Rowwise sub-storage on Hopper has _data populated. Force a stale + # _transpose and invalidate flag to mimic the regression scenario. + if out._rowwise_storage._transpose is None: + # Set up a fake stale transpose (non-None, marked invalid by + # the mismatching shape would catch nothing, so just plant + # a tensor and clear the invalid flag to "valid"). + out._rowwise_storage._transpose = torch.zeros_like(out._rowwise_storage._data).t() + out._rowwise_storage._transpose_invalid = False + stale_transpose_id = id(out._rowwise_storage._transpose) + + # Drive a real all-gather round trip via the protocol + sharded_tensors, metadata = param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + out2, _ = param.fsdp_post_all_gather(sharded_tensors, metadata, param.dtype, out=out) + + # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned + # columnwise=False, so update_usage(rowwise=True, columnwise=False) + # must clear the stale _transpose (preventing the silent + # stale-cache regression on Hopper). + assert out2._rowwise_storage._transpose is None or ( + out2._rowwise_storage._transpose_invalid + and id(out2._rowwise_storage._transpose) != stale_transpose_id + ), "Stale _transpose was not invalidated after fsdp_post_all_gather" diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py new file mode 100644 index 0000000000..422b3fd975 --- /dev/null +++ b/tests/pytorch/test_identity_quantizer.py @@ -0,0 +1,1867 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for IdentityQuantizer (high-precision passthrough) and its use as a +per-direction component of HybridQuantizer to express mixed forward/backward +precision via the CustomRecipe + qfactory machinery. Scoped to single-GPU +TE GEMM modules. +""" + +import io + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + HybridQuantizer, + HybridQuantizedTensor, + IdentityQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) +from transformer_engine.pytorch.tensor.identity_tensor import IdentityTensor +from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( + IdentityTensorStorage, +) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR = ( + "Columnwise-only per-tensor FP8 quantization is not implemented" +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + + +# ── Module-level qfactories (picklable / autocast-friendly) ────────── + + +def identity_all_factory(role): # pylint: disable=unused-argument + """Whole layer in high precision: Identity for every slot.""" + return IdentityQuantizer() + + +def _fp8_cs(fp8_dtype=tex.DType.kFloat8E4M3): + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _mxfp8(fp8_dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def _float8_blockwise(fp8_dtype=tex.DType.kFloat8E4M3): + return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) + + +def _nvfp4(): + return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + + +_HYBRID_IDENTITY_FORMATS = [ + pytest.param( + "fp8_current", + marks=pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}"), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + ), + pytest.param( + "float8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), + ), +] + +_HYBRID_IDENTITY_RECOMPUTE_FORMATS = [ + pytest.param( + "fp8_current", + marks=pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}"), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), + ), +] + + +def _format_quantizer(format_name): + if format_name == "fp8_current": + return _fp8_cs(tex.DType.kFloat8E4M3) + if format_name == "mxfp8": + return _mxfp8(tex.DType.kFloat8E4M3) + if format_name == "float8_blockwise": + return _float8_blockwise(tex.DType.kFloat8E4M3) + if format_name == "nvfp4": + return _nvfp4() + raise ValueError(format_name) + + +def _hybrid_quantized_fwd_identity_bwd_factory(format_name): + def qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + + return qfactory + + +def fwd_hp_bwd_fp8_factory(role): + """High-precision forward, FP8 backward (per-direction via hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _fp8_cs(tex.DType.kFloat8E5M2) + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + ) + + +def fwd_fp8_bwd_hp_factory(role): + """FP8 forward, high-precision backward (per-direction via hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ) + + +def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): + """FP8 forward, high-precision backward from dequantized fprop values.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + +def hybrid_all_identity_factory(role): + """All directions high precision, expressed through the hybrid container. + + weight / input / output -> Hybrid(Identity, Identity); grad -> Identity. + Exercises the HybridQuantizedTensor path with Identity sub-storages in both + directions (distinct from the non-hybrid whole-layer-HP path). + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + + +def fp8_fwd_factory(role): + """Plain FP8 current scaling for every slot (E4M3 fwd, E5M2 grad). + + Used with ``backward_override="high_precision"`` as the reference that the + per-direction Identity machinery (``fwd_fp8_bwd_hp_factory``) must reproduce + bitwise: same FP8 forward, high-precision backward. + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _fp8_cs(tex.DType.kFloat8E5M2) + return _fp8_cs(tex.DType.kFloat8E4M3) + + +def _offload_roundtrip(tensor): + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + tid = state.push_tensor(tensor) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + reloaded = state.pop_tensor(tid) + torch.cuda.synchronize() + try: + return reloaded + finally: + state.release_all_memory() + + +# ── Unit tests ─────────────────────────────────────────────────────── + + +class TestIdentityQuantizerUnit: + """IdentityQuantizer / IdentityTensorStorage basic behavior.""" + + def test_quantize_returns_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + out = IdentityQuantizer()(x) + assert isinstance(out, IdentityTensor) + + def test_internal_returns_storage(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + out = q(x) + assert isinstance(out, IdentityTensorStorage) + assert not isinstance(out, IdentityTensor) + + def test_make_empty_internal_returns_storage(self): + q = IdentityQuantizer() + q.internal = True + + out = q.make_empty((8, 16), dtype=torch.bfloat16, device="cuda") + + assert isinstance(out, IdentityTensorStorage) + assert not isinstance(out, IdentityTensor) + assert out.size() == torch.Size((8, 16)) + assert out.dequantize().dtype == torch.bfloat16 + + def test_grouped_split_all_identity_uses_plain_tensor_views(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_non_hybrid, + ) + + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + m_splits = [3, 5] + quantizers = [IdentityQuantizer(), IdentityQuantizer()] + + out = _split_quantize_non_hybrid(x, m_splits, quantizers, activation_dtype=torch.bfloat16) + + assert all(isinstance(t, torch.Tensor) for t in out) + assert not any(isinstance(t, IdentityTensorStorage) for t in out) + for actual, expected in zip(out, torch.split(x, m_splits)): + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + cast_quantizers = [ + IdentityQuantizer(dtype=torch.float32), + IdentityQuantizer(dtype=torch.float32), + ] + cast_out = _split_quantize_non_hybrid( + x, m_splits, cast_quantizers, activation_dtype=torch.bfloat16 + ) + assert all(isinstance(t, IdentityTensorStorage) for t in cast_out) + for actual, expected in zip(cast_out, torch.split(x, m_splits)): + dequantized = actual.dequantize() + assert dequantized.dtype == torch.float32 + torch.testing.assert_close( + dequantized, + expected.to(torch.float32), + rtol=0.0, + atol=0.0, + ) + + def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + cases = [ + [IdentityQuantizer(), _mxfp8(tex.DType.kFloat8E4M3)], + [ + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ), + HybridQuantizer( + rowwise_quantizer=_mxfp8(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ), + ], + ] + + for quantizers in cases: + with pytest.raises(ValueError, match="mix Identity-backed and quantized"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + from transformer_engine.pytorch.module.grouped_linear import _split_quantize_hybrid + + calls = [] + + def fake_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation=False): + calls.append(disable_bulk_allocation) + return [ + quantizer(tensor_part) + for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) + ] + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", fake_split_quantize) + monkeypatch.setattr( + grouped_linear, + "_supports_native_split_quantize", + lambda quantizer: True, + ) + x = torch.randn(8, 16, dtype=torch.bfloat16) + m_splits = [3, 5] + quantizers = [ + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + for _ in m_splits + ] + + out = _split_quantize_hybrid( + x, + m_splits, + quantizers, + disable_bulk_allocation=True, + ) + + assert calls == [True, True] + assert len(out) == len(m_splits) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_grouped_linear_cpu_offload_disables_bulk_allocation_for_hybrid_input( + self, monkeypatch + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + class StopAfterFlagCapture(RuntimeError): + pass + + def qfactory(role): + if role is not None and role.module_type == "grouped_linear": + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + ) + return _fp8_cs(tex.DType.kFloat8E4M3) + + calls = [] + + def fake_split_quantize_hybrid( + tensor, m_splits, quantizers, *, disable_bulk_allocation=False, **kwargs + ): + del tensor, m_splits, quantizers, kwargs + calls.append(disable_bulk_allocation) + raise StopAfterFlagCapture("captured hybrid split kwargs") + + monkeypatch.setattr(grouped_linear, "is_cpu_offload_enabled", lambda: True) + monkeypatch.setattr(grouped_linear, "_split_quantize_hybrid", fake_split_quantize_hybrid) + + model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + + with pytest.raises(StopAfterFlagCapture): + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + model(x, m_splits=m_splits) + + assert calls == [True] + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_grouped_linear_rejects_mixed_identity_weight_quantizers(self): + weight_count = 0 + + def qfactory(role): + nonlocal weight_count + if role is not None and role.module_type == "grouped_linear": + if role.tensor_type == "weight": + weight_count += 1 + if weight_count == 1: + return IdentityQuantizer() + return _mxfp8(tex.DType.kFloat8E4M3) + return _mxfp8(tex.DType.kFloat8E4M3) + + model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="mix Identity-backed and quantized"): + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + model(x, m_splits=m_splits) + + def test_identity_contiguous_preserves_wrapper_and_values(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16).t() + t = IdentityQuantizer()(x) + + out = t.contiguous() + + assert isinstance(out, IdentityTensor) + assert out.is_contiguous() + torch.testing.assert_close(out.dequantize(), x.contiguous(), rtol=0.0, atol=0.0) + + def test_hybrid_identity_contiguous_preserves_wrapper_and_values(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + + out = t.contiguous() + + assert out is t + assert isinstance(out, HybridQuantizedTensor) + assert isinstance(out.rowwise_sub_storage, IdentityTensor) + assert isinstance(out.columnwise_sub_storage, IdentityTensor) + torch.testing.assert_close(out.dequantize(), x, rtol=0.0, atol=0.0) + + def test_hybrid_identity_cpu_preserves_nested_storage_types(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + + out = t.cpu() + + assert isinstance(out, HybridQuantizedTensor) + assert out.device.type == "cpu" + assert isinstance(out.rowwise_sub_storage, IdentityTensor) + assert isinstance(out.columnwise_sub_storage, IdentityTensor) + assert out.rowwise_sub_storage.device.type == "cpu" + assert out.columnwise_sub_storage.device.type == "cpu" + torch.testing.assert_close(out.dequantize(), x.cpu(), rtol=0.0, atol=0.0) + assert len(out.get_data_tensors()) == 4 + out.copy_(torch.ones_like(x, device="cpu")) + torch.testing.assert_close( + out.dequantize(), torch.ones_like(x, device="cpu"), rtol=0.0, atol=0.0 + ) + + def test_hybrid_quantizer_copy_preserves_parent_flags(self): + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + q.set_usage(rowwise=True, columnwise=False) + q.internal = True + q.optimize_for_gemm = True + + out = q.copy() + + assert isinstance(out, HybridQuantizer) + assert out is not q + assert out.rowwise_quantizer is not q.rowwise_quantizer + assert out.columnwise_quantizer is not q.columnwise_quantizer + assert out.rowwise_usage is True + assert out.columnwise_usage is False + assert out.internal is True + assert out.optimize_for_gemm is True + assert out.rowwise_quantizer.rowwise_usage is True + assert out.rowwise_quantizer.columnwise_usage is False + assert out.columnwise_quantizer.rowwise_usage is False + assert out.columnwise_quantizer.columnwise_usage is True + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_te_ops_basic_linear_accepts_hybrid_identity_quantized_weight(self): + import transformer_engine.pytorch.ops as te_ops + + def qfactory(role): # pylint: disable=unused-argument + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + + torch.manual_seed(1701) + ref = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) + custom_recipe = CustomRecipe(qfactory=qfactory) + torch.manual_seed(1702) + with te.quantized_model_init(enabled=True, recipe=custom_recipe): + test = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + test.weight.copy_(ref.weight) + + torch.manual_seed(1703) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + x_ref = x.detach().clone().requires_grad_(True) + x_test = x.detach().clone().requires_grad_(True) + + y_ref = ref(x_ref) + with te.autocast(enabled=True, recipe=custom_recipe): + y_test = test(x_test) + + torch.manual_seed(1704) + grad_output = torch.randn_like(y_ref) + y_ref.backward(grad_output) + y_test.backward(grad_output) + + assert isinstance(test.weight, HybridQuantizedTensor) + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0.0, atol=0.0) + torch.testing.assert_close(test.weight.grad, ref.weight.grad, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(lambda role: IdentityQuantizer(), id="identity"), + pytest.param( + lambda role: HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ), + id="hybrid_identity", + ), + ], + ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactory): + import transformer_engine.pytorch.ops as te_ops + + ref = te_ops.GELU() + test = te_ops.Sequential(te_ops.Quantize(forward=True), te_ops.GELU()) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + x_ref = x.detach().clone().requires_grad_(True) + x_test = x.detach().clone().requires_grad_(True) + + y_ref = ref(x_ref) + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + y_test = test(x_test) + + torch.manual_seed(1801) + grad_output = torch.randn_like(y_ref) + y_ref.backward(grad_output) + y_test.backward(grad_output) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0.0, atol=0.0) + + def test_hybrid_fsdp_rejects_storage_only_sub_storages(self): + row_quantizer = IdentityQuantizer() + col_quantizer = IdentityQuantizer() + row_quantizer.internal = True + col_quantizer.internal = True + q = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + ) + t = q(torch.randn(8, 16, device="cuda", dtype=torch.bfloat16)) + + with pytest.raises(NotImplementedError, match="storage-only rowwise sub-storage"): + t.fsdp_pre_all_gather( + mesh=None, + orig_size=t.shape, + contiguous_orig_stride=t.stride(), + module=None, + mp_policy=None, + ) + + def test_hybrid_quantizer_rejects_nested_quantizer_requests(self): + from transformer_engine.pytorch.quantization import DelayedScalingRequest + + with pytest.raises(TypeError, match="does not support nested QuantizerRequest"): + HybridQuantizer( + rowwise_quantizer=DelayedScalingRequest(), + columnwise_quantizer=IdentityQuantizer(), + ) + + def test_fp8_dpa_rejects_identity_quantizer_with_type_error(self): + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + META_DO, + META_DP, + META_DQKV, + META_O, + META_QKV, + META_S, + ) + + n_fwd = max(META_QKV, META_S, META_O) + 1 + n_bwd = max(META_DO, META_DP, META_DQKV) + 1 + quantizers = { + "scaling_fwd": [IdentityQuantizer() for _ in range(n_fwd)], + "scaling_bwd": [IdentityQuantizer() for _ in range(n_bwd)], + } + + with pytest.raises(TypeError, match="FP8 attention requires FP8-compatible quantizers"): + dpa_utils.get_attention_quantizers(True, quantizers) + + def test_dequantize_bitwise_identical(self): + x = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16) + out = IdentityQuantizer()(x) + assert torch.equal(out.dequantize(), x) + + def test_dtype_cast(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.float32) + out = IdentityQuantizer(dtype=torch.bfloat16)(x) + dequantized = out.dequantize() + assert dequantized.dtype == torch.bfloat16 + torch.testing.assert_close(dequantized, x.to(torch.bfloat16), rtol=0.0, atol=0.0) + + def test_make_empty_honors_configured_dtype(self): + tensor = IdentityQuantizer(dtype=torch.float32).make_empty( + (4, 8), + dtype=torch.bfloat16, + device="cuda", + ) + + assert tensor.dtype == torch.float32 + assert tensor._hp_data.dtype == torch.float32 + assert tensor.dequantize().dtype == torch.float32 + + def test_update_quantized_synchronizes_dtype(self): + dst = IdentityQuantizer().make_empty( + (4, 8), + dtype=torch.bfloat16, + device="cuda", + ) + quantizer = IdentityQuantizer(dtype=torch.float32) + + quantizer.update_quantized(torch.ones_like(dst._hp_data), dst) + + assert dst.dtype == torch.float32 + assert dst._hp_data.dtype == torch.float32 + assert dst.dequantize().dtype == torch.float32 + + @pytest.mark.parametrize("noop", [0, 1]) + def test_update_quantized_honors_noop(self, noop): + quantizer = IdentityQuantizer() + dst = quantizer(torch.full((4, 8), 3.0, device="cuda")) + src = torch.full((4, 8), 7.0, device="cuda") + + quantizer.update_quantized( + src, dst, noop_flag=torch.tensor(noop, dtype=torch.float32, device="cuda") + ) + + expected = 3.0 if noop else 7.0 + torch.testing.assert_close(dst.dequantize(), torch.full_like(src, expected)) + + @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") + def test_hybrid_update_honors_shared_noop(self): + quantizer = HybridQuantizer( + rowwise_quantizer=_fp8_cs(), + columnwise_quantizer=IdentityQuantizer(), + ) + old = torch.full((32, 32), 3.0, dtype=torch.bfloat16, device="cuda") + dst = quantizer(old) + expected_row = dst._rowwise_storage.dequantize().clone() + + quantizer.update_quantized( + torch.full_like(old, 7.0), + dst, + noop_flag=torch.ones(1, dtype=torch.float32, device="cuda"), + ) + + torch.testing.assert_close(dst._rowwise_storage.dequantize(), expected_row) + torch.testing.assert_close(dst._columnwise_storage.dequantize(), old) + + def test_hybrid_new_zeros_preserves_identity_dtypes(self): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(dtype=torch.bfloat16), + columnwise_quantizer=IdentityQuantizer(dtype=torch.float32), + ) + tensor = quantizer(torch.ones(4, 8, dtype=torch.bfloat16, device="cuda")) + + zeros = tensor.new_zeros(tensor.shape) + + assert zeros.rowwise_sub_storage._hp_data.dtype == torch.bfloat16 + assert zeros.columnwise_sub_storage._hp_data.dtype == torch.float32 + + def test_storage_dequantize_defaults_to_nominal_dtype(self): + payload = torch.randn(4, 8, dtype=torch.float32) + storage = IdentityTensorStorage( + hp_data=payload, + fake_dtype=torch.bfloat16, + ) + + default = storage.dequantize() + assert default.dtype == torch.bfloat16 + torch.testing.assert_close(default, payload.to(torch.bfloat16), rtol=0.0, atol=0.0) + + explicit = storage.dequantize(dtype=torch.float16) + assert explicit.dtype == torch.float16 + torch.testing.assert_close(explicit, payload.to(torch.float16), rtol=0.0, atol=0.0) + + same_dtype_storage = IdentityTensorStorage( + hp_data=payload, + fake_dtype=torch.float32, + ) + assert same_dtype_storage.dequantize() is payload + + def test_update_usage_is_noop(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + st = q(x) + st.update_usage(rowwise_usage=False, columnwise_usage=True) + assert torch.equal(st.dequantize(), x) + assert st.get_usages() == {"rowwise": True, "columnwise": True} + + def test_save_restore_roundtrip(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + st = q(x) + tensors, _ = st.prepare_for_saving() + assert st._hp_data is None + leftover = st.restore_from_saved(tensors) + assert leftover == [] + assert torch.equal(st.dequantize(), x) + + def test_update_quantized_inplace(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + st = q.make_empty((4, 8), dtype=torch.bfloat16, device="cuda") + q.update_quantized(x, st) + assert torch.equal(st.dequantize(), x) + + def test_tensor_ops_preserve_identity_and_values(self): + x = torch.arange(24, device="cuda", dtype=torch.bfloat16).reshape(6, 4) + t = IdentityQuantizer()(x) + + view = t.view(3, 8) + assert isinstance(view, IdentityTensor) + torch.testing.assert_close(view.dequantize(), x.view(3, 8), rtol=0.0, atol=0.0) + + pieces = torch.split(t, 2, dim=0) + assert all(isinstance(piece, IdentityTensor) for piece in pieces) + for piece, ref in zip(pieces, torch.split(x, 2, dim=0)): + torch.testing.assert_close(piece.dequantize(), ref, rtol=0.0, atol=0.0) + + sliced = t[1:5:2] + assert isinstance(sliced, IdentityTensor) + torch.testing.assert_close(sliced.dequantize(), x[1:5:2], rtol=0.0, atol=0.0) + + strided = torch.as_strided(t, (3, 4), (4, 1), 4) + assert isinstance(strided, IdentityTensor) + torch.testing.assert_close(strided.dequantize(), x[1:4], rtol=0.0, atol=0.0) + + cloned = torch.clone(t) + assert isinstance(cloned, IdentityTensor) + torch.testing.assert_close(cloned.dequantize(), x, rtol=0.0, atol=0.0) + + zeros = t.new_zeros((2, 3)) + assert isinstance(zeros, IdentityTensor) + torch.testing.assert_close( + zeros.dequantize(), + torch.zeros((2, 3), device="cuda", dtype=x.dtype), + rtol=0.0, + atol=0.0, + ) + + dst = IdentityQuantizer().make_empty(x.shape, dtype=x.dtype, device="cuda") + dst.copy_(t) + torch.testing.assert_close(dst.dequantize(), x, rtol=0.0, atol=0.0) + + def test_view_ops_preserve_offsets_strides_and_aliasing(self): + base = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor = IdentityQuantizer()(base) + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 2, 6, 1) + omitted_offset = torch.ops.aten.as_strided.default( + sliced, + [2, 4], + [4, 1], + ) + explicit_offset = torch.ops.aten.as_strided.default( + sliced, + [2, 4], + [4, 1], + 0, + ) + torch.testing.assert_close(omitted_offset.dequantize(), base[2:4], rtol=0.0, atol=0.0) + torch.testing.assert_close(explicit_offset.dequantize(), base[:2], rtol=0.0, atol=0.0) + assert sliced.storage_offset() == base[2:].storage_offset() + assert omitted_offset.storage_offset() == base[2:].storage_offset() + assert explicit_offset.storage_offset() == 0 + + noncontiguous = torch.ops.aten.slice.Tensor(tensor, 1, 0, 4, 2) + assert noncontiguous.stride() == base[:, ::2].stride() + assert noncontiguous.storage_offset() == base[:, ::2].storage_offset() + replacement = torch.full_like(base[:, ::2], -3) + noncontiguous.copy_(replacement) + torch.testing.assert_close(base[:, ::2], replacement, rtol=0.0, atol=0.0) + + pieces = torch.split(tensor, 2, dim=0) + piece_value = torch.full_like(base[:2], 7) + pieces[0].copy_(piece_value) + torch.testing.assert_close(base[:2], piece_value, rtol=0.0, atol=0.0) + + def test_detach_and_clone_preserve_view_metadata(self): + base = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor = IdentityQuantizer()(base) + + offset_view = torch.ops.aten.slice.Tensor(tensor, 0, 2, 6, 1) + detached = offset_view.detach() + assert detached.stride() == detached._hp_data.stride() == offset_view.stride() + assert detached.storage_offset() == detached._hp_data.storage_offset() == 8 + replacement = torch.full_like(base[2:6], -5) + detached.copy_(replacement) + torch.testing.assert_close(base[2:6], replacement, rtol=0.0, atol=0.0) + + base2 = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor2 = IdentityQuantizer()(base2) + noncontiguous = torch.ops.aten.slice.Tensor(tensor2, 1, 0, 4, 2) + detached_noncontiguous = noncontiguous.detach() + assert detached_noncontiguous.stride() == detached_noncontiguous._hp_data.stride() + assert detached_noncontiguous.storage_offset() == 0 + + transposed = torch.ops.aten.as_strided.default( + tensor2, + [4, 6], + [1, 4], + ) + cloned = transposed.clone() + assert cloned.stride() == cloned._hp_data.stride() + assert cloned.storage_offset() == cloned._hp_data.storage_offset() + torch.testing.assert_close(cloned.dequantize(), base2.as_strided((4, 6), (1, 4))) + original_snapshot = base2.clone() + cloned.copy_(torch.zeros_like(cloned.dequantize())) + torch.testing.assert_close(base2, original_snapshot, rtol=0.0, atol=0.0) + + def test_copy_from_plain_and_identity_sources(self): + dst_data = torch.zeros(3, 4, device="cuda", dtype=torch.float32) + dst = IdentityQuantizer()(dst_data) + + plain_src = torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + assert dst.copy_(plain_src) is dst + torch.testing.assert_close(dst.dequantize(), plain_src, rtol=0.0, atol=0.0) + + identity_src_data = plain_src.neg() + identity_src = IdentityQuantizer()(identity_src_data) + assert dst.copy_(identity_src) is dst + torch.testing.assert_close(dst.dequantize(), identity_src_data, rtol=0.0, atol=0.0) + + def test_fsdp_pre_post_all_gather_roundtrip(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + sharded_tensors, metadata = t.fsdp_pre_all_gather( + mesh=None, orig_size=t.shape, contiguous_orig_stride=None, module=None, mp_policy=None + ) + gathered, outputs = t.fsdp_post_all_gather(sharded_tensors, metadata, t.dtype, out=None) + assert isinstance(gathered, IdentityTensor) + assert outputs is sharded_tensors + torch.testing.assert_close(gathered.dequantize(), x, rtol=0.0, atol=0.0) + + reuse, _ = t.fsdp_post_all_gather(sharded_tensors, metadata, t.dtype, out=gathered) + assert reuse is gathered + torch.testing.assert_close(reuse.dequantize(), x, rtol=0.0, atol=0.0) + + def test_torch_weights_only_load_preserves_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, IdentityTensor) + assert isinstance(loaded._quantizer, IdentityQuantizer) + torch.testing.assert_close(loaded.dequantize(), x, rtol=0.0, atol=0.0) + + def test_torch_weights_only_load_preserves_hybrid_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, HybridQuantizedTensor) + assert isinstance(loaded._quantizer, HybridQuantizer) + assert isinstance(loaded._rowwise_storage, IdentityTensor) + assert isinstance(loaded._columnwise_storage, IdentityTensor) + torch.testing.assert_close(loaded.dequantize(), x, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_torch_weights_only_load_preserves_hybrid_mxfp8_identity_tensor(self): + x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_mxfp8(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + expected = t.dequantize() + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, HybridQuantizedTensor) + assert isinstance(loaded._quantizer, HybridQuantizer) + assert isinstance(loaded._columnwise_storage, IdentityTensor) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_cpu_offload_roundtrip_identity_exact(self): + x = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + + reloaded = _offload_roundtrip(t) + + assert isinstance(reloaded, IdentityTensor) + torch.testing.assert_close(reloaded.dequantize(), x, rtol=0.0, atol=0.0) + + def test_replace_raw_data_preserves_identity_values(self): + from transformer_engine.pytorch.tensor.utils import replace_raw_data + + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + new_raw = torch.empty_like(x) + replace_raw_data(t, new_raw) + assert t._hp_data is new_raw + torch.testing.assert_close(t.dequantize(), x, rtol=0.0, atol=0.0) + + def test_quantize_master_weights_identity_exact_nonzero_offset(self): + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing, + quantize_master_weights, + ) + + group = _ensure_single_rank_dp_group() + q = IdentityQuantizer() + weight = q.make_empty((4, 8), dtype=torch.bfloat16, device="cuda") + original = torch.randn_like(weight.dequantize()) + q.update_quantized(original, weight) + + master_full = torch.randn(4, 8, device="cuda", dtype=torch.float32) + start_offset = master_full.numel() // 2 + master_shard = master_full.reshape(-1)[start_offset:].contiguous() + + quantize_master_weights([weight], [master_shard], [start_offset], group=group) + post_all_gather_processing([weight]) + + expected = original.clone() + expected.reshape(-1)[start_offset:] = master_shard.to(torch.bfloat16) + torch.testing.assert_close(weight.dequantize(), expected, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_quantize_master_weights_hybrid_identity_fp8_current(self): + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing, + quantize_master_weights, + ) + + group = _ensure_single_rank_dp_group() + recipe = CustomRecipe(qfactory=fwd_hp_bwd_fp8_factory) + torch.manual_seed(123) + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(32, 32, bias=False, params_dtype=torch.bfloat16).cuda() + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor) + assert isinstance(weight._rowwise_storage, IdentityTensorStorage) + + master = torch.randn_like(weight.dequantize(dtype=torch.float32)).reshape(-1).contiguous() + quantize_master_weights([weight], [master], [0], group=group) + post_all_gather_processing([weight]) + + expected = master.to(torch.bfloat16) + row_deq = weight._rowwise_storage.dequantize().reshape(-1) + col_deq = weight._columnwise_storage.dequantize(dtype=torch.float32).reshape(-1) + torch.testing.assert_close(row_deq, expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(col_deq, master, rtol=0.125, atol=0.1) + + +# ── te.Linear integration ──────────────────────────────────────────── + + +def _make_linears(in_f, out_f, seed=1234, dtype=torch.bfloat16, bias=True): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + ref = te.Linear(in_f, out_f, bias=bias, params_dtype=dtype).cuda() + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + test = te.Linear(in_f, out_f, bias=bias, params_dtype=dtype).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + return ref, test + + +def _rel_l2_error(actual, reference): + """Relative L2-norm error ``||actual - reference|| / ||reference||``. + + The right metric for comparing a quantized result to a high-precision + reference: element-wise ``rtol`` is meaningless here because reference grads + contain near-zero entries (relative error on ~1e-12 values explodes), while + the aggregate norm error reflects the true quantization noise. + """ + a = actual.float() + b = reference.float() + return (a - b).norm().item() / (b.norm().item() + 1e-12) + + +def _ensure_single_rank_dp_group(): + import pathlib + import tempfile + + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _fwd_bwd(model, x, recipe=None): + x = x.clone().detach().requires_grad_(True) + if recipe is not None: + with te.autocast(enabled=True, recipe=recipe): + y = model(x) + else: + y = model(x) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +def _fwd_bwd_checkpoint(model, x, recipe, use_reentrant): + x = x.clone().detach().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + if use_reentrant is None: + y = model(x) + else: + y = te.checkpoint(model, x, use_reentrant=use_reentrant) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +_IDENTITY_MODULE_NAMES = ( + "Linear", + "LayerNormLinear", + "LayerNormMLP", + "GroupedLinear", + "TransformerLayer", +) + + +def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16, bias=True): + hidden_size = 64 + ffn_hidden_size = 128 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + if module_name == "Linear": + return te.Linear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "LayerNormMLP": + return te.LayerNormMLP(hidden_size, ffn_hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "GroupedLinear": + return te.GroupedLinear(2, hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "TransformerLayer": + return te.TransformerLayer( + hidden_size, + ffn_hidden_size, + 4, + hidden_dropout=0.0, + attention_dropout=0.0, + bias=bias, + params_dtype=dtype, + ).cuda() + raise ValueError(module_name) + + +def _make_identity_module_pair(module_name, seed=1234, dtype=torch.bfloat16, bias=True): + ref = _make_identity_module(module_name, seed=seed, dtype=dtype, bias=bias) + test = _make_identity_module(module_name, seed=seed + 1, dtype=dtype, bias=bias) + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + return ref, test + + +def _identity_module_input(module_name): + torch.manual_seed(7) + if module_name == "TransformerLayer": + return torch.randn(4, 2, 64, device="cuda", dtype=torch.bfloat16) + return torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + + +def _identity_module_forward(module_name, module, x): + if module_name == "GroupedLinear": + m_splits = torch.tensor([8, 8], device="cuda", dtype=torch.int32) + return module(x, m_splits=m_splits) + return module(x) + + +def _fwd_bwd_module(module_name, model, x, recipe=None): + x = x.clone().detach().requires_grad_(True) + if recipe is not None: + with te.autocast(enabled=True, recipe=recipe): + y = _identity_module_forward(module_name, model, x) + else: + y = _identity_module_forward(module_name, model, x) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestIdentityTEModuleCoverage: + """All-Identity recipes should route every TE module through HP-compatible paths.""" + + @pytest.mark.parametrize("module_name", _IDENTITY_MODULE_NAMES) + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(identity_all_factory, id="plain_identity"), + pytest.param(hybrid_all_identity_factory, id="hybrid_identity"), + ], + ) + def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): + # Keep the GEMM topology identical. Identity-backed grad-output uses an + # unfused bgrad path, while the plain BF16 path may fuse bgrad with + # wgrad. Those mathematically equivalent kernels need not accumulate + # wgrad in the same order on every architecture. + ref, test = _make_identity_module_pair(module_name, seed=7300, bias=False) + x = _identity_module_input(module_name) + recipe = CustomRecipe(qfactory=qfactory) + + y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd_module(module_name, test, x, recipe=recipe) + + # Identity is a high-precision passthrough. Entering the recipe context + # must not change any module's BF16 kernel selection or arithmetic. + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_grouped_linear_mxfp8_forward_identity_backward_matches_override(self): + def mxfp8_all_factory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + def run(model, x, recipe): + x = x.detach().clone().requires_grad_(True) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + with te.autocast(enabled=True, recipe=recipe): + y = model(x, m_splits=m_splits) + torch.manual_seed(9001) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + torch.manual_seed(8300) + ref = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + torch.manual_seed(8301) + test = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + + torch.manual_seed(8302) + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + y_bo, dx_bo, wg_bo = run( + ref, + x, + CustomRecipe(qfactory=mxfp8_all_factory, backward_override="high_precision"), + ) + y_id, dx_id, wg_id = run( + test, + x, + CustomRecipe(qfactory=_hybrid_quantized_fwd_identity_bwd_factory("mxfp8")), + ) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + +class TestIdentityHybridFormatProtocols: + SHAPE = (256, 256) + OFFLOAD_SHAPE = (1024, 1024) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_FORMATS) + def test_save_restore_keeps_identity_direction_exact(self, format_name): + torch.manual_seed(401) + x = torch.randn(*self.SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + hybrid = q.quantize(x) + expected_row = hybrid._rowwise_storage.dequantize().clone() + expected_col = x.clone() + + tensors, obj = hybrid.prepare_for_saving() + leftover = obj.restore_from_saved(tensors) + + assert leftover == [] + assert isinstance(hybrid._columnwise_storage, IdentityTensorStorage) + torch.testing.assert_close( + hybrid._columnwise_storage.dequantize(), expected_col, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + hybrid._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 + ) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_FORMATS) + def test_cpu_offload_keeps_identity_direction_exact(self, format_name): + torch.manual_seed(402) + x = torch.randn(*self.OFFLOAD_SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + hybrid = q.quantize(x) + expected_row = hybrid._rowwise_storage.dequantize().clone() + + reloaded = _offload_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert isinstance(reloaded._columnwise_storage, IdentityTensorStorage) + torch.testing.assert_close(reloaded._columnwise_storage.dequantize(), x, rtol=0.0, atol=0.0) + torch.testing.assert_close( + reloaded._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 + ) + + @pytest.mark.parametrize("format_name", ["mxfp8", "float8_blockwise", "nvfp4"]) + def test_quantize_master_weights_per_block_hybrid_identity_rejected(self, format_name): + if format_name == "mxfp8" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + if format_name == "float8_blockwise" and not fp8_block_scaling_available: + pytest.skip(f"Float8Blockwise: {reason_for_no_fp8_block_scaling}") + if format_name == "nvfp4" and not (fp8_available and nvfp4_available): + pytest.skip(f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}") + + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + x = torch.randn(*self.SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + weight = q.quantize(x) + master = torch.randn_like(x, dtype=torch.float32).reshape(-1).contiguous() + + with pytest.raises(NotImplementedError, match="HybridQuantizer"): + quantize_master_weights([weight], [master], [0], group=group) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_RECOMPUTE_FORMATS) + @pytest.mark.parametrize("use_reentrant", [True, False]) + def test_activation_recompute_matches_no_checkpoint(self, format_name, use_reentrant): + recipe = CustomRecipe(qfactory=_hybrid_quantized_fwd_identity_bwd_factory(format_name)) + ref, test = _make_linears(128, 128, seed=440) + torch.manual_seed(441) + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16) + + y_ref, dx_ref, wg_ref = _fwd_bwd_checkpoint(ref, x, recipe, use_reentrant=None) + y_test, dx_test, wg_test = _fwd_bwd_checkpoint(test, x, recipe, use_reentrant=use_reentrant) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_test) == len(wg_ref) + for g_test, g_ref in zip(wg_test, wg_ref): + torch.testing.assert_close(g_test, g_ref, rtol=0.0, atol=0.0) + + +_ZOO_DEQUANTIZED_MODULES = [ + pytest.param("Linear", id="linear"), + pytest.param("LayerNormLinear", id="layernorm_linear"), + pytest.param("GroupedLinear", id="grouped_linear"), + pytest.param( + "LayerNormMLP", + marks=pytest.mark.xfail( + reason="LayerNormMLP does not support built-in backward_override=dequantized", + raises=AssertionError, + strict=True, + ), + id="layernorm_mlp", + ), + pytest.param("LayerNormLinearLinear", id="layernorm_linear_linear"), +] + +_ZOO_DEQUANTIZED_CASES = [ + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + id="mxfp8", + ), + pytest.param( + "nvfp4_row_scaled", + marks=pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}"), + id="nvfp4_row_scaled", + ), +] + + +def _make_zoo_dequantized_module(module_name, *, save_original_input=False): + hidden_size = 128 + if module_name == "Linear": + return te.Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ).cuda() + if module_name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, hidden_size, params_dtype=torch.bfloat16).cuda() + if module_name == "GroupedLinear": + return te.GroupedLinear( + 2, + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ).cuda() + if module_name == "LayerNormMLP": + return te.LayerNormMLP( + hidden_size, + 2 * hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + if module_name == "LayerNormLinearLinear": + return torch.nn.Sequential( + te.LayerNormLinear(hidden_size, hidden_size, params_dtype=torch.bfloat16), + te.Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ), + ).cuda() + raise ValueError(module_name) + + +def _make_zoo_dequantized_module_pair(module_name): + torch.manual_seed(9400) + torch.cuda.manual_seed(9400) + ref = _make_zoo_dequantized_module(module_name, save_original_input=False) + torch.manual_seed(9401) + torch.cuda.manual_seed(9401) + test = _make_zoo_dequantized_module(module_name, save_original_input=True) + test.load_state_dict(ref.state_dict()) + return ref, test + + +def _zoo_dequantized_input(module_name): + torch.manual_seed(9402) + batch = 128 if module_name == "GroupedLinear" else 64 + return torch.randn(batch, 128, device="cuda", dtype=torch.bfloat16) + + +def _zoo_dequantized_forward(module_name, module, inp): + if module_name == "GroupedLinear": + m_splits = torch.tensor([64, 64], device="cuda", dtype=torch.int32) + return module(inp, m_splits=m_splits) + return module(inp) + + +def _fwd_bwd_zoo_dequantized_module(module_name, module, inp, recipe): + inp = inp.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + out = _zoo_dequantized_forward(module_name, module, inp) + torch.manual_seed(9403) + grad = torch.randn_like(out) + out.backward(grad) + param_grads = { + name: param.grad.detach().clone() + for name, param in module.named_parameters() + if param.grad is not None + } + return out.detach().clone(), inp.grad.detach().clone(), param_grads + + +def _mxfp8_all_qfactory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + +def _zoo_dequantized_qfactory(case_name): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory, + ) + + if case_name == "mxfp8": + return mxfp8_fwd_dequantized_bwd_quantizer_factory + if case_name == "nvfp4_row_scaled": + return nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory + raise ValueError(case_name) + + +def _zoo_dequantized_recipes(case_name): + from transformer_engine.common import recipe as te_recipe + + if case_name == "mxfp8": + return ( + CustomRecipe(qfactory=_mxfp8_all_qfactory, backward_override="dequantized"), + CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)), + ) + if case_name == "nvfp4_row_scaled": + ref_recipe = te_recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) + ref_recipe.fp4_quant_fwd_inp = te_recipe.QParams() + ref_recipe.fp4_quant_fwd_weight = te_recipe.QParams() + ref_recipe.fp4_quant_bwd_grad = te_recipe.QParams() + return ( + ref_recipe, + CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)), + ) + raise ValueError(case_name) + + +def _assert_zoo_layernorm_mlp_role_semantics(case_name): + module = _make_zoo_dequantized_module("LayerNormMLP") + qfactory = _zoo_dequantized_qfactory(case_name) + + fwd_roles = module.get_quantizer_roles(fwd=True, num_quantizers=6) + fwd_gemm_roles = [ + role for role in fwd_roles if role is not None and role.tensor_type in ("input", "weight") + ] + assert len(fwd_gemm_roles) == 5 + for role in fwd_gemm_roles: + quantizer = qfactory(role) + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) + if case_name == "mxfp8": + assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) + else: + assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) + assert quantizer.rowwise_quantizer.row_scaled_nvfp4 == (role.tensor_type == "input") + + bwd_roles = module.get_quantizer_roles(fwd=False, num_quantizers=4) + grad_output_roles = [ + role for role in bwd_roles if role is not None and role.tensor_type == "grad_output" + ] + assert len(grad_output_roles) == 3 + for role in grad_output_roles: + assert isinstance(qfactory(role), IdentityQuantizer) + + +class TestZooDequantizedBackwardFactoryModuleCoverage: + """Zoo dequantized-backward recipes should match base recipes across TE modules.""" + + @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) + @pytest.mark.parametrize("module_name", _ZOO_DEQUANTIZED_MODULES) + def test_matches_base_recipe_bitwise(self, case_name, module_name): + ref, test = _make_zoo_dequantized_module_pair(module_name) + inp = _zoo_dequantized_input(module_name) + ref_recipe, qfactory_recipe = _zoo_dequantized_recipes(case_name) + + y_ref, dx_ref, grads_ref = _fwd_bwd_zoo_dequantized_module( + module_name, ref, inp, ref_recipe + ) + y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( + module_name, test, inp, qfactory_recipe + ) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) + assert grads_test.keys() == grads_ref.keys() + for name, grad_ref in grads_ref.items(): + torch.testing.assert_close( + grads_test[name], + grad_ref, + rtol=0.0, + atol=0.0, + msg=f"{case_name}/{module_name} grad mismatch for {name}", + ) + + @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) + def test_layernorm_mlp_runs_and_uses_dequantized_backward_roles(self, case_name): + _assert_zoo_layernorm_mlp_role_semantics(case_name) + + module = _make_zoo_dequantized_module("LayerNormMLP") + inp = _zoo_dequantized_input("LayerNormMLP") + recipe = CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)) + out, dx, grads = _fwd_bwd_zoo_dequantized_module("LayerNormMLP", module, inp, recipe) + + assert torch.isfinite(out).all() + assert torch.isfinite(dx).all() + expected_grads = {name for name, param in module.named_parameters() if param.requires_grad} + assert grads.keys() == expected_grads + for grad in grads.values(): + assert torch.isfinite(grad).all() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestIdentityLinear: + """End-to-end te.Linear with Identity-based recipes.""" + + IN_F = 128 + OUT_F = 128 + BATCH = 64 + + def _input(self): + torch.manual_seed(7) + return torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(identity_all_factory, id="plain_identity"), + pytest.param(hybrid_all_identity_factory, id="hybrid_identity"), + ], + ) + def test_identity_bias_path_matches_bf16_bitwise(self, qfactory): + """The unfused Identity bgrad path must preserve exact bias semantics. + + Weight-gradient bitwise parity is covered with ``bias=False`` so both + sides use the same GEMM topology. This bias-enabled test separately + locks output, input-gradient, and bias-gradient parity to zero tolerance. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=True) + x = self._input() + + y_ref, dx_ref, _ = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, _ = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=qfactory)) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(test.bias.grad, ref.bias.grad, rtol=0.0, atol=0.0) + + def test_whole_layer_hp_matches_bf16_bitwise(self): + """Identity for every slot => all GEMMs high precision => bitwise-equal + to a plain BF16 te.Linear (no autocast).""" + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=identity_all_factory)) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fwd_hp_bwd_fp8_forward_bitwise(self): + """High-precision forward must be bitwise-equal to BF16 forward; the + backward runs in FP8 (finite, close to BF16 within a loose tolerance).""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_h, dx_h, wg_h = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_hp_bwd_fp8_factory)) + + # Forward is high precision -> bitwise equal. + torch.testing.assert_close(y_h, y_ref, rtol=0.0, atol=0.0) + # Backward is FP8 (E4M3 weight col, E5M2 grad) -> relative L2 error vs the + # BF16 reference reflects pure FP8 quant noise. Measured: dgrad ~5.7e-2, + # weight-grad ~5.8e-2 (E4M3 ~6% step). Bound 7e-2 keeps a small margin. + assert torch.isfinite(dx_h).all() + assert _rel_l2_error(dx_h, dx_ref) < 7e-2 + for g, g_ref in zip(wg_h, wg_ref): + assert torch.isfinite(g).all() + if g.dim() == 1: + # Bias grad = sum(dY) is computed in high precision (dY is bitwise + # identical since the forward is bitwise), so it must match exactly. + torch.testing.assert_close(g, g_ref, rtol=0.0, atol=0.0) + else: + assert _rel_l2_error(g, g_ref) < 7e-2 + + def test_fwd_fp8_bwd_hp_runs_and_backward_high_precision(self): + """FP8 forward + high-precision backward. Forward differs from BF16 + (quantized), backward GEMMs run in high precision.""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, _ = _fwd_bwd(ref, x, recipe=None) + y_q, dx_q, wg_q = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory)) + + # Forward is FP8 (E4M3) -> relative L2 error vs BF16 is the quant noise. + # Measured ~3.7e-2; bound 5e-2. + assert torch.isfinite(y_q).all() + assert _rel_l2_error(y_q, y_ref) < 5e-2 + # Backward GEMMs run in high precision. dgrad differs from the BF16 + # reference only because the FP8 forward perturbs dY; measured ~1.2e-2, + # bound 3e-2 (the bitwise HP-backward guarantee is locked by the + # backward_override equivalence test below). + assert torch.isfinite(dx_q).all() + assert _rel_l2_error(dx_q, dx_ref) < 3e-2 + for g in wg_q: + assert torch.isfinite(g).all() + + def test_hybrid_all_identity_matches_bf16_bitwise(self): + """All-Identity through the *hybrid* container must be bitwise-equal to a + plain BF16 te.Linear. Complements the non-hybrid whole-layer-HP test: this + exercises HybridQuantizedTensor with Identity sub-storages in both + directions and the per-operand unwrap of every GEMM.""" + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=hybrid_all_identity_factory) + ) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_identity_reproduces_backward_override_high_precision_bitwise(self): + """The per-direction Identity machinery must reproduce + ``backward_override="high_precision"`` **bitwise**. + + Both runs quantize the forward to the same FP8 (current scaling) and run + the backward in high precision against the original operands. The Identity + path expresses this per-tensor (weight/input = Hybrid(row=FP8, col=Identity), + grad = Identity); the reference uses the global ``backward_override`` knob. + Identical forward FP8 + identical original HP backward operands => bitwise. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="high_precision"), + ) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory)) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + def test_identity_reproduces_backward_override_dequantized_bitwise(self): + """Per-direction Identity must reproduce ``backward_override=dequantized``. + + Both runs quantize forward to the same FP8 values. The reference saves + rowwise fprop payloads and dequantizes them for high-precision backward; + the hybrid path stores Identity columnwise data sourced from the rowwise + dequantized value and uses Identity grad tensors. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="dequantized"), + ) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_rowwise_dequantized_hp_factory) + ) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + def test_identity_matches_bf16_multistep_training_bitwise(self): + """Multi-step SGD: an all-Identity recipe must track a plain BF16 + te.Linear bitwise across optimizer steps (no drift from workspace caching + or any hidden state).""" + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) + opt_test = torch.optim.SGD(test.parameters(), lr=0.1) + recipe = CustomRecipe(qfactory=identity_all_factory) + + for step in range(4): + torch.manual_seed(1000 + step) + x = torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(2000 + step) + target = torch.randn(self.BATCH, self.OUT_F, device="cuda", dtype=torch.bfloat16) + + opt_ref.zero_grad() + y_ref = ref(x) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + opt_ref.step() + + opt_test.zero_grad() + with te.autocast(enabled=True, recipe=recipe): + y_test = test(x) + loss_test = torch.nn.functional.mse_loss(y_test, target) + loss_test.backward() + opt_test.step() + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + torch.testing.assert_close(p_test, p_ref, rtol=0.0, atol=0.0, msg=f"step {step}") + + def test_quantized_model_init_identity_matches_bf16_bitwise(self): + """Persistent Identity params from quantized_model_init should match BF16 exactly.""" + torch.manual_seed(314) + ref = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(2718) + with te.quantized_model_init(enabled=True, recipe=recipe): + test = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + assert isinstance(p_test, IdentityTensor) + p_test.copy_(p_ref) + + x = self._input() + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=recipe) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_quantized_model_init_identity_training_loss_decreases_bitwise(self): + """All-Identity quantized params train like BF16 and loss decreases.""" + torch.manual_seed(777) + ref = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(888) + with te.quantized_model_init(enabled=True, recipe=recipe): + test = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + assert isinstance(p_test, IdentityTensor) + p_test.copy_(p_ref) + + torch.manual_seed(909) + x = torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + target = torch.zeros(self.BATCH, self.OUT_F, device="cuda", dtype=torch.bfloat16) + opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) + opt_test = torch.optim.SGD(test.parameters(), lr=0.1) + losses_ref = [] + losses_test = [] + + for _ in range(5): + opt_ref.zero_grad() + y_ref = ref(x) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + opt_ref.step() + losses_ref.append(loss_ref.detach().clone()) + + opt_test.zero_grad() + with te.autocast(enabled=True, recipe=recipe): + y_test = test(x) + loss_test = torch.nn.functional.mse_loss(y_test, target) + loss_test.backward() + opt_test.step() + losses_test.append(loss_test.detach().clone()) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + torch.testing.assert_close(p_test.dequantize(), p_ref, rtol=0.0, atol=0.0) + + assert all( + losses_ref[i + 1].item() < losses_ref[i].item() for i in range(len(losses_ref) - 1) + ), f"BF16 loss did not strictly decrease: {[x.item() for x in losses_ref]}" + for loss_test, loss_ref in zip(losses_test, losses_ref): + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("use_reentrant", [True, False]) + def test_identity_activation_recompute_matches_bf16_bitwise(self, use_reentrant): + """All-Identity recompute should be exactly the BF16 no-checkpoint path.""" + ref, test = _make_linears(self.IN_F, self.OUT_F, seed=4242, bias=False) + recipe = CustomRecipe(qfactory=identity_all_factory) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd_checkpoint(test, x, recipe, use_reentrant=use_reentrant) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_quantized_model_init_identity_state_dict_save_load_exact(self): + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(5151) + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + + torch.manual_seed(5152) + x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + out_before = model(x) + + buffer = io.BytesIO() + torch.save(model.state_dict(), buffer) + buffer.seek(0) + + with te.quantized_model_init(enabled=True, recipe=recipe): + model2 = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + model2.load_state_dict(torch.load(buffer, weights_only=True)) + + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + out_after = model2(x) + + assert isinstance(model2.weight, IdentityTensor) + torch.testing.assert_close(out_after, out_before, rtol=0.0, atol=0.0) + + def test_load_bf16_state_dict_into_identity_model_exact(self): + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(6161) + ref = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + + model.load_state_dict(ref.state_dict()) + assert isinstance(model.weight, IdentityTensor) + torch.testing.assert_close(model.weight.dequantize(), ref.weight, rtol=0.0, atol=0.0) + + torch.manual_seed(6162) + x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + out_ref = ref(x) + with te.autocast(enabled=True, recipe=recipe): + out_id = model(x) + + torch.testing.assert_close(out_id, out_ref, rtol=0.0, atol=0.0) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2f77ecd66..a07d54d22b 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -20,6 +20,7 @@ Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, + MXFP8TensorStorage, NVFP4Tensor, QuantizedTensor, ) @@ -670,6 +671,31 @@ def test_cpu_dequantize( assert y_cpu.shape == ref_cpu.shape torch.testing.assert_close(y_cpu, ref_cpu, rtol=0, atol=0) + def test_mxfp8_fsdp_extract_keeps_partial_columnwise_scale_block(self) -> None: + """FSDP extraction must retain the scale for a partial final 32-row block.""" + rows, cols = 48, 64 + columnwise_scale_inv = torch.arange(4 * 128, dtype=torch.float32).reshape(4, 128) + storage = MXFP8TensorStorage( + rowwise_data=None, + rowwise_scale_inv=None, + columnwise_data=torch.zeros((rows, cols), dtype=torch.uint8), + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=te.DType.kFloat8E4M3, + quantizer=None, + with_gemm_swizzled_scales=False, + fake_dtype=torch.bfloat16, + ) + + buffers, metadata = storage.fsdp_extract_buffers() + + assert metadata == {"field_names": ("_columnwise_data", "_columnwise_scale_inv")} + assert len(buffers) == 2 + assert buffers[0] is storage._columnwise_data + extracted_scale_inv = buffers[1] + assert extracted_scale_inv is not None + assert extracted_scale_inv.shape == (2, 128) + torch.testing.assert_close(extracted_scale_inv, columnwise_scale_inv[:2], rtol=0, atol=0) + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("dim", [0, 1]) def test_chunk( diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8c209ace15..a89ddba917 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -666,6 +666,14 @@ class CustomRecipe(Recipe): `high_precision` keeps original high-precision operands for backward, and `dequantized` dequantizes saved operands to the active high-precision compute dtype (e.g. BF16/FP16/FP32) for backward. + quantization_alignment : int, default = 128 + Conservative recipe-wide fallback used by automatic padding for grouped operations. + This must be at least the largest alignment required by any quantizer + that ``qfactory`` may return for a grouped operation, across all roles + and module names. The default of 128 safely supports all current TE + formats. It can be lowered when the factory's full output space is known; + for example, a factory restricted to MXFP8 may use 32. + Automatic padding reads this value without invoking ``qfactory``. """ qfactory: Callable[..., Any] @@ -678,15 +686,19 @@ class CustomRecipe(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + quantization_alignment: int = 128 def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + if self.quantization_alignment <= 0: + raise ValueError("CustomRecipe quantization_alignment must be positive.") def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " - f"backward_override={self.backward_override}" + f"backward_override={self.backward_override}, " + f"quantization_alignment={self.quantization_alignment}" ) diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index 3a8536587c..02d64bcfff 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -161,13 +161,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) static_assert(std::is_same::value); const CType scale_inv = 1.0f / block_tile_scale; - size_t row_idx = tile_id_y; - size_t col_idx = tile_id_x; - tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + if (tile_scales_inv_c != nullptr) { + size_t row_idx = tile_id_y; + size_t col_idx = tile_id_x; + tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + } if constexpr (kReturnTranspose) { - row_idx = tile_id_x; - col_idx = tile_id_y; + size_t row_idx = tile_id_x; + size_t col_idx = tile_id_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -189,7 +191,9 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) thrd_tile_out_trans[j].data.elt[i] = scaled_elt; } } - tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length); + if (output_c != nullptr) { + tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length); + } } // Step 4: store transpose into shared memory @@ -388,13 +392,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose static_assert(std::is_same::value); const CType scale_inv = 1.0f / block_tile_scale; - size_t row_idx = tile_id_y; - size_t col_idx = tile_id_x; - tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + if (tile_scales_inv_c != nullptr) { + size_t row_idx = tile_id_y; + size_t col_idx = tile_id_x; + tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + } if constexpr (kReturnTranspose) { - row_idx = tile_id_x; - col_idx = tile_id_y; + size_t row_idx = tile_id_x; + size_t col_idx = tile_id_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -433,8 +439,10 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose thrd_tile_out_trans[j].data.elt[i] = scaled_elt; } } - tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0, - thread_tile_ncols); + if (output_c != nullptr) { + tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0, + thread_tile_ncols); + } } if constexpr (kReturnTranspose) { @@ -492,19 +500,26 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor "with MXFP8, which requires using power of two scaling factors."); } - NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); + const bool return_identity = output.dptr != nullptr; + if (return_identity) { + NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); + } + NVTE_CHECK(return_identity || return_transpose, + "At least one of rowwise or columnwise output must be requested."); const size_t row_length = input.shape.size() > 0 ? input.shape.back() : 1; size_t num_rows = 1; for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) { num_rows *= input.shape.at(i); } - NVTE_CHECK(scale_inv.shape.size() == 2, "scale_inv must have 2 dimensions."); - - size_t scale_k = scale_inv.shape[1]; - - const size_t scale_stride_x = 1; - const size_t scale_stride_y = scale_k; + size_t scale_k = 0; + const size_t scale_stride_x = return_identity ? 1 : 0; + size_t scale_stride_y = 0; + if (return_identity) { + NVTE_CHECK(scale_inv.shape.size() == 2, "scale_inv must have 2 dimensions."); + scale_k = scale_inv.shape[1]; + scale_stride_y = scale_k; + } size_t scale_t_stride_x = 0; size_t scale_t_stride_y = 0; @@ -522,7 +537,9 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor ") and output_t (shape=", output_t.shape, ") have incompatible dims."); } } - NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); + if (return_identity) { + NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); + } NVTE_CHECK(scale_inv_t.shape.size() == 2, "scale_inv_t must have 2 dimensions."); @@ -530,6 +547,8 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor scale_t_stride_y = scale_inv_t.shape[1]; } + const auto out_dtype = return_identity ? output.dtype : output_t.dtype; + const size_t num_blocks_x = DIVUP(row_length, BLOCK_TILE_DIM); const size_t num_blocks_y = DIVUP(num_rows, BLOCK_TILE_DIM); @@ -537,7 +556,7 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor input.dtype, InputType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output.dtype, OutputType, + out_dtype, OutputType, TRANSFORMER_ENGINE_SWITCH_CONDITION( return_transpose, kReturnTranspose, diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 06db28ee27..4c6b7fc67a 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -93,6 +93,12 @@ from transformer_engine.pytorch.tensor import MXFP8Tensor from transformer_engine.pytorch.tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor import NVFP4Tensor +from transformer_engine.pytorch.tensor import HybridQuantizer +from transformer_engine.pytorch.tensor import HybridQuantizedTensorStorage +from transformer_engine.pytorch.tensor import IdentityQuantizer +from transformer_engine.pytorch.tensor import IdentityTensorStorage +from transformer_engine.pytorch.tensor import HybridQuantizedTensor +from transformer_engine.pytorch.tensor import IdentityTensor from transformer_engine.pytorch.tensor.float8_tensor import ( _make_float8_tensor_in_reduce_ex, ) @@ -105,6 +111,12 @@ from transformer_engine.pytorch.tensor.float8_blockwise_tensor import ( _make_float8_blockwise_tensor_in_reduce_ex, ) +from transformer_engine.pytorch.tensor.hybrid_tensor import ( + _make_hybrid_quantized_tensor_in_reduce_ex, +) +from transformer_engine.pytorch.tensor.identity_tensor import ( + _make_identity_tensor_in_reduce_ex, +) try: torch._dynamo.config.error_on_nested_jit_trace = False @@ -129,6 +141,8 @@ MXFP8TensorStorage, NVFP4TensorStorage, Float8BlockwiseQTensorStorage, + HybridQuantizedTensorStorage, + IdentityTensorStorage, # Quantizer types embedded in metadata Quantizer, Float8Quantizer, @@ -136,6 +150,8 @@ MXFP8Quantizer, NVFP4Quantizer, Float8BlockQuantizer, + HybridQuantizer, + IdentityQuantizer, # Python IntEnum used as Quantizer.dtype. DType, # pybind11 enum used as Quantizer.dtype. @@ -146,6 +162,8 @@ _make_mxfp8_tensor_in_reduce_ex, _make_nvfp4_tensor_in_reduce_ex, _make_float8_blockwise_tensor_in_reduce_ex, + _make_hybrid_quantized_tensor_in_reduce_ex, + _make_identity_tensor_in_reduce_ex, ] ) except (ImportError, AttributeError): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 4cc4cab1b8..9ae0df880a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -85,6 +85,75 @@ "_alibi_bias_require_update": False, } + +def _infer_custom_dpa_local_recipes( + fp8_recipe: Recipe, + fp8_meta: Dict[str, Any], + quantizers: Dict[str, Any], +) -> Optional[List[Recipe]]: + """Infer native-equivalent DPA recipe labels for CustomRecipe control-flow. + + CustomRecipe owns quantizer construction, but DPA backend selection and a + few fused-attention branches still dispatch on recipe predicates. Attach + local recipe labels that match the qfactory DPA quantizer family while + keeping the actual qfactory-created quantizers untouched. + """ + try: + qkv_quantizer = quantizers["scaling_fwd"][dpa_utils.META_QKV] + except (KeyError, IndexError, TypeError): + return None + + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + if isinstance(qkv_quantizer, MXFP8Quantizer): + return [ + MXFP8BlockScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + ] + + def _delayed_scaling_recipe() -> Optional[DelayedScaling]: + fwd_state = fp8_meta.get("scaling_fwd") + ds_recipe = getattr(fwd_state, "_inner_delayed_scaling_recipe", None) + if ds_recipe is None: + return None + return DelayedScaling( + fp8_format=ds_recipe.fp8_format, + margin=ds_recipe.margin, + amax_history_len=ds_recipe.amax_history_len, + amax_compute_algo=ds_recipe.amax_compute_algo, + scaling_factor_compute_algo=ds_recipe.scaling_factor_compute_algo, + reduce_amax=ds_recipe.reduce_amax, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + ds_recipe = _delayed_scaling_recipe() + if ds_recipe is not None: + return [ + Float8CurrentScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ), + ds_recipe, + ] + + if isinstance(qkv_quantizer, Float8Quantizer): + ds_recipe = _delayed_scaling_recipe() + if ds_recipe is not None: + return [ds_recipe] + + return None + + """ This feature is **experimental** and subject to change. @@ -499,6 +568,14 @@ def __init__( ) -> None: super().__init__(name=name) + # Cache the native recipe labels inferred from custom DPA quantizers. + # ``init_fp8_metadata`` runs on every forward, while the quantizers only + # change when their recipe state is rebuilt. + self._custom_dpa_local_recipes_cache_key: Optional[Tuple[Any, ...]] = None + self._custom_dpa_local_recipes_cache: Optional[List[Recipe]] = None + self._qkv_capabilities_quantizer: Optional[Any] = None + self._qkv_capabilities_cache: Optional[Tuple[bool, bool]] = None + self.logger = logging.getLogger("DotProductAttention") self.logger.setLevel(attn_log._log_level) if not self.logger.hasHandlers(): @@ -719,6 +796,26 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): super().init_fp8_metadata(num_gemms=num_gemms) + fwd_quantizers = self.quantizers.get("scaling_fwd", ()) + cache_key = ( + id(self.fp8_meta.get("scaling_fwd")), + tuple(id(quantizer) for quantizer in fwd_quantizers), + fp8_recipe.fp8_format, + fp8_recipe.fp8_dpa, + fp8_recipe.fp8_mha, + ) + if cache_key != self._custom_dpa_local_recipes_cache_key: + self._custom_dpa_local_recipes_cache = _infer_custom_dpa_local_recipes( + fp8_recipe, self.fp8_meta, self.quantizers + ) + self._custom_dpa_local_recipes_cache_key = cache_key + + if self._custom_dpa_local_recipes_cache is None: + # Do not leave labels from an earlier supported quantizer + # family attached after a rebuild to an unsupported family. + self.fp8_meta.pop("local_recipes", None) + else: + self.fp8_meta["local_recipes"] = self._custom_dpa_local_recipes_cache return # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to @@ -925,6 +1022,57 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Clear cached workspaces as they were created with the old recipe/quantizer type self._fp8_workspaces.clear() + def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: + """Return MHA boundary capabilities from the canonical QKV quantizer. + + The returned flags are ``(float8_current_scaling, mxfp8_scaling)``. + """ + self.init_fp8_metadata(num_gemms=3) + try: + qkv_quantizer = self.quantizers["scaling_fwd"][dpa_utils.META_QKV] + except (KeyError, IndexError, TypeError) as exc: + role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=self.name or "", + ) + raise RuntimeError( + f"DotProductAttention did not materialize the canonical QKV quantizer for {role}." + ) from exc + + if qkv_quantizer is self._qkv_capabilities_quantizer: + assert self._qkv_capabilities_cache is not None + return self._qkv_capabilities_cache + + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + capabilities = (True, False) + elif isinstance(qkv_quantizer, MXFP8Quantizer): + capabilities = (False, True) + elif isinstance(qkv_quantizer, Float8Quantizer): + capabilities = (False, False) + else: + capabilities = None + + if capabilities is not None: + self._qkv_capabilities_quantizer = qkv_quantizer + self._qkv_capabilities_cache = capabilities + return capabilities + + role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=self.name or "", + ) + raise TypeError( + f"Unsupported CustomRecipe quantizer for {role}: {type(qkv_quantizer).__name__}." + ) + def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> None: """Override to allow multiple recipes. Init scales and amaxes for fwd | bwd.""" if isinstance(recipe, Recipe) and recipe.custom(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6c47d0f1cb..52a9afc07b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2688,14 +2688,15 @@ def get_attention_quantizers(fp8, quantizers): ]: if _q is None and _name in _allow_none: continue - assert isinstance(_q, _fp8_types), ( - "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " - f"but {_name} quantizer is {type(_q).__name__}. " - "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " - "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " - "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " - "(boundary slots like O output and dQKV grad-input)." - ) + if not isinstance(_q, _fp8_types): + raise TypeError( + "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " + f"but {_name} quantizer is {type(_q).__name__}. " + "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " + "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " + "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " + "(boundary slots like O output and dQKV grad-input)." + ) return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 0b73daf667..f87365bf7c 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -500,11 +500,11 @@ def _update_output_quantizer_roles( 1. ``qkv_fp8_output`` — **QKV linear → DPA (fwd)**: the QKV linear's ``output_quantizer_role`` is told its consumer is DPA. - 2. ``proj_fp8_grad`` — **Proj linear ← DPA (bwd)**: proj's - ``grad_input_quantizer_role`` is told its producer is DPA. + 2. ``proj_fp8_grad`` — **Proj linear → DPA (bwd)**: proj's + ``grad_input_quantizer_role`` is told its consumer is DPA. 3. ``dpa_fp8_output`` — **DPA → Proj linear (fwd)**: DPA's ``output_quantizer_role`` is told its consumer is the proj linear. - 4. ``dpa_fp8_output`` — **DPA ← QKV linear (bwd)**: DPA's + 4. ``dpa_fp8_output`` — **DPA → QKV linear (bwd)**: DPA's ``grad_input_quantizer_role`` is told its consumer is QKV linear. When a flag is ``False`` the corresponding role is reset to ``None`` @@ -530,7 +530,7 @@ def _update_output_quantizer_roles( self.query_layer.output_quantizer_role = qkv_output_role self.key_value.output_quantizer_role = qkv_output_role - # ── Boundary 2 (bwd): Proj grad-input ← produced by DPA ────────── + # ── Boundary 2 (bwd): Proj grad-input (dO) → consumed by DPA ───── proj_grad_input_role = ( QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name) if proj_fp8_grad @@ -870,12 +870,27 @@ def forward( # ====================== fp8 = FP8GlobalStateManager.is_fp8_enabled() + custom_recipe = False if _dpa_fp8_recipe == "": fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + custom_recipe = fp8_recipe.custom() fp8_dpa = fp8_recipe.fp8_dpa fp8_mha = fp8_recipe.fp8_mha float8_current_scaling = fp8_recipe.float8_current_scaling() mxfp8_scaling = fp8_recipe.mxfp8() + if fp8 and custom_recipe and fp8_mha: + # Wire every boundary this CustomRecipe may quantize before + # DPA materializes its recipe state. Some quantizer families + # disable the corresponding output below, but pre-wiring avoids + # rebuilding that state after inspecting its canonical QKV slot. + self._update_output_quantizer_roles( + rotary_pos_emb is None, + True, + True, + ) + float8_current_scaling, mxfp8_scaling = ( + self.core_attention.get_qkv_quantization_capabilities() + ) else: fp8_dpa = _dpa_fp8_recipe_dpa fp8_mha = _dpa_fp8_recipe_mha @@ -897,12 +912,18 @@ def forward( # DPA: produce FP8 output to take advantage of O amax from DPA; Projection Gemm can take FP8 or F16 inputs # 1. FP8DS/FP8CS recipe: produce FP8 output # 2. MXFP8 recipe: produce F16 output; again, due to quantization dimensions mismatch - dpa_fp8_output = fp8 and (fp8_dpa or fp8_mha) and not mxfp8_scaling + # For CustomRecipe, fp8_dpa only controls DPA-internal quantization. + # External MHA boundary tensors become FP8 only when fp8_mha is enabled. + dpa_fp8_output_enabled = fp8_mha if custom_recipe else (fp8_dpa or fp8_mha) + dpa_fp8_output = fp8 and dpa_fp8_output_enabled and not mxfp8_scaling # Projection Gemm: match DPA output except # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling - self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Custom fp8_mha boundaries were wired before DPA recipe-state setup so + # querying its canonical QKV quantizer cannot trigger a second build. + if not (fp8 and custom_recipe and fp8_mha and _dpa_fp8_recipe == ""): + self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) # Packed pass-through to DotProductAttention: the fused QKV/KV projection # already produces one packed buffer, which DPA accepts directly via its diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index d5446e7b61..b4a8b21b00 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -4,7 +4,7 @@ """Python interface for GEMM extensions""" -from typing import Iterable, Optional, Tuple, Union, List +from typing import Iterable, Literal, Optional, Tuple, Union, List import os import functools import torch @@ -12,16 +12,21 @@ from ..constants import TE_DType, DType from ..utils import get_sm_count, _empty_tensor -from ..quantized_tensor import Quantizer +from ..quantized_tensor import QuantizedTensorStorage, Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer +from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage +from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer - __all__ = [ "general_gemm", "general_grouped_gemm", @@ -103,6 +108,67 @@ def _nvfp4_row_scaled_gemm_inputs( ) +_NATIVE_GEMM_INPUT_STORAGES = ( + Float8TensorStorage, + MXFP8TensorStorage, + Float8BlockwiseQTensorStorage, + NVFP4TensorStorage, +) + + +def _unwrap_tensor( + tensor: Union[torch.Tensor, QuantizedTensorStorage], + usage: Literal["rowwise", "columnwise"], +) -> Union[torch.Tensor, QuantizedTensorStorage]: + """Prepare a tensor for native or custom GEMM dispatch.""" + if usage not in ("rowwise", "columnwise"): + raise ValueError(f"Unsupported GEMM tensor usage ({usage})") + + # Plain PyTorch tensor + if not isinstance(tensor, QuantizedTensorStorage): + return tensor + + # Select the direction of a hybrid tensor, then process its sub-storage. + if isinstance(tensor, HybridQuantizedTensorStorage): + sub_storage = ( + tensor.rowwise_sub_storage if usage == "rowwise" else tensor.columnwise_sub_storage + ) + return _unwrap_tensor(sub_storage, usage) + + # Preserve custom tensors for custom_gemm dispatch. + if is_custom(tensor): + return tensor + + # Quantized tensor formats with native GEMM support + if isinstance(tensor, _NATIVE_GEMM_INPUT_STORAGES): + return tensor + + # Fall back to high-precision GEMM for other quantized tensor formats. + return tensor.dequantize() + + +_NATIVE_GEMM_OUTPUT_QUANTIZERS = ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, +) + + +def _validate_native_gemm_output_quantizer(quantization_params): + """Validate that the native C++ GEMM path can convert an output quantizer.""" + if quantization_params is not None and not isinstance( + quantization_params, _NATIVE_GEMM_OUTPUT_QUANTIZERS + ): + raise NotImplementedError( + f"{type(quantization_params).__name__} is not supported as a native " + "GEMM output quantizer. " + "Return a TE-native quantizer for output/grad_input roles or disable " + "quantized GEMM output for this boundary." + ) + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -129,6 +195,9 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" + A = _unwrap_tensor(A, "rowwise" if transa else "columnwise") + B = _unwrap_tensor(B, "columnwise" if transb else "rowwise") + alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) workspace = get_cublas_workspace(A.device.index, ub is not None, False) @@ -174,6 +243,8 @@ def general_gemm( A = A.get_tensor(not transa) B = B.get_tensor(transb) + _validate_native_gemm_output_quantizer(quantization_params) + # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] @@ -299,6 +370,9 @@ def general_grouped_gemm( transa = layout[0] == "T" transb = layout[1] == "T" + A = [_unwrap_tensor(a, "rowwise" if transa else "columnwise") for a in A] + B = [_unwrap_tensor(b, "columnwise" if transb else "rowwise") for b in B] + empty_tensor = _empty_tensor() empty_tensors = [empty_tensor] * num_gemms diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 17237fa9b8..dbc27eee36 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -17,6 +17,25 @@ namespace transformer_engine::pytorch { namespace { +/*! @brief Reject unsupported columnwise-only per-tensor FP8 quantization + * + * Per-tensor FP8 uses the legacy cast-transpose kernel whenever columnwise + * output is requested. That kernel requires both rowwise and columnwise + * output buffers, so a transpose-only output cannot be produced without an + * otherwise-unused rowwise allocation and write. Keep this limitation + * explicit until the kernel supports independently selecting its outputs. + */ +void check_per_tensor_fp8_quantize_output(const TensorWrapper& output) { + const auto rowwise_data = output.get_rowwise_data(); + const auto columnwise_data = output.get_columnwise_data(); + if (rowwise_data.data_ptr == nullptr && columnwise_data.data_ptr != nullptr) { + PyErr_SetString(PyExc_NotImplementedError, + "Columnwise-only per-tensor FP8 quantization is not implemented; " + "the cast-transpose kernel requires rowwise output storage."); + throw py::error_already_set(); + } +} + /*! @brief Resolve an optional device to a concrete CUDA device * * If no device is provided, uses the current CUDA device. @@ -615,6 +634,7 @@ void Float8Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, if (input.numel() == 0) { return; } + check_per_tensor_fp8_quantize_output(out); QuantizationConfigWrapper quant_config; if (noop_flag) { quant_config.set_noop_tensor(noop_flag->data()); @@ -975,6 +995,7 @@ void Float8CurrentScalingQuantizer::quantize_impl(const TensorWrapper& input, Te out.set_scale(nullptr, DType::kFloat32, out.defaultShape); return; } + check_per_tensor_fp8_quantize_output(out); // Quantization configs QuantizationConfigWrapper quant_config; diff --git a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py similarity index 82% rename from transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py rename to transformer_engine/pytorch/custom_recipes/quantization_factory_base.py index 45febaf413..e5125fd92f 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py @@ -3,18 +3,18 @@ # See LICENSE for license information. """ -Quantizer factory examples using real silicon quantizers. +Quantizer factory examples using native Transformer Engine quantizers. -Each factory below replicates the behaviour of built-in TE recipe but via the -``CustomRecipe`` + ``qfactory`` interface. This is useful when you want to -start from a known-good recipe and then selectively override quantizer settings -for specific layers / tensor types. +For TE ``Linear`` and ``GroupedLinear`` roles, each factory below mirrors the +nominal defaults of a built-in recipe through the ``CustomRecipe`` + +``qfactory`` interface. This provides a built-in-equivalent starting point for +selectively overriding quantizer settings for specific layers or tensor types. Usage (any factory):: from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) @@ -32,6 +32,24 @@ from ..constants import DType +def high_precision_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "IdentityQuantizer": + """Factory that runs all GEMMs in high precision (no quantization). + + Returns an :class:`IdentityQuantizer` for every slot, so no tensor is + quantized. This is the simplest base factory and a good starting point to + branch from: keep most roles in high precision and selectively override the + ones you want to quantize. + + Dispatch logic: + * every role -> ``IdentityQuantizer`` (no quantization) + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + return IdentityQuantizer() + + def delayed_scaling_quantizer_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": @@ -69,8 +87,8 @@ def current_scaling_quantizer_factory( return Float8CurrentScalingQuantizer( fp8_dtype=fp8_dtype, device=torch.device("cuda"), - force_pow_2_scales=False, # constrain scale to powers of 2 - amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + force_pow_2_scales=False, + amax_epsilon=0.0, ) @@ -113,7 +131,7 @@ def float8_block_scaling_quantizer_factory( fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True, - amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + amax_epsilon=0.0, force_pow_2_scales=True, block_scaling_dim=block_scaling_dim, # 1 = 1D (1×128), 2 = 2D (128×128) ) @@ -156,8 +174,6 @@ def nvfp4_quantizer_factory( if is_grad: return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=True, with_post_rht_amax=True, with_2d_quantization=False, @@ -168,8 +184,6 @@ def nvfp4_quantizer_factory( # For input and unknown roles return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=True, with_post_rht_amax=True, with_2d_quantization=False, diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py deleted file mode 100644 index e88adbe4cc..0000000000 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py +++ /dev/null @@ -1,270 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -""" -Quantizer factory examples. - -Demonstrates how to use the ``CustomRecipe`` + ``qfactory`` interface to apply -*different* quantization recipes to different module/tensor types/instances within the same model. - -Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_mxfp8_grouped_linear_factory, - nvfp4_linear_fp8_dpa_factory, - nvfp4_linear_mxfp8_dpa_factory, - ) - - # Mixed module types: NVFP4 for Linear, MXFP8 for GroupedLinear - recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_grouped_linear_factory) - with autocast(recipe=recipe): - output = model(input) - - # NVFP4 for Linear, FP8 current-scaling + delayed-scaling for DPA - recipe = CustomRecipe(qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True) - with autocast(recipe=recipe): - output = model(input) - - # NVFP4 for Linear, MXFP8 for DPA - recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) - with autocast(recipe=recipe): - output = model(input) -""" - -from __future__ import annotations - -from typing import Optional - -from transformer_engine.pytorch.quantization import QuantizerRole -from ..constants import DType - - -def nvfp4_linear_mxfp8_grouped_linear_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``GroupedLinear``. - - Dispatch logic: - * ``role.module_type == "grouped_linear"`` -> MXFP8 (E4M3, block-32) - * everything else (``"linear"`` or unknown) -> NVFP4 (E2M1) - - NVFP4 settings follow the built-in ``NVFP4BlockScaling`` defaults: - * Weights: 2D quantization (16x16), no RHT, no stochastic rounding - * Inputs: 1D quantization, RHT enabled, no stochastic rounding - * Grads: 1D quantization, RHT enabled, stochastic rounding enabled - """ - is_grouped_linear = role is not None and role.module_type == "grouped_linear" - - if is_grouped_linear: - return _make_mxfp8_quantizer() - - return _make_nvfp4_quantizer(role) - - -def _make_mxfp8_quantizer(): - """Return an MXFP8 quantizer with default settings (E4M3, block-32, E8M0 scales).""" - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - return MXFP8Quantizer( - fp8_dtype=DType.kFloat8E4M3, - ) - - -def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): - """Return an NVFP4 quantizer configured per tensor role. - - Mirrors :class:`NVFP4BlockScaling` recipe defaults. - """ - from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - - is_linear = role is not None and role.module_type == "linear" - is_weight = is_linear and role.tensor_type == "weight" - is_grad = is_linear and role.tensor_type == "grad_output" - - if is_weight: - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - with_rht=False, - with_post_rht_amax=False, - with_2d_quantization=True, - stochastic_rounding=False, - with_random_sign_mask=True, - ) - - if is_grad: - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, - with_rht=True, - with_post_rht_amax=True, - with_2d_quantization=False, - stochastic_rounding=True, - with_random_sign_mask=True, - ) - - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, - with_rht=True, - with_post_rht_amax=True, - with_2d_quantization=False, - stochastic_rounding=False, - with_random_sign_mask=True, - ) - - -def nvfp4_linear_fp8_dpa_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, mixed FP8 for ``DotProductAttention``. - - This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` - to combine NVFP4 quantization for linear layers with FP8 attention. - - DPA tensor types (``role.module_type == "dpa"``): - - =========== ============================================================ - tensor_type Description - =========== ============================================================ - ``"qkv"`` Query, Key, Value inputs to the first attention GEMM - ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"o"`` Attention output (O = S·V) - ``"do"`` Gradient of the attention output (dO), backward input - ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - ``"dqkv"`` Gradient flowing back to Q, K, V - =========== ============================================================ - - Dispatch logic: - * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` - -> FP8 delayed scaling (stateful amax tracking) - * ``role.module_type == "dpa"`` (QKV, dO) - -> FP8 current scaling (E4M3) - * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) - -> FP8 current scaling placeholder. The fused attention kernel requires - FP8-compatible quantizers in all DPA slots, even when the output is - produced in BF16 (``fp8_mha=False``). DPA emits these hint-only roles - (with empty ``module_type`` and ``tensor_type``) when the downstream - consumer is unknown. - * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) - -> NVFP4 (E2M1), configured per tensor role - - Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_fp8_dpa_factory, - ) - - recipe = CustomRecipe( - qfactory=nvfp4_linear_fp8_dpa_factory, - fp8_dpa=True, - ) - with autocast(recipe=recipe): - output = model(input) - """ - from transformer_engine.pytorch.quantization import DelayedScalingRequest - from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer - - is_dpa = role is not None and role.module_type == "dpa" - is_softmax_or_dp = is_dpa and role.tensor_type in ("s", "dp") - - if is_softmax_or_dp: - return DelayedScalingRequest() - - if is_dpa: - return Float8CurrentScalingQuantizer( - fp8_dtype=DType.kFloat8E4M3, - device="cuda", - ) - - # DPA boundary slots (O output / dQKV grad-input): the fused attention - # kernel only supports FP8 quantizers here, regardless of the linear recipe. - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) - ) - if is_dpa_boundary: - return Float8CurrentScalingQuantizer( - fp8_dtype=DType.kFloat8E4M3, - device="cuda", - ) - - return _make_nvfp4_quantizer(role) - - -def nvfp4_linear_mxfp8_dpa_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``DotProductAttention``. - - Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from - :mod:`transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention` - (see the recipe-combination table at the top of that module). With - ``CustomRecipe`` the per-tensor decision is made directly here, so the - ``NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling"`` env override that the - built-in recipes would otherwise need is unnecessary. - - DPA tensor types (``role.module_type == "dpa"``): - - =========== ============================================================ - tensor_type Description - =========== ============================================================ - ``"qkv"`` Query, Key, Value inputs to the first attention GEMM - ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"o"`` Attention output (O = S·V) - ``"do"`` Gradient of the attention output (dO), backward input - ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - ``"dqkv"`` Gradient flowing back to Q, K, V - =========== ============================================================ - - Dispatch logic: - * ``role.module_type == "dpa"`` -> MXFP8 (E4M3, block-32) - The MXFP8 fused-attention kernel handles the S/dP slots - internally, so any quantizer returned for those roles is later - nulled out by ``get_attention_quantizers``. Returning MXFP8 is - the simplest valid choice. - * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in - ``role.name``) -> MXFP8 placeholder. The fused attention kernel - requires FP8-compatible quantizers in all DPA slots. - * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) - -> NVFP4 (E2M1), configured per tensor role. - - Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_mxfp8_dpa_factory, - ) - - recipe = CustomRecipe( - qfactory=nvfp4_linear_mxfp8_dpa_factory, - fp8_dpa=True, - ) - with autocast(recipe=recipe): - output = model(input) - """ - is_dpa = role is not None and role.module_type == "dpa" - if is_dpa: - return _make_mxfp8_quantizer() - - # DPA boundary slots (O output / dQKV grad-input): emitted by DPA with - # empty `module_type` and a `name` like ".dpa_output". The fused - # attention kernel requires an FP8-compatible quantizer here even when - # the downstream consumer is unknown. - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) - ) - if is_dpa_boundary: - return _make_mxfp8_quantizer() - - return _make_nvfp4_quantizer(role) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py new file mode 100644 index 0000000000..8b23897050 --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -0,0 +1,478 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Quantizer factory zoo. + +A collection of composed/mixed-recipe factories built on top of the single-format +building blocks. They demonstrate how to use the ``CustomRecipe`` + ``qfactory`` +interface to apply *different* quantization recipes to different +module/tensor types/instances within the same model, and how to use +``HybridQuantizer`` when rowwise and columnwise tensor directions should use +different representations or sources. + +Within the Linear/GroupedLinear and RL-oriented families, examples are roughly +ordered by increasingly aggressive forward quantization. This is an +organizational convention, not an expected accuracy or performance ranking. +These factories demonstrate what can be composed; they are not necessarily +tuned for end-to-end performance. + +When the columnwise/backward representation is higher precision than +the rowwise/forward representation, consider +``columnwise_source="rowwise_dequantized"`` as a starting point. It constructs +backward operands from the value obtained during forward quantization, +improving forward/backward representation consistency. +Note: dequantization does not recover information discarded during forward quantization. + +Organization: + * Linear / grouped-linear recipes (pre-training). Favor more precision + on the forward pass. + * RL-oriented recipes: Favor more precision in backward GEMMs. + * Linear + attention recipes: factories that also cover ``DotProductAttention`` + roles and require ``CustomRecipe(..., fp8_dpa=True)``. + +.. warning:: + + Use these with caution. These are **not** official, supported recipes + provided by Transformer Engine -- they are illustrative examples meant to + inspire your own experiments, not drop-in production defaults. Most include + a motivating rationale in their per-factory docstrings, but they have not + been broadly validated for accuracy, convergence, or performance across + models and hardware. Treat them as starting points: benchmark and verify on + your own workload before relying on any of them. + +Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, + nvfp4_linear_mxfp8_dpa_factory, + ) + + # Linear-only recipe (no attention quantization): the qfactory is the only knob. + recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) + with autocast(recipe=recipe): + output = model(input) + + # Recipe that also quantizes DotProductAttention: set ``fp8_dpa=True`` so the + # attention GEMMs request quantizers from the factory (DPA roles) too. + recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) + with autocast(recipe=recipe): + output = model(input) + + # The other factories in this module follow the same two patterns; see their + # docstrings for the exact per-role dispatch. +""" + +from __future__ import annotations + +from typing import Optional + +from transformer_engine.pytorch.quantization import QuantizerRole +from ..constants import DType +from .quantization_factory_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory + +# ----------------------------------------------------------------------------- +# Linear / GroupedLinear Recipes (pre-training) +# ----------------------------------------------------------------------------- + + +def high_precision_fwd_mxfp8_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: high-precision forward, MXFP8 backward. + + Dispatch logic: + * ``grad_output`` -> MXFP8 (E4M3, block-32) + * everything else -> ``Hybrid(rowwise=IdentityQuantizer, columnwise=MXFP8)`` + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "grad_output": + return mxfp8_quantizer_factory(role) + + # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=mxfp8_quantizer_factory(role), + ) + + +def mxfp8_fwd_nvfp4_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: MXFP8 forward, NVFP4 backward. + + Per-GEMM format consumption: + * fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` + * dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` + * wgrad: ``input.col(NVFP4) x grad_output.col(NVFP4)`` + + The NVFP4 side mirrors :func:`nvfp4_quantizer_factory` role semantics: + weights use 2D scaling, activations/grads use the base 1D RHT/SR settings. + ``HybridQuantizer`` then pins each sub-quantizer to the direction that is + actually consumed. + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + if is_linear and role.tensor_type == "grad_output": + return nvfp4_quantizer_factory(role) + return mxfp8_quantizer_factory(role) + + +def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): + """NVFP4 quantizer without RHT, stochastic rounding, or 2D scaling.""" + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + return NVFP4Quantizer( + fp4_dtype=DType.kFloat4E2M1, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, + ) + + +def nvfp4_1d_double_quantized_weight_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 recipe with 1D weight double quantization. + + Dispatch logic: + * ``linear`` / ``grouped_linear`` ``weight`` -> + ``Hybrid(rowwise=plain 1D NVFP4, columnwise=plain 1D NVFP4, + columnwise_source="rowwise_dequantized")`` + * everything else -> :func:`nvfp4_quantizer_factory` + + ``W.T`` is quantized from the dequantized forward/rowwise ``W`` value + instead of directly from the original high-precision weight. In + ``HybridQuantizer`` terms, that source choice is expressed with + ``columnwise_source="rowwise_dequantized"``. + + All non-weight roles keep the standard NVFP4 factory behavior, including RHT + for inputs and stochastic rounding for gradients. The weight override uses + plain 1D NVFP4 in both directions: no RHT, stochastic rounding, row-scaled + activations, or 2D weight scaling. + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + return nvfp4_quantizer_factory(role) + + +# ----------------------------------------------------------------------------- +# RL-Oriented Recipes +# ----------------------------------------------------------------------------- + + +def mxfp8_fwd_dequantized_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: MXFP8 forward, high-precision dequantized backward. + + This expresses the linear/grouped-linear equivalent of + ``backward_override="dequantized"`` through per-direction quantizers: + + * ``input`` / ``weight`` -> + ``Hybrid(rowwise=MXFP8, columnwise=Identity, columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> MXFP8 + + Backward GEMMs use high-precision operands dequantized from the MXFP8 + forward payload, avoiding gradient quantization. + + This recipe targets RL-style training use cases and is motivated by + NVIDIA/TransformerEngine#2644, where ``backward_override="dequantized"`` + was introduced: + https://github.com/NVIDIA/TransformerEngine/pull/2644 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return mxfp8_quantizer_factory(role) + + +def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, dequantized-source MXFP8 backward. + + This RL-related recipe is inspired by the Composer 2 MoE grouped-GEMM + recipe described in arXiv:2603.24477. + + Derived from the report: Composer 2 describes row-scaled NVFP4 for the MoE + forward pass and standard MXFP8 for the MoE backward pass. This factory maps + that format split onto ``GroupedLinear`` roles. + + Assumed here: regular non-MoE ``Linear`` layers use the MXFP8 fallback. The + public report does not specify the precision used for non-MoE linears. + + Dispatch logic: + + * ``GroupedLinear`` ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=MXFP8, + columnwise_source="rowwise_dequantized")`` + * ``GroupedLinear`` ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=MXFP8, + columnwise_source="rowwise_dequantized")`` + * regular ``Linear`` -> MXFP8 + * ``grad_output`` -> MXFP8 + * everything else -> MXFP8 + + Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid + RHT, stochastic rounding, and 2D scaling. This sample assumes MXFP8 + backward operands are quantized from the dequantized rowwise NVFP4 forward + value, expressed with ``columnwise_source="rowwise_dequantized"``. To + quantize MXFP8 backward operands from the original high-precision tensor + instead, construct the same ``HybridQuantizer`` while omitting + ``columnwise_source="rowwise_dequantized"``. + + Composer 2 Technical Report: + https://arxiv.org/abs/2603.24477 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + 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=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_source="rowwise_dequantized", + ) + if is_grouped_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_source="rowwise_dequantized", + ) + if is_grouped_linear and role.tensor_type == "grad_output": + return mxfp8_quantizer_factory(role) + if is_linear: + return mxfp8_quantizer_factory(role) + return mxfp8_quantizer_factory(role) + + +def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, dequantized backward. + + This expresses a linear/grouped-linear variant of + ``NVFP4BlockScaling(row_scaled_activation=True, + backward_override="dequantized")`` through per-direction quantizers: + + * ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> plain NVFP4 + + Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, + stochastic rounding, and 2D scaling. + + This recipe targets RL-style training use cases and builds on + NVIDIA/TransformerEngine#2931, which introduced row-scaled NVFP4: + https://github.com/NVIDIA/TransformerEngine/pull/2931 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return _plain_nvfp4_quantizer() + + +# ----------------------------------------------------------------------------- +# Linear + Attention Recipes +# ----------------------------------------------------------------------------- + + +def nvfp4_linear_fp8_dpa_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 for ``Linear``, FP8 for ``DotProductAttention``. + + This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` + to combine NVFP4 quantization for linear layers with FP8 attention. + + DPA-owned tensor types (``role.module_type == "dpa"``): + + =========== ============================================================ + tensor_type Description + =========== ============================================================ + ``"qkv"`` Query, Key, Value inputs to the first attention GEMM + ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM + ``"do"`` Gradient of the attention output (dO), backward input + ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward + =========== ============================================================ + + Dispatch logic: + * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` + -> FP8 delayed scaling (``Format.HYBRID``, most_recent, history length 1) + * other DPA roles + -> FP8 current scaling (``Format.HYBRID``: E4M3 fwd, E5M2 bwd) + * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) + -> FP8 current scaling placeholder. The fused attention kernel requires + FP8-compatible quantizers in all DPA slots, even when the output is + produced in BF16 (``fp8_mha=False``). DPA emits these hint-only roles + (with empty ``module_type`` and ``tensor_type``) when the downstream + consumer is unknown. + * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) + -> NVFP4 (E2M1), configured per tensor role + + Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + """ + from transformer_engine.common.recipe import Format + from transformer_engine.pytorch.quantization import DelayedScalingRequest + from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + # Native NVFP4 + FP8 attention uses delayed scaling for S/dP. + if is_dpa and role.tensor_type in ("s", "dp"): + return DelayedScalingRequest( + fp8_format=Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + reduce_amax=True, + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + return nvfp4_quantizer_factory(role) + + +def nvfp4_linear_mxfp8_dpa_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``DotProductAttention``. + + Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from + :mod:`transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention` + (see the recipe-combination table at the top of that module). With + ``CustomRecipe`` the per-tensor decision is made directly here, so the + ``NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling"`` env override that the + built-in recipes would otherwise need is unnecessary. + + DPA-owned tensor types (``role.module_type == "dpa"``): + + =========== ============================================================ + tensor_type Description + =========== ============================================================ + ``"qkv"`` Query, Key, Value inputs to the first attention GEMM + ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM + ``"do"`` Gradient of the attention output (dO), backward input + ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward + =========== ============================================================ + + Dispatch logic: + * ``role.module_type == "dpa"`` -> MXFP8 + (``Format.HYBRID``: E4M3 fwd, E5M2 bwd) + The MXFP8 fused-attention kernel handles the S/dP slots + internally, so any quantizer returned for those roles is later + nulled out by ``get_attention_quantizers``. Returning MXFP8 is + the simplest valid choice. + * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in + ``role.name``) -> MXFP8 placeholder. The fused attention kernel + requires FP8-compatible quantizers in all DPA slots. + * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) + -> NVFP4 (E2M1), configured per tensor role. + + Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_linear_mxfp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + """ + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + return nvfp4_quantizer_factory(role) diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 9b914e2d6c..634bdd83d1 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -13,6 +13,7 @@ from .. import cpp_extensions as tex from ..constants import TE_DType from ..export import is_in_onnx_export_mode +from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method @@ -31,6 +32,36 @@ def set_quantizer_amax_reduction_group(quantizer, amax_reduction_group) -> None: target.amax_reduction_group = amax_reduction_group +def set_quantizer_usage_for_wgrad_all_gather(quantizer) -> None: + """Configure an all-gather output for consumption by wgrad.""" + if quantizer is None: + return + + target = getattr(quantizer, "parent_quantizer", quantizer) + + # Hybrid currently gathers in high precision, then quantizes the full + # result, so request the columnwise representation consumed by wgrad. + if isinstance(target, HybridQuantizer): + target.set_usage(rowwise=False, columnwise=True) + elif target.supports_only_rowwise_all_gather(): + # Per-tensor FP8 gathers rowwise data and synthesizes its transpose. + target.set_usage(rowwise=True, columnwise=False) + else: + target.set_usage(rowwise=False, columnwise=True) + + +def can_reconstruct_wgrad_input_from_original(quantizer) -> bool: + """Whether wgrad input can be reconstructed from a saved original tensor.""" + target = getattr(quantizer, "parent_quantizer", quantizer) + if target is None: + target = quantizer + if isinstance(target, HybridQuantizer): + if target.columnwise_source == "original": + return True + return target.rowwise_quantizer.is_requantization_safe() + return target.is_requantization_safe() + + def _get_normalization_func(normalization: str, forward: bool): fwd_normalization_funcs = { "LayerNorm": tex.layernorm_fwd, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 7c11e635c6..563836ee36 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -52,9 +52,12 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from ..utils import ( is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, @@ -752,6 +755,14 @@ def _is_weight_workspace_valid( return False if quantizer.columnwise_usage and workspace._columnwise_data is None: return False + elif isinstance(workspace, HybridQuantizedTensorStorage): + # Workspace cached under one flag setting (e.g. inference with + # ``columnwise=False``) becomes stale when the next call needs the + # missing direction; invalidate so a fresh workspace is built. + if quantizer.rowwise_usage and workspace._rowwise_storage is None: + return False + if quantizer.columnwise_usage and workspace._columnwise_storage is None: + return False if isinstance(workspace, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer): return False return True @@ -1049,6 +1060,8 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Return early if recipe state matches recipe if self.fp8_meta_tensors_initialized: recipe_state = self.fp8_meta[fp8_meta_tensor_key] + # TODO(#3157): Match built-in recipes by full config, not just RecipeState type, so + # same-class mid-training changes rebuild quantizers/workspaces correctly. if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState): self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd) return @@ -1065,6 +1078,9 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): + # TODO(#3157): Compare CustomRecipe/qfactory config here. qfactory changes made + # mid-training on the same recipe object currently do not take effect because + # stale quantizers are reused. if recipe_state.recipe is recipe: return @@ -1709,8 +1725,12 @@ def grad_output_preprocess( ): grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0) else: - if isinstance(quantizer, Float8BlockQuantizer): - # unfuse bgrad for now until cast_transpose + dgrad calculation is ready for Float8BlockQuantizer. + if isinstance( + quantizer, (Float8BlockQuantizer, HybridQuantizer, IdentityQuantizer) + ): + # Float8BlockQuantizer: unfused until cast_transpose + dgrad is ready. + # HybridQuantizer: tex.bgrad_quantize doesn't recognize hybrid quantizers. + # IdentityQuantizer: high-precision passthrough; bgrad computed in HP. grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0) else: grad_bias, grad_output = tex.bgrad_quantize(grad_output, quantizer) @@ -1777,8 +1797,11 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: raise RuntimeError("Weight quantizer has not been initialized") quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) quantizer.internal = False + # HybridQuantizer is included so its current-scaling / NVFP4 + # sub-quantizers get the same cross-shard amax reduction as the + # vanilla path (no-op for block-scaled sub-quantizers like MXFP8). if is_dtensor and isinstance( - quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer) + quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer, HybridQuantizer) ): device_mesh = dtensor_param.device_mesh amax_reduction_group = ( diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index 8ac49c9bae..5c46230b47 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -80,9 +80,10 @@ class Fp8Padding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - the alignment size for the input tensor. If not provided, the alignment size will - be determined by the FP8/FP4 recipe (32 for MXFP8/NVFP4 and 16 for others) in the first - forward pass. + Alignment size for each grouped input. If not provided, it is + determined from the active recipe on the first forward pass: + 32 for MXFP8, 128 for NVFP4, 16 for other built-in recipes, + and ``CustomRecipe.quantization_alignment`` for custom recipes. """ def __init__( diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index c5d396837f..d34605e195 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -78,9 +78,10 @@ class Fp8Unpadding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - The alignment size for the input tensor. If not provided, the alignment size will - be automatically determined based on the FP8/FP4 recipe in the first forward pass: - 32 for MXFP8 or NVFP4, otherwise 16. + Alignment size used to pad each grouped input. If not provided, + it is determined from the active recipe on the first forward + pass: 32 for MXFP8, 128 for NVFP4, 16 for other built-in + recipes, and ``CustomRecipe.quantization_alignment`` for custom recipes. """ def __init__( diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 143fee7ba2..3f3963c549 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -4,7 +4,7 @@ """GroupedLinear API""" -from typing import Union, Optional, Callable, Tuple, List +from typing import Union, Optional, Callable, Tuple, List, Sequence from itertools import chain import os import warnings @@ -28,7 +28,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ._common import WeightGradStore +from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( divide, @@ -65,6 +65,8 @@ Float8BlockQuantizer, Float8CurrentScalingQuantizer, Float8Quantizer, + HybridQuantizer, + IdentityQuantizer, MXFP8Quantizer, NVFP4Quantizer, ) @@ -77,6 +79,368 @@ from ...debug.pytorch.debug_quantization import DebugQuantizer from ...debug.pytorch.debug_state import TEDebugState + +_NATIVE_SPLIT_QUANTIZER_TYPES = frozenset( + { + Float8Quantizer, + Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, + } +) + + +def _supports_native_split_quantize(quantizer): + """Whether ``tex.split_quantize`` has an exact converter for this quantizer.""" + return type(quantizer) in _NATIVE_SPLIT_QUANTIZER_TYPES + + +def _uses_identity_quantizer(quantizer): + """Whether a quantizer, including a hybrid sub-quantizer, is Identity-backed.""" + if quantizer is None: + return False + if isinstance(quantizer, IdentityQuantizer): + return True + if isinstance(quantizer, HybridQuantizer): + return _uses_identity_quantizer(quantizer.rowwise_quantizer) or _uses_identity_quantizer( + quantizer.columnwise_quantizer + ) + return False + + +def _identity_quantizer_signature(quantizer): + """Identity usage per GEMM direction: (rowwise, columnwise).""" + if isinstance(quantizer, HybridQuantizer): + return ( + _uses_identity_quantizer(quantizer.rowwise_quantizer), + _uses_identity_quantizer(quantizer.columnwise_quantizer), + ) + identity = isinstance(quantizer, IdentityQuantizer) + return (identity, identity) + + +_DYNAMIC_QUANTIZER_SIGNATURE_FIELDS = frozenset( + { + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + } +) + + +def _backend_quantizer_signature(quantizer): + """Return backend configuration that grouped kernels require to be uniform.""" + if quantizer is None: + return None + + # Identity is not registered as a torch.compile value quantizer, but its + # dtype changes the grouped GEMM input type and therefore must be uniform. + if isinstance(quantizer, IdentityQuantizer): + return (type(quantizer), (("dtype", quantizer.dtype),)) + + fields = quantizer._value_fields() + if fields is None: + # Delayed-scaling Float8Quantizer carries per-expert scale/amax tensors, + # which are intentionally different, but its emitted FP8 dtype is a + # group-wide backend choice. Other unregistered/custom quantizers retain + # the conservative exact-family behavior until they expose value fields. + fields = ("dtype",) if isinstance(quantizer, Float8Quantizer) else () + + config = [] + for name in fields: + if name in _DYNAMIC_QUANTIZER_SIGNATURE_FIELDS: + continue + value = getattr(quantizer, name) + if name == "dtype": + value = int(value) + config.append((name, value)) + return (type(quantizer), tuple(config)) + + +def _validate_backend_match(reference, quantizer, operand_name, direction, expert_index): + """Validate one expert against the group's reference backend.""" + if type(quantizer) is not type(reference): + raise ValueError( + f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend" + f" families across experts: expert 0 uses {type(reference).__name__}, but expert" + f" {expert_index} uses {type(quantizer).__name__}. Grouped operands require one" + " quantizer family per direction." + ) + reference_signature = _backend_quantizer_signature(reference) + quantizer_signature = _backend_quantizer_signature(quantizer) + if quantizer_signature != reference_signature: + raise ValueError( + f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend" + f" configurations across experts: expert 0 uses {reference_signature}, but expert" + f" {expert_index} uses {quantizer_signature}. Grouped operands require the same" + " backend-relevant configuration per direction." + ) + + +def _validate_grouped_quantizer_list(quantizers, *, operand_name="operand") -> None: + """Validate one grouped operand once when its quantizer generation changes.""" + if not quantizers: + return + + reference = quantizers[0] + reference_is_hybrid = isinstance(reference, HybridQuantizer) + reference_identity = _identity_quantizer_signature(reference) + + for expert_index, quantizer in enumerate(quantizers[1:], start=1): + if (quantizer is None) != (reference is None): + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix None and concrete quantizers" + f" across experts: expert 0 is {type(reference).__name__}, but expert" + f" {expert_index} is {type(quantizer).__name__}." + ) + if reference is None: + continue + + quantizer_is_hybrid = isinstance(quantizer, HybridQuantizer) + if quantizer_is_hybrid != reference_is_hybrid: + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix HybridQuantizer and non-hybrid" + f" quantizers across experts: expert 0 is {type(reference).__name__}, but expert" + f" {expert_index} is {type(quantizer).__name__}." + ) + + identity = _identity_quantizer_signature(quantizer) + if identity != reference_identity: + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix Identity-backed and quantized" + f" directions across experts: expert 0 uses {reference_identity}, but expert" + f" {expert_index} uses {identity}." + ) + + if reference_is_hybrid: + _validate_backend_match( + reference.rowwise_quantizer, + quantizer.rowwise_quantizer, + operand_name, + "rowwise", + expert_index, + ) + _validate_backend_match( + reference.columnwise_quantizer, + quantizer.columnwise_quantizer, + operand_name, + "columnwise", + expert_index, + ) + if quantizer.columnwise_source != reference.columnwise_source: + raise ValueError( + f"GroupedLinear {operand_name} HybridQuantizer list has mixed columnwise" + " source policies across experts: expert 0 uses" + f" {reference.columnwise_source!r}, but expert {expert_index} uses" + f" {quantizer.columnwise_source!r}." + ) + else: + _validate_backend_match( + reference, + quantizer, + operand_name, + "plain", + expert_index, + ) + + +def _split_quantize_non_hybrid( + tensor, + m_splits, + quantizers, + activation_dtype, + *, + disable_bulk_allocation=False, + allow_identity_views=True, +): + """Split and quantize one homogeneous, non-Hybrid quantizer list.""" + reference = quantizers[0] + if _supports_native_split_quantize(reference): + return tex.split_quantize( + tensor, + m_splits, + quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + + tensor = cast_if_needed(tensor, activation_dtype) + if ( + allow_identity_views + and type(reference) is IdentityQuantizer + and (reference.dtype is None or reference.dtype == activation_dtype) + ): + return torch.split(tensor, m_splits) + + return [ + quantizer(tensor_part) if quantizer is not None else tensor_part + for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) + ] + + +def _split_quantize_hybrid( + tensor, + m_splits, + quantizers, + *, + disable_bulk_allocation=False, +): + """Grouped split+quantize for an all-hybrid, generation-validated operand.""" + from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage + + reference = quantizers[0] + rowwise_enabled = reference.rowwise_usage + columnwise_enabled = reference.columnwise_usage + columnwise_source = reference.columnwise_source + rowwise_quantizers = [quantizer.rowwise_quantizer for quantizer in quantizers] + columnwise_quantizers = [quantizer.columnwise_quantizer for quantizer in quantizers] + + needs_rowwise_result = rowwise_enabled or ( + columnwise_enabled and columnwise_source == "rowwise_dequantized" + ) + row_results = ( + _split_quantize_non_hybrid( + tensor, + m_splits, + rowwise_quantizers, + tensor.dtype, + disable_bulk_allocation=disable_bulk_allocation, + allow_identity_views=False, + ) + if needs_rowwise_result + else [None] * len(quantizers) + ) + + columnwise_src = tensor + if columnwise_enabled and columnwise_source == "rowwise_dequantized": + # Assemble the exact grouped row results in split order. NVFP4 padding + # and scale layout can differ from independently quantizing each split. + columnwise_src = torch.cat( + [result.dequantize(dtype=tensor.dtype) for result in row_results], + dim=0, + ) + col_results = ( + _split_quantize_non_hybrid( + columnwise_src, + m_splits, + columnwise_quantizers, + tensor.dtype, + disable_bulk_allocation=disable_bulk_allocation, + allow_identity_views=False, + ) + if columnwise_enabled + else [None] * len(quantizers) + ) + + return [ + HybridStorage( + rowwise_storage=row if rowwise_enabled else None, + columnwise_storage=col, + quantizer=q, + fake_dtype=tensor.dtype, + ) + for row, col, q in zip( + row_results, + col_results, + quantizers, + ) + ] + + +def _split_quantize( + tensor: torch.Tensor, + split_sizes: List[int], + with_quantized_output: bool, + quantizers: Optional[List[Quantizer]], + dtype: torch.dtype, + with_debug_quantizers: bool, + disable_bulk_allocation: bool, +) -> Sequence[Union[torch.Tensor, QuantizedTensorStorage]]: + """Split a tensor and quantize each part if needed.""" + if not with_quantized_output: + return torch.split(cast_if_needed(tensor, dtype), split_sizes) + + if quantizers is None or quantizers[0] is None: + raise ValueError("Quantizers are required for quantized split output") + + if with_debug_quantizers: + return DebugQuantizer.multi_tensor_quantize(tensor, quantizers, split_sizes, dtype) + + reference = quantizers[0] + if isinstance(reference, HybridQuantizer): + return _split_quantize_hybrid( + tensor, + split_sizes, + quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + + return _split_quantize_non_hybrid( + tensor, + split_sizes, + quantizers, + dtype, + disable_bulk_allocation=disable_bulk_allocation, + ) + + +def _split_quantize_and_bias( + tensor: torch.Tensor, + split_sizes: List[int], + *, + fp8: bool, + debug: bool, + quantizers: Optional[List[Quantizer]], + dtype: torch.dtype, + use_bias: bool, + recipe: Recipe, + disable_bulk_allocation: bool, +) -> Tuple[ + Sequence[Union[torch.Tensor, QuantizedTensorStorage]], + List[Optional[torch.Tensor]], +]: + """Split grad output, quantize if needed, and compute unfused bias gradients.""" + num_splits = len(split_sizes) + grad_biases = [None] * num_splits + reference = quantizers[0] + identity = _uses_identity_quantizer(reference) + hybrid = isinstance(reference, HybridQuantizer) and not identity + + use_native_bgrad_quantize = ( + fp8 + and not debug + and not hybrid + and use_bias + and not identity + and (recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8()) + ) + if use_native_bgrad_quantize: + outputs = [None] * num_splits + for i, tensor_part in enumerate(torch.split(tensor, split_sizes)): + grad_biases[i], outputs[i] = tex.bgrad_quantize(tensor_part, quantizers[i]) + return outputs, grad_biases + + with_quantized_output = fp8 or debug + if with_quantized_output and (use_bias or debug): + for i, tensor_part in enumerate(torch.split(tensor, split_sizes)): + grad_biases[i] = tensor_part.sum(dim=0) + + # Preserve the existing CPU-offload policy: only Hybrid split-quantize + # disables bulk allocation in backward. + disable_bulk_allocation = disable_bulk_allocation if hybrid else False + outputs = _split_quantize( + tensor, + split_sizes, + with_quantized_output=with_quantized_output, + quantizers=quantizers, + dtype=dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=disable_bulk_allocation, + ) + return outputs, grad_biases + + __all__ = ["GroupedLinear"] @@ -513,6 +877,8 @@ def forward( cache_weight, skip_fp8_weight_update, save_original_input, + delayed_scaling_input_quantizer, + unsafe_requantization_input_quantizer, debug, ) = non_tensor_args if fp8: @@ -535,20 +901,34 @@ def forward( if is_dist_weight: weights = materialize_weight_for_forward(weights) - # Configure quantizers - if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): - if FP8GlobalStateManager.get_fp8_recipe().custom(): - # Custom recipe factory may produce DS quantizers unknown to caller. - # TODO(negvet): fix on Megatron side — guard should also exclude 'custom', or - # better: check at runtime whether quantizers are DS-based. + backward_needs_input = is_grad_enabled and weight_requires_grad + if backward_override is None and save_original_input and backward_needs_input: + if delayed_scaling_input_quantizer is not None: + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError( + "DelayedScaling recipe is not supported with save_original_input" + ) + + # Megatron-Core may enable this automatically to reuse an activation + # already retained by an upstream operation. The resolved quantizer + # generation is classified once in ``_validate_quantizer_generation``. + if save_original_input and unsafe_requantization_input_quantizer is not None: warnings.warn( - "save_original_input is incompatible with delayed-scaling quantizers " - "(Float8Quantizer). Disabling save_original_input for this module.", + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({unsafe_requantization_input_quantizer}).", stacklevel=2, ) save_original_input = False - else: - raise ValueError("DelayedScaling recipe is not supported with save_original_input") + + # Configure quantizers if input_quantizers[0] is not None: for input_quantizer in input_quantizers: input_quantizer.set_usage( @@ -630,23 +1010,18 @@ def forward( m_splits = m_splits.tolist() inp_view = inp.reshape(-1, in_features) - inputmats: list - if fp8 and not debug: - # Disable bulk allocation when CPU offloading is active: offloading skips small - # tensors (like scales), but bulk allocation shares storage across all tensors, - # so if scales can't be offloaded, nothing in the group can be offloaded. - inputmats = tex.split_quantize( - inp_view, - m_splits, - input_quantizers, - disable_bulk_allocation=cpu_offloading, - ) - elif debug: - inputmats = DebugQuantizer.multi_tensor_quantize( - inp_view, input_quantizers, m_splits, activation_dtype - ) - else: - inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) + # Disable bulk allocation when CPU offloading is active: offloading skips small + # tensors (like scales), but bulk allocation shares storage across all tensors, + # so if scales can't be offloaded, nothing in the group can be offloaded. + inputmats = _split_quantize( + inp_view, + m_splits, + with_quantized_output=fp8 or debug, + quantizers=input_quantizers, + dtype=activation_dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=cpu_offloading, + ) if cpu_offloading: start_offload(*inputmats) @@ -1060,52 +1435,27 @@ def backward( # Preprocess grad output grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) - grad_output = [None] * ctx.num_gemms - grad_biases = [None] * ctx.num_gemms - if ctx.fp8 and not ctx.debug: - if ctx.use_bias: - grad_output_mats = torch.split(grad_output_view, ctx.m_splits) - recipe = ctx.fp8_recipe - if recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8(): - # Fused bias grad + quantize kernel - for i in range(ctx.num_gemms): - grad_biases[i], grad_output[i] = tex.bgrad_quantize( - grad_output_mats[i], - ctx.grad_output_quantizers[i], - ) - else: - # Unfused bias grad and multi-tensor quantize - for i in range(ctx.num_gemms): - grad_biases[i] = grad_output_mats[i].sum(dim=0) - grad_output = tex.split_quantize( - grad_output_view, - ctx.m_splits, - ctx.grad_output_quantizers, - ) - else: - # Multi-tensor quantize - grad_output = tex.split_quantize( - grad_output_view, - ctx.m_splits, - ctx.grad_output_quantizers, + grad_output_reference = ctx.grad_output_quantizers[0] + if ctx.fp8 and isinstance(grad_output_reference, HybridQuantizer): + # Usage is a runtime decision, not part of generation validation. + # Apply it uniformly so dispatch can read the first parent without + # rescanning every expert. + for grad_output_quantizer in ctx.grad_output_quantizers: + grad_output_quantizer.set_usage( + rowwise=ctx.requires_dgrad, + columnwise=ctx.weights_requires_grad, ) - elif ctx.debug: - grad_output_mats = torch.split(grad_output_view, ctx.m_splits) - for i in range(ctx.num_gemms): - grad_biases[i] = grad_output_mats[i].sum(dim=0) - grad_output = DebugQuantizer.multi_tensor_quantize( - grad_output_view, - ctx.grad_output_quantizers, - ctx.m_splits, - ctx.activation_dtype, - ) - else: - # Only split grad output. Grad bias is fused with - # wgrad GEMM. - grad_output = torch.split( - cast_if_needed(grad_output_view, ctx.activation_dtype), - ctx.m_splits, - ) + grad_output, grad_biases = _split_quantize_and_bias( + grad_output_view, + ctx.m_splits, + fp8=ctx.fp8, + debug=ctx.debug, + quantizers=ctx.grad_output_quantizers, + dtype=ctx.activation_dtype, + use_bias=ctx.use_bias, + recipe=ctx.fp8_recipe, + disable_bulk_allocation=ctx.cpu_offloading, + ) if is_dist_weight: accumulate_wgrad_into_param_main_grad = False @@ -1206,20 +1556,15 @@ def backward( input_quantizer.set_usage(rowwise=True, columnwise=True) else: input_quantizer.set_usage(rowwise=False, columnwise=True) - inputmats: list - if ctx.fp8 and not ctx.debug: - inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) - elif ctx.debug: - inputmats = DebugQuantizer.multi_tensor_quantize( - inp_view, - ctx.input_quantizers, - ctx.m_splits, - ctx.activation_dtype, - ) - else: - inputmats = torch.split( - cast_if_needed(inp_view, ctx.activation_dtype), ctx.m_splits - ) + inputmats = _split_quantize( + inp_view, + ctx.m_splits, + with_quantized_output=ctx.fp8 or ctx.debug, + quantizers=ctx.input_quantizers, + dtype=ctx.activation_dtype, + with_debug_quantizers=ctx.debug, + disable_bulk_allocation=ctx.cpu_offloading, + ) elif ctx.backward_override == "dequantized": inputmats_dequant = [] for inputmat in inputmats: @@ -1366,7 +1711,8 @@ class GroupedLinear(TransformerEngineBaseModule): If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. - Cannot work with FP8 DelayedScaling recipe. + Requires input quantizers that can safely reproduce their results from the + original input. Cannot work with FP8 DelayedScaling recipe. single_grouped_weight : bool, default = False If set to ``True``, grouped weights are stored as a single grouped parameter instead of one parameter per GEMM. @@ -1451,6 +1797,9 @@ def __init__( "fwd": 3, "bwd": 2, } + self._validated_quantizer_generations = {} + self._delayed_scaling_input_quantizer = None + self._unsafe_requantization_input_quantizer = None if tp_group is None: self.tp_size = tp_size @@ -1539,6 +1888,57 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) + self._validate_quantizer_generation(fwd) + + def _validate_quantizer_generation(self, fwd: bool) -> None: + """Validate grouped-kernel invariants once per quantizer generation.""" + # Recipe state replaces this list object only when it constructs a new + # quantizer generation. The O(1) identity guard keeps validation off the + # steady-state forward path. Record a generation only after all of its + # operand roles pass, so a failed recipe transition is retried. + meta_key = "scaling_fwd" if fwd else "scaling_bwd" + generation = self.quantizers.get(meta_key) + if generation is None: + return + if self._validated_quantizer_generations.get(meta_key) is generation: + return + + if fwd: + stride = self._num_fp8_tensors_per_gemm["fwd"] + input_quantizers = tuple( + generation[self._offsets["input"] + i * stride] for i in range(self.num_gemms) + ) + weight_quantizers = tuple( + generation[self._offsets["weight"] + i * stride] for i in range(self.num_gemms) + ) + _validate_grouped_quantizer_list(input_quantizers, operand_name="input") + _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight") + delayed_scaling_input_quantizer = next( + (q for q in input_quantizers if isinstance(q, Float8Quantizer)), + None, + ) + unsafe_requantization_input_quantizer = next( + ( + q + for q in input_quantizers + if q is not None and not can_reconstruct_wgrad_input_from_original(q) + ), + None, + ) + self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer + self._unsafe_requantization_input_quantizer = unsafe_requantization_input_quantizer + else: + stride = self._num_fp8_tensors_per_gemm["bwd"] + grad_output_quantizers = tuple( + generation[self._offsets["grad_output"] + i * stride] for i in range(self.num_gemms) + ) + _validate_grouped_quantizer_list( + grad_output_quantizers, + operand_name="grad_output", + ) + + self._validated_quantizer_generations[meta_key] = generation + def get_quantizer_roles( self, *, @@ -1575,6 +1975,21 @@ def make_grouped_weights(self, defer_init=False) -> None: return weight_quantizers = self._get_weight_quantizers() + # TODO(#3158): Support Identity/Hybrid single grouped weights. + unsupported_quantizers = tuple( + type(quantizer).__name__ + for quantizer in weight_quantizers + if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)) + ) + if unsupported_quantizers: + quantizer_names = ", ".join(dict.fromkeys(unsupported_quantizers)) + raise NotImplementedError( + "GroupedLinear(single_grouped_weight=True) does not support " + f"{quantizer_names} weight quantizers yet. Set " + "single_grouped_weight=False or unset " + "NVTE_GROUPED_LINEAR_SINGLE_PARAM. See #3158." + ) + recipe = ( weight_quantizers[0]._get_compatible_recipe() if weight_quantizers and weight_quantizers[0] is not None @@ -1943,6 +2358,8 @@ def forward( cache_weight, skip_fp8_weight_update, self.save_original_input, + self._delayed_scaling_input_quantizer, + self._unsafe_requantization_input_quantizer, debug, ) out, new_workspaces = linear_fn( @@ -2090,6 +2507,14 @@ def _get_weight_quantizers(self) -> List[Quantizer]: return weight_quantizers def _get_quantizers(self): + if self.fp8: + # Normally validated while installing recipe metadata. Keep this + # O(1) generation guard so failed transitions cannot reuse stale + # validation state if base metadata takes an early return on retry. + self._validate_quantizer_generation(True) + if torch.is_grad_enabled(): + self._validate_quantizer_generation(False) + weight_quantizers = self._get_weight_quantizers() input_quantizers, output_quantizers = ( [None] * self.num_gemms, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a588e21a0c..561e813348 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -69,6 +69,7 @@ apply_normalization, noop_cat, set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, ) from ..quantized_tensor import ( @@ -80,6 +81,8 @@ ) from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -236,6 +239,8 @@ def forward( # Avoid quantized norm kernel if norm output will be returned # or if a gather of ln_out must be in high precision. custom = is_custom(input_quantizer) + hybrid = isinstance(input_quantizer, HybridQuantizer) + identity = isinstance(input_quantizer, IdentityQuantizer) with_quantized_norm = ( fp8 and not debug @@ -243,6 +248,8 @@ def forward( and not return_layernorm_output_gathered and backward_override is None and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() + and not hybrid + and not identity ) # Apply normalization @@ -767,12 +774,7 @@ def backward( quantizer = None if ctx.input_quantizer is not None and ctx.fp8: quantizer = ctx.input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if ctx.ub_bulk_dgrad: ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, @@ -929,7 +931,7 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.grad_output_quantizer is not None: - ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(ctx.grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 19e20d775c..6a2ad31461 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -74,7 +74,14 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer -from ._common import apply_normalization, set_quantizer_amax_reduction_group, WeightGradStore +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer +from ._common import ( + apply_normalization, + set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, + WeightGradStore, +) from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -413,12 +420,16 @@ def _forward( # for debug: : layernorm output = High precision to enable processing of this norm custom = is_custom(fc1_input_quantizer) + hybrid = isinstance(fc1_input_quantizer, HybridQuantizer) + identity = isinstance(fc1_input_quantizer, IdentityQuantizer) with_quantized_norm = ( fp8 and not debug and not return_layernorm_output and not return_layernorm_output_gathered and not custom + and not hybrid + and not identity ) # Apply normalization @@ -1171,12 +1182,7 @@ def backward( quantizer = None if ctx.fp8 or ctx.debug: quantizer = ctx.fc1_input_quantizer - if isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if ctx.ub_bulk_dgrad: ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( @@ -1294,7 +1300,7 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.fc2_grad_output_quantizer is not None: - ctx.fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(ctx.fc2_grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, @@ -1471,7 +1477,10 @@ def fc2_wgrad_gemm( if ctx.fp8: # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now if ( - isinstance(ctx.fc1_grad_output_quantizer, Float8BlockQuantizer) + isinstance( + ctx.fc1_grad_output_quantizer, + (Float8BlockQuantizer, IdentityQuantizer), + ) or ctx.fp8_recipe.custom() ): fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) @@ -2491,7 +2500,7 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): rowwise=True, columnwise=isinstance( fc2_input_quantizer, - (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer), + (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer, HybridQuantizer), ), ) fc2_input_quantizer.internal = True diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 6b0f941bc4..56622db5e6 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -30,7 +30,13 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ._common import noop_cat, set_quantizer_amax_reduction_group, WeightGradStore +from ._common import ( + can_reconstruct_wgrad_input_from_original, + noop_cat, + set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, + WeightGradStore, +) from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( cast_if_needed, @@ -294,10 +300,38 @@ def _linear_forward_impl( debug = args.debug backward_override = args.backward_override is_fsdp2 = args.is_fsdp2 + backward_needs_input = is_grad_enabled and weight.requires_grad if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": save_original_input = False + if ( + backward_override is None + and save_original_input + and backward_needs_input + and input_quantizer is not None + ): + # Megatron-Core enables this automatically for attention output projections + # to reuse the high-precision DPA output saved by attention backward. Validate + # the resolved quantizer since this is not necessarily an informed user opt-in. + if isinstance(input_quantizer, Float8Quantizer): + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError("DelayedScaling recipe is not supported with save_original_input") + elif not can_reconstruct_wgrad_input_from_original(input_quantizer): + warnings.warn( + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({input_quantizer}).", + stacklevel=2, + ) + save_original_input = False # NVTX label for profiling nvtx_label = "transformer_engine._Linear.forward" @@ -310,7 +344,6 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -341,10 +374,6 @@ def _linear_forward_impl( own_quantized_input = False if fp8: assert_dim_for_fp8_exec(inputmat, weight) - if save_original_input: - assert not isinstance( - input_quantizer, Float8Quantizer - ), "DelayedScaling recipe is not supported with save_original_input" if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor @@ -951,12 +980,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. quantizer = None if bwd_args.fp8 or bwd_args.debug: quantizer = input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if bwd_args.ub_bulk_dgrad: inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, @@ -1128,7 +1152,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. and bwd_args.ub_obj_gradout.with_cublasmp() ): if grad_output_quantizer is not None: - grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, bwd_args.tp_group, @@ -1531,7 +1555,8 @@ class Linear(TransformerEngineBaseModule): If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. - Cannot work with FP8 DelayedScaling recipe. + Requires an input quantizer that can safely reproduce its result from the + original input. Cannot work with FP8 DelayedScaling recipe. """ def __init__( diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 607346ce30..f39115c4c6 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -13,11 +13,36 @@ from transformer_engine_torch import FP8TensorMeta from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager +from ..quantized_tensor import QuantizedTensorStorage, Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..tensor.float8_tensor import Float8Tensor -from ..quantized_tensor import QuantizedTensorStorage from ..utils import canonicalize_dtype +def get_fused_normalization_quantizer( + quantizer: Optional[Quantizer], +) -> Optional[Quantizer]: + """Return a quantizer supported by fused normalization kernels.""" + if isinstance( + quantizer, + ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, + ), + ): + return quantizer + return None + + def validate_or_alloc_output( buffer: Optional[torch.Tensor], shape: tuple[int, ...] | list[int], diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 3fda5145c6..3372952add 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -24,7 +24,11 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize +from .._common import ( + get_fused_normalization_quantizer, + maybe_autocast_dtype, + maybe_dequantize, +) class LayerNorm(BasicOperation): @@ -183,6 +187,9 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) + # Fall back to a high-precision output when fused quantization is unsupported. + next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + # Check tensor dims weight = self.weight weight_dims = tuple(weight.size()) diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 1d8d8be971..96851c8db4 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -24,7 +24,11 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize +from .._common import ( + get_fused_normalization_quantizer, + maybe_autocast_dtype, + maybe_dequantize, +) class RMSNorm(BasicOperation): @@ -166,6 +170,9 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) + # Fall back to a high-precision output when fused quantization is unsupported. + next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + # Check tensor dims weight = self.weight weight_dims = tuple(weight.size()) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f1bffb122f..07a3b80483 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -260,7 +260,17 @@ def get_default_recipe() -> Recipe: def get_align_size_for_quantization(recipe: Recipe) -> int: - """Get the alignment size for quantization.""" + """Get the alignment used to pad grouped quantized operations. + + Built-in recipes use their format requirement. Custom recipes use their + declarative ``quantization_alignment`` contract rather than invoking + ``qfactory``, since a factory may be stateful or role-dependent. + """ + # TODO(#3158): Prefer module/role-specific alignment derived from canonical + # cached quantizers when that context is available. Keep the recipe-wide + # alignment as the conservative fallback for context-free callers. + if recipe.custom(): + return recipe.quantization_alignment if recipe.mxfp8(): return 32 if recipe.nvfp4(): @@ -1894,18 +1904,17 @@ def make_quantizers(self) -> list: ) roles = [QuantizerRole() for _ in range(self.num_quantizers)] - # qfactory must return a Quantizer or QuantizerRequest for every slot. - # None is not a valid return value — it would silently disable quantization - # for that tensor, risking hard-to-detect performance regressions. - # TODO(negvet): Introduce an explicit IdentityQuantizer for intentional no-op - # quantization. Until then, None is rejected. + # qfactory returns one quantizer-like object per slot; use + # ``IdentityQuantizer`` for intentional high-precision passthrough. raw = [qfactory(roles[i]) for i in range(self.num_quantizers)] for i, q in enumerate(raw): if q is None: raise ValueError( f"CustomRecipe qfactory returned None for slot {i} " f"(role={roles[i]}). Every slot must return a Quantizer " - "instance or a QuantizerRequest." + "instance or a QuantizerRequest. For an intentional no-op " + "(high-precision / unquantized) slot, return an " + "IdentityQuantizer instead of None." ) # -- Delayed scaling sub-state -- diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index c4194e7f52..11224516bb 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -133,6 +133,65 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: f"{self.__class__.__name__} class does not implement copy_from_storage function" ) + # ── FSDP2 buffer protocol ─────────────────────────────────────── + # + # These three methods decouple FSDP2 all-gather buffer extraction from + # format-specific padding/layout tricks. `HybridQuantizedTensor` uses them + # to aggregate buffers from its two sub-storages without knowing each + # sub-storage's internal field layout. + # + # Contract: + # * ``fsdp_buffer_fields`` returns an ordered tuple of attribute names + # on *self* that hold the tensor buffers that must be all-gathered. + # Scalars/metadata that only need broadcasting (e.g. per-tensor FP8 + # ``_scale_inv``) are NOT listed here — they travel via the hook's + # metadata tuple instead. + # * ``fsdp_extract_buffers`` returns ``(buffers, reassembly_meta)``. + # The default implementation reads the fields as-is. Sub-storages with + # gather-time padding (MXFP8 block scales) override this to strip the + # padding before gather. + # * ``fsdp_assign_gathered`` writes the gathered buffers back into the + # storage's fields. Sub-storages with gather-time padding override + # this to re-apply the padding before assignment. + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Ordered attribute names holding tensor buffers gathered by FSDP2.""" + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement fsdp_buffer_fields" + ) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Return ``(buffers, reassembly_meta)`` for FSDP2 all-gather. + + Default implementation reads the fields named by ``fsdp_buffer_fields`` + verbatim. Override when the on-disk layout differs from the + gather-ready layout (e.g. MXFP8 block scales carry alignment padding). + """ + names = self.fsdp_buffer_fields() + buffers = tuple(getattr(self, name) for name in names) + return buffers, {"field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered buffers into the fields named in ``meta``. + + Override when the gather-ready layout needs a format-specific transform + (e.g. MXFP8 scales must be padded back to ``[128, 4]`` / ``[4, 128]``). + """ + names = meta["field_names"] + if len(names) != len(gathered): + raise RuntimeError( + f"{type(self).__name__}.fsdp_assign_gathered got " + f"{len(gathered)} buffers for {len(names)} fields" + ) + for name, buf in zip(names, gathered): + setattr(self, name, buf) + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], @@ -394,6 +453,14 @@ def supports_only_rowwise_all_gather(self) -> bool: """Returns True if the quantizer supports only rowwise all-gather""" return False + def is_requantization_safe(self) -> bool: + """Whether repeated quantization of the same input reproduces the same value. + + Stateful or stochastic quantizers should return ``False``. This lets callers + decide whether a quantized value may be discarded and reconstructed later. + """ + return False + def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-argument """Whether tensor supports quantized all-gather @@ -501,6 +568,7 @@ def __new__( requires_grad: bool = False, device: Optional[torch.device] = None, stride: Optional[Iterable[int]] = None, + storage_offset: int = 0, ): if fake_dtype is not None and fake_dtype != dtype: raise ValueError(f"fake_dtype ({fake_dtype}) does not match dtype ({dtype})") @@ -513,7 +581,7 @@ def __new__( cls, shape, strides=stride, - storage_offset=0, + storage_offset=storage_offset, dtype=dtype, layout=torch.strided, requires_grad=requires_grad, diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 426c656d47..c3355b6c62 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -19,11 +19,15 @@ from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .storage.nvfp4_tensor_storage import NVFP4TensorStorage from .storage.grouped_tensor_storage import GroupedTensorStorage +from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .grouped_tensor import GroupedTensor +from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer +from .identity_tensor import IdentityTensor, IdentityQuantizer +from .storage.identity_tensor_storage import IdentityTensorStorage from .utils import cast_master_weights_to_fp8, replace_raw_data __all__ = [ @@ -33,18 +37,24 @@ "MXFP8Quantizer", "Float8BlockQuantizer", "NVFP4Quantizer", + "HybridQuantizer", + "IdentityQuantizer", "QuantizedTensorStorage", "Float8TensorStorage", "MXFP8TensorStorage", "Float8BlockwiseQTensorStorage", "NVFP4TensorStorage", "GroupedTensorStorage", + "HybridQuantizedTensorStorage", + "IdentityTensorStorage", "QuantizedTensor", "Float8Tensor", + "IdentityTensor", "MXFP8Tensor", "Float8BlockwiseQTensor", "NVFP4Tensor", "GroupedTensor", + "HybridQuantizedTensor", "prepare_for_saving", "restore_from_saved", "restore_from_func_ctx", @@ -97,5 +107,9 @@ def get_all_tensor_types(): NVFP4TensorStorage, GroupedTensor, GroupedTensorStorage, + HybridQuantizedTensor, + HybridQuantizedTensorStorage, + IdentityTensor, + IdentityTensorStorage, ] return all_tensor_types diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 1b08039dda..10672bbcfb 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -9,13 +9,66 @@ """ from __future__ import annotations -from typing import Callable, Optional, Tuple, Any, Dict, TYPE_CHECKING +from typing import Callable, Optional, Tuple, Any, Dict, Iterable, TYPE_CHECKING import torch if TYPE_CHECKING: from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +def _resolve_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: + """Resolve a requested view shape with PyTorch-compatible semantics. + + The concrete-integer path avoids constructing a temporary meta tensor. If + either shape contains symbolic dimensions, retain the previous meta-tensor + path so that PyTorch remains responsible for symbolic shape handling. + """ + input_shape = tuple(input_shape) + shape = tuple(shape) + if len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)): + shape = tuple(shape[0]) + + # Avoid comparisons that specialize or guard SymInts. The meta fallback is + # also useful for preserving PyTorch's type checking of non-integer dims. + if any(not isinstance(dim, int) or isinstance(dim, bool) for dim in (*input_shape, *shape)): + return torch.empty(input_shape, device="meta").view(shape).shape + + input_numel = 1 + for dim in input_shape: + input_numel *= dim + + inferred_dim = None + known_numel = 1 + for index, dim in enumerate(shape): + if dim == -1: + if inferred_dim is not None: + raise RuntimeError("only one dimension can be inferred") + inferred_dim = index + elif dim < 0: + raise RuntimeError( + f"invalid shape dimension {dim} at index {index} of shape {list(shape)}" + ) + else: + known_numel *= dim + + resolved_shape = list(shape) + if inferred_dim is not None: + if known_numel == 0: + if input_numel == 0: + raise RuntimeError( + f"cannot reshape tensor of 0 elements into shape {list(shape)} because " + "the unspecified dimension size -1 can be any value and is ambiguous" + ) + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + if input_numel % known_numel != 0: + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + resolved_shape[inferred_dim] = input_numel // known_numel + elif known_numel != input_numel: + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + + return torch.Size(resolved_shape) + + class _QuantizeFunc(torch.autograd.Function): """Quantize tensor""" diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 09fde86f17..6d26ba384e 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -122,6 +122,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def is_requantization_safe(self) -> bool: + """Block-FP8 scales are derived deterministically from each input.""" + return True + def get_scale_shape(self, shape: Iterable[int], columnwise: bool) -> Tuple[int, int]: """Scaling tensor shape. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 2e0491d49a..075561a30e 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -19,7 +19,11 @@ from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer from ..dynamo import register_value_opaque_quantizer -from ._quantization_helpers import _IdentityFunc, safe_quantized_repr +from ._quantization_helpers import ( + _IdentityFunc, + _resolve_view_shape, + safe_quantized_repr, +) from ..constants import dist_group_type, DType aten = torch.ops.aten @@ -38,6 +42,14 @@ } +def _columnwise_shape_for(rowwise_shape: Iterable[int]) -> torch.Size: + """Physical columnwise FP8 shape for a logical rowwise shape.""" + shape = torch.Size(rowwise_shape) + if len(shape) == 0: + return shape + return torch.Size((shape[-1], *shape[:-1])) + + class Float8Quantizer(Quantizer): """Builder class for FP8 tensors with per-tensor delayed scaling @@ -387,6 +399,10 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + def is_requantization_safe(self) -> bool: + """Current scaling is derived deterministically from each input.""" + return True + register_value_opaque_quantizer(Float8CurrentScalingQuantizer) @@ -479,8 +495,12 @@ def detach(self) -> Float8Tensor: def clone(self) -> Float8Tensor: # pylint: disable=missing-function-docstring - assert self._data is not None - data = self._data.detach().clone() + # ``_data`` may be None for columnwise-only sub-storages of a + # HybridQuantizedTensor on architectures without native non-TN FP8 + # GEMM (Hopper / L40), where columnwise-only Float8 allocates + # ``_transpose`` instead of ``_data``. On Blackwell+ the C++ + # override keeps ``_data`` populated even in columnwise-only mode. + data = self._data.detach().clone() if self._data is not None else None data_transpose = None if self._transpose is not None: data_transpose = self._transpose.detach().clone() @@ -537,6 +557,12 @@ def _reset_caches(self) -> None: Set transpose cache as invalid. Should be called after any in-place operation. """ + if self._data is None and self._transpose is not None: + # Columnwise-only Float8 tensors on Hopper / L40 store their only + # live FP8 payload in _transpose. Treat it as primary storage, not + # as a derived cache that can be invalidated. + self._transpose_invalid = False + return self._transpose_invalid = True def remove_caches(self) -> None: @@ -581,24 +607,34 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.view.default: tensor = args[0] data = tensor._data - out_data = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - out_shape = out_data.size() + out_data = None + if data is not None: + out_data = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + out_shape = out_data.size() + else: + out_shape = _resolve_view_shape(tensor.shape, args[1:]) + out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if ( - out_transpose_shape[0] != out_shape[-1] - or out_transpose_shape[1:] != out_shape[:-1] - ): + view_shape_for_transpose = _columnwise_shape_for(out_shape) + if out_transpose.shape != view_shape_for_transpose: + if data is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data is only supported " + "when the requested shape preserves the columnwise layout" + ) out_transpose = None else: - view_shape_for_transpose = [out_shape[-1]] + list(out_shape[:-1]) out_transpose = out_transpose.view(*view_shape_for_transpose) + if data is None and out_transpose is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data requires a valid columnwise buffer" + ) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, @@ -614,17 +650,22 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func in (aten.slice.Tensor, aten.select.int): tensor = args[0] data = tensor._data - data_slice = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) + data_slice = None + if data is not None: + data_slice = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) transpose_slice = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose - ndim = data.dim() + ndim = tensor.dim() + if ndim == 0: + return super().__torch_dispatch__(func, types, args, kwargs) dim = args[1] if len(args) > 1 else 0 + dim %= ndim t_dim = 0 if dim == ndim - 1 else dim + 1 transpose_slice = transpose.__torch_dispatch__( func, @@ -632,34 +673,55 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, t_dim] + list(args[2:]), kwargs, ) + if func == aten.select.int and dim == ndim - 1 and transpose_slice.dim() > 1: + transpose_slice = transpose_slice.movedim(-1, 0).contiguous() + + if data_slice is not None: + out_shape = data_slice.shape + else: + logical = torch.empty(tensor.shape, device="meta") + out_shape = func(logical, *args[1:], **(kwargs or {})).shape + if transpose_slice is None: + raise RuntimeError( + "Float8Tensor slice/select requires rowwise or columnwise data" + ) + expected_transpose_shape = _columnwise_shape_for(out_shape) + if transpose_slice.shape != expected_transpose_shape: + raise RuntimeError( + "Float8Tensor slice/select produced incompatible columnwise storage: " + f"expected {tuple(expected_transpose_shape)}, got " + f"{tuple(transpose_slice.shape)}" + ) return Float8Tensor.make_like( tensor, data=data_slice, data_transpose=transpose_slice, - shape=data_slice.shape, + shape=out_shape, ) # Related to FSDP2 if func == aten.split.Tensor: tensor = args[0] data = tensor._data - func_out = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - t_func_out = [None] * len(func_out) - # Compute corresponding split of the transpose cache if available + # _data may be None for columnwise-only sub-storages (hybrid quantization) + if data is not None: + func_out = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + else: + func_out = None + + t_func_out = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose - ndim = data.dim() - # Figure out the original split dim + ndim = tensor.dim() if "dim" in kwargs: dim_to_split = kwargs["dim"] else: dim_to_split = args[2] if len(args) > 2 else 0 - # Dimension along which transpose needs to be split t_dim = 0 if dim_to_split == ndim - 1 else dim_to_split + 1 t_func_out = transpose.__torch_dispatch__( func, @@ -667,12 +729,30 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, args[1], t_dim], kwargs, ) + + ref_out = func_out if func_out is not None else t_func_out + if ref_out is None: + return super().__torch_dispatch__(func, types, args, kwargs) + + num_splits = len(ref_out) + if func_out is None: + func_out = [None] * num_splits + if t_func_out is None: + t_func_out = [None] * num_splits + outs = [ Float8Tensor.make_like( tensor, data=split_tensor, data_transpose=split_transpose_tensor, - shape=split_tensor.shape, + shape=( + split_tensor.shape + if split_tensor is not None + else ( + *split_transpose_tensor.shape[1:], + split_transpose_tensor.shape[0], + ) + ), ) for split_tensor, split_transpose_tensor in zip(func_out, t_func_out) ] @@ -682,12 +762,18 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # create fresh new tensor with zeros. tensor = args[0] data = tensor._data - func_out = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) + storage_kwargs = dict(kwargs or {}) + output_dtype = storage_kwargs.pop("dtype", None) or tensor.dtype + storage_kwargs.pop("layout", None) + storage_kwargs.pop("requires_grad", None) + func_out = None + if data is not None: + func_out = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + storage_kwargs, + ) func_transposed_out = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose @@ -697,25 +783,69 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): func, types, [transpose, t_shape] + list(args[2:]), - kwargs, + storage_kwargs, ) + if func_out is None and func_transposed_out is None: + raise RuntimeError("Float8Tensor.new_zeros requires rowwise or columnwise data") scale_inv = tensor._scale_inv.detach().clone() + reference = func_out if func_out is not None else func_transposed_out + if scale_inv.device != reference.device: + scale_inv = scale_inv.to(reference.device) quantizer = tensor._quantizer # Deep-copied in constructor out_tensor = Float8Tensor( data=func_out, - shape=func_out.shape, - dtype=tensor.dtype, + shape=torch.Size(args[1]), + dtype=output_dtype, fp8_dtype=tensor._fp8_dtype, fp8_scale_inv=scale_inv, data_transpose=func_transposed_out, quantizer=quantizer, - device=tensor.device, + device=reference.device, ) return out_tensor if func == torch.ops.aten.as_strided.default: tensor = args[0] data = tensor._data + if data is None: + size = torch.Size(args[1]) + stride = tuple(args[2]) + storage_offset = (kwargs or {}).get( + "storage_offset", + args[3] if len(args) > 3 else tensor.storage_offset(), + ) + # A contiguous 2D row shard maps to a column slice in the + # persistent transpose and can remain a Float8Tensor. + has_valid_transpose = ( + tensor._transpose is not None and not tensor._transpose_invalid + ) + is_contiguous_row_shard = ( + tensor.dim() == 2 + and tensor.shape[1] > 0 + and len(size) == 2 + and size[1] == tensor.shape[1] + and stride == tuple(tensor.stride()) + ) + if ( + has_valid_transpose + and is_contiguous_row_shard + and storage_offset % tensor.shape[1] == 0 + ): + row_start = storage_offset // tensor.shape[1] + row_end = row_start + size[0] + if 0 <= row_start and row_end <= tensor.shape[0]: + transpose_out = tensor._transpose[:, row_start:row_end] + return Float8Tensor.make_like( + tensor, + data=None, + data_transpose=transpose_out, + shape=size, + ) + + # Arbitrary logical strides are not generally affine views of + # transposed storage. Fall back explicitly to high precision. + return func(tensor.dequantize(), *args[1:], **(kwargs or {})) + # Apply as_strided to the primary uint8 data func_out = data.__torch_dispatch__( func, @@ -927,9 +1057,19 @@ def __reduce_ex__(self, protocol: int) -> tuple: with ``torch.serialization.add_safe_globals`` to keep ``torch.load(weights_only=True)`` compatibility. """ + data_transpose = None + if self._data is None and self._transpose is not None and not self._transpose_invalid: + data_transpose = self._transpose return ( _make_float8_tensor_in_reduce_ex, - (self._data, self._fp8_dtype, self._scale_inv, self.dtype, self.shape), + ( + self._data, + self._fp8_dtype, + self._scale_inv, + self.dtype, + self.shape, + data_transpose, + ), ) @classmethod @@ -1025,11 +1165,12 @@ def _set_data(self, tensor: torch.Tensor) -> None: def _make_float8_tensor_in_reduce_ex( - data: torch.Tensor, + data: Optional[torch.Tensor], fp8_dtype: DType, fp8_scale_inv: torch.Tensor, dtype: torch.dtype, shape: torch.Size, + data_transpose: Optional[torch.Tensor] = None, ) -> Float8Tensor: """Reconstruct a ``Float8Tensor`` from its ``__reduce_ex__`` payload.""" return Float8Tensor( @@ -1038,7 +1179,12 @@ def _make_float8_tensor_in_reduce_ex( fp8_scale_inv=fp8_scale_inv, dtype=dtype, shape=shape, - device=data.device if data is not None else None, + data_transpose=data_transpose, + device=( + data.device + if data is not None + else data_transpose.device if data_transpose is not None else None + ), ) @@ -1059,16 +1205,28 @@ def forward( ctx.shape = tensor.shape if shape is None: return tensor.detach() - out_data = tensor._data.view(*shape) - out_shape = out_data.size() + out_data = None + if tensor._data is not None: + out_data = tensor._data.view(*shape) + out_shape = out_data.size() + else: + out_shape = _resolve_view_shape(tensor.shape, shape) out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if out_transpose_shape[0] != out_shape[-1] or out_transpose_shape[1:] != out_shape[:-1]: + view_shape_for_transpose = _columnwise_shape_for(out_shape) + if out_transpose.shape != view_shape_for_transpose: + if tensor._data is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data is only supported " + "when the requested shape preserves the columnwise layout" + ) out_transpose = None else: - view_shape_for_transpose = [shape[-1]] + list(shape[:-1]) out_transpose = out_transpose.view(*view_shape_for_transpose) + if tensor._data is None and out_transpose is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data requires a valid columnwise buffer" + ) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py new file mode 100644 index 0000000000..8f6051c5f8 --- /dev/null +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -0,0 +1,1169 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tensor and quantizer classes for composed rowwise and columnwise representations.""" + +from __future__ import annotations +from typing import Any, Dict, Iterable, Literal, Optional, Tuple, Union + +import torch + +from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from .storage.identity_tensor_storage import IdentityTensorStorage +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer + +aten = torch.ops.aten + + +class HybridQuantizer(Quantizer): + """Quantizer that composes rowwise and columnwise representations. + + When both representations are requested, applies ``rowwise_quantizer`` to + produce the rowwise representation and ``columnwise_quantizer`` to produce + the columnwise representation. The results are wrapped in a + ``HybridQuantizedTensor``. The children may use the same or different + Transformer Engine formats, or custom ``Quantizer`` implementations. + + Parameters + ---------- + rowwise_quantizer : Quantizer + Quantizer for the rowwise direction (e.g. MXFP8Quantizer). + columnwise_quantizer : Quantizer + Quantizer for the columnwise direction (e.g. NVFP4Quantizer). + columnwise_source : {"original", "rowwise_dequantized"}, default = "original" + Source tensor for columnwise quantization. ``"original"`` quantizes + columnwise directly from the input tensor. ``"rowwise_dequantized"`` + quantizes rowwise first, dequantizes the rowwise result, then uses that + value as the columnwise source. + + Notes + ----- + Rowwise and columnwise describe storage and GEMM orientations. Whether a + representation is consumed in the forward or backward pass depends on the + tensor's role in the operation. + + ``HybridQuantizer`` pins each sub-quantizer to its designated direction by + mutating its usage flags, so it takes ownership of the supplied quantizer + instances. The rowwise and columnwise quantizers must be distinct objects. + If both directions need shared state, construct two quantizer instances that + reference the same external state object. + + Each ``HybridQuantizer`` must receive its own sub-quantizer instances; do not + reuse a sub-quantizer instance across multiple ``HybridQuantizer`` objects. + + Example + ------- + MXFP8 rowwise data plus a high-precision columnwise representation derived + from the rowwise value can be expressed as:: + + HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer, + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + In a ``CustomRecipe`` factory this can be combined with role-based + specialization. For example, a factory can return a ``HybridQuantizer`` only + for ``role.tensor_type == "weight"`` and use regular quantizers for inputs + and gradients. See + ``custom_recipes.quantization_factory_zoo.nvfp4_1d_double_quantized_weight_quantizer_factory`` + for a weight-only double-quantization example. + + """ + + _COLUMNWISE_SOURCES = ("original", "rowwise_dequantized") + + rowwise_quantizer: Quantizer + columnwise_quantizer: Quantizer + columnwise_source: Literal["original", "rowwise_dequantized"] + + def __init__( + self, + *, + rowwise_quantizer: Quantizer, + columnwise_quantizer: Quantizer, + columnwise_source: Literal["original", "rowwise_dequantized"] = "original", + ) -> None: + super().__init__(rowwise=True, columnwise=True) + from transformer_engine.pytorch.quantization import QuantizerRequest # local import + + for role, quantizer in ( + ("rowwise", rowwise_quantizer), + ("columnwise", columnwise_quantizer), + ): + if isinstance(quantizer, QuantizerRequest): + # TODO(#3158): Support delayed-scaling requests inside hybrid sub-quantizers. + raise TypeError( + "HybridQuantizer does not support nested QuantizerRequest " + f"objects yet; got {type(quantizer).__name__} for the {role} " + "direction. Delayed scaling in CustomRecipe is currently " + "supported only when the qfactory returns DelayedScalingRequest " + "as a top-level slot. Resolving delayed-scaling requests inside " + "HybridQuantizer is future work; pass a concrete Quantizer " + "instance instead." + ) + if not isinstance(quantizer, Quantizer): + raise TypeError( + "HybridQuantizer requires concrete Quantizer instances for " + f"both directions, but the {role} argument is " + f"{type(quantizer).__name__}." + ) + if rowwise_quantizer is columnwise_quantizer: + raise ValueError( + "HybridQuantizer requires distinct rowwise and columnwise quantizer" + " instances. If both directions need shared state, construct two" + " quantizer objects that reference the same shared state." + ) + if columnwise_source not in self._COLUMNWISE_SOURCES: + raise ValueError( + "HybridQuantizer columnwise_source must be one of " + f"{self._COLUMNWISE_SOURCES}, got {columnwise_source!r}." + ) + self.rowwise_quantizer = rowwise_quantizer + self.columnwise_quantizer = columnwise_quantizer + self.columnwise_source = columnwise_source + + # Pin each sub-quantizer to its designated direction + self.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) + self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f"rowwise_usage={self.rowwise_usage}, " + f"columnwise_usage={self.columnwise_usage}, " + f"columnwise_source={self.columnwise_source!r}, " + f"internal={self.internal}, " + ")" + ) + + def copy(self) -> "HybridQuantizer": + """Create a shallow copy, preserving parent and sub-quantizer state.""" + quantizer = HybridQuantizer( + rowwise_quantizer=self.rowwise_quantizer.copy(), + columnwise_quantizer=self.columnwise_quantizer.copy(), + columnwise_source=self.columnwise_source, + ) + quantizer.set_usage( + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + @property + def with_amax_reduction(self) -> bool: + """Whether either sub-quantizer has cross-rank amax reduction enabled.""" + return getattr(self.rowwise_quantizer, "with_amax_reduction", False) or getattr( + self.columnwise_quantizer, "with_amax_reduction", False + ) + + @with_amax_reduction.setter + def with_amax_reduction(self, value: bool) -> None: + # Set on the HybridQuantizer by module / FSDP2 code, but read by the C++ + # kernel off the sub-quantizer that runs -- hence forwarded, not stored here. + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + if hasattr(sub, "with_amax_reduction"): + sub.with_amax_reduction = value + + @property + def amax_reduction_group(self): + """Amax-reduction group of the sub-quantizers, or ``None`` if unset.""" + result = None + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + group = getattr(sub, "amax_reduction_group", None) + if group is None: + continue + if result is None: + result = group + elif group is not result: + raise RuntimeError( + "HybridQuantizer sub-quantizers have inconsistent amax_reduction_group values." + ) + return result + + @amax_reduction_group.setter + def amax_reduction_group(self, value) -> None: + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + if hasattr(sub, "amax_reduction_group"): + sub.amax_reduction_group = value + + def _columnwise_src_from_rowwise( + self, + tensor: torch.Tensor, + rowwise_result: Optional[Any], + ) -> torch.Tensor: + if rowwise_result is None: + rowwise_result = self.rowwise_quantizer.quantize(tensor) + return rowwise_result.dequantize(dtype=tensor.dtype) + + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: + # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers + # are pinned to one direction in ``__init__``; the parent flag decides + # whether to invoke them. + rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None + columnwise_src = tensor + if self.columnwise_usage and self.columnwise_source == "rowwise_dequantized": + columnwise_src = self._columnwise_src_from_rowwise(tensor, rowwise_result) + columnwise_result = ( + self.columnwise_quantizer.quantize(columnwise_src) if self.columnwise_usage else None + ) + + if self.internal: + return HybridQuantizedTensorStorage( + rowwise_storage=rowwise_result, + columnwise_storage=columnwise_result, + quantizer=self, + fake_dtype=tensor.dtype, + ) + + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=rowwise_result, + columnwise_storage=columnwise_result, + quantizer=self, + ) + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, + pin_memory: bool = False, + ) -> Union["HybridQuantizedTensor", HybridQuantizedTensorStorage]: + # Mirror ``quantize_impl``: invoke each sub-quantizer with its own + # ``internal`` setting (no toggle), so the produced sub-storages have + # the same type that ``quantize_impl`` would produce via + # ``sub_quantizer.quantize(tensor)``. + rowwise_empty = ( + self.rowwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory + ) + if self.rowwise_usage + else None + ) + columnwise_empty = ( + self.columnwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory + ) + if self.columnwise_usage + else None + ) + + if self.internal: + return HybridQuantizedTensorStorage( + rowwise_storage=rowwise_empty, + columnwise_storage=columnwise_empty, + quantizer=self, + fake_dtype=dtype, + ) + + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + requires_grad=requires_grad, + device=device, + rowwise_storage=rowwise_empty, + columnwise_storage=columnwise_empty, + quantizer=self, + ) + + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, + ) -> QuantizedTensorStorage: + """Re-quantize sub-storages of a hybrid tensor in-place. + + Each direction is refreshed only when the parent usage flag is set + **and** the corresponding sub-storage exists. + """ + if not isinstance(dst, HybridQuantizedTensorStorage): + raise ValueError( + "HybridQuantizer can only update HybridQuantizedTensorStorage, got" + f" {type(dst).__name__}" + ) + rowwise_result_for_columnwise = None + if self.rowwise_usage and dst._rowwise_storage is not None: + self.rowwise_quantizer.update_quantized(src, dst._rowwise_storage, noop_flag=noop_flag) + rowwise_result_for_columnwise = dst._rowwise_storage + if self.columnwise_usage and dst._columnwise_storage is not None: + columnwise_src = src + if self.columnwise_source == "rowwise_dequantized": + columnwise_src = self._columnwise_src_from_rowwise( + src, rowwise_result_for_columnwise + ) + self.columnwise_quantizer.update_quantized( + columnwise_src, dst._columnwise_storage, noop_flag=noop_flag + ) + return dst + + def supports_only_rowwise_all_gather(self) -> bool: + """Whether all-gather requires a rowwise-dequantizable source. + + Hybrid tensors currently use the high-precision fallback, which + dequantizes the local shard before communication. Preserve rowwise + data when required by the rowwise sub-quantizer or when the + columnwise sub-quantizer is NVFP4, whose columnwise-only storage + cannot be dequantized. + """ + if self.rowwise_quantizer.supports_only_rowwise_all_gather(): + return True + # Local import avoids a circular dependency chain + # (nvfp4_tensor → quantized_tensor → hybrid_tensor at module import). + from .nvfp4_tensor import NVFP4Quantizer # noqa: PLC0415 + + return isinstance(self.columnwise_quantizer, NVFP4Quantizer) + + def is_requantization_safe(self) -> bool: + """Whether repeated quantization reproduces all requested representations.""" + if self.rowwise_usage and not self.rowwise_quantizer.is_requantization_safe(): + return False + if self.columnwise_usage and not self.columnwise_quantizer.is_requantization_safe(): + return False + if ( + self.columnwise_usage + and self.columnwise_source == "rowwise_dequantized" + and not self.rowwise_quantizer.is_requantization_safe() + ): + return False + return True + + def _get_compatible_recipe(self): + # HybridQuantizer is only reachable via CustomRecipe (the qfactory + # returns HybridQuantizer per role). Checking that the autocast recipe + # is also CustomRecipe catches the obvious mismatch (e.g. hybrid + # quantized_model_init + built-in MXFP8BlockScaling autocast). + # We trust that users who write a CustomRecipe know what they're doing + # with regard to per-operand scaling mode compatibility. + # + # TODO(#3158): validate per-operand scaling-mode compatibility at + # recipe-build time instead of at cuBLAS-dispatch time. Concretely: + # 1. Walk the qfactory outputs for a given module_type (``linear``, + # ``grouped_linear``, ``dpa``) — call the factory for each + # ``QuantizerRole.tensor_type`` the module uses. + # 2. Extract the scaling_mode of each sub-quantizer: + # weight_row, weight_col (from HybridQuantizer) + # input_row, input_col (from HybridQuantizer) + # grad_output_row, grad_output_col (plain quantizer OR + # HybridQuantizer) + # 3. Assert the three GEMM pairs share a scaling_mode each: + # fprop TN: weight_row == input_row (FormatA) + # dgrad NN: weight_col == grad_output_row (FormatB) + # wgrad NT: input_col == grad_output_col (FormatC) + # Mismatches raise ``ValueError`` naming the offending slots, e.g. + # "dgrad GEMM: weight columnwise format (MXFP8) does not match + # grad_output rowwise format (NVFP4)". + # 4. Blocked on `semantic_quantizer_roles` / PR #2620 for the + # ``QuantizerRole`` dataclass — the factory signature is role- + # aware only on that branch. + from transformer_engine.common.recipe import CustomRecipe # avoid circular import + + return CustomRecipe + + +class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): + """Tensor holding independently produced rowwise and columnwise representations. + + The tensor presents as having a standard logical dtype, but + internally stores representations produced by the two sub-quantizers. + These may use the same format, different formats, or a high-precision + representation such as ``IdentityTensorStorage``. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Nominal tensor datatype. + rowwise_storage : QuantizedTensorStorage, optional + Sub-storage for the rowwise representation. + columnwise_storage : QuantizedTensorStorage, optional + Sub-storage for the columnwise representation. + quantizer : HybridQuantizer + Parent hybrid quantizer that owns the rowwise and columnwise sub-quantizers. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + + """ + + def __new__( + cls, + *args, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + quantizer: HybridQuantizer, + **kwargs, + ): + instance = super().__new__( + cls, + *args, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + quantizer=quantizer, + **kwargs, + ) + return instance + + def __repr__(self, *, tensor_contents=None): + row_type = ( + type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + ) + col_type = ( + type(self._columnwise_storage).__name__ + if self._columnwise_storage is not None + else "None" + ) + return ( + f"HybridQuantizedTensor(rowwise={row_type}, columnwise={col_type}, dtype={self.dtype})" + ) + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if dtype is None: + dtype = self.dtype + return HybridQuantizedTensorStorage.dequantize(self, dtype=dtype) + + def detach(self) -> HybridQuantizedTensor: + """Return a new HybridQuantizedTensor with cloned sub-storage wrappers. + + Each sub-storage is re-wrapped via its own ``make_like`` so the + new hybrid tensor has independent sub-storage objects that share + the *underlying* buffer tensors with ``self``. This is required for + the cpu_offload_v2 pattern at ``cpu_offload.py:378-382``:: + + tensor_copy = tensor.detach() + saved_tensors, _ = tensor_copy.prepare_for_saving() # nulls fields + + If ``detach()`` merely shared sub-storage objects, the + ``prepare_for_saving`` call above would null out fields on the + original ``tensor`` too (since both hybrids would point at the same + sub-storage Python objects), and subsequent operations — even a + bare ``.device`` read during ``_check_if_offload`` for a follow-up + ``push_tensor`` on the same original — would crash with + `` has no data!``. + """ + row = None + if self._rowwise_storage is not None: + row_cls = type(self._rowwise_storage) + if hasattr(row_cls, "make_like"): + row = row_cls.make_like(self._rowwise_storage) + else: + raise NotImplementedError( + "HybridQuantizedTensor.detach() does not support storage-only " + f"rowwise sub-storage {row_cls.__name__}" + ) + col = None + if self._columnwise_storage is not None: + col_cls = type(self._columnwise_storage) + if hasattr(col_cls, "make_like"): + col = col_cls.make_like(self._columnwise_storage) + else: + raise NotImplementedError( + "HybridQuantizedTensor.detach() does not support storage-only " + f"columnwise sub-storage {col_cls.__name__}" + ) + return HybridQuantizedTensor( + shape=self.shape, + dtype=self.dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=self._quantizer, + ) + + def get_metadata(self) -> Dict[str, Any]: + return HybridQuantizedTensorStorage.get_metadata(self) + + @staticmethod + def _move_metadata_value( + value: Any, + *, + target_device: torch.device, + non_blocking: bool, + pin_memory: bool, + ) -> Any: + if isinstance(value, torch.Tensor): + value = value.to(device=target_device, non_blocking=non_blocking) + if pin_memory and target_device.type == "cpu": + value = value.pin_memory() + return value + + @classmethod + def _move_sub_storage( + cls, + sub_storage: Optional[QuantizedTensorStorage], + *, + target_device: torch.device, + non_blocking: bool, + pin_memory: bool, + ) -> Optional[QuantizedTensorStorage]: + if sub_storage is None: + return None + metadata = { + key: cls._move_metadata_value( + value, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + for key, value in sub_storage.get_metadata().items() + } + if isinstance(sub_storage, QuantizedTensor): + metadata.update( + { + "shape": sub_storage.shape, + "dtype": sub_storage.dtype, + "requires_grad": sub_storage.requires_grad, + "device": target_device, + } + ) + return type(sub_storage)(**metadata) + + def contiguous( + self, + memory_format: torch.memory_format = torch.contiguous_format, + ) -> "HybridQuantizedTensor": + """Return a HybridQuantizedTensor with contiguous sub-storages.""" + + def _contiguous_sub( + role: str, + sub_storage: Optional[QuantizedTensorStorage], + ) -> Optional[QuantizedTensorStorage]: + if sub_storage is None: + return None + if not isinstance(sub_storage, torch.Tensor): + raise ValueError( + "HybridQuantizedTensor.contiguous does not support storage-only " + f"{role} sub-storage {type(sub_storage).__name__}. This path is " + "only supported for tensor sub-storages." + ) + try: + return sub_storage.contiguous(memory_format=memory_format) + except (NotImplementedError, ValueError) as err: + raise ValueError( + "HybridQuantizedTensor.contiguous could not make the " + f"{role} sub-storage {type(sub_storage).__name__} contiguous " + f"with memory_format={memory_format}." + ) from err + + row = _contiguous_sub("rowwise", self._rowwise_storage) + col = _contiguous_sub("columnwise", self._columnwise_storage) + if row is self._rowwise_storage and col is self._columnwise_storage: + return self + return HybridQuantizedTensor( + shape=self.shape, + dtype=self.dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=self.device, + ) + + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling. + + Without this, the default ``torch.Tensor.__reduce_ex__`` rebuilds + the parameter as a plain ``torch.Tensor``, dropping the + sub-storages and per-tensor scale state. DCP then reloads the + parameter via ``aten.copy_(dst, plain_tensor)`` which routes to + ``dst.quantize_(plain_tensor)`` — re-quantizing dequantized data + loses precision. + + Mirrors the per-format ``__reduce_ex__`` on ``Float8Tensor``, + ``MXFP8Tensor``, ``NVFP4Tensor``, and ``Float8BlockwiseQTensor``. + Each sub-storage is itself pickled via its own ``__reduce_ex__`` + (preserving FP8 bytes + ``_scale_inv``); the quantizers travel as + regular Python objects and must therefore be picklable + themselves. + """ + return ( + _make_hybrid_quantized_tensor_in_reduce_ex, + ( + self._rowwise_storage, + self._columnwise_storage, + self._quantizer, + self.dtype, + self.shape, + ), + ) + + # ── FSDP2 protocol ────────────────────────────────────────────── + + def fsdp_pre_all_gather( # pylint: disable=unused-argument + self, mesh, orig_size, contiguous_orig_stride, module, mp_policy + ): + """Extract plain tensor buffers from both sub-storages for FSDP2 all-gather. + + Always send both directions. This gives a stable buffer count/shape + across forward and backward, at the cost of gathering the unused + direction each pass. No requantization, no BF16 fallback. + + Buffer extraction is delegated to each sub-storage's + :meth:`QuantizedTensorStorage.fsdp_extract_buffers`, which strips any + format-specific padding (e.g. MXFP8 block-scale alignment) before the + gather so concatenation along dim-0 is well-defined. + + TODO(#3158): bandwidth optimization — pack both directions into a + single flat buffer sized ``max(row_bytes, col_bytes)`` (not + ``row_bytes + col_bytes``) to halve comm volume for asymmetric format + pairs. Planned implementation: a new per-sub-storage + ``fsdp_pack_into(flat_buffer, offset, meta)`` helper that layouts + both directions back-to-back with offsets stored in the metadata + tuple; ``fsdp_post_all_gather`` would slice the gathered flat buffer + using those offsets. + """ + # Mirror ``Float8Tensor.fsdp_pre_all_gather``: enable cross-shard amax + # reduction so the post-optimizer re-quantization of the sharded weight + # keeps one shared scale across shards (no-op for sub-quantizers without + # amax reduction, e.g. MXFP8). + if mesh is not None: + self._quantizer.amax_reduction_group = mesh.get_group() + self._quantizer.with_amax_reduction = True + + # Quick, targeted error for sub-storages whose FSDP2 support isn't + # implemented yet (e.g. NVFP4). Without this, users hit + # NotImplementedError from deep inside fsdp_extract_buffers with a + # generic message. + for role, sub in ( + ("rowwise", self._rowwise_storage), + ("columnwise", self._columnwise_storage), + ): + if sub is None: + continue + if not isinstance(sub, QuantizedTensor): + raise NotImplementedError( + "Hybrid FSDP2 all-gather does not support storage-only " + f"{role} sub-storage {type(sub).__name__}. This usually means " + "a HybridQuantizer sub-quantizer had internal=True; use " + "tensor sub-storages for Hybrid FSDP2 or disable Hybrid FSDP2 " + "for this parameter." + ) + try: + sub.fsdp_buffer_fields() + except NotImplementedError as err: + raise NotImplementedError( + "Hybrid FSDP2 all-gather is not supported for a " + f"{type(sub).__name__} {role} sub-storage: it does not " + "implement fsdp_buffer_fields. " + "NVFP4 sub-storages need packed-FP4 dim-0 alignment, " + "columnwise dequantization and RHT-cache handling before " + "they can be gathered. Use a supported sub-quantizer " + "(Float8CurrentScaling, MXFP8, Float8Block) or run without " + "FSDP2." + ) from err + + row_buffers: Tuple[Optional[torch.Tensor], ...] = () + col_buffers: Tuple[Optional[torch.Tensor], ...] = () + row_meta: Optional[Dict[str, Any]] = None + col_meta: Optional[Dict[str, Any]] = None + if self._rowwise_storage is not None: + row_buffers, row_meta = self._rowwise_storage.fsdp_extract_buffers() + if self._columnwise_storage is not None: + col_buffers, col_meta = self._columnwise_storage.fsdp_extract_buffers() + + sharded_tensors = row_buffers + col_buffers + + metadata = ( + len(row_buffers), + row_meta, + col_meta, + self._rowwise_storage, # original sharded sub-storage (for make_like on iter-1) + self._columnwise_storage, + self._quantizer, + ) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[HybridQuantizedTensor] = None, + ): + """Reconstruct HybridQuantizedTensor from all-gathered buffers. + + On iteration 1 (``out=None``): clone each sub-storage via + :meth:`make_like` from the sharded original, then delegate the + gathered-buffer writeback (and any format-specific re-padding) to + :meth:`QuantizedTensorStorage.fsdp_assign_gathered`. + On iteration 2+ (``out=prev``): delegate directly to the existing + sub-storages' ``fsdp_assign_gathered``. + """ + ( + n_row_buffers, + row_meta, + col_meta, + orig_row_sub, + orig_col_sub, + hybrid_quantizer, + ) = metadata + + row_quantizer = hybrid_quantizer.rowwise_quantizer + col_quantizer = hybrid_quantizer.columnwise_quantizer + + row_gathered = all_gather_outputs[:n_row_buffers] + col_gathered = all_gather_outputs[n_row_buffers:] + + def _infer_shape(gathered_buffers): + for buf in gathered_buffers: + if buf is not None: + return buf.shape + return None + + # ``update_usage`` after gathered writeback mirrors what vanilla + # ``Float8Tensor.fsdp_post_all_gather`` does + # — invalidates any stale ``_transpose`` cache on Float8 sub-storages + # and recreates the transpose on non-Hopper architectures where the + # FP8 cuBLAS path requires it. No-op on Blackwell. The flags come + # from the sub-quantizer's pinned direction (set by + # ``HybridQuantizer.__init__``), so we honor whatever the inner + # quantizer thinks its direction is. + def _sync_usage(sub_storage, sub_quantizer): + if sub_storage is None or sub_quantizer is None: + return + sub_storage.update_usage( + rowwise_usage=sub_quantizer.rowwise_usage, + columnwise_usage=sub_quantizer.columnwise_usage, + ) + + if out is not None: + # Iteration 2+: in-place field update on existing sub-storages + if out._rowwise_storage is not None and row_meta is not None: + out._rowwise_storage.fsdp_assign_gathered(row_gathered, row_meta) + _sync_usage(out._rowwise_storage, out._quantizer.rowwise_quantizer) + if out._columnwise_storage is not None and col_meta is not None: + out._columnwise_storage.fsdp_assign_gathered(col_gathered, col_meta) + _sync_usage(out._columnwise_storage, out._quantizer.columnwise_quantizer) + else: + # First iteration: clone the original sharded sub-storages via make_like, + # then write gathered (full-size) buffers via each sub-storage's own + # fsdp_assign_gathered so padding is re-applied where applicable. + row_sub = None + if orig_row_sub is not None and isinstance(orig_row_sub, QuantizedTensor): + gathered_shape = _infer_shape(row_gathered) + row_sub = type(orig_row_sub).make_like(orig_row_sub, shape=gathered_shape) + if row_meta is not None: + row_sub.fsdp_assign_gathered(row_gathered, row_meta) + _sync_usage(row_sub, row_quantizer) + + col_sub = None + if orig_col_sub is not None and isinstance(orig_col_sub, QuantizedTensor): + gathered_shape = _infer_shape(col_gathered) + col_sub = type(orig_col_sub).make_like(orig_col_sub, shape=gathered_shape) + if col_meta is not None: + col_sub.fsdp_assign_gathered(col_gathered, col_meta) + _sync_usage(col_sub, col_quantizer) + + ref_sub = row_sub if row_sub is not None else col_sub + out = HybridQuantizedTensor( + shape=( + ref_sub.shape + if ref_sub is not None + else _infer_shape(row_gathered + col_gathered) + ), + dtype=param_dtype, + rowwise_storage=row_sub, + columnwise_storage=col_sub, + quantizer=hybrid_quantizer, + ) + + return out, all_gather_outputs + + @classmethod + def _delegate_reshape_op(cls, func, tensor, args, kwargs): + """Delegate a shape-altering op (slice, as_strided) to each sub-storage. + + Returns a new ``HybridQuantizedTensor`` when every non-None sub-storage + returns a ``QuantizedTensorStorage`` of the same kind (i.e. real + op support, as Float8Tensor provides via its own + ``__torch_dispatch__``). Returns ``None`` when any sub-storage + dequantized to a plain ``torch.Tensor`` (i.e. the sub-storage does not + support this op — MXFP8Tensor / Float8BlockwiseQTensor fall through + that way for real slicing today). On ``None`` the caller should defer + to ``super().__torch_dispatch__`` for a consistent BF16 fallback. + """ + + def _delegate(sub): + if sub is None: + return None + return func(sub, *args[1:], **kwargs) + + row_out = _delegate(tensor._rowwise_storage) + col_out = _delegate(tensor._columnwise_storage) + + row_ok = row_out is None or isinstance(row_out, QuantizedTensorStorage) + col_ok = col_out is None or isinstance(col_out, QuantizedTensorStorage) + if not (row_ok and col_ok): + return None + if row_out is None and col_out is None: + return None + + ref = row_out if row_out is not None else col_out + return HybridQuantizedTensor( + shape=ref.shape, + dtype=tensor.dtype, + rowwise_storage=row_out, + columnwise_storage=col_out, + quantizer=tensor._quantizer, + ) + + @staticmethod + def _new_zeroed_sub_storage( + sub_storage: QuantizedTensorStorage, + shape: torch.Size, + *, + dtype: torch.dtype, + device: torch.device, + pin_memory: bool, + ) -> QuantizedTensorStorage: + """Allocate and initialize a zero sub-storage without touching live state.""" + + try: + quantizer = sub_storage._get_quantizer().copy() + out = quantizer.make_empty( + shape, + dtype=dtype, + device=device, + requires_grad=False, + pin_memory=pin_memory, + ) + except NotImplementedError as exc: + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros cannot construct zero-initialized " + f"{type(sub_storage).__name__} storage" + ) from exc + + if not isinstance(out, type(sub_storage)): + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros did not preserve sub-storage format: " + f"expected {type(sub_storage).__name__}, got {type(out).__name__}" + ) + + buffers, storage = out.prepare_for_saving() + try: + with torch.no_grad(): + for buffer in buffers: + if buffer is not None: + buffer.zero_() + finally: + leftover = storage.restore_from_saved(buffers) + if leftover: + raise RuntimeError( + f"{type(out).__name__}.restore_from_saved did not consume all buffers" + ) + return out + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if kwargs is None: + kwargs = {} + + if func == aten.detach.default: + return args[0].detach() + + if func == aten._to_copy.default: + tensor = args[0] + kw = dict(kwargs) if kwargs else {} + dtype = kw.get("dtype", None) + if dtype is None or dtype == tensor.dtype: + target_device = torch.device(kw.get("device", tensor.device) or tensor.device) + pin_memory = bool(kw.get("pin_memory", False)) + non_blocking = bool(kw.get("non_blocking", False)) + row = cls._move_sub_storage( + tensor._rowwise_storage, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + col = cls._move_sub_storage( + tensor._columnwise_storage, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=tensor._quantizer, + requires_grad=tensor.requires_grad, + device=target_device, + ) + + # ── FSDP2: view ────────────────────────────────────────────── + if func == aten.view.default: + tensor = args[0] + shape = args[1] + # Identity view fast-path (FSDP2 reset_sharded_param issues a view + # to the param's own shape). The columnwise sub-storage's own shape + # is transposed relative to the hybrid for some formats (e.g. a 2D + # block-scaled Float8BlockwiseQTensor has shape (N, M) for an + # (M, N) weight). Forwarding the hybrid's row-major shape to it + # would be a spurious last-2-dims change, which 2D block scaling + # cannot represent and so dequantizes the sub-storage to a plain + # tensor -- breaking the FSDP2 sub-storage protocol later. Preserve + # the sub-storages as-is, mirroring the as_strided / slice no-op + # fast paths below. + if list(shape) == list(tensor.shape): + return HybridQuantizedTensor.make_like(tensor) + row_view = None + col_view = None + if tensor._rowwise_storage is not None: + row_view = tensor._rowwise_storage.view(shape) + if tensor._columnwise_storage is not None: + col_view = tensor._columnwise_storage.view(shape) + return HybridQuantizedTensor( + shape=shape, + dtype=tensor.dtype, + rowwise_storage=row_view, + columnwise_storage=col_view, + quantizer=tensor._quantizer, + ) + + # ── FSDP2: split ───────────────────────────────────────────── + if func == aten.split.Tensor: + tensor = args[0] + split_size = args[1] + dim = kwargs.get("dim", args[2] if len(args) > 2 else 0) + + if dim != 0: + return super().__torch_dispatch__(func, types, args, kwargs) + + row_pieces = ( + torch.split(tensor._rowwise_storage, split_size, dim=dim) + if tensor._rowwise_storage is not None + else None + ) + col_pieces = ( + torch.split(tensor._columnwise_storage, split_size, dim=dim) + if tensor._columnwise_storage is not None + else None + ) + + if row_pieces is None and col_pieces is None: + return super().__torch_dispatch__(func, types, args, kwargs) + + # TODO(#3158): Support Hybrid sub-storages that fall back to a + # high-precision tensor for an unquantizable local shard. + for direction, pieces in ( + ("rowwise", row_pieces), + ("columnwise", col_pieces), + ): + if pieces is None: + continue + for piece in pieces: + if not isinstance(piece, QuantizedTensor): + raise NotImplementedError( + "HybridQuantizedTensor split produced a high-precision " + f"{type(piece).__name__} for the {direction} sub-storage " + f"with local shape {tuple(piece.shape)}. Hybrid FSDP2 does " + "not support high-precision fallback children. MXFP8 FSDP " + "shards must have a first dimension divisible by 32; adjust " + "the sharding topology or disable Hybrid MXFP8 for this " + "parameter. See #3158." + ) + + num_pieces = len(row_pieces) if row_pieces is not None else len(col_pieces) + return [ + HybridQuantizedTensor( + shape=(row_pieces[i].shape if row_pieces is not None else col_pieces[i].shape), + dtype=tensor.dtype, + rowwise_storage=row_pieces[i] if row_pieces is not None else None, + columnwise_storage=col_pieces[i] if col_pieces is not None else None, + quantizer=tensor._quantizer, + ) + for i in range(num_pieces) + ] + + # ── FSDP2: as_strided / slice ──────────────────────────────── + # Fast path for no-op (common during FSDP2 reset_sharded_param); + # otherwise delegate per sub-storage so we inherit each sub-storage's + # own support level. Float8Tensor implements real slicing/as_strided + # via `_data.__torch_dispatch__`; MXFP8Tensor and Float8BlockwiseQTensor + # handle only the no-op case and fall through to dequantize for real + # ops (matching their vanilla FSDP2 behaviour). If any sub-storage + # returns a plain torch.Tensor (dequantized), we can't rewrap into a + # hybrid so we fall through to super() for a consistent BF16 fallback. + if func == aten.as_strided.default: + tensor = args[0] + shape = args[1] + strides = args[2] + storage_offset = kwargs.get("storage_offset", args[3] if len(args) > 3 else None) + if storage_offset is None: + storage_offset = tensor.storage_offset() + if ( + tuple(shape) == tuple(tensor.size()) + and tuple(strides) == tuple(tensor.stride()) + and storage_offset == tensor.storage_offset() + ): + return HybridQuantizedTensor.make_like(tensor) + out = cls._delegate_reshape_op(func, tensor, args, kwargs) + if out is not None: + return out + return super().__torch_dispatch__(func, types, args, kwargs) + + if func == aten.slice.Tensor: + tensor = args[0] + dim = args[1] + start = args[2] + end = args[3] + step = args[4] if len(args) > 4 else 1 + if start == 0 and end == tensor.size(dim) and step == 1: + return HybridQuantizedTensor.make_like(tensor) + out = cls._delegate_reshape_op(func, tensor, args, kwargs) + if out is not None: + return out + return super().__torch_dispatch__(func, types, args, kwargs) + + # ── FSDP2: copy_ ───────────────────────────────────────────── + # Fast path for hybrid-to-hybrid (FSDP2 fills buffer allocated via + # new_zeros/make_empty). Other src types (e.g. a BF16 master weight + # during checkpoint load) fall through to QuantizedTensor's base + # dispatch which routes to ``dst.quantize_(src)``. + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(dst, HybridQuantizedTensor) and isinstance(src, HybridQuantizedTensor): + dst_usages = dst.get_usages() + src_usages = src.get_usages() + if dst_usages != src_usages: + raise NotImplementedError( + "HybridQuantizedTensor.copy_ requires matching rowwise/columnwise " + f"usages, but source has {src_usages} and destination has " + f"{dst_usages}. Copy from a high-precision tensor instead." + ) + + if dst._rowwise_storage is not None and src._rowwise_storage is not None: + aten.copy_.default(dst._rowwise_storage, src._rowwise_storage) + if dst._columnwise_storage is not None and src._columnwise_storage is not None: + aten.copy_.default(dst._columnwise_storage, src._columnwise_storage) + return dst + + # ── FSDP2: new_zeros ───────────────────────────────────────── + if func == aten.new_zeros.default: + tensor = args[0] + new_shape = torch.Size(args[1]) + if tensor._rowwise_storage is None and tensor._columnwise_storage is None: + raise RuntimeError( + "HybridQuantizedTensor.new_zeros requires at least one present sub-storage" + ) + dtype = kwargs.get("dtype") or tensor.dtype + device = torch.device(kwargs.get("device") or tensor.device) + layout = kwargs.get("layout") + pin_memory = bool(kwargs.get("pin_memory")) + if layout is not None and layout != torch.strided: + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros only supports torch.strided layout" + ) + + supported_quantized_dtypes = (torch.float32, torch.float16, torch.bfloat16) + for direction, sub_storage in ( + ("rowwise", tensor._rowwise_storage), + ("columnwise", tensor._columnwise_storage), + ): + if sub_storage is None or isinstance(sub_storage, IdentityTensorStorage): + continue + if dtype not in supported_quantized_dtypes: + raise TypeError( + "HybridQuantizedTensor.new_zeros only supports float32, float16, or " + f"bfloat16 with a {direction} {type(sub_storage).__name__}; got {dtype}." + ) + + row = ( + cls._new_zeroed_sub_storage( + tensor._rowwise_storage, + new_shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + if tensor._rowwise_storage is not None + else None + ) + col = ( + cls._new_zeroed_sub_storage( + tensor._columnwise_storage, + new_shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + if tensor._columnwise_storage is not None + else None + ) + + # The source's parent usage flags are mutable and may no longer + # describe its surviving sub-storages. Give the result an isolated + # parent whose dynamic usage matches what was actually allocated; + # otherwise a later copy_ from a plain tensor can silently skip a + # present direction. Bind its children to copies of the newly + # allocated storage quantizers so their format/state stays coherent. + quantizer = tensor._quantizer.copy() + if row is not None: + quantizer.rowwise_quantizer = row._get_quantizer().copy() + quantizer.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) + if col is not None: + quantizer.columnwise_quantizer = col._get_quantizer().copy() + quantizer.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + quantizer.set_usage( + rowwise=row is not None, + columnwise=col is not None, + ) + + return HybridQuantizedTensor( + shape=new_shape, + dtype=dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=quantizer, + requires_grad=False, + device=device, + ) + + # ── FSDP2: clone ───────────────────────────────────────────── + if func == aten.clone.default: + tensor = args[0] + row_clone = ( + torch.clone(tensor._rowwise_storage) + if tensor._rowwise_storage is not None + else None + ) + col_clone = ( + torch.clone(tensor._columnwise_storage) + if tensor._columnwise_storage is not None + else None + ) + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=row_clone, + columnwise_storage=col_clone, + quantizer=tensor._quantizer, + ) + + return super().__torch_dispatch__(func, types, args, kwargs) + + +def _make_hybrid_quantized_tensor_in_reduce_ex( + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + quantizer: HybridQuantizer, + dtype: torch.dtype, + shape: torch.Size, +) -> HybridQuantizedTensor: + """Reconstruct a ``HybridQuantizedTensor`` from its ``__reduce_ex__`` payload.""" + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + quantizer=quantizer, + ) diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py new file mode 100644 index 0000000000..9fb980a755 --- /dev/null +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -0,0 +1,355 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""High-precision passthrough quantizer and tensor. + +``IdentityQuantizer`` stores a tensor directly without a low-precision encoding, +preserving the input dtype by default. It exists so the ``CustomRecipe`` + +``qfactory`` machinery can express *high-precision* tensors and, composed inside +a :class:`HybridQuantizer`, *high-precision directions* without scattering +``None``/``isinstance`` special-cases across the modules. Here high precision +means the held compute dtype, typically BF16, FP16, or FP32; it does not mean +FP32 specifically. +""" + +from __future__ import annotations +from typing import Any, Iterable, Optional, Tuple, Union + +import torch +from torch.ops import aten + +from .storage.identity_tensor_storage import IdentityTensorStorage +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer + + +class IdentityQuantizer(Quantizer): + """Quantizer that produces a high-precision passthrough representation. + + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) + holding the tensor directly, without a low-precision encoding. + ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it + in the held dtype. + + Parameters + ---------- + dtype : torch.dtype, optional + If set, the held tensor is cast to this dtype on quantize. ``None`` + (default) keeps the input's dtype. + rowwise, columnwise : bool + Usage flags (kept for interface compatibility; the single + high-precision buffer serves both directions). + """ + + def __init__( + self, + *, + dtype: Optional[torch.dtype] = None, + rowwise: bool = True, + columnwise: bool = True, + ) -> None: + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.dtype = dtype + + def copy(self) -> "IdentityQuantizer": + """Create shallow copy.""" + quantizer = IdentityQuantizer( + dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def _maybe_cast(self, tensor: torch.Tensor) -> torch.Tensor: + # Detach so the held buffer is plain "data" with no autograd graph edge, + # mirroring the real quantizers (whose quantize kernels emit fresh, + # non-differentiable tensors). Autograd connectivity for the *quantize* + # op is provided separately by ``_QuantizeFunc`` in ``Quantizer.quantize``; + # the surrounding TE module Function computes dgrad/wgrad manually. Without + # the detach the produced tensor aliases a grad-requiring input (e.g. the + # weight workspace returned across the module Function boundary), which + # creates a spurious empty grad edge. + out = tensor.detach() + if self.dtype is not None and out.dtype != self.dtype: + return out.to(self.dtype) + return out + + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensorStorage: + data = self._maybe_cast(tensor) + if self.internal: + return IdentityTensorStorage( + hp_data=data, + fake_dtype=data.dtype, + quantizer=self, + ) + # requires_grad=False: this is the quantized "data" tensor. Autograd + # connectivity is provided by ``_QuantizeFunc`` in ``Quantizer.quantize`` + # (mirrors the real quantizers, which return non-differentiable data). + return IdentityTensor( + data.shape, + data.dtype, + hp_data=data, + quantizer=self, + requires_grad=False, + device=data.device, + ) + + def is_requantization_safe(self) -> bool: + """Identity quantization is deterministic.""" + return True + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, + pin_memory: bool = False, + ) -> Union["IdentityTensor", IdentityTensorStorage]: + if device is None: + device = torch.device("cuda") + device = torch.device(device) + data_dtype = self.dtype if self.dtype is not None else dtype + data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory) + if self.internal: + return IdentityTensorStorage( + hp_data=data, + fake_dtype=data_dtype, + quantizer=self, + ) + return IdentityTensor( + data.shape, + data_dtype, + hp_data=data, + quantizer=self, + requires_grad=requires_grad, + device=device, + ) + + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, + ) -> QuantizedTensorStorage: + if not isinstance(dst, IdentityTensorStorage): + raise ValueError( + f"IdentityQuantizer can only update IdentityTensorStorage, got {type(dst).__name__}" + ) + data = self._maybe_cast(src) + if ( + dst._hp_data is not None + and dst._hp_data.shape == data.shape + and dst._hp_data.dtype == data.dtype + and dst._hp_data.device == data.device + ): + if noop_flag is None: + dst._hp_data.copy_(data) + else: + torch.where(noop_flag == 0, data, dst._hp_data, out=dst._hp_data) + else: + if noop_flag is not None and noop_flag.item() != 0: + return dst + dst._hp_data = data.detach() + dst._dtype = data.dtype + return dst + + def calibrate(self, tensor: torch.Tensor) -> None: + # No state to calibrate. + return + + def _get_compatible_recipe(self): + # Only reachable via CustomRecipe (qfactory returns IdentityQuantizer). + from transformer_engine.common.recipe import CustomRecipe # avoid circular import + + return CustomRecipe + + +class IdentityTensor(IdentityTensorStorage, QuantizedTensor): + """High-precision passthrough tensor produced by :class:`IdentityQuantizer`. + + Presents as a standard tensor of its nominal dtype; internally it just + holds data directly in that dtype, without a low-precision encoding. + """ + + def __repr__(self, *, tensor_contents=None): + return f"IdentityTensor(dtype={self.dtype}, data={self._hp_data})" + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + return IdentityTensorStorage.dequantize(self, dtype=dtype) + + def view(self, *shape) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + flat_shape = shape[0] if len(shape) == 1 and not isinstance(shape[0], int) else shape + return self._wrap_data_view(self._hp_data.view(*flat_shape)) + + def detach(self) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + return self._wrap_data_view(self._hp_data.detach(), requires_grad=False) + + def clone(self) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + data = self._hp_data.detach().clone() if self._hp_data is not None else None + return IdentityTensor( + self.shape, + self.dtype, + hp_data=data, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=self.device, + stride=data.stride(), + storage_offset=data.storage_offset(), + ) + + def contiguous( + self, + memory_format: torch.memory_format = torch.contiguous_format, + ) -> "IdentityTensor": + """Return an IdentityTensor with contiguous high-precision storage.""" + if self._hp_data is not None and self._hp_data.is_contiguous(memory_format=memory_format): + return self + return self._wrap_data_view(self._hp_data.contiguous(memory_format=memory_format)) + + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling that preserves the high-precision payload.""" + return ( + _make_identity_tensor_in_reduce_ex, + (self._hp_data, self._quantizer, self.dtype, self.shape), + ) + + def fsdp_pre_all_gather( # pylint: disable=unused-argument + self, mesh, orig_size, contiguous_orig_stride, module, mp_policy + ): + """Extract the high-precision buffer for FSDP2 all-gather.""" + return (self._hp_data,), (self._quantizer,) + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional["IdentityTensor"] = None, + ): + """Rebuild IdentityTensor from the gathered high-precision buffer.""" + (data,) = all_gather_outputs + (quantizer,) = metadata + logical_dtype = ( + quantizer.dtype + if quantizer is not None and quantizer.dtype is not None + else param_dtype + ) + if data.dtype != logical_dtype: + raise RuntimeError( + "IdentityTensor FSDP payload dtype does not match its logical dtype: " + f"payload={data.dtype}, logical={logical_dtype}." + ) + if out is not None: + out._hp_data = data + out._dtype = logical_dtype + else: + out = IdentityTensor( + shape=data.shape, + dtype=logical_dtype, + hp_data=data, + quantizer=quantizer, + requires_grad=False, + device=data.device, + ) + return out, all_gather_outputs + + def _wrap_data_view( + self, data: torch.Tensor, *, requires_grad: Optional[bool] = None + ) -> "IdentityTensor": + requires_grad = self.requires_grad if requires_grad is None else requires_grad + return IdentityTensor( + shape=data.shape, + dtype=self.dtype, + hp_data=data, + quantizer=self._quantizer, + requires_grad=requires_grad, + device=data.device, + stride=data.stride(), + storage_offset=data.storage_offset(), + ) + + @classmethod + def _delegate_view_op(cls, func, tensor, args, kwargs): + """Apply an alias-preserving view op to the held tensor and rewrap it.""" + + result = func(tensor._hp_data, *args[1:], **kwargs) + + def _wrap(value): + if isinstance(value, torch.Tensor): + return tensor._wrap_data_view(value) + return value + + if isinstance(result, tuple): + return tuple(_wrap(value) for value in result) + if isinstance(result, list): + return [_wrap(value) for value in result] + return _wrap(result) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if kwargs is None: + kwargs = {} + + if func == aten.detach.default: + return args[0].detach() + + if func == aten.clone.default: + return args[0].clone() + + if func in ( + aten.view.default, + aten.split.Tensor, + aten.as_strided.default, + aten.slice.Tensor, + ): + # Preserve optional arguments exactly; omitted as_strided offset + # means reuse the input view's current storage offset, not zero. + return cls._delegate_view_op(func, args[0], args, kwargs) + + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(dst, IdentityTensor): + src_data = src._hp_data if isinstance(src, IdentityTensor) else src + dst._hp_data.copy_(src_data, *args[2:], **kwargs) + return dst + + if func == aten.new_zeros.default: + tensor = args[0] + new_shape = args[1] + if tensor._quantizer is not None: + out = tensor._quantizer.make_empty( + new_shape, + dtype=kwargs.get("dtype") or tensor.dtype, + device=kwargs.get("device") or tensor.device, + pin_memory=bool(kwargs.get("pin_memory", False)), + ) + out._hp_data.zero_() + return out + + return super().__torch_dispatch__(func, types, args, kwargs) + + +def _make_identity_tensor_in_reduce_ex( + hp_data: torch.Tensor, + quantizer: Optional[Quantizer], + dtype: torch.dtype, + shape: torch.Size, +) -> IdentityTensor: + """Reconstruct an ``IdentityTensor`` from its ``__reduce_ex__`` payload.""" + return IdentityTensor( + shape=shape, + dtype=dtype, + hp_data=hp_data, + quantizer=quantizer, + requires_grad=False, + device=hp_data.device if hp_data is not None else None, + ) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 3045662216..f30108aecc 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -180,6 +180,10 @@ def onnx_dequantize(self, tensor: Union[MXFP8TensorStorage, MXFP8Tensor]) -> tor def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return MXFP8BlockScaling + def is_requantization_safe(self) -> bool: + """MXFP8 block scales are derived deterministically from each input.""" + return True + register_value_opaque_quantizer(MXFP8Quantizer) @@ -288,8 +292,10 @@ def detach(self) -> MXFP8Tensor: def clone(self) -> MXFP8Tensor: # pylint: disable=missing-function-docstring - assert self._rowwise_data is not None - rowwise_data = self._rowwise_data.detach().clone() + # _rowwise_data may be None for columnwise-only sub-storages (hybrid quantization) + rowwise_data = ( + self._rowwise_data.detach().clone() if self._rowwise_data is not None else None + ) columnwise_data = None if self._columnwise_data is not None: columnwise_data = self._columnwise_data.detach().clone() @@ -403,67 +409,73 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ): return super().__torch_dispatch__(func, types, args, kwargs) - out_data = [] - for data in [tensor._rowwise_data, tensor._columnwise_data]: - func_out = ( - data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - if data is not None - else None + def _split_data(data): + if data is None: + return None + return data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, ) - out_data.append(func_out) - - scale_invs = [tensor._rowwise_scale_inv, tensor._columnwise_scale_inv] - split_sizes_for_scale = [split_size, split_size // MXFP8_BLOCK_SCALING_SIZE] - # Padding requirements: rowwise dim0 should be divisble by 128, columnwise dim0 should be divisble by 4 - padding_multiples = [128, 4] - for scale_inv, scale_split_size, pad_multiple in zip( - scale_invs, split_sizes_for_scale, padding_multiples - ): - scale_inv_out = ( + + row_data_splits = _split_data(tensor._rowwise_data) + col_data_splits = _split_data(tensor._columnwise_data) + + def _split_scale_inv(scale_inv, scale_split_size, pad_multiple): + if scale_inv is None: + return None + scale_inv_out = list( scale_inv.__torch_dispatch__( func, types, [scale_inv, scale_split_size] + list(args[2:]), kwargs, ) - if scale_inv is not None - else None ) - scale_inv_out = list(scale_inv_out) if scale_inv_out is not None else None - # Pad scale_inv_out to be a multiple of pad_multiple - if scale_inv_out is not None: - for idx, split_scale_inv_out in enumerate(scale_inv_out): - current_shape = split_scale_inv_out.shape - pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple - if pad_dim0 > 0: - scale_inv_out[idx] = torch.nn.functional.pad( - split_scale_inv_out, (0, 0, 0, pad_dim0) - ) - out_data.append(scale_inv_out) + for idx, split_scale_inv_out in enumerate(scale_inv_out): + current_shape = split_scale_inv_out.shape + pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple + if pad_dim0 > 0: + scale_inv_out[idx] = torch.nn.functional.pad( + split_scale_inv_out, (0, 0, 0, pad_dim0) + ) + return scale_inv_out + + row_scale_splits = _split_scale_inv( + tensor._rowwise_scale_inv, + split_size, + 128, + ) + col_scale_splits = _split_scale_inv( + tensor._columnwise_scale_inv, + split_size // MXFP8_BLOCK_SCALING_SIZE, + 4, + ) + + ref_splits = row_data_splits if row_data_splits is not None else col_data_splits + num_splits = len(ref_splits) return [ MXFP8Tensor( shape=( - splitted_tensor_data[0].size() - if splitted_tensor_data[0] is not None - else splitted_tensor_data[1].size() + row_data_splits[i].size() + if row_data_splits is not None + else col_data_splits[i].size() ), dtype=tensor.dtype, - rowwise_data=splitted_tensor_data[0], - rowwise_scale_inv=splitted_tensor_data[2], - columnwise_data=splitted_tensor_data[1], - columnwise_scale_inv=splitted_tensor_data[3], + rowwise_data=row_data_splits[i] if row_data_splits is not None else None, + rowwise_scale_inv=row_scale_splits[i] if row_scale_splits is not None else None, + columnwise_data=col_data_splits[i] if col_data_splits is not None else None, + columnwise_scale_inv=( + col_scale_splits[i] if col_scale_splits is not None else None + ), quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, with_gemm_swizzled_scales=False, device=tensor.device, ) - for splitted_tensor_data in zip(*out_data) + for i in range(num_splits) ] if func == torch.ops.aten.as_strided.default: @@ -554,6 +566,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): device=tensor.device, ) + if func == torch.ops.aten.clone.default: + return cls.clone(args[0]) + # Default case return super().__torch_dispatch__(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ccf06ac166..3f8a11d2e1 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -254,6 +254,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def is_requantization_safe(self) -> bool: + """NVFP4 quantization is replay-safe unless stochastic rounding is enabled.""" + return not self.stochastic_rounding + def is_quantizable(self, inp: torch.Tensor) -> bool: """Returns whether or not given inp can be quantized""" if self.row_scaled_nvfp4: diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 464a7c3b23..8ff1a0dd1f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -17,7 +17,7 @@ from ...constants import TE_DType_To_Torch, DType -from ...utils import _empty_tensor +from ...utils import _empty_tensor, round_up_to_nearest_multiple class _FromFloat8BlockwiseFunc(torch.autograd.Function): @@ -511,3 +511,154 @@ def get_usages(self) -> Dict[str, bool]: "rowwise": self._rowwise_data is not None, "columnwise": self._columnwise_data is not None, } + + # ── FSDP2 sub-storage buffer protocol ──────────────────────────── + # + # Float8Block stores columnwise data N-major (transposed) for the GEMM, so + # it cannot be dim-0 all-gathered directly. Each direction is made + # self-contained: the columnwise direction fp8-transposes its own data to + # M-major for the gather and back on assign, using only its own buffers (no + # dependency on a rowwise sibling, which in a hybrid tensor may be a + # different format). Block-scale GEMM alignment padding (round-up-to-4) is + # stripped before the gather and re-applied after. Only 2D block scaling is + # supported -- the 1D scale layout has M in dim1, incompatible with FSDP2's + # dim-0 all-gather. + + _FSDP_BLOCK_LEN = 128 + + def _fsdp_logical_mn(self) -> Tuple[int, int]: + """Flattened ``(M, N)`` of this sub-storage's logical shape.""" + shape = self.size() + last_dim = shape[-1] if len(shape) > 0 else 1 + leading = 1 + for dim in shape[:-1]: + leading *= dim + return leading, last_dim + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for Float8 block scaling (2D scaling only).""" + if not self._is_2D_scaled: + raise NotImplementedError( + "FSDP2 for Float8BlockwiseQTensor requires 2D block scaling " + "(block_scaling_dim=2). 1D block scaling is not supported because " + "its scale layout has M in dim1, which is incompatible with FSDP2 " + "dim-0 all-gather." + ) + fields = [] + if self._rowwise_data is not None: + fields.extend(("_rowwise_data", "_rowwise_scale_inv")) + if self._columnwise_data is not None: + fields.extend(("_columnwise_data", "_columnwise_scale_inv")) + return tuple(fields) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract M-major, alignment-stripped buffers for dim-0 all-gather. + + Rowwise data is already M-major; columnwise data is N-major and is + fp8-transposed to M-major here (and transposed back in + :meth:`fsdp_assign_gathered`). The block-scale round-up-to-4 alignment + padding is stripped so dim-0 concatenation across shards is well-defined. + """ + names = self.fsdp_buffer_fields() + block_len = self._FSDP_BLOCK_LEN + m, n = self._fsdp_logical_mn() + if m % block_len != 0: + raise RuntimeError( + "FSDP2 cannot all-gather a 2-D Float8BlockwiseQTensor whose " + f"local flattened M dimension ({m}) is not a multiple of {block_len}; " + "the shard boundary splits a block-scale tile. Choose aligned shard " + "boundaries or use a supported non-blockwise recipe." + ) + m_tiles = (m + block_len - 1) // block_len + last_tiles = (n + block_len - 1) // block_len + + if self._rowwise_data is not None: + # Rowwise scale is (m_tiles, round_up(last_tiles, 4)); m_tiles sits in + # dim-0 (sharded/gathered) unpadded, the round-up padding is on dim-1 + # (not sharded). Strip dim-1 to the compact tile count. + scale = self._rowwise_scale_inv + if scale is not None and scale.size(1) > last_tiles: + scale = scale[:, :last_tiles].contiguous() + buffers = (self._rowwise_data, scale) + direction = "rowwise" + else: + # Columnwise data is N-major (N, M); transpose to M-major (M, N). + col_data = self._columnwise_data + if not col_data.is_contiguous(): + col_data = col_data.contiguous() + data_m = tex.fp8_transpose(col_data, self._fp8_dtype, out=None) + # Columnwise scale is (last_tiles, round_up(m_tiles, 4)); transpose to + # (round_up(m_tiles, 4), last_tiles) and strip dim-0 to m_tiles so the + # gathered (dim-0) axis is the M-tiles, matching the rowwise layout. + scale = self._columnwise_scale_inv.transpose(0, 1).contiguous() + if scale.size(0) > m_tiles: + scale = scale[:m_tiles].contiguous() + buffers = (data_m, scale) + direction = "columnwise" + + return buffers, {"direction": direction, "field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered buffers back, re-applying transpose + scale padding. + + Inverse of :meth:`fsdp_extract_buffers`: rowwise re-pads the scale's + last-dim alignment; columnwise transposes the M-major gathered data back + to N-major and re-pads/transposes the scale to the GEMM scale layout + produced by ``get_scale_shape(..., columnwise=True)``. + """ + block_len = self._FSDP_BLOCK_LEN + direction = meta["direction"] + data, scale = gathered + if direction not in ("rowwise", "columnwise"): + raise RuntimeError(f"Invalid Float8Block FSDP gather direction: {direction!r}") + if data is None or scale is None: + raise RuntimeError( + "Float8Block FSDP gathered data and scale buffers must both be present." + ) + + m_full = 1 + for dim in data.shape[:-1]: + m_full *= dim + last_dim = data.size(-1) if data.dim() > 0 else 1 + expected_scale_shape = ( + (m_full + block_len - 1) // block_len, + (last_dim + block_len - 1) // block_len, + ) + if tuple(scale.shape) != expected_scale_shape: + raise RuntimeError( + "Float8Block FSDP gathered scale geometry does not match gathered data: " + f"got scale shape {tuple(scale.shape)}, expected {expected_scale_shape} " + f"for gathered data shape {tuple(data.shape)}." + ) + + if direction == "rowwise": + last_dim = data.size(-1) + last_tiles = (last_dim + block_len - 1) // block_len + if scale is not None: + pad = round_up_to_nearest_multiple(last_tiles, 4) - last_tiles + if pad > 0: + scale = torch.nn.functional.pad(scale, (0, pad)) + self._rowwise_data = data + self._rowwise_scale_inv = scale + return + + # Columnwise: gathered data is M-major (M_full, N); transpose to N-major. + data_m = data if data.is_contiguous() else data.contiguous() + self._columnwise_data = tex.fp8_transpose(data_m, self._fp8_dtype, out=None) + m_full = 1 + for dim in data.shape[:-1]: + m_full *= dim + m_tiles_full = (m_full + block_len - 1) // block_len + # Gathered scale is compact (m_tiles_full, last_tiles); transpose to + # (last_tiles, m_tiles_full) and re-pad the M-tile dim to multiple of 4. + scale_t = scale.transpose(0, 1).contiguous() + pad = round_up_to_nearest_multiple(m_tiles_full, 4) - m_tiles_full + if pad > 0: + scale_t = torch.nn.functional.pad(scale_t, (0, pad)) + self._columnwise_scale_inv = scale_t.contiguous() diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 374d0e1e72..37b669f707 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -12,7 +12,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from .._quantization_helpers import safe_quantized_repr +from .._quantization_helpers import _resolve_view_shape, safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -45,7 +45,23 @@ def forward( # Cast from FP8 return tex.dequantize(tensor, te_dtype) - raise NotImplementedError("Casting back from the transpose not implemented yet!") + if tensor._transpose is not None and not tensor._transpose_invalid: + # A columnwise-only tensor stores the logical last dimension first + # in its physical FP8 buffer. Dequantize that buffer as ordinary + # FP8 data, then restore logical row-major dimension order. + transpose_tensor = Float8TensorStorage( + data=tensor._transpose, + data_transpose=None, + fp8_scale_inv=tensor._scale_inv, + fp8_dtype=tensor._fp8_dtype, + fake_dtype=tensor._dtype, + ) + output = _FromFloat8Func.forward(None, transpose_tensor, dtype) + if output.dim() > 0: + output = output.movedim(0, -1).contiguous() + return output + + raise RuntimeError("Float8TensorStorage has neither rowwise nor columnwise data") @staticmethod def backward( @@ -175,12 +191,15 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: dtype = self._dtype return _FromFloat8Func.forward(None, self, dtype) - def size(self, *args, **kwargs): + def size(self, dim: Optional[int] = None) -> Union[torch.Size, int]: # pylint: disable=missing-function-docstring if self._data is not None: - return self._data.size(*args, **kwargs) - size = self._transpose.size(*args, **kwargs) - return torch.Size([size[-1], math.prod(size[:-1])]) + return self._data.size(dim) + transpose_size = self._transpose.size() + size = torch.Size([transpose_size[-1], math.prod(transpose_size[:-1])]) + if dim is None: + return size + return size[dim] @property def device(self): @@ -193,12 +212,31 @@ def device(self): def view(self, shape: torch.Size): # pylint: disable=missing-function-docstring - out_data = self._data.view(shape) + out_data = self._data.view(shape) if self._data is not None else None + if out_data is not None: + out_shape = out_data.size() + else: + out_shape = _resolve_view_shape(self.size(), shape) out_transpose = None if self._transpose_invalid else self._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if out_transpose_shape[0] != shape[-1] or out_transpose_shape[1:] != shape[:-1]: + if len(out_shape) == 0: + view_shape_for_transpose = out_shape + else: + view_shape_for_transpose = torch.Size((out_shape[-1], *out_shape[:-1])) + if out_transpose.shape != view_shape_for_transpose: + if self._data is None: + raise NotImplementedError( + "Float8TensorStorage view with columnwise-only data is only " + "supported when the requested shape preserves the columnwise layout" + ) out_transpose = None + else: + out_transpose = out_transpose.view(*view_shape_for_transpose) + if self._data is None and out_transpose is None: + raise NotImplementedError( + "Float8TensorStorage view with columnwise-only data requires a valid " + "columnwise buffer" + ) return Float8TensorStorage( data=out_data, @@ -224,6 +262,10 @@ def __repr__(self): def _create_transpose(self): """Update FP8 transpose cache""" data = self._data + # Columnwise-only Float8Tensors (e.g. hybrid quantization sub-storages) + # have _data=None — nothing to transpose. + if data is None: + return if not data.is_contiguous(): data = data.contiguous() self._transpose = tex.fp8_transpose(data, self._fp8_dtype, out=self._transpose) @@ -277,3 +319,90 @@ def get_usages(self) -> Dict[str, bool]: else: usages["columnwise"] = self._transpose is not None and not self._transpose_invalid return usages + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for per-tensor FP8. + + ``_scale_inv`` is a per-tensor scalar; it travels through the hook's + metadata tuple (mirroring :meth:`Float8Tensor.fsdp_pre_all_gather`). + + Direction-aware: a vanilla Float8Tensor parameter has ``_data`` + populated, but a columnwise-only sub-storage (used inside + ``HybridQuantizedTensor`` on Hopper / L40 where non-TN FP8 GEMM is + not natively supported) holds its quantized data in ``_transpose`` + instead. Returning ``("_data",)`` unconditionally would have + ``fsdp_extract_buffers`` produce ``(None,)`` and FSDP2 would + all-gather a ``None`` tensor. + + The per-sub-storage direction is fixed at construction (pinned by + ``HybridQuantizer.__init__`` via ``set_usage``), so this check is + stable across iterations even though it inspects the current + field state. + """ + if self._data is not None: + return ("_data",) + if self._transpose is not None: + return ("_transpose",) + # Degenerate: fully empty storage. Fall back to ``_data`` so the + # base ``fsdp_extract_buffers`` returns ``(None,)`` — same surface + # the caller would have seen pre-direction-aware logic. + return ("_data",) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract dim-0 gather-safe buffers for per-tensor Float8. + + Rowwise ``_data`` is already M-major. A transpose-only Hopper/L40 + payload is physically ``[N, *M]`` and cannot be gathered directly + because FSDP2 grows physical dimension 0. Move N back to the logical + last dimension and make the transport buffer contiguous so FSDP2 sees + ``[*M, N]``. :meth:`fsdp_assign_gathered` restores the persistent + columnwise layout after the collective. + """ + names = self.fsdp_buffer_fields() + if names == ("_transpose",): + if self._transpose is None or self._transpose_invalid: + raise RuntimeError( + "Float8TensorStorage cannot extract an invalid columnwise transpose" + ) + transport = self._transpose + if transport.dim() > 0: + transport = transport.movedim(0, -1).contiguous() + return (transport,), { + "field_names": names, + "transport_layout": "columnwise_m_major", + } + buffers = tuple(getattr(self, name) for name in names) + return buffers, {"field_names": names, "transport_layout": "native"} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Restore gathered Float8 buffers to their persistent layout.""" + transport_layout = meta.get("transport_layout", "native") + if transport_layout == "columnwise_m_major": + if meta.get("field_names") != ("_transpose",) or len(gathered) != 1: + raise RuntimeError( + "Float8TensorStorage got invalid metadata for columnwise FSDP transport" + ) + (transport,) = gathered + if transport is None: + raise RuntimeError( + "Float8TensorStorage got a None columnwise FSDP transport buffer" + ) + self._data = None + if transport.dim() > 0: + transport = transport.movedim(-1, 0).contiguous() + self._transpose = transport + self._transpose_invalid = False + return + if transport_layout != "native": + raise RuntimeError( + f"Float8TensorStorage got unknown FSDP transport layout {transport_layout!r}" + ) + super().fsdp_assign_gathered(gathered, meta) + if "_transpose" in meta["field_names"]: + self._transpose_invalid = False diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 7ede75943a..3473024c03 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -654,6 +654,18 @@ def make_grouped_tensor( A GroupedTensor. """ + if quantizer is not None: + # TODO(#3158): Support Identity/Hybrid packed grouped storage. + from ..hybrid_tensor import HybridQuantizer + from ..identity_tensor import IdentityQuantizer + + if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)): + raise NotImplementedError( + "GroupedTensorStorage does not support IdentityQuantizer or " + "HybridQuantizer yet. Use separate tensors, or set " + "GroupedLinear(single_grouped_weight=False). See #3158." + ) + # Set device if device is None: device = torch.cuda.current_device() diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py new file mode 100644 index 0000000000..822bb342fe --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -0,0 +1,232 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Storage class for composed rowwise and columnwise representations.""" + +from __future__ import annotations +from collections.abc import Iterable +from typing import Any, Dict, Optional, Tuple + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer + + +class HybridQuantizedTensorStorage(QuantizedTensorStorage): + """Storage that composes rowwise and columnwise sub-storages. + + One sub-storage provides the rowwise representation and the other provides + the columnwise representation. Either may be absent when that direction is + not needed. The sub-storages may use different formats or include an + high-precision representation such as ``IdentityTensorStorage``. + + """ + + _rowwise_storage: Optional[QuantizedTensorStorage] + _columnwise_storage: Optional[QuantizedTensorStorage] + _quantizer: Quantizer + + def __new__( + cls, + *args, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + quantizer: Quantizer, + fake_dtype: Optional[torch.dtype] = None, + **kwargs, + ): + if quantizer is None or not isinstance(quantizer, Quantizer): + raise TypeError( + "HybridQuantizedTensorStorage requires a parent HybridQuantizer; " + f"got {type(quantizer).__name__}." + ) + if not hasattr(quantizer, "rowwise_quantizer") or not hasattr( + quantizer, "columnwise_quantizer" + ): + raise TypeError( + "HybridQuantizedTensorStorage requires a parent HybridQuantizer " + "with rowwise_quantizer and columnwise_quantizer attributes; " + f"got {type(quantizer).__name__}." + ) + + if cls is HybridQuantizedTensorStorage: + instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) + + instance._rowwise_storage = rowwise_storage + instance._columnwise_storage = columnwise_storage + instance._quantizer = quantizer.copy() + return instance + + @property + def rowwise_sub_storage(self) -> Optional[QuantizedTensorStorage]: + """The sub-storage providing rowwise quantized data.""" + return self._rowwise_storage + + @property + def columnwise_sub_storage(self) -> Optional[QuantizedTensorStorage]: + """The sub-storage providing columnwise quantized data.""" + return self._columnwise_storage + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + # A storage object cannot reconstruct a representation that has already + # been dropped. Validate every requested direction before mutating + # either one so mixed drop/enable requests are atomic. + if rowwise_usage and self._rowwise_storage is None: + raise RuntimeError( + "Requested rowwise usage, but HybridQuantizedTensorStorage " + "has no rowwise sub-storage" + ) + if columnwise_usage and self._columnwise_storage is None: + raise RuntimeError( + "Requested columnwise usage, but HybridQuantizedTensorStorage " + "has no columnwise sub-storage" + ) + if rowwise_usage is not None and not rowwise_usage: + self._rowwise_storage = None + if columnwise_usage is not None and not columnwise_usage: + self._columnwise_storage = None + + def clear(self): + """Deallocate both sub-storages' buffers. + + Delegates to each sub-storage's own ``clear()``; no-op when a + sub-storage is ``None`` (columnwise-only or rowwise-only hybrid). + + Used by ``cpu_offload_v1`` after the offloader has taken its own + reference to the extracted buffers, to release the GPU-resident + originals. + """ + if self._rowwise_storage is not None: + self._rowwise_storage.clear() + if self._columnwise_storage is not None: + self._columnwise_storage.clear() + + def get_usages(self) -> Dict[str, bool]: + return { + "rowwise": self._rowwise_storage is not None, + "columnwise": self._columnwise_storage is not None, + } + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], HybridQuantizedTensorStorage]: + tensors = [] + if self._rowwise_storage is not None: + row_tensors, _ = self._rowwise_storage.prepare_for_saving() + tensors.extend(row_tensors) + if self._columnwise_storage is not None: + col_tensors, _ = self._columnwise_storage.prepare_for_saving() + tensors.extend(col_tensors) + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + if self._rowwise_storage is not None: + tensors = self._rowwise_storage.restore_from_saved(tensors) + if self._columnwise_storage is not None: + tensors = self._columnwise_storage.restore_from_saved(tensors) + return tensors + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """Dequantize using the first available sub-storage.""" + if dtype is None: + dtype = self._dtype + # TODO: choose the best native sub-storage for dequantization, preferring # pylint: disable=fixme + # identity/high-precision and FP8-ish formats over FP4 when both exist. + if self._rowwise_storage is not None: + return self._rowwise_storage.dequantize(dtype=dtype) + if self._columnwise_storage is not None: + return self._columnwise_storage.dequantize(dtype=dtype) + raise RuntimeError("HybridQuantizedTensorStorage has no data to dequantize") + + def get_data_tensors(self): + """Return raw data tensors from both available sub-storages.""" + row_tensors = () + col_tensors = () + if self._rowwise_storage is not None: + result = self._rowwise_storage.get_data_tensors() + row_tensors = result if isinstance(result, tuple) else (result,) + if self._columnwise_storage is not None: + result = self._columnwise_storage.get_data_tensors() + col_tensors = result if isinstance(result, tuple) else (result,) + return row_tensors + col_tensors + + def size(self, *args, **kwargs): + """Return the logical size from the first available sub-storage.""" + if self._rowwise_storage is not None: + return self._rowwise_storage.size(*args, **kwargs) + if self._columnwise_storage is not None: + return self._columnwise_storage.size(*args, **kwargs) + raise RuntimeError("HybridQuantizedTensorStorage has no data") + + @property + def device(self): + """Return the device from the first available sub-storage.""" + if self._rowwise_storage is not None: + return self._rowwise_storage.device + if self._columnwise_storage is not None: + return self._columnwise_storage.device + raise RuntimeError("HybridQuantizedTensorStorage has no data") + + def view(self, *shape): + """View delegates to each sub-storage. Used by FSDP2 reset_sharded_param. + + Identity views are handled without forwarding a reshape to the + sub-storages: the columnwise sub-storage's own shape is transposed + relative to the hybrid for some formats (e.g. a 2D block-scaled + Float8BlockwiseQTensor has shape ``(N, M)`` for an ``(M, N)`` weight), + so forwarding the hybrid's row-major shape would be a spurious + last-2-dims change that dequantizes it to a plain tensor. + """ + flat_shape = shape[0] if len(shape) == 1 and isinstance(shape[0], Iterable) else shape + if list(flat_shape) == list(self.size()): + return HybridQuantizedTensorStorage( + rowwise_storage=self._rowwise_storage, + columnwise_storage=self._columnwise_storage, + quantizer=self._quantizer, + fake_dtype=self._dtype, + ) + row_view = self._rowwise_storage.view(*shape) if self._rowwise_storage is not None else None + col_view = ( + self._columnwise_storage.view(*shape) if self._columnwise_storage is not None else None + ) + return HybridQuantizedTensorStorage( + rowwise_storage=row_view, + columnwise_storage=col_view, + quantizer=self._quantizer, + fake_dtype=self._dtype, + ) + + def get_metadata(self) -> Dict[str, Any]: + """Return constructor metadata for make_like and serialization paths.""" + return { + "rowwise_storage": self._rowwise_storage, + "columnwise_storage": self._columnwise_storage, + "quantizer": self._quantizer, + "fake_dtype": self._dtype, + } + + def __repr__(self): + row_type = ( + type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + ) + col_type = ( + type(self._columnwise_storage).__name__ + if self._columnwise_storage is not None + else "None" + ) + return ( + "HybridQuantizedTensorStorage(" + f"rowwise={row_type}, " + f"columnwise={col_type}, " + f"dtype={self._dtype})" + ) diff --git a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py new file mode 100644 index 0000000000..c1dea51097 --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py @@ -0,0 +1,152 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Storage class for a high-precision IdentityTensor representation.""" + +from __future__ import annotations +from typing import Any, Dict, Optional, Tuple + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...utils import _empty_tensor + + +class IdentityTensorStorage(QuantizedTensorStorage): + """Passthrough storage that holds a high-precision tensor directly. + + Produced by :class:`IdentityQuantizer`. It implements the + ``QuantizedTensorStorage`` interface so it can flow through the same + module / GEMM / save-for-backward / FSDP machinery as the real quantized + storages, but it uses no low-precision encoding. ``general_gemm`` + materializes it as a plain tensor in the held dtype. + + The data is direction-agnostic -- the same tensor serves both the rowwise + and columnwise directions (the GEMM transposes via its layout flags), so a + single buffer is stored. This is what lets a ``HybridQuantizer`` mix one + quantized direction with one high-precision direction. + """ + + _hp_data: Optional[torch.Tensor] + + def __new__( + cls, + *args, + hp_data: Optional[torch.Tensor], + fake_dtype: Optional[torch.dtype] = None, + quantizer: Optional[Quantizer] = None, + **kwargs, + ): + if cls is IdentityTensorStorage: + instance = object.__new__(cls) + if fake_dtype is not None: + instance._dtype = fake_dtype + elif hp_data is not None: + instance._dtype = hp_data.dtype + else: + instance._dtype = torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) + instance._hp_data = hp_data + instance._quantizer = quantizer.copy() if quantizer is not None else None + return instance + + def clear(self): + """Deallocate the held tensor's memory.""" + if self._hp_data is not None: + self._hp_data.data = _empty_tensor() + + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data from another IdentityTensorStorage.""" + if not isinstance(src, IdentityTensorStorage): + raise TypeError("copy_from_storage expects IdentityTensorStorage") + if self._hp_data is not None and src._hp_data is not None: + self._hp_data.copy_(src._hp_data) + + def get_metadata(self) -> Dict[str, Any]: + """Get this tensor's metadata.""" + return { + "hp_data": self._hp_data, + "quantizer": self._quantizer, + "fake_dtype": self._dtype, + } + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], "IdentityTensorStorage"]: + """Prepare the tensor base for saving for backward.""" + tensors = [self._hp_data] + self._hp_data = None + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + """Restore the held tensor from the saved tensors list.""" + self._hp_data = tensors[0] + return tensors[1:] + + def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = True): + """Get this tensor's data. The single HP buffer serves both directions.""" + if rowwise_data and columnwise_data: + return self._hp_data, None + if rowwise_data: + return self._hp_data + if columnwise_data: + return self._hp_data + raise ValueError("No data to get, both rowwise_data and columnwise_data are False") + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """Return the held high-precision tensor, casting when requested.""" + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data to dequantize") + if dtype is None: + dtype = self._dtype + if self._hp_data.dtype != dtype: + return self._hp_data.to(dtype) + return self._hp_data + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + """No-op: the single high-precision buffer serves both directions.""" + # High-precision data is not direction-specific, so there is nothing + # to drop or synthesize. Honor the request only insofar as keeping the + # buffer (a request to drop both would leave no data, which is invalid). + + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the tensor.""" + has_data = self._hp_data is not None + return {"rowwise": has_data, "columnwise": has_data} + + def size(self, *args, **kwargs): + # pylint: disable=missing-function-docstring + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + return self._hp_data.size(*args, **kwargs) + + @property + def device(self): + """Return the device of the held tensor.""" + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data!") + return self._hp_data.device + + def view(self, *shape): + # pylint: disable=missing-function-docstring + flat_shape = shape[0] if len(shape) == 1 and not isinstance(shape[0], int) else shape + return IdentityTensorStorage( + hp_data=self._hp_data.view(*flat_shape) if self._hp_data is not None else None, + fake_dtype=self._dtype, + quantizer=self._quantizer, + ) + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Field gathered by FSDP2 for the high-precision passthrough.""" + return ("_hp_data",) + + def __repr__(self): + return f"IdentityTensorStorage(dtype={self._dtype}, data={self._hp_data})" diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 606ac9e74b..da5d6c4e6b 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -15,7 +15,11 @@ from ...quantized_tensor import QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr -from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType +from ...constants import ( + TE_DType as torch_to_transformer_engine_dtype, + MXFP8_BLOCK_SCALING_SIZE, + DType, +) from ...utils import _empty_tensor @@ -322,3 +326,77 @@ def get_usages(self) -> Dict[str, bool]: "rowwise": self._rowwise_data is not None, "columnwise": self._columnwise_data is not None, } + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for MXFP8. + + Block scales are per-block and direction-specific — each direction + gathers both its data buffer and its scale-inv buffer. ``None``-valued + directions (e.g. a columnwise-only sub-storage in hybrid quantization) + are excluded so the gather tuple only contains real tensors. + """ + fields = [] + if self._rowwise_data is not None: + fields.extend(("_rowwise_data", "_rowwise_scale_inv")) + if self._columnwise_data is not None: + fields.extend(("_columnwise_data", "_columnwise_scale_inv")) + return tuple(fields) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract MXFP8 buffers, unpadding block-scale alignment before gather. + + MXFP8 kernels require scale-inv tensors aligned to ``[128, 4]`` + (rowwise) and ``[4, 128]`` (columnwise). That padding is attached to + the local shard but would produce misaligned concatenation under + FSDP2's dim-0 all-gather. Strip it here and re-apply in + :meth:`fsdp_assign_gathered`. + """ + if self._with_gemm_swizzled_scales: + raise NotImplementedError( + "FSDP2 is only supported for MXFP8Tensors with compact scales" + ) + names = self.fsdp_buffer_fields() + buffers = [] + shape = self.size() + flattened_in_shape0 = math.prod(shape[:-1]) + for name in names: + t = getattr(self, name) + if name == "_rowwise_scale_inv" and t is not None: + if t.size(0) != flattened_in_shape0: + t = t[:flattened_in_shape0] + elif name == "_columnwise_scale_inv" and t is not None: + expected = math.ceil(flattened_in_shape0 / MXFP8_BLOCK_SCALING_SIZE) + if t.size(0) != expected: + t = t[:expected] + buffers.append(t) + return tuple(buffers), {"field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered MXFP8 buffers back, re-padding block scales. + + Inverse of :meth:`fsdp_extract_buffers`: the gathered scale-inv tensors + are padded back up to ``[128, 4]`` / ``[4, 128]`` alignment before + being assigned to the storage. + """ + names = meta["field_names"] + if len(names) != len(gathered): + raise RuntimeError( + "MXFP8TensorStorage.fsdp_assign_gathered got " + f"{len(gathered)} buffers for {len(names)} fields" + ) + for name, buf in zip(names, gathered): + if buf is not None and name == "_rowwise_scale_inv": + pad = (128 - buf.size(0) % 128) % 128 + if pad > 0: + buf = torch.nn.functional.pad(buf, (0, 0, 0, pad)) + elif buf is not None and name == "_columnwise_scale_inv": + pad = (4 - buf.size(0) % 4) % 4 + if pad > 0: + buf = torch.nn.functional.pad(buf, (0, 0, 0, pad)) + setattr(self, name, buf) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 1789a0c98e..e35d57b363 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -19,6 +19,9 @@ from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer +from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer +from .identity_tensor import IdentityQuantizer +from .storage.identity_tensor_storage import IdentityTensorStorage from ..optimizers.multi_tensor_apply import multi_tensor_applier from ..utils import is_non_tn_fp8_gemm_supported from ..constants import NVFP4_BLOCK_SCALING_SIZE, DType @@ -65,12 +68,139 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): new_raw_data.detach().copy_(old_rowwise) tensor._rowwise_data = new_raw_data del old_rowwise + elif isinstance(tensor, IdentityTensorStorage): + old_raw_data = tensor._hp_data + if old_raw_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + if old_raw_data.dtype != new_raw_data.dtype: + raise ValueError( + "The data types of raw data don't match: " + f"old dtype={old_raw_data.dtype}, new dtype={new_raw_data.dtype}" + ) + new_raw_data.detach().copy_(old_raw_data) + tensor._hp_data = new_raw_data + del old_raw_data elif isinstance(tensor, MXFP8Tensor): raise NotImplementedError("replace_raw_data for MXFP8Tensor is not supported yet") + elif isinstance(tensor, HybridQuantizedTensor): + # The distopt all-gather buffer routes at the rowwise sub-storage only; + # the columnwise sub-storage is refreshed each iteration via + # ``HybridQuantizer.update_quantized``. The underlying call delegates + # to the rowwise sub-storage's own ``replace_raw_data`` (which may + # raise for sub-storage types that don't implement it). + if tensor._rowwise_storage is None: + raise NotImplementedError( + "replace_raw_data for HybridQuantizedTensor without a rowwise " + "sub-storage is not supported." + ) + replace_raw_data(tensor._rowwise_storage, new_raw_data) else: raise ValueError(f"replace_raw_data for {type(tensor)} is not supported yet") +def _is_float8_transpose_only(tensor: QuantizedTensor) -> bool: + """Whether a Float8 tensor stores its live payload only in _transpose.""" + return ( + isinstance(tensor, Float8Tensor) + and tensor._data is None + and tensor._transpose is not None + and not tensor._transpose_invalid + ) + + +def _validate_flat_fragment( + model_weight: QuantizedTensor, master_weight: torch.Tensor, start_offset +): + """Validate a flat logical shard and return its exclusive end offset.""" + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") + end_offset = start_offset + master_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) + return end_offset + + +def _cast_master_weight_to_rowwise_fp8_bytes( + master_weight: torch.Tensor, + model_weight: Float8Tensor, + quantizer: Float8Quantizer, +) -> torch.Tensor: + """Cast a flat master shard to row-major FP8 bytes using ``quantizer`` scale state.""" + rowwise_quantizer = Float8Quantizer( + scale=quantizer.scale, + amax=quantizer.amax, + fp8_dtype=quantizer.dtype, + rowwise=True, + columnwise=False, + ) + raw = torch.empty((1, master_weight.numel()), dtype=torch.uint8, device=model_weight.device) + temp = rowwise_quantizer.create_tensor_from_data(raw, model_weight.dtype) + rowwise_quantizer.update_quantized(master_weight.reshape(1, -1), temp) + if temp._data is None: + raise RuntimeError("Expected rowwise Float8 temporary to populate _data") + return temp._data.reshape(-1) + + +def _update_transpose_only_float8_flat_fragment( + model_weight: QuantizedTensor, + master_weight: torch.Tensor, + start_offset, + quantizer: Float8Quantizer, +) -> bool: + """Update a logical flat shard in a transpose-only Float8 tensor. + + Hopper / L40 columnwise-only Float8 sub-storages keep their live FP8 + bytes in ``_transpose`` with physical shape ``[K, rows]`` for a logical + ``[rows, K]`` tensor. A row-major logical shard is not contiguous in that + storage, so cast the shard once and scatter the resulting FP8 bytes by + logical row into the transpose buffer. + """ + if not _is_float8_transpose_only(model_weight): + return False + + _validate_flat_fragment(model_weight, master_weight, start_offset) + numel = master_weight.numel() + if numel == 0: + return True + + shape = tuple(model_weight.shape) + if len(shape) == 0: + raise ValueError("Float8 scalar transpose-only flat update is not supported") + logical_cols = int(shape[-1]) + if logical_cols <= 0 or model_weight.numel() % logical_cols != 0: + raise ValueError(f"Invalid Float8 logical shape for transpose-only update: {shape}") + logical_rows = model_weight.numel() // logical_cols + + transpose = model_weight._transpose + if transpose.numel() != model_weight.numel(): + raise ValueError( + "Float8 transpose-only storage has unexpected numel: " + f"transpose={transpose.numel()}, logical={model_weight.numel()}" + ) + transpose_2d = transpose.reshape(logical_cols, logical_rows) + fp8_bytes = _cast_master_weight_to_rowwise_fp8_bytes(master_weight, model_weight, quantizer) + + remaining = numel + logical_offset = start_offset + src_offset = 0 + while remaining > 0: + row = logical_offset // logical_cols + col = logical_offset % logical_cols + n = min(remaining, logical_cols - col) + transpose_2d[col : col + n, row].copy_(fp8_bytes[src_offset : src_offset + n]) + logical_offset += n + src_offset += n + remaining -= n + + model_weight._transpose_invalid = False + return True + + def quantize_master_weights( model_weights, master_weights, @@ -111,12 +241,29 @@ def quantize_master_weights( blockwise_scaling_params = [] mxfp8_scaling_params = [] nvfp4_params = [] + identity_params = [] if fsdp_shard_model_weights is None: use_fsdp_shard_model_weights = False fsdp_shard_model_weights = [None] * len(model_weights) else: use_fsdp_shard_model_weights = True + # Validate the entire batch before clearing initialization state or + # populating any per-format work buckets. + for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( + model_weights, master_weights, start_offsets, fsdp_shard_model_weights + ): + _validate_per_tensor_fp8_fsdp_hopper_policy( + model_weight, + fsdp_shard_model_weight, + ) + if isinstance(model_weight, HybridQuantizedTensor): + _validate_hybrid_partial_master_policy( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + ) for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights @@ -165,6 +312,20 @@ def quantize_master_weights( mxfp8_scaling_params.append( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) + elif isinstance(quantizer, IdentityQuantizer): + identity_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, HybridQuantizer): + _route_hybrid_to_buckets( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + delayed_scaling_params=delayed_scaling_params, + current_scaling_params=current_scaling_params, + identity_params=identity_params, + ) else: raise ValueError(f"quantize_master_weights for {type(quantizer)} is not supported yet") @@ -179,6 +340,8 @@ def quantize_master_weights( _cast_master_weights_to_fp8_mxfp8_scaling(mxfp8_scaling_params, *extra_args) if len(nvfp4_params) > 0: _cast_master_weights_to_nvfp4_2d(nvfp4_params, *extra_args) + if len(identity_params) > 0: + _cast_master_weights_to_identity(identity_params, *extra_args) def cast_master_weights_to_fp8( @@ -257,6 +420,10 @@ def _cast_master_weights_to_fp8_delayed_scaling( # master_weight may be smaller than model_weight because it could be distributed across # multiple ranks. So we need to create a dummy weight using the raw data from model_weight. if not use_fsdp_shard_model_weights: + if _update_transpose_only_float8_flat_fragment( + model_weight, master_weight, start_offset, quantizer + ): + continue shard_model_weight_raw = model_weight._data.view(-1)[start_offset:end_offset] shard_model_weight_fp8 = quantizer.create_tensor_from_data( shard_model_weight_raw.view(1, -1), @@ -402,13 +569,17 @@ def _cast_master_weights_to_fp8_current_scaling( # Cast master weight to FP8 end_offset = start_offset + master_weight.numel() - if not use_fsdp_shard_model_weights: - model_weight_fragment = model_weight.reshape(-1)[start_offset:end_offset] quantizer = Float8Quantizer( scale=scale, amax=torch.Tensor(), fp8_dtype=model_weight._fp8_dtype, ) + if not use_fsdp_shard_model_weights: + if _update_transpose_only_float8_flat_fragment( + model_weight, master_weight, start_offset, quantizer + ): + continue + model_weight_fragment = model_weight.reshape(-1)[start_offset:end_offset] if use_fsdp_shard_model_weights and not isinstance(model_weight_fragment, Float8Tensor): # NOTE: The fsdp shard model weight may be a unit8 tensor instead of # a float8 tensor. We should handle this situation properly. @@ -794,6 +965,53 @@ def _cast_master_weights_to_nvfp4_2d( ) +def _identity_storage_data(tensor): + if not isinstance(tensor, IdentityTensorStorage): + raise TypeError(f"Expected IdentityTensorStorage, got {type(tensor).__name__}") + if tensor._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + return tensor._hp_data + + +def _cast_master_weights_to_identity( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): + del group, manual_post_all_gather_processing + + for model_weight, master_weight, start_offset, model_weight_fragment in params: + if master_weight is None: + continue + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") + end_offset = start_offset + master_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) + + if use_fsdp_shard_model_weights: + target = model_weight_fragment + if target is None: + raise RuntimeError("FSDP shard model weight is required for Identity writeback") + if isinstance(target, IdentityTensorStorage): + target_flat = _identity_storage_data(target).reshape(-1) + else: + target_flat = target.reshape(-1) + target_slice = target_flat[: master_weight.numel()] + else: + target_slice = _identity_storage_data(model_weight).reshape(-1)[start_offset:end_offset] + + if target_slice.numel() != master_weight.numel(): + raise ValueError( + f"Identity target slice has {target_slice.numel()} elements, " + f"but master_weight has {master_weight.numel()}" + ) + target_slice.copy_(master_weight.reshape(-1)) + + def _cast_master_weights_to_fp8_mxfp8_scaling( params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False ): # pylint: disable=unused-argument @@ -933,6 +1151,252 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( ) +# --------------------------------------------------------------------------------------------- +# HybridQuantizer helpers for `quantize_master_weights` / `post_all_gather_processing`. +# +# Dispatch is per-direction: `_route_hybrid_to_buckets` iterates over both sub-storages +# of a `HybridQuantizedTensor` and routes each one independently into the per-format +# bucket matching its own sub-quantizer type. Row and col make their own decisions and +# can mix any pair of currently-supported sub-quantizers. +# +# Supported (per-tensor Float8 or Identity sub-quantizers, any direction): +# - Float8Quantizer (delayed scaling) +# - Float8CurrentScalingQuantizer (current scaling) +# - IdentityQuantizer (high-precision passthrough) +# +# Per-tensor Float8 works because `_cast_master_weights_to_fp8_{delayed,current}_scaling` +# accept any Float8Tensor (single direction is fine — each entry is one Float8Tensor +# with its own `_scale_inv` and the helper writes that one entry's `_data`). Each +# hybrid sub-storage IS a single-direction Float8Tensor, so we route them as two +# independent entries (into the same bucket for same-format, or into different +# buckets for cross-format Float8 — e.g. delayed row + current col). +# The FSDP-sharded columnwise-only representation on Hopper is the exception: +# neither per-tensor cast helper can safely flatten its transpose-only payload, +# so `_validate_per_tensor_fp8_fsdp_hopper_policy` rejects it before mutation. +# +# Identity routes to an exact copy bucket. Single-direction hybrid (only one +# sub-storage populated) routes the present direction only. Both-None hybrids +# raise ValueError. Per-block sub-quantizers still hit their per-direction TODO. +# +# Not supported (raise NotImplementedError per-direction + TODO): +# +# - MXFP8Quantizer as a hybrid sub-quantizer (any direction) +# TODO(#3158, hybrid-mxfp8-distopt): the distopt cast kernels +# (`tex.mxfp8_scaling_compute_partial_amax`, `tex.mxfp8_scaling_partial_cast`) +# are bidirectional — both rowwise and colwise outputs required — so they +# cannot ingest a single-direction hybrid sub-storage. (Unrelated to the +# regular `tex.quantize` kernel used by forward/backward, which natively +# supports single-direction output.) Unblocker: add single-direction +# variants of the two distopt kernels, then route hybrid sub-storages +# per-direction into `mxfp8_scaling_params` matching the Float8 path above. +# Also unlocks cross-format MXFP8 row + col. +# +# - NVFP4Quantizer as a hybrid sub-quantizer (any direction) +# TODO(#3158, hybrid-nvfp4-distopt): load-bearing blocker is the kernel assertion +# `return_identity || !use_2d_quantization` in +# `quantize_transpose_vector_blockwise_fp4.cu`, which rejects exactly the +# columnwise-only 2D configuration that `HybridQuantizer.__init__` produces +# for the col sub-quantizer. Blocks hybrid 2D NVFP4 weight construction at +# `quantized_model_init` time. 1D NVFP4 is unaffected. The assertion is an +# explicitly-marked unwritten code path, not an algorithmic limit (see the +# kernel author's note above the early-return guard). +# +# Secondary blocker (gated on the kernel fix): the distopt helper +# `_cast_master_weights_to_nvfp4_2d` writes only `_rowwise_data` and relies +# on per-tensor post-AG `_create_columnwise()` — for hybrid, the columnwise +# data needs to land in a SEPARATE col sub-storage, so the post-AG branch +# must be made hybrid-aware (derive `col_sub._columnwise_data` from +# `row_sub`'s gathered rowwise). +# +# - Float8BlockQuantizer as a hybrid sub-quantizer +# TODO(#3158, hybrid-fp8-blockwise): same shape as the NVFP4 secondary blocker — +# `_cast_master_weights_to_fp8_blockwise_scaling` writes only `_rowwise_data` +# with per-tensor post-AG `_create_columnwise()` that doesn't reach hybrid's +# separate col sub-storage. Unlike NVFP4, there is no kernel-level +# construction blocker (the Block FP8 kernel natively supports +# columnwise-only mode), so hybrid Block FP8 weights construct fine via the +# non-distopt FusedAdam path today; only the sharded-master distopt cast +# path is blocked. Unblocker is a Python-side hybrid-aware post-AG branch; +# no C++ work needed. +# +# --------------------------------------------------------------------------------------------- + + +def _validate_per_tensor_fp8_fsdp_hopper_policy( + model_weight, + fsdp_shard_model_weight, +): + """Reject unsupported transpose-only per-tensor FP8 FSDP updates on Hopper. + + The FSDP shard path can receive a ``Float8Tensor`` and pass it directly to + a per-tensor FP8 cast helper. On Hopper, both delayed and current scaling + use a transpose-only representation for columnwise-only tensors that the + sharded update path cannot flatten safely. Reject the whole batch before + any scale, cache, or storage mutation instead. + """ + if fsdp_shard_model_weight is None or is_non_tn_fp8_gemm_supported(): + return + + if isinstance(model_weight, HybridQuantizedTensor): + quantizer = model_weight._get_quantizer() + candidates = ( + (model_weight._rowwise_storage, quantizer.rowwise_quantizer), + (model_weight._columnwise_storage, quantizer.columnwise_quantizer), + ) + else: + candidates = ((model_weight, model_weight._get_quantizer()),) + + for storage, quantizer in candidates: + if ( + storage is not None + and isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)) + and _is_float8_transpose_only(storage) + ): + raise NotImplementedError( + "Columnwise-only per-tensor FP8 quantization is not implemented for " + "FSDP updates on this architecture." + ) + + +def _validate_hybrid_partial_master_policy( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, +): + """Reject unsupported Hybrid column-source updates before distopt mutation.""" + row_sub = model_weight._rowwise_storage + col_sub = model_weight._columnwise_storage + quantizer = model_weight._get_quantizer() + + if col_sub is not None and quantizer.columnwise_source == "rowwise_dequantized": + raise NotImplementedError( + "quantize_master_weights does not support HybridQuantizer with a live " + "columnwise representation and " + "columnwise_source='rowwise_dequantized'. The column must be derived " + "after the rowwise update/all-gather. See #3158." + ) + + if master_weight is None or row_sub is None or col_sub is None: + return + + shard_has_both_directions = ( + isinstance(fsdp_shard_model_weight, HybridQuantizedTensor) + and fsdp_shard_model_weight._rowwise_storage is not None + and fsdp_shard_model_weight._columnwise_storage is not None + ) + if shard_has_both_directions: + return + + end_offset = _validate_flat_fragment(model_weight, master_weight, start_offset) + if start_offset == 0 and end_offset == model_weight.numel(): + return + + raise ValueError( + "quantize_master_weights cannot update a two-direction " + "HybridQuantizedTensor from a partial master shard when " + "columnwise_source='original': the one-payload distributed-optimizer " + "path cannot preserve an independently quantized columnwise value. " + "Provide full-master data or retain only the rowwise representation." + ) + + +def _route_hybrid_to_buckets( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + *, + delayed_scaling_params, + current_scaling_params, + identity_params, +): + """Decompose a `HybridQuantizedTensor` into per-direction entries and route each + into the appropriate per-format bucket used by `quantize_master_weights`. + + Per-direction dispatch: each sub-storage routes independently based on its + own sub-quantizer type. Per-tensor Float8 sub-quantizers (delayed and/or + current scaling) are supported in any combination per direction; single- + direction hybrid (one sub-storage dropped via ``update_usage``) is also + supported. See the TODO block above this helper for the per-block-format + rejection rationale and unblocker shapes. + """ + row_sub = model_weight._rowwise_storage + col_sub = model_weight._columnwise_storage + sub_q_row = model_weight._quantizer.rowwise_quantizer + sub_q_col = model_weight._quantizer.columnwise_quantizer + + if row_sub is None and col_sub is None: + raise ValueError( + "quantize_master_weights called on HybridQuantizedTensor with both " + "rowwise and columnwise sub-storages dropped (via update_usage). " + "Nothing to cast — this is most likely a caller bug." + ) + + # Per-direction routing: each (sub_storage, sub_quantizer) pair selects its + # own bucket based on the sub-quantizer's type. Directions that have been + # dropped via ``update_usage`` are silently skipped. + for direction, sub_storage, sub_q in ( + ("rowwise", row_sub, sub_q_row), + ("columnwise", col_sub, sub_q_col), + ): + if sub_storage is None: + continue + shard_fragment = fsdp_shard_model_weight + if shard_fragment is not None and isinstance(shard_fragment, HybridQuantizedTensor): + shard_fragment = ( + shard_fragment._rowwise_storage + if direction == "rowwise" + else shard_fragment._columnwise_storage + ) + entry = (sub_storage, master_weight, start_offset, shard_fragment) + if isinstance(sub_q, Float8Quantizer): + # Delayed scaling: the per-format helper iterates entries + # independently and does a per-DP amax all-reduce across the bucket. + delayed_scaling_params.append(entry) + elif isinstance(sub_q, Float8CurrentScalingQuantizer): + current_scaling_params.append(entry) + elif isinstance(sub_q, IdentityQuantizer): + identity_params.append(entry) + elif isinstance(sub_q, MXFP8Quantizer): + # TODO(#3158, hybrid-mxfp8-distopt): the distopt cast kernels are + # bidirectional, so a single-direction hybrid sub-storage cannot be + # fed in. See top-of-file TODO block for the unblocker (single- + # direction variants of the two distopt kernels). + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with MXFP8Quantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for the unblocker shape." + ) + elif isinstance(sub_q, NVFP4Quantizer): + # TODO(#3158, hybrid-nvfp4-distopt): load-bearing blocker is the kernel + # assertion that rejects columnwise-only 2D NVFP4 — which is + # exactly what hybrid's col sub-quantizer pin produces. Secondary + # blocker (gated on the kernel fix) is the per-tensor post-AG + # `_create_columnwise()` not reaching hybrid's separate col + # sub-storage. See top-of-file TODO block for details. + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with NVFP4Quantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for details." + ) + elif isinstance(sub_q, Float8BlockQuantizer): + # Pending hybrid-fp8-blockwise work (#3158): same shape as the NVFP4 + # secondary blocker (and only that one — no kernel-level construction + # blocker for Block FP8). Python-side post-AG fix. See the + # _route_hybrid_to_buckets design note above for details. + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with Float8BlockQuantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for details." + ) + else: + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with " + f"{type(sub_q).__name__} {direction} sub-quantizer is not supported yet." + ) + + def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Tensor]]): """ Post-processing after all-gather for weights in distributed optimizer. @@ -941,6 +1405,9 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten - Plain pytorch tensor: noop. For NVFP4 tensors, uses batched multi-tensor processing to reduce CPU overhead. + + For `HybridQuantizedTensor`, recurses per-direction so each present + sub-storage runs its native post-processing. Identity sub-storages are no-op. """ if not isinstance(model_weights, list): model_weights = [model_weights] @@ -963,6 +1430,12 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten elif isinstance(model_weight, MXFP8Tensor): # MXFP8 scaling: no need to do anything. pass + elif isinstance(model_weight, IdentityTensorStorage): + pass + elif isinstance(model_weight, HybridQuantizedTensor): + for sub in (model_weight._rowwise_storage, model_weight._columnwise_storage): + if sub is not None: + post_all_gather_processing(sub) elif isinstance(model_weight, QuantizedTensor): raise ValueError(f"post_processing for {type(model_weight)} is not supported")