From 8f1fd24cba18234eee0022c8004d6e6929edca52 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Fri, 7 Aug 2026 14:37:17 +0000 Subject: [PATCH] [PyTorch] Fix deferred initialization in te.ops Ops replace their params when materializing them on the first fuser forward pass, which invalidated the fuser's cached params and ops.Linear's parameter aliases. Backward then failed with a meta-vs-cuda device mismatch. - Re-cache the basic ops' params after first-forward materialization - Run first-forward initialization on the top-level ops, before the cached params are used to pick which ops require backward - Sync ops.Linear weight/bias aliases with its basic ops - Add deferred initialization tests for the affected ops Fixes #3322 Signed-off-by: Alp Dener --- tests/pytorch/test_fusible_ops.py | 134 +++++++++++++++++++++++ transformer_engine/pytorch/ops/fuser.py | 44 ++++++-- transformer_engine/pytorch/ops/linear.py | 21 ++++ 3 files changed, 190 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..8c9109f9cb 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -649,6 +649,140 @@ def test_pyt_autocast( assert x.grad.dtype == model_dtype assert op.weight.grad.dtype == model_dtype + @pytest.mark.parametrize( + "op_type", + ( + "basic_linear", + "bias", + "layer_norm", + "rmsnorm", + "grouped_linear", + "grouped_linear_single_param", + "linear", + "sequential", + ), + ) + @pytest.mark.parametrize("dtype", _dtypes) + def test_deferred_param_init( + self, + monkeypatch, + *, + op_type: str, + size: int = 32, + dtype: torch.dtype, + device: torch.device = "cuda", + ) -> None: + """Test ops constructed on the meta device + + Ops replace their params when materializing them on the first + forward pass. See + https://github.com/NVIDIA/TransformerEngine/issues/3322. + + """ + + # Construct operation on meta device + kwargs = {"device": "meta", "dtype": dtype} + num_groups = 2 + if op_type == "basic_linear": + op = te_ops.BasicLinear(size, size, **kwargs) + elif op_type == "bias": + op = te_ops.Bias(size, **kwargs) + elif op_type == "layer_norm": + op = te_ops.LayerNorm(size, **kwargs) + elif op_type == "rmsnorm": + op = te_ops.RMSNorm(size, **kwargs) + elif op_type == "grouped_linear": + op = te_ops.GroupedLinear(num_groups, size, size, bias=True, **kwargs) + elif op_type == "grouped_linear_single_param": + # Materializing params also changes how many there are + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + op = te_ops.GroupedLinear( + num_groups, + size, + size, + bias=True, + single_grouped_weight=True, + single_grouped_bias=True, + **kwargs, + ) + elif op_type == "linear": + op = te_ops.Linear(size, size, bias=True, **kwargs) + elif op_type == "sequential": + op = te_ops.Sequential( + te_ops.LayerNorm(size, **kwargs), + te_ops.Linear(size, size, bias=True, **kwargs), + ) + else: + raise ValueError(f"Unsupported op type ({op_type})") + for param in op.parameters(): + assert param.device.type == "meta" + + # Forward and backward pass + in_shape = (size, size) + extra_inputs = [] + if op_type.startswith("grouped_linear"): + in_shape = (size * num_groups, size) + extra_inputs.append(torch.tensor([size] * num_groups, dtype=torch.int, device=device)) + x = torch.randn(in_shape, dtype=dtype, device=device, requires_grad=True) + y = op(x, *extra_inputs) + y.backward(torch.randn_like(y)) + + # Check that params have been materialized + params = dict(op.named_parameters()) + assert params + for name, param in params.items(): + assert param.device.type == device, f"{name} was not materialized on {device}" + assert param.grad is not None, f"{name} did not get a grad" + assert param.grad.device.type == device, f"{name} got a grad on {param.grad.device}" + assert x.grad is not None and x.grad.device.type == device + + # Check that fused ops do not alias stale params + for module in op.modules(): + if isinstance(module, te_ops.Linear): + assert module.weight is module.basic_ops[module._linear_idx].weight + assert module.bias is module.basic_ops[module._bias_idx].bias + + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantized_weight", (False, True)) + def test_deferred_param_init_quantized( + self, + *, + size: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + quantization: Optional[str], + quantized_weight: bool, + ) -> None: + """Test quantized op constructed on the meta device""" + + # Skip invalid configurations + in_shape = (size, size) + if quantization is None: + pytest.skip("Quantization scheme is not specified") + maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + + # Construct operation on meta device + recipe = make_recipe(quantization) + with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): + op = te_ops.Linear(size, size, bias=True, device="meta", dtype=dtype) + + # Forward and backward pass + x = torch.randn(in_shape, dtype=dtype, device=device, requires_grad=True) + with te.autocast(recipe=recipe): + y = op(x) + y.backward(torch.randn_like(y)) + + # Check that params have been materialized + assert op.weight is op.basic_ops[op._linear_idx].weight + assert op.bias is op.basic_ops[op._bias_idx].bias + if quantized_weight: + assert isinstance(op.weight, QuantizedTensor) + for name, param in op.named_parameters(): + assert param.device.type == device, f"{name} was not materialized on {device}" + assert param.grad is not None, f"{name} did not get a grad" + assert param.grad.device.type == device, f"{name} got a grad on {param.grad.device}" + assert x.grad is not None and x.grad.device.type == device + class TestBasicOps: """Tests for individual operations""" diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..d858aeb146 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -330,9 +330,12 @@ def __init__( ops: list[FusibleOperation], ) -> None: + # Ops before flattening out fused ops + self._ops: list[FusibleOperation] = list(ops) + # Get list of basic operations basic_ops = [] - for op in ops: + for op in self._ops: if op.is_fused_op: basic_ops.extend(op.basic_ops) else: @@ -354,7 +357,19 @@ def __init__( self.backward_override = None self._last_amax_history_len = 0 - # Flatten list of parameters + # Flattened list of params, populated on the first forward pass + self._basic_op_params: list[list[torch.nn.Parameter]] + self._basic_op_num_params: list[int] + self._flat_basic_op_params: list[torch.nn.Parameter] + + def _cache_parameters(self) -> None: + """Cache the basic ops' params + + Walking the ops' params is expensive, so it is kept out of the + steady-state forward pass. Must be called again whenever an op + may have replaced its params. + + """ self._basic_op_params = [list(op.parameters()) for op in self._basic_ops] self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) @@ -423,6 +438,20 @@ def maybe_fuse_ops( ): """Attempt to fuse operations if neccesary""" + # Initialize ops before the first forward pass + # Note: Recipe state must be reset first since params may be + # quantized. Ops are initialized top-down so fused ops can sync + # params aliased from their basic ops. + is_first_forward = self.recipe_type is None + if is_first_forward: + for op in self._basic_ops: + op.reset_recipe_state(recipe=recipe) + for op in self._ops: + op.pre_first_fuser_forward() + + # Cache params now that ops have initialized them + self._cache_parameters() + # Determine which basic ops require backward if not is_grad_enabled: first_op_requiring_backward = self._num_basic_ops @@ -437,7 +466,7 @@ def maybe_fuse_ops( break # Early exit if fusion parameters haven't changed - need_reset = False + need_reset = is_first_forward recipe_type = type(recipe) backward_override = recipe.backward_override if recipe is not None else None fusion_params = (recipe_type, first_op_requiring_backward, backward_override) @@ -459,13 +488,10 @@ def maybe_fuse_ops( return # Reset recipe state - for op in self._basic_ops: - op.reset_recipe_state(recipe=recipe) - - # Check if this is the first iteration - if self.recipe_type is None: + # Note: Already done above on the first forward pass. + if not is_first_forward: for op in self._basic_ops: - op.pre_first_fuser_forward() + op.reset_recipe_state(recipe=recipe) # Apply joint forward-backward fusions first joint_ops = OperationFuser._apply_fusions( diff --git a/transformer_engine/pytorch/ops/linear.py b/transformer_engine/pytorch/ops/linear.py index c6ca4786b8..4e49a5a0e0 100644 --- a/transformer_engine/pytorch/ops/linear.py +++ b/transformer_engine/pytorch/ops/linear.py @@ -166,6 +166,27 @@ def register_parameter(self, name: str, param: Optional[torch.nn.Parameter]) -> elif name == "bias" and self._bias_idx is not None: self.basic_ops[self._bias_idx].bias = param + def _sync_parameters(self) -> None: + """Update the params registered in this op to match the basic ops + + Only re-registers changed params since this may run on every + forward pass. + + """ + weight = self.basic_ops[self._linear_idx].weight + if weight is not self._parameters["weight"]: + self.register_parameter("weight", weight) + if self._bias_idx is not None: + bias = self.basic_ops[self._bias_idx].bias + if bias is not self._parameters["bias"]: + self.register_parameter("bias", bias) + + def pre_first_fuser_forward(self) -> None: + super().pre_first_fuser_forward() + + # Basic ops replace their params during deferred initialization + self._sync_parameters() + def state_dict(self, *, prefix: str = "", **kwargs) -> dict[str, Any]: """Save state""" state_dict = super().state_dict(prefix=prefix, **kwargs)