diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index e24fff9049..8c24b0bdc5 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -228,6 +228,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) @@ -1086,6 +1113,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/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 7c19dc6536..9c0359a3c0 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -650,6 +650,8 @@ def test_sanity_grouped_linear( loss = out.sum() loss.backward() assert out.shape == (num_tokens, ffn_hidden_size) + if single_param: + assert te_grouped_linear.weight.grad is not None @pytest.mark.parametrize("dtype", param_types) @@ -1154,6 +1156,158 @@ 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() + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("fuse_wgrad_accumulation", (False, True)) +def test_grouped_linear_single_param_legacy_wgrad_matches_discrete( + monkeypatch, + fuse_wgrad_accumulation, +): + """Single-grouped and discrete MXFP8 weights produce identical legacy wgrads.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0") + dtype = torch.bfloat16 + num_gemms = 2 + + def make_module(single_grouped_weight): + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + with quantized_model_init(enabled=True, recipe=recipe.MXFP8BlockScaling()): + return GroupedLinear( + num_gemms=num_gemms, + in_features=32, + out_features=64, + bias=False, + params_dtype=dtype, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + single_grouped_weight=single_grouped_weight, + ).cuda() + + discrete_module = make_module(False) + grouped_module = make_module(True) + + if fuse_wgrad_accumulation: + for index in range(num_gemms): + weight = getattr(discrete_module, f"weight{index}") + weight.main_grad = torch.zeros( + tuple(weight.shape), dtype=weight.dtype, device=weight.device + ) + weight.grad_added_to_main_grad = False + grouped_module.weight.main_grad = torch.zeros( + tuple(grouped_module.weight.shape), + dtype=grouped_module.weight.dtype, + device=grouped_module.weight.device, + ) + grouped_module.weight.grad_added_to_main_grad = False + + torch.manual_seed(5678) + torch.cuda.manual_seed(5678) + inp = torch.randn(64, 32, dtype=dtype, device="cuda") + grad_output = torch.randn(64, 64, dtype=dtype, device="cuda") + + def run_backward(module): + module_input = inp.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=recipe.MXFP8BlockScaling()): + output = module(module_input, [32, 32]) + output.backward(grad_output) + + run_backward(discrete_module) + run_backward(grouped_module) + + if fuse_wgrad_accumulation: + discrete_wgrads = [ + getattr(discrete_module, f"weight{i}").main_grad for i in range(num_gemms) + ] + grouped_wgrad = grouped_module.weight.main_grad + assert grouped_module.weight.grad_added_to_main_grad + else: + discrete_wgrads = [getattr(discrete_module, f"weight{i}").grad for i in range(num_gemms)] + grouped_wgrad = grouped_module.weight.grad + + assert all(wgrad is not None for wgrad in discrete_wgrads) + assert grouped_wgrad is not None + discrete_wgrad = torch.stack(discrete_wgrads) + torch.testing.assert_close(grouped_wgrad, discrete_wgrad, rtol=0, atol=0) + + 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 82ee7953c1..0afcaf6fef 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 @@ -68,6 +71,66 @@ __all__ = ["GroupedLinear"] +class _GroupedWeightAutogradBridge(torch.autograd.Function): + """Connect cached per-GEMM views to their registered grouped parameter. + + The legacy GroupedLinear autograd function consumes cached member tensors. Those + tensors are storage views, not registered parameters, so using them directly drops + the autograd edge to the grouped parent. This bridge makes the parent the autograd + input while preserving the member objects required by the grouped GEMM kernels. + """ + + @staticmethod + def forward(ctx, grouped_weight: torch.Tensor, fuse_wgrad_accumulation: bool): + # Backward only needs the parent's identity/attributes, not saved tensor values; + # use a weakref to avoid extending its lifetime or creating a reference cycle. + ctx.grouped_weight_ref = weakref.ref(grouped_weight) + ctx.fuse_wgrad_accumulation = bool(fuse_wgrad_accumulation) + members = grouped_weight.quantized_tensors + if members is None: + members = grouped_weight.split_into_quantized_tensors() + return tuple(members) + + @staticmethod + def backward(ctx, *member_grads): + grouped_weight = ctx.grouped_weight_ref() + if grouped_weight is None: + raise RuntimeError("Grouped weight was released before backward completed") + + if ctx.fuse_wgrad_accumulation: + # GroupedLinear has already written each wgrad into a view of the parent's + # main_grad. Trigger the parent's accumulator/DDP hook with TE's cached + # dummy wgrad and tell MCore not to add that dummy to main_grad. + grouped_grad = None + if hasattr(grouped_weight, "grad_added_to_main_grad"): + grouped_weight.grad_added_to_main_grad = True + grouped_grad = get_dummy_wgrad( + list(grouped_weight.shape), + grouped_weight.dtype, + zero=getattr(grouped_weight, "zero_out_wgrad", False), + ) + else: + members = grouped_weight.quantized_tensors + if members is None: + members = grouped_weight.split_into_quantized_tensors() + grouped_grad = torch.stack( + [ + ( + grad + if grad is not None + else torch.zeros( + tuple(member.shape), + dtype=grouped_weight.dtype, + device=grouped_weight.device, + ) + ) + for member, grad in zip(members, member_grads) + ], + dim=0, + ) + return grouped_grad, None + + class _GroupedLinear(torch.autograd.Function): """GroupedLinear semi-top level module Calls custom cuda extensions. @@ -1467,6 +1530,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, @@ -1490,9 +1566,20 @@ 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) + for member in grouped_parameter.quantized_tensors: + member.requires_grad_(grouped_parameter.requires_grad) + 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"], @@ -1904,6 +1991,28 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage if weight_tensors is None: # TODO(ksivaman): Remove this after GEMM integration. weight_tensors = grouped_weight.split_into_quantized_tensors() + if torch.is_grad_enabled() and grouped_weight.requires_grad: + weight_tensors = list( + _GroupedWeightAutogradBridge.apply(grouped_weight, self.fuse_wgrad_accumulation) + ) + # Quantized tensor subclasses preserve the cached wrapper's original + # requires_grad bit when a custom autograd Function returns an alias. + # The grouped parent is trainable, so make that state explicit on the + # bridged member outputs consumed by legacy GroupedLinear autograd. + for weight in weight_tensors: + weight.requires_grad_(True) + if self.fuse_wgrad_accumulation: + grouped_main_grad = getattr(grouped_weight, "main_grad", None) + for index, weight in enumerate(weight_tensors): + if grouped_main_grad is not None: + weight.main_grad = grouped_main_grad[index] + for attr in ( + "grad_added_to_main_grad", + "overwrite_main_grad", + "zero_out_wgrad", + ): + if hasattr(grouped_weight, attr): + setattr(weight, attr, getattr(grouped_weight, attr)) else: weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5c96d4658e..f08002c580 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -213,12 +213,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 diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index da3aa13156..dd3a005c46 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -949,7 +949,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: @@ -981,7 +987,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: