[PyTorch] Fix deferred initialization in fusible ops

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"""

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(

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)