From 12406c993de14bd92a88a6219e5ff698d8f057ad Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Tue, 21 Jul 2026 11:11:09 +0200 Subject: [PATCH 1/5] [PyTorch] Preserve grouped weight initialization metadata Signed-off-by: CarlosGomes98 --- tests/pytorch/test_sanity.py | 79 +++++++++++++++++++ transformer_engine/pytorch/module/base.py | 34 +++++--- .../pytorch/module/grouped_linear.py | 27 ++++++- 3 files changed, 127 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 7c19dc6536..dff94cef9d 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1154,6 +1154,85 @@ def test_quantized_model_init_high_precision_init_val(): ), "clear_high_precision_init_val() not work" +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_grouped_linear_single_param_preserves_high_precision_init(monkeypatch): + """Grouped MXFP8 and discrete weights produce identical FP32 master initialization.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + num_gemms = 3 + + def make_module(single_grouped_weight): + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + with quantized_model_init( + enabled=True, + recipe=recipe.MXFP8BlockScaling(), + preserve_high_precision_init_val=True, + ): + return GroupedLinear( + num_gemms=num_gemms, + in_features=32, + out_features=64, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=single_grouped_weight, + ).cuda() + + discrete_module = make_module(False) + grouped_module = make_module(True) + discrete_init_vals = [ + getattr(discrete_module, f"weight{i}").get_high_precision_init_val() + for i in range(num_gemms) + ] + expected_grouped_init = torch.stack(discrete_init_vals, dim=0) + + grouped_weight = grouped_module.weight + assert hasattr(grouped_weight, "get_high_precision_init_val") + assert hasattr(grouped_weight, "clear_high_precision_init_val") + grouped_init = grouped_weight.get_high_precision_init_val() + assert grouped_init.device.type == "cpu" + assert grouped_init.shape == grouped_weight.shape + torch.testing.assert_close( + grouped_init, + expected_grouped_init, + rtol=0, + atol=0, + ) + + # This is the exact layout consumed when constructing the FP32 optimizer master. + discrete_master = expected_grouped_init.float() + grouped_master = grouped_init.float() + torch.testing.assert_close(grouped_master, discrete_master, rtol=0, atol=0) + + grouped_weight.clear_high_precision_init_val() + assert grouped_weight.get_high_precision_init_val() is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_grouped_linear_rejects_partial_high_precision_init(monkeypatch): + """Packing fails rather than mixing preserved and dequantized initialization.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + with quantized_model_init( + enabled=True, + recipe=recipe.MXFP8BlockScaling(), + preserve_high_precision_init_val=True, + ): + module = GroupedLinear( + num_gemms=2, + in_features=32, + out_features=64, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=False, + ).cuda() + + module.weight0.clear_high_precision_init_val() + with pytest.raises( + RuntimeError, + match="inconsistent high-precision initialization state", + ): + module.make_grouped_weights() + + def test_sanity_checkpointing_on_callables(): """Test that TE checkpointing works correctly on callable modules.""" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5eab45b7ac..c237220672 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -87,6 +87,27 @@ layers_atomic_ring_exchange = [] +def _get_high_precision_init_val(parameter: torch.Tensor) -> Optional[torch.Tensor]: + """Return temporary pre-quantization initialization stored on a parameter.""" + return getattr(parameter, "_high_precision_init_val", None) + + +def _clear_high_precision_init_val(parameter: torch.Tensor) -> None: + """Release temporary pre-quantization initialization stored on a parameter.""" + if hasattr(parameter, "_high_precision_init_val"): + del parameter._high_precision_init_val + + +def _attach_high_precision_init_val( + parameter: torch.Tensor, + high_precision_init_val: torch.Tensor, +) -> None: + """Attach TE's temporary high-precision initialization contract to a parameter.""" + parameter._high_precision_init_val = high_precision_init_val + parameter.get_high_precision_init_val = MethodType(_get_high_precision_init_val, parameter) + parameter.clear_high_precision_init_val = MethodType(_clear_high_precision_init_val, parameter) + + def is_ub_initialized() -> bool: """Whether the Userbuffers communicators have been initialized.""" return _ub_initialized @@ -1787,21 +1808,10 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: # should call `clear_high_precision_init_val` to remove it after master weight # is initialized. - def get(self): - if hasattr(self, "_high_precision_init_val"): - return self._high_precision_init_val - return None - - def clear(self): - if hasattr(self, "_high_precision_init_val"): - del self._high_precision_init_val - # DTensor.from_local() does not preserve object identity, # so attach to the DTensor's local tensor when applicable. target = dtensor_param._local_tensor if is_dtensor else param - target._high_precision_init_val = high_precision_init_val - target.get_high_precision_init_val = MethodType(get, target) - target.clear_high_precision_init_val = MethodType(clear, target) + _attach_high_precision_init_val(target, high_precision_init_val) if not is_dtensor: self.module_setattr(name, param) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index a65ee3b5c3..b244386f67 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -27,6 +27,9 @@ _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, + _attach_high_precision_init_val, + _clear_high_precision_init_val, + _get_high_precision_init_val, ) from ._common import WeightGradStore from ..quantization import FP8GlobalStateManager, QuantizerRole @@ -1522,6 +1525,19 @@ def make_grouped_weights(self, defer_init=False) -> None: weights = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] + # TE preserves the original BF16/FP16 initialization on each quantized + # parameter so distributed optimizers can construct lossless FP32 masters. + # Packing the parameters must transfer those values to the new registered + # grouped parameter; otherwise its master is initialized by dequantizing + # MXFP8 and starts from a different value than the discrete-weight layout. + high_precision_init_vals = [_get_high_precision_init_val(weight) for weight in weights] + if any(value is not None for value in high_precision_init_vals) and not all( + value is not None for value in high_precision_init_vals + ): + raise RuntimeError( + "Grouped weights have inconsistent high-precision initialization state" + ) + # Create the weight storage. grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=self.num_gemms, @@ -1545,9 +1561,18 @@ def make_grouped_weights(self, defer_init=False) -> None: and (weight_quantizers[0] is None or not weight_quantizers[0].internal) ): raise RuntimeError("Found internal quantizer with `single_grouped_weight=True`.") + grouped_parameter = torch.nn.Parameter(grouped_weights) + if all(value is not None for value in high_precision_init_vals): + _attach_high_precision_init_val( + grouped_parameter, + torch.stack(high_precision_init_vals, dim=0), + ) + for weight in weights: + _clear_high_precision_init_val(weight) + self.register_parameter( "weight", - torch.nn.Parameter(grouped_weights), + grouped_parameter, init_fn=self.init_method, get_rng_state_tracker=self.get_rng_state_tracker, fp8_meta_index=self._offsets["weight"], From 99b3f378b6aade8e253e431d92bd9f10a0e0a5f8 Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Tue, 21 Jul 2026 11:11:57 +0200 Subject: [PATCH 2/5] [PyTorch] Mark late grouped weights for delayed wgrad Signed-off-by: CarlosGomes98 --- tests/pytorch/test_grouped_mlp.py | 27 +++++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 23 +++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 76550db5f8..a3d91cea70 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -232,6 +232,33 @@ def make_reference_and_test_tensors( class TestGroupedLinearOp: """Tests for advanced features with grouped linear basic op""" + def test_meta_single_grouped_weight_with_delayed_wgrad(self, monkeypatch) -> None: + """A deferred op shell must not access its grouped parent before it is attached.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + op = te.ops.GroupedLinear( + 2, + 16, + 16, + device="meta", + dtype=torch.bfloat16, + single_grouped_weight=True, + delay_wgrad_compute=True, + ) + + assert op._parameters.get("weight") is None + assert op.weight0.device.type == "meta" + + # Mirror an external caller attaching the grouped parent after constructing + # the parameterless shell, then verify delayed-wgrad metadata is applied. + grouped_weight = torch.nn.Parameter(torch.empty(2, 16, 16, device="meta")) + assert not hasattr(grouped_weight, "skip_backward_post_hook") + op.register_parameter("weight", grouped_weight) + for group_idx in range(op.num_groups): + op.register_parameter(f"weight{group_idx}", None) + + assert grouped_weight.skip_backward_post_hook + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5ef0fa4339..18b5322ebf 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -220,12 +220,33 @@ def __init__( self._apply_delay_wgrad_param_hooks() + def register_parameter( + self, + name: str, + param: Optional[torch.nn.Parameter], + ) -> None: + """Register a parameter and apply delayed-wgrad metadata when needed.""" + super().register_parameter(name, param) + + # A single-grouped-weight op may be constructed on the meta device and + # receive its grouped parent later. Mark that parent at attachment time, + # before DDP/FSDP inspects the parameter to install backward hooks. + if name == "weight" and param is not None and getattr(self, "single_grouped_weight", False): + wgrad_store = getattr(self, "wgrad_store", None) + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + param.skip_backward_post_hook = True + def _apply_delay_wgrad_param_hooks(self) -> None: """Set ``skip_backward_post_hook`` on weights when delaying wgrad (bias uses main backward).""" if not self.wgrad_store.delay_wgrad_compute(): return if self.single_grouped_weight: - self.weight.skip_backward_post_hook = True + # A meta-device op may be created as a parameterless shell and have its + # grouped parent attached after construction. In that case there is no + # ``weight`` parameter to mark yet. + weight = self._parameters.get("weight") + if weight is not None: + weight.skip_backward_post_hook = True else: for group_idx in range(self.num_groups): getattr(self, f"weight{group_idx}").skip_backward_post_hook = True From aa4e2336eef78a1fe1f9017a842bc7849173d201 Mon Sep 17 00:00:00 2001 From: CarlosGomes98 Date: Tue, 21 Jul 2026 11:12:26 +0200 Subject: [PATCH 3/5] [PyTorch] Preserve grouped MXFP8 columnwise usage Signed-off-by: CarlosGomes98 --- tests/pytorch/test_grouped_mlp.py | 58 +++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 13 ++++- 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index a3d91cea70..cadfd6c1f6 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1131,6 +1131,64 @@ def test_grouped_mlp_fp16( activation=activation, ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_single_grouped_weight_eval_preserves_columnwise_usage( + self, + *, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + group_size: int = 2, + hidden_size: int = 128, + ) -> None: + """Eager eval must not drop storage still used by captured training dgrad.""" + + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + split_sizes = torch.full( + (group_size,), + 128, + dtype=torch.int64, + device=device, + ) + num_tokens = int(split_sizes.sum()) + recipe = make_recipe("mxfp8") + + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + device=device, + dtype=dtype, + single_grouped_weight=True, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + device=device, + dtype=dtype, + single_grouped_weight=True, + ) + module = te.ops.Sequential( + fc1, + te.ops.ScaledSwiGLU(glu_interleave_size=32), + fc2, + ) + + x = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, requires_grad=True) + probs = torch.ones(num_tokens, device=device, dtype=dtype) + with te.autocast(enabled=True, recipe=recipe): + module(x, split_sizes, probs, split_sizes) + assert fc1.weight.quantizer.columnwise_usage + assert fc2.weight.quantizer.columnwise_usage + + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + module(x.detach(), split_sizes, probs, split_sizes) + assert fc1.weight.quantizer.columnwise_usage + assert fc2.weight.quantizer.columnwise_usage + @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 31189af09c..f9ec62221c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -973,7 +973,13 @@ def fuser_forward( "FC1 expected GroupedTensor weight with single_grouped_weight=True." ) if fc1_op.weight.quantizer is not None: - fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + # Full-iteration CUDA graph replay does not rerun this Python metadata + # update. Remember that a grad-enabled forward requested columnwise + # storage so eager eval cannot drop buffers still used by captured dgrad. + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=input_requires_grad or fc1_weight_quantizer.columnwise_usage, + ) fc1_op.weight.quantizer = fc1_weight_quantizer grouped_fc1_weight = fc1_op.weight else: @@ -1005,7 +1011,10 @@ def fuser_forward( "FC2 expected GroupedTensor weight with single_grouped_weight=True." ) if fc2_op.weight.quantizer is not None: - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=input_requires_grad or fc2_weight_quantizer.columnwise_usage, + ) fc2_op.weight.quantizer = fc2_weight_quantizer grouped_fc2_weight = fc2_op.weight else: From 0e91ea5abd636b08c00798d01e78113cd36d02e9 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 23 Jul 2026 11:21:34 -0700 Subject: [PATCH 4/5] Fix bug where grouped MLP used same quantizer in grouped linear op and weight param Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- .../pytorch/ops/fused/grouped_mlp.py | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index f9ec62221c..c4b1359e2b 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -973,14 +973,13 @@ def fuser_forward( "FC1 expected GroupedTensor weight with single_grouped_weight=True." ) if fc1_op.weight.quantizer is not None: - # Full-iteration CUDA graph replay does not rerun this Python metadata - # update. Remember that a grad-enabled forward requested columnwise - # storage so eager eval cannot drop buffers still used by captured dgrad. - fc1_weight_quantizer.set_usage( - rowwise=True, - columnwise=input_requires_grad or fc1_weight_quantizer.columnwise_usage, - ) - fc1_op.weight.quantizer = fc1_weight_quantizer + # Update param quantizer + # Note: Only add usages, never drop. The same weight may be used in multiple steps. + fc1_weight_quantizer = fc1_op.weight.quantizer + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=True if input_requires_grad else None, + ) grouped_fc1_weight = fc1_op.weight else: if fc1_op.weight.rowwise_data is None: @@ -1011,11 +1010,11 @@ def fuser_forward( "FC2 expected GroupedTensor weight with single_grouped_weight=True." ) if fc2_op.weight.quantizer is not None: - fc2_weight_quantizer.set_usage( - rowwise=True, - columnwise=input_requires_grad or fc2_weight_quantizer.columnwise_usage, - ) - fc2_op.weight.quantizer = fc2_weight_quantizer + fc2_weight_quantizer = fc2_op.weight.quantizer + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=True if requires_grad else None, + ) grouped_fc2_weight = fc2_op.weight else: if fc2_op.weight.rowwise_data is None: From 2c7db67a7fa0b174b1ec32b3a54ec56dee5f225e Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 23 Jul 2026 14:46:39 -0700 Subject: [PATCH 5/5] Fix inconsistent line endings GitHub inserted CRLFs. Microsoft... Signed-off-by: Tim Moon --- .../pytorch/ops/fused/grouped_mlp.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index c4b1359e2b..ef97567e95 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -973,13 +973,13 @@ def fuser_forward( "FC1 expected GroupedTensor weight with single_grouped_weight=True." ) if fc1_op.weight.quantizer is not None: - # Update param quantizer - # Note: Only add usages, never drop. The same weight may be used in multiple steps. - fc1_weight_quantizer = fc1_op.weight.quantizer - fc1_weight_quantizer.set_usage( - rowwise=True, - columnwise=True if input_requires_grad else None, - ) + # Update param quantizer + # Note: Only add usages, never drop. The same weight may be used in multiple steps. + fc1_weight_quantizer = fc1_op.weight.quantizer + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=True if input_requires_grad else None, + ) grouped_fc1_weight = fc1_op.weight else: if fc1_op.weight.rowwise_data is None: @@ -1010,11 +1010,11 @@ def fuser_forward( "FC2 expected GroupedTensor weight with single_grouped_weight=True." ) if fc2_op.weight.quantizer is not None: - fc2_weight_quantizer = fc2_op.weight.quantizer - fc2_weight_quantizer.set_usage( - rowwise=True, - columnwise=True if requires_grad else None, - ) + fc2_weight_quantizer = fc2_op.weight.quantizer + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=True if requires_grad else None, + ) grouped_fc2_weight = fc2_op.weight else: if fc2_op.weight.rowwise_data is None: