diff --git a/tests/pytorch/test_weight_swizzle_in_layers.py b/tests/pytorch/test_weight_swizzle_in_layers.py index b9f19f2fd9..ff961c5083 100644 --- a/tests/pytorch/test_weight_swizzle_in_layers.py +++ b/tests/pytorch/test_weight_swizzle_in_layers.py @@ -21,19 +21,25 @@ # skipped individually when the hardware/recipe is unavailable. _SWIZZLING_RECIPES = [ pytest.param( - MXFP8BlockScaling, + "mxfp8", marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), - id="mxfp8", ), pytest.param( - NVFP4BlockScaling, + "nvfp4-2d", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + ), + pytest.param( + "nvfp4-1d", marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), - id="nvfp4", ), ] _LAYER_TYPES = ["Linear", "LayerNormLinear", "LayerNormMLP", "GroupedLinear"] +# 1024 is 128-aligned (NVFP4 swizzle fusion eligible). 1056 is valid MXFP8/NVFP4 +# (divisible by 32) but not 128-aligned, so NVFP4 keeps compact cached scales. +_WEIGHT_SHAPES = [1024, 1056] + def _make_module(layer_type, in_features, out_features, device, num_gemms=1): common = dict(bias=True, params_dtype=torch.bfloat16) @@ -65,12 +71,17 @@ def _grouped_m_splits(layer_type, batch, num_gemms): return [batch // num_gemms] * num_gemms -def _make_recipe(recipe_cls): +def _make_recipe(recipe_name): """Instantiate a recipe with run-to-run nondeterminism disabled where it exists (NVFP4 stochastic rounding); MXFP8 has none.""" - if recipe_cls is NVFP4BlockScaling: - return recipe_cls(disable_stochastic_rounding=True) - return recipe_cls() + if recipe_name == "mxfp8": + return MXFP8BlockScaling() + if recipe_name in ("nvfp4-2d", "nvfp4-1d"): + return NVFP4BlockScaling( + disable_stochastic_rounding=True, + disable_2d_quantization=recipe_name == "nvfp4-1d", + ) + raise ValueError(f"unknown recipe {recipe_name}") def _forward_backward(module, x, is_first_microbatch, recipe, m_splits): @@ -110,21 +121,16 @@ def _run_step(module, x, is_first_microbatch, recipe, m_splits): @pytest.mark.parametrize("layer_type", _LAYER_TYPES) -@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES) -def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls): - """Caching the quantized weight across microbatches must pre-swizzle its - scales: the weight quantizer enables ``optimize_for_gemm`` and the cached - workspace carries ``_with_gemm_swizzled_scales``. - - Generic across every swizzling recipe (MXFP8, NVFP4) and module type. - """ +@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES) +def test_weight_swizzling_with_workspace_caching(layer_type, recipe_name): + """Cached weights use preswizzled scales only when direct fusion is supported.""" torch.manual_seed(1234) device = "cuda" in_features, out_features = 1024, 1024 batch = 512 num_gemms = 2 if layer_type == "GroupedLinear" else 1 m_splits = _grouped_m_splits(layer_type, batch, num_gemms) - recipe = recipe_cls() + recipe = _make_recipe(recipe_name) module = _make_module(layer_type, in_features, out_features, device, num_gemms) x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True) @@ -132,11 +138,12 @@ def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls): # is_first_microbatch=True caches the quantized weight. weight_quantizers = _forward_backward(module, x, True, recipe, m_splits) + expect_preswizzle = recipe_name != "nvfp4-1d" for weight_quantizer in weight_quantizers: assert weight_quantizer is not None assert ( - weight_quantizer.optimize_for_gemm is True - ), f"optimize_for_gemm must be enabled for cached {layer_type} weights" + weight_quantizer.optimize_for_gemm is expect_preswizzle + ), f"unexpected optimize_for_gemm for cached {layer_type} weights" workspaces = module._fp8_workspaces assert len(workspaces) == _expected_num_weights( @@ -144,13 +151,13 @@ def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls): ), f"unexpected cached weight workspace count for {layer_type}: {len(workspaces)}" for name, ws in workspaces.items(): assert ( - getattr(ws, "_with_gemm_swizzled_scales", False) is True - ), f"cached weight workspace {name!r} scales were not pre-swizzled" + getattr(ws, "_with_gemm_swizzled_scales", False) is expect_preswizzle + ), f"unexpected scale layout for cached weight workspace {name!r}" @pytest.mark.parametrize("layer_type", _LAYER_TYPES) -@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES) -def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls): +@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES) +def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_name): """With quantized_model_init the weight parameter is itself quantized and is all-gathered (FSDP2) / optimizer-updated in its unswizzled layout, so the eager-swizzle optimization must stay off: the weight quantizer must keep @@ -164,7 +171,7 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls): batch = 512 num_gemms = 2 if layer_type == "GroupedLinear" else 1 m_splits = _grouped_m_splits(layer_type, batch, num_gemms) - recipe = recipe_cls() + recipe = _make_recipe(recipe_name) with te.quantized_model_init(enabled=True, recipe=recipe): module = _make_module(layer_type, in_features, out_features, device, num_gemms) @@ -191,8 +198,35 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls): @pytest.mark.parametrize("layer_type", _LAYER_TYPES) -@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES) -def test_weight_caching_matches_no_caching(layer_type, recipe_cls): +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +def test_weight_optimize_for_gemm_disabled_without_swizzle_fusion(layer_type): + """NVFP4 weights that cannot use in-kernel swizzle fusion must keep compact cached scales.""" + torch.manual_seed(1234) + device = "cuda" + # Valid for NVFP4 quantization, but not aligned to 128x128 swizzle tiles. + in_features, out_features = 1056, 1056 + batch = 512 + num_gemms = 2 if layer_type == "GroupedLinear" else 1 + m_splits = _grouped_m_splits(layer_type, batch, num_gemms) + recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + + module = _make_module(layer_type, in_features, out_features, device, num_gemms) + x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True) + + weight_quantizers = _forward_backward(module, x, True, recipe, m_splits) + + for weight_quantizer in weight_quantizers: + assert weight_quantizer is not None + assert weight_quantizer.optimize_for_gemm is False + + for _, ws in module._fp8_workspaces.items(): + assert getattr(ws, "_with_gemm_swizzled_scales", False) is False + + +@pytest.mark.parametrize("layer_type", _LAYER_TYPES) +@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES) +@pytest.mark.parametrize("features", _WEIGHT_SHAPES) +def test_weight_caching_matches_no_caching(layer_type, recipe_name, features): """Caching the quantized (pre-swizzled) weight across microbatches must be numerically identical to the uncached flow that re-quantizes the weight every microbatch. Verified per microbatch for fprop output, dgrad and wgrad, across @@ -205,12 +239,12 @@ def test_weight_caching_matches_no_caching(layer_type, recipe_cls): """ torch.manual_seed(1234) device = "cuda" - in_features, out_features = 1024, 1024 + in_features, out_features = features, features batch = 512 microbatches = 4 num_gemms = 2 if layer_type == "GroupedLinear" else 1 m_splits = _grouped_m_splits(layer_type, batch, num_gemms) - recipe = _make_recipe(recipe_cls) + recipe = _make_recipe(recipe_name) # Identical modules: one drives the cached path, one the uncached path. cached = _make_module(layer_type, in_features, out_features, device, num_gemms) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5eab45b7ac..7c11e635c6 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -58,6 +58,7 @@ from ..utils import ( is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, + get_device_compute_capability, get_nvtx_range_context, nvtx_range_push, nvtx_range_pop, @@ -1207,6 +1208,40 @@ def _get_weight_quantizers(self) -> List[Quantizer]: f"{self.__class__.__name__} class does not implement _get_weight_quantizers function" ) + def _enable_weight_preswizzle( + self, + quantizer: Quantizer, + weight: torch.Tensor, + ) -> bool: + """Whether to fuse scale-factor swizzling into weight quantization. + + When enabled, scales are preswizzled during quantization instead of lazily + inside every GEMM. Disabled when primary weights are already quantized + (dequant and FSDP2 all-gather expect the unswizzled layout). For NVFP4, + enabled only for shapes/architectures where the fused swizzle+quantize + kernel is supported. Weight quantization always uses the single-tensor + kernel (including GroupedLinear's cached weights), so NVFP4 RHT + eligibility uses that kernel's 64-row alignment. + """ + if self.primary_weights_in_fp8: + return False + if isinstance(quantizer, MXFP8Quantizer): + return True + if isinstance(quantizer, NVFP4Quantizer): + rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1] + arch_supported = get_device_compute_capability() >= (10, 0) + if quantizer.with_rht: + return arch_supported and rows % 64 == 0 and cols % 128 == 0 + return ( + arch_supported + and quantizer.with_2d_quantization + and not quantizer.row_scaled_nvfp4 + and not quantizer.nvfp4_use_4over6 + and rows % 128 == 0 + and cols % 128 == 0 + ) + return False + def init_fp8_meta_tensors(self, recipe: Recipe) -> None: """Init scales and amaxes.""" self.set_meta_tensor(True, recipe) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 32963c07ca..143fee7ba2 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1900,6 +1900,13 @@ def forward( grad_weight_quantizers, grad_output_quantizers, ) = quantizers + if not debug and weight_quantizers[0] is not None: + # Experts share shape and recipe settings: compute once and broadcast. + optimize_for_gemm = self._enable_weight_preswizzle( + weight_quantizers[0], weight_tensors[0] + ) + for q in weight_quantizers: + q.optimize_for_gemm = optimize_for_gemm if is_grad_enabled: linear_fn = _GroupedLinear.apply @@ -2078,16 +2085,8 @@ def _get_weight_quantizers(self) -> List[Quantizer]: ] for i in range(self.num_gemms) ] - # Preswizzle the weights during quantization instead of lazily inside every GEMM. - # This wont work when primay weights are in fp8 because of 2 reasons - # 1. optimizer step updates would need to dequantize the weights. But swizzled weights - # currently dont support dequantization. - # 2. For FSDP2, quantized weight all-gather would need to be done in the - # unswizzled layout. for i in range(self.num_gemms): weight_quantizers[i].internal = not self.primary_weights_in_fp8 - if not self.primary_weights_in_fp8: - weight_quantizers[i].optimize_for_gemm = True return weight_quantizers def _get_quantizers(self): diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index cffc5ba1e6..a588e21a0c 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1747,6 +1747,10 @@ def forward( grad_weight_quantizer, grad_output_quantizer, ) = quantizers + if weight_quantizer is not None and not debug: + weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + weight_quantizer, weight_tensor + ) if is_grad_enabled: fwd_fn = _LayerNormLinear.apply @@ -1994,12 +1998,4 @@ def _get_weight_quantizers(self) -> List[Quantizer]: return [None] weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] weight_quantizer.internal = True - # Preswizzle the weights during quantization instead of lazily inside every GEMM. - # This wont work when primay weights are in fp8 because of 2 reasons - # 1. optimizer step updates would need to dequantize the weights. But swizzled weights - # currently dont support dequantization. - # 2. For FSDP2, quantized weight all-gather would need to be done in the - # unswizzled layout. - if not self.primary_weights_in_fp8: - weight_quantizer.optimize_for_gemm = True return [weight_quantizer] diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ad629c3f11..19e20d775c 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2337,6 +2337,15 @@ def forward( fc1_weight, fc2_weight = self._get_weight_tensors() fc1_bias = self.fc1_bias if self.use_bias else None fc2_bias = self.fc2_bias if self.use_bias else None + if not debug: + if fc1_weight_quantizer is not None: + fc1_weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + fc1_weight_quantizer, fc1_weight + ) + if fc2_weight_quantizer is not None: + fc2_weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + fc2_weight_quantizer, fc2_weight + ) if not self.fp8: if isinstance(fc1_weight, Float8Tensor): fc1_weight = fc1_weight.dequantize() @@ -2713,17 +2722,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]: fc1_weight_quantizer.internal = True fc2_weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_WEIGHT] fc2_weight_quantizer.internal = True - # Weight scale factors must be GEMM-swizzled before cuBLAS/CUTLASS can - # consume them. Pre-swizzle once at quantize time (persisted on the cached - # workspace when the weight is cached) instead of lazily inside every GEMM. - # No-op for recipes whose scales don't need swizzling (e.g. per-tensor FP8). - # The exception is quantized_model_init, where the weight parameter is - # itself quantized and gets all-gathered (FSDP2) and optimizer-updated in - # its unswizzled layout; swizzling would break the all-gather and the - # dequantize-on-update, so leave the quantizer state untouched. - if not self.primary_weights_in_fp8: - fc1_weight_quantizer.optimize_for_gemm = True - fc2_weight_quantizer.optimize_for_gemm = True return [fc1_weight_quantizer, fc2_weight_quantizer] def backward_dw(self): diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index cef6d3d533..6b0f941bc4 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1918,6 +1918,10 @@ def forward( grad_weight_quantizer, grad_output_quantizer, ) = quantizers + if weight_quantizer is not None and not debug: + weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + weight_quantizer, weight_tensor + ) if is_grad_enabled: linear_fn = _Linear.apply @@ -2205,12 +2209,4 @@ def _get_weight_quantizers(self) -> List[Quantizer]: return [None] weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] weight_quantizer.internal = True - # Preswizzle the weights during quantization instead of lazily inside every GEMM. - # This wont work when primay weights are in fp8 because of 2 reasons - # 1. optimizer step updates would need to dequantize the weights. But swizzled weights - # currently dont support dequantization. - # 2. For FSDP2, quantized weight all-gather would need to be done in the - # unswizzled layout. - if not self.primary_weights_in_fp8: - weight_quantizer.optimize_for_gemm = True return [weight_quantizer]