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)