diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 395476dbd1d..d2ceb7c7f33 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1475,7 +1475,9 @@ def test_mha_accuracy(dtype, bs, model, mask_type): assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype]) -def _test_granular_accuracy(block, bs, dtype, config, delay_wgrad_compute=False, recipe=None): +def _test_granular_accuracy( + block, bs, dtype, config, delay_wgrad_compute=False, recipe=None, grad_output=None +): reset_rng_states() fp8 = recipe is not None if fp8: @@ -1493,8 +1495,11 @@ def _test_granular_accuracy(block, bs, dtype, config, delay_wgrad_compute=False, out = block(inp_hidden_states) if isinstance(out, (List, Tuple)): out = out[0] - loss = out.sum() - loss.backward() + if grad_output is None: + loss = out.sum() + loss.backward() + else: + out.backward(grad_output) if delay_wgrad_compute: block.backward_dw() @@ -2058,15 +2063,35 @@ def test_layernorm_accuracy(dtype, bs, model, eps, zero_centered_gamma): assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype]) -@pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("bs", batch_sizes) -@pytest.mark.parametrize("model", ["small"]) -@pytest.mark.parametrize("normalization", all_normalizations) -@pytest.mark.parametrize("zero_centered_gamma", all_boolean) -@pytest.mark.parametrize("return_bias", all_boolean) -@pytest.mark.parametrize("bias", all_boolean) +@pytest.mark.parametrize( + "dtype,bs,model,normalization,zero_centered_gamma,return_bias,bias,split_storage", + [ + (dtype, bs, "small", normalization, zero_centered_gamma, return_bias, bias, None) + for dtype in param_types + for bs in batch_sizes + for normalization in all_normalizations + for zero_centered_gamma in all_boolean + for return_bias in all_boolean + for bias in all_boolean + ] + + [ + pytest.param( + torch.bfloat16, + 1, + "small", + "LayerNorm", + False, + False, + True, + storage, + id=f"split-{storage}-storage", + marks=pytest.mark.skipif(not is_bf16_available(), reason="BF16 is not supported"), + ) + for storage in ("shared", "separate") + ], +) def test_layernorm_linear_accuracy( - dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias + dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias, split_storage ): config = model_configs[model] @@ -2081,8 +2106,25 @@ def test_layernorm_linear_accuracy( return_bias=return_bias, bias=bias, device="cuda", + parameters_split=( + dict(q=2 * config.hidden_size, k=config.hidden_size, v=config.hidden_size) + if split_storage + else None + ), ) + te_module = te_ln_linear.te_module + weights = [getattr(te_module, name) for name in te_module.weight_names] + biases = [getattr(te_module, name) for name in te_module.bias_names] if bias else [] + if split_storage: + for parts in (weights, biases): + if split_storage == "separate": + for part in parts: + part.data = part.detach().clone() + assert len({part.untyped_storage().data_ptr() for part in parts}) == ( + len(parts) if split_storage == "separate" else 1 + ) + torch_ln_linear = ( TorchLayerNormLinear( config.hidden_size, @@ -2105,12 +2147,26 @@ def test_layernorm_linear_accuracy( torch_ln_linear.layernorm.bias = Parameter( te_ln_linear.te_module.layer_norm_bias.clone() ) - torch_ln_linear.linear.weight = Parameter(te_ln_linear.te_module.weight.clone()) + torch_ln_linear.linear.weight = Parameter(torch.cat(weights)) if bias: - torch_ln_linear.linear.bias = Parameter(te_ln_linear.te_module.bias.clone()) + torch_ln_linear.linear.bias = Parameter(torch.cat(biases)) - te_outputs = _test_granular_accuracy(te_ln_linear, bs, dtype, config) - torch_outputs = _test_granular_accuracy(torch_ln_linear, bs, dtype, config) + grad_output = ( + torch.randn(config.max_seqlen_q, bs, 4 * config.hidden_size, dtype=dtype, device="cuda") + if split_storage + else None + ) + te_outputs = _test_granular_accuracy(te_ln_linear, bs, dtype, config, grad_output=grad_output) + torch_outputs = _test_granular_accuracy( + torch_ln_linear, bs, dtype, config, grad_output=grad_output + ) + if split_storage: + torch_outputs = [ + *torch_outputs[:-2], + *torch_outputs[-2].split(te_module.parameter_split_sizes), + *torch_outputs[-1].split(te_module.parameter_split_sizes), + ] + assert len(te_outputs) == len(torch_outputs) atol = { torch.float32: 2.5e-4, diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6868de7ef18..06270955a81 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1933,47 +1933,232 @@ def test_to_tensor_spec_quantized(factory, shape): # --------------------------------------------------------------------------- -@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") -@pytest.mark.parametrize("compile_mode", _compile_modes) -@pytest.mark.parametrize( - "fp8_recipe", - [None, *_all_recipes], - ids=lambda r: "bf16" if r is None else type(r).__name__, -) -def test_te_linear_compiles(fp8_recipe, compile_mode): - """ - torch.compile(fullgraph=True) of ``te.Linear`` under every built-in - recipe (plus the bf16-only baseline with no autocast), for both the default - backend and ``mode="reduce-overhead"`` (CUDA-graph trees). - """ - dtype = torch.bfloat16 - device = "cuda" +def _check_module_outputs_and_grads( + model, + fp8_recipe=None, + compile_mode=None, + *, + grads="all", + check_storage=False, + recompute_workspace=False, +): + """Compare module outputs and gradients through parameter storage changes.""" + torch._dynamo.reset() + training = grads != "none" + dtype = next(model.parameters()).dtype + in_features = model.in_features + tolerances = ( + dict(atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + if len(model.weight_names) == 1 + else dtype_tols(dtype) + ) + bias_mode = "returned" if model.return_bias else "fused" if model.use_bias else "none" + for name, parameter in model.named_parameters(): + parameter.requires_grad_( + grads == "all" + or (grads == "mixed" and name not in ("k_weight", "v_bias")) + or (grads == "weight" and name == "v_weight") + ) + reference = copy.deepcopy(model) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01) + quant_recipe = copy.deepcopy(fp8_recipe) + if quant_recipe is not None and quant_recipe.nvfp4(): + quant_recipe.disable_stochastic_rounding = True - # FP8 GEMMs require leading dimensions divisible by 16. - model = te.Linear(64, 32, params_dtype=dtype, device=device) + autocast_enabled = False - def fn(inp): - if fp8_recipe is None: - return model(inp) - with te.autocast(recipe=fp8_recipe): - return model(inp) + def run(inp, current=model): + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled): + if quant_recipe is None: + out = current(inp) + else: + with te.autocast(recipe=quant_recipe): + out = current(inp) + return out[0] + out[1] if bias_mode == "returned" else out + + def reference_parameters(): + weight = torch.cat([getattr(reference, name) for name in reference.weight_names]) + bias = ( + torch.cat([getattr(reference, name) for name in reference.bias_names]) + if reference.use_bias + else None + ) + return weight, bias, None, None + + reference._get_weight_and_bias = reference_parameters + + if compile_mode == "default" and check_storage: + inp = torch.randn(32, in_features, device="cuda", dtype=dtype, requires_grad=True) + out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp) + with torch.no_grad(): + model.v_weight.add_(0.01) + reference.v_weight.add_(0.01) + with pytest.raises(RuntimeError, match="modified by an inplace operation"): + out.sum().backward() - torch._dynamo.reset() if compile_mode == "reduce-overhead": _cudagraph_warmup( - fn, - torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), - backward=True, + run, + torch.randn( + 32, + in_features, + device="cuda", + dtype=dtype, + requires_grad=grads in ("all", "mixed", "input"), + ), + backward=training, ) model.zero_grad(set_to_none=True) - compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run + check_cats = ( + check_storage + and compile_mode != "reduce-overhead" + and dtype == torch.bfloat16 + and fp8_recipe is None + and bias_mode == "fused" + ) + mutations = ( + ["none", "optimizer", "data", "parameter", "assign", "dtype", "mixed_dtype"] + if check_storage + else ["none"] + ) + if recompute_workspace: + mutations.append("fsdp2") + for mutation in mutations: + autocast_enabled = mutation == "mixed_dtype" + for current in (model, reference): + if mutation == "data": + current.k_weight.data = current.k_weight.detach().clone() + 0.01 + current.k_bias.data = current.k_bias.detach().clone() + 0.01 + elif mutation == "parameter": + current.q_weight = torch.nn.Parameter( + current.q_weight.detach().clone() + 0.01, + requires_grad=current.q_weight.requires_grad, + ) + elif mutation == "assign": + state = {name: value.clone() for name, value in current.state_dict().items()} + current.load_state_dict(state, assign=True) + elif mutation == "dtype": + current.to(torch.float64).to(dtype) + elif mutation == "mixed_dtype": + current.k_weight.data = current.k_weight.float() + if current.use_bias: + current.k_bias.data = current.k_bias.float() + elif mutation == "fsdp2": + current._is_fsdp2 = True + if mutation == "optimizer" and training: + optimizer.step() + ref_optimizer.step() + for iteration in range(3): + torch.compiler.cudagraph_mark_step_begin() + inp = torch.randn( + 32, + in_features, + device="cuda", + dtype=dtype, + requires_grad=grads in ("all", "mixed", "input"), + ) + ref_inp = inp.detach().clone().requires_grad_(inp.requires_grad) + model.zero_grad(set_to_none=True) + reference.zero_grad(set_to_none=True) + grad = torch.randn(32, model.out_features, device="cuda", dtype=dtype) + rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None + profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data") + with ( + torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) + if profile + else contextlib.nullcontext() + ) as prof: + with torch.set_grad_enabled(training): + out = actual_forward(inp) + if training: + out.backward(grad) + if profile: + cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat") + expected_cats = ( + 0 + if mutation in ("none", "optimizer") + else (4 if compile_mode and training else 2) + ) + assert cats == expected_cats + if rng_state is not None: + torch.cuda.set_rng_state(rng_state) + with torch.set_grad_enabled(training): + ref_out = run(ref_inp, reference) + if training: + ref_out.backward(grad) + torch.testing.assert_close(out, ref_out, **tolerances) + for actual, expected in zip( + (inp, *model.parameters()), (ref_inp, *reference.parameters()) + ): + if not actual.requires_grad: + assert actual.grad is None + else: + assert actual.grad is not None + torch.testing.assert_close(actual.grad, expected.grad, **tolerances) - # Iterate a few times so reduce-overhead actually replays a captured graph. - n_iters = 3 if compile_mode == "reduce-overhead" else 1 + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize( + "fp8_recipe,compile_mode,dtype,parameters_split,bias_mode,grads", + [ + (q, mode, torch.bfloat16, None, "fused", "all") + for q in (None, *_all_recipes) + for mode in _compile_modes + ] + + [ + (None, "default", torch.float32, ("q", "k", "v"), "fused", "all"), + (None, "default", torch.float16, dict(q=128, k=64, v=64), "fused", "all"), + (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "fused", "mixed"), + (None, "reduce-overhead", torch.bfloat16, ("q", "k", "v"), "fused", "all"), + (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "none", "all"), + (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "returned", "all"), + ] + + [ + (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "fused", grads) + for grads in ("none", "input", "weight") + ] + + [ + (q, "default", torch.bfloat16, dict(q=128, k=64, v=64), "fused", "all") + for q in _all_recipes + ] + + [ + ( + recipe.Float8CurrentScaling(backward_override=override), + "default", + torch.bfloat16, + dict(q=128, k=64, v=64), + "fused", + "all", + ) + for override in ("high_precision", "dequantized") + if fp8_available + ], +) +def test_te_linear_compiles(fp8_recipe, compile_mode, dtype, parameters_split, bias_mode, grads): + """Check compiled Linear outputs and gradients, including split parameter updates.""" + device = "cuda" + model = te.Linear( + 128 if parameters_split else 64, + (256 if isinstance(parameters_split, dict) else 384) if parameters_split else 32, + params_dtype=dtype, + device=device, + parameters_split=parameters_split, + bias=bias_mode != "none", + return_bias=bias_mode == "returned", + ) with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): - for _ in range(n_iters): - base = torch.randn(32, 64, dtype=dtype, device=device) - _assert_close_eager_compiled(fn, compiled, model, base) + _check_module_outputs_and_grads( + model, + fp8_recipe, + compile_mode, + grads=grads, + check_storage=grads == "mixed", + recompute_workspace=bool(parameters_split) + and isinstance(fp8_recipe, recipe.Float8CurrentScaling) + and fp8_recipe.backward_override is None, + ) @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") @@ -2238,39 +2423,56 @@ def fn(inp): @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") -@pytest.mark.parametrize("compile_mode", _compile_modes) -@pytest.mark.parametrize("deferred_backward", [False, True]) @pytest.mark.parametrize( - "fp8_recipe", + "fp8_recipe,compile_mode,deferred_backward,parameters_split", [ pytest.param( - recipe.Float8CurrentScaling(), - marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), - ), - pytest.param( - recipe.Float8BlockScaling(), - marks=pytest.mark.skipif( - not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling + q, + mode, + deferred, + None, + marks=pytest.mark.skipif(not available, reason=reason), + id=f"{recipe_id(q)}-{mode}-{deferred}", + ) + for q, available, reason in ( + (recipe.Float8CurrentScaling(), fp8_available, reason_for_no_fp8), + ( + recipe.Float8BlockScaling(), + fp8_block_scaling_available, + reason_for_no_fp8_block_scaling, ), - ), - pytest.param( - recipe.MXFP8BlockScaling(), - marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), - ), + (recipe.MXFP8BlockScaling(), mxfp8_available, reason_for_no_mxfp8), + ( + recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), + nvfp4_available, + reason_for_no_nvfp4, + ), + ) + for mode in _compile_modes + for deferred in (False, True) + ] + + [ pytest.param( - recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), - marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), - ), + recipe.Float8CurrentScaling(), + "default", + False, + dict(q=64, k=32, v=32), + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + id="split-cache", + ) ], - ids=recipe_id, ) def test_te_linear_compile_is_first_microbatch( - fp8_recipe, compile_mode, deferred_backward, monkeypatch + fp8_recipe, compile_mode, deferred_backward, parameters_split, monkeypatch ): """Reuse cached weights across microbatches and refresh after optimizer updates.""" dtype, device = torch.bfloat16, "cuda" - model = te.Linear(128, 128, params_dtype=dtype, device=device) - ref_model = te.Linear(128, 128, params_dtype=dtype, device=device) + model = te.Linear( + 128, 128, params_dtype=dtype, device=device, parameters_split=parameters_split + ) + ref_model = te.Linear( + 128, 128, params_dtype=dtype, device=device, parameters_split=parameters_split + ) ref_model.load_state_dict(model.state_dict()) def fn(inp, is_first): diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py new file mode 100644 index 00000000000..b0f1a315aa8 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/concat.py @@ -0,0 +1,62 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Deferred row concatenation for tensors consumed inside a custom op.""" + +import math +from typing import List, Sequence, Tuple, Union + +import torch + +from .tensor_spec import TensorSpec + + +def concat_shape(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> Tuple[int, ...]: + """Shape after row concatenation.""" + return (sum(t.shape[0] for t in parts), *parts[0].shape[1:]) + + +def concat_to_spec(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> TensorSpec: + """Describe the concatenation without accessing its storage.""" + first = parts[0] + dtype = first.dtype + for tensor in parts[1:]: + dtype = torch.promote_types(dtype, tensor.dtype) + return TensorSpec( + shape=concat_shape(parts), + dtype=dtype, + device=first.device, + requires_grad=any(t.requires_grad for t in parts), + ) + + +def materialize_concat(value: List[torch.Tensor]) -> torch.Tensor: + """Build a checked view, or copy parts whose storage is not adjacent.""" + if len(value) == 1: + return value[0] + first = value[0] + if first.shape[0] == 0: + return torch.cat(value) + storage = first.untyped_storage() + offset = first.storage_offset() + for tensor in value: + if ( + tensor.dtype != first.dtype + or tensor.device != first.device + or tensor.is_neg() + or tensor.is_conj() + ): + return torch.cat(value) + if ( + tensor.shape[1:] != first.shape[1:] + or not tensor.is_contiguous() + or tensor.untyped_storage().data_ptr() != storage.data_ptr() + or tensor.storage_offset() != offset + ): + return torch.cat(value) + offset += tensor.numel() + if offset * first.element_size() > storage.nbytes(): + return torch.cat(value) + shape = concat_shape(value) + return first.as_strided(shape, (math.prod(shape[1:]), *first.stride()[1:])) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 1ed84017ecd..8d391941ec4 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -37,6 +37,7 @@ * ``TENSOR_OR_QUANTIZED`` -- a field that may be a plain tensor, a bare quantized storage, or ``None``: three slots (the tensor, its flat inner buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. + * ``TENSOR_LIST`` -- a ``List[Tensor]`` field: one ``Tensor[]`` slot. * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes, quantizers -- value-opaque constants baked into the graph -- and nested collections of them), gathered into one shared ``OpaqueValueBundle`` slot. @@ -396,6 +397,7 @@ class _FieldKind(Enum): """How one dataclass field crosses the custom-op boundary.""" TENSOR = "tensor" # one ``Tensor`` / ``Tensor?`` slot + TENSOR_LIST = "tensor_list" # one ``Tensor[]`` / ``Tensor[]?`` slot TENSOR_OR_QUANTIZED = "tensor_or_quantized" # 3 slots: tensor / inner / meta PROCESS_GROUP = "process_group" # c10d group name inside the shared bundle SIMPLE = "simple" # value carried verbatim inside the shared bundle @@ -413,8 +415,7 @@ class _TensorOrQuantizedKind(Enum): _TQ_KIND_KEY = "__kind__" _SIMPLE_META_SLOT = "_simple_meta" -# Matched by exact member set, so a bare quantized annotation or an accidental -# extra union member is rejected rather than silently taken as tensor-or-quantized. +# Tensor unions are matched by exact member sets. _TQ_MEMBERS = frozenset(get_args(TensorOrQuantized)) @@ -439,8 +440,13 @@ class _FieldPlan: slots: Tuple[_SlotSpec, ...] +def _is_tensor_list_annot(annot: Any) -> bool: + """Whether an annotation is a list of tensors.""" + return get_origin(annot) is list and get_args(annot) == (torch.Tensor,) + + def _is_tensor_storage_union(annot: Any) -> bool: - """Whether ``annot`` is exactly the tensor-or-quantized union.""" + """Whether a tensor union can use the tensor / buffers / metadata slots.""" if not _is_union(annot): return False members = frozenset(a for a in get_args(annot) if a is not type(None)) @@ -488,6 +494,9 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan: ) return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots) stripped, is_optional = _strip_optional(annot) + if _is_tensor_list_annot(stripped): + slot = _SlotSpec(name, "Tensor[]?" if is_optional else "Tensor[]") + return _FieldPlan(name, _FieldKind.TENSOR_LIST, (slot,)) if stripped is torch.Tensor: slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor") return _FieldPlan(name, _FieldKind.TENSOR, (slot,)) @@ -582,7 +591,7 @@ def tensor_field_names(self) -> Tuple[str, ...]: return tuple( f.name for f in self.fields - if f.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED) + if f.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_LIST, _FieldKind.TENSOR_OR_QUANTIZED) ) def tensor_or_quantized_offsets(self) -> List[int]: @@ -608,7 +617,11 @@ def resolve_grad_targets(self, input_tensors_for_grad: Sequence[str]) -> List[in index: Dict[str, int] = {} offset = 0 for field in self.fields: - if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED): + if field.kind in ( + _FieldKind.TENSOR, + _FieldKind.TENSOR_LIST, + _FieldKind.TENSOR_OR_QUANTIZED, + ): # The gradient flows to the group's first slot -- for # tensor-or-quantized that is the ``Tensor?`` slot, the one # autograd sees the (subclass) tensor in. @@ -631,7 +644,7 @@ def pack(self, obj: Any) -> Dict[str, Any]: for field in self.fields: value = getattr(obj, field.name, None) match field.kind: - case _FieldKind.TENSOR: + case _FieldKind.TENSOR | _FieldKind.TENSOR_LIST: slots[field.slots[0].name] = value case _FieldKind.TENSOR_OR_QUANTIZED: _pack_tensor_or_quantized(field, value, slots) @@ -665,7 +678,7 @@ def unpack(self, slots: Dict[str, Any]) -> Any: bundle = slots.get(_SIMPLE_META_SLOT) for field in self.fields: match field.kind: - case _FieldKind.TENSOR: + case _FieldKind.TENSOR | _FieldKind.TENSOR_LIST: kwargs[field.name] = slots[field.slots[0].name] case _FieldKind.TENSOR_OR_QUANTIZED: kwargs[field.name] = _unpack_tensor_or_quantized(field, slots) @@ -741,7 +754,9 @@ def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any: overrides: Dict[str, Any] = {} for name in tensor_field_names: value = getattr(obj, name, None) - if value is not None and not isinstance(value, TensorSpec): + if isinstance(value, list): + overrides[name] = [to_tensor_spec(tensor) for tensor in value] + elif value is not None and not isinstance(value, TensorSpec): overrides[name] = to_tensor_spec(value) if not overrides: return obj @@ -1030,10 +1045,13 @@ def _register_autograd_for_op( the plan on ``ctx`` so backward can slice its grads per user output. """ bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields) + fwd_slot_indices = {name: i for i, name in enumerate(fwd_plan.slot_names)} def _setup_context(ctx, inputs, output): ctx.fwd_tensor_list_lengths = { - i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) + name: len(value) + for name, value in zip(fwd_plan.slot_names, inputs) + if isinstance(value, list) } fwd_obj = fwd_plan.unpack(dict(zip(fwd_plan.slot_names, inputs))) spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names()) @@ -1076,14 +1094,25 @@ def _autograd_backward(ctx, *grad_outputs): kwargs = bwd_plan.pack(bwd_obj) bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] + if hasattr(bwd_obj, "postprocess_gradients"): + grads = bwd_obj.postprocess_gradients(grads) + if len(grads) != len(grad_targets): + raise RuntimeError("Expected one gradient per input_tensors_for_grad entry") ctx.backward_objects = None # One grad per input schema slot: default None, but a ``Tensor[]`` slot # (always recorded in ``fwd_tensor_list_lengths``) needs a # list-shaped no-grad of matching length. out: List[Any] = [None] * len(fwd_plan.slot_names) - for pos, length in ctx.fwd_tensor_list_lengths.items(): - out[pos] = [None] * length + for name, length in ctx.fwd_tensor_list_lengths.items(): + out[fwd_slot_indices[name]] = [None] * length for pos, g in zip(grad_targets, grads): + name = fwd_plan.slot_names[pos] + if name in ctx.fwd_tensor_list_lengths: + if g is not None: + if not isinstance(g, list) or len(g) != ctx.fwd_tensor_list_lengths[name]: + raise RuntimeError("Expected one gradient per tensor-list element") + out[pos] = g + continue if g is not None: shape = ctx.grad_input_shapes.get(pos) if shape is not None and g.shape != shape: @@ -1354,9 +1383,10 @@ def register_custom_op_with_autograd( * ``setup_context(bwd_obj, fwd_args, user_outputs, ctx_attrs, saved) -> tensors_to_save`` -- populate ``bwd_obj`` from forward state; return the tensors to persist across the boundary. - * ``bwd_impl(bwd_args) -> grads`` -- exactly one grad per - ``input_tensors_for_grad`` entry, in that order (``None`` for a - non-differentiable input). + * ``bwd_impl(bwd_args) -> grads`` -- flat tensor grads (or ``None``). + Without a ``bwd_arg_type.postprocess_gradients(grads)`` hook, return one + grad per ``input_tensors_for_grad`` entry. With the hook, it maps the + flat results to those entries outside the op, including tensor-list grads. * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning :class:`TensorSpec` grads. * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook; skipped if @@ -1430,7 +1460,9 @@ def _register_custom_op_with_autograd_impl( tags=fwd_tags, ) bwd_qualname = f"{_TE_OP_NAMESPACE}::{op_name}_backward_base" - num_grad_inputs = len(input_tensors_for_grad) + num_grad_inputs = ( + None if hasattr(bwd_arg_type, "postprocess_gradients") else len(input_tensors_for_grad) + ) bwd_op = _register_op( name=f"{op_name}_backward", arg_type=bwd_arg_type, diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 3d04882698a..6dbad7fad9b 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -6,14 +6,16 @@ import dataclasses import queue -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Any, Callable, List, Optional, Sequence, Tuple, Union import torch from .. import cpp_extensions as tex from ..constants import TE_DType from ..distributed import in_fp8_activation_recompute_phase +from ..dynamo import TensorOrQuantized from ..export import is_in_onnx_export_mode +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage from ..quantization import FP8GlobalStateManager from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method @@ -239,6 +241,27 @@ def noop_cat( return _NoopCatFunc.apply(dim, *tensors) +def prepare_parameter( + tensors: Sequence[TensorOrQuantized], *, defer: bool = False +) -> Tuple[Optional[TensorOrQuantized], Optional[List[torch.Tensor]]]: + """Return a parameter operand or its deferred parts.""" + if len(tensors) == 1: + return tensors[0], None + if ( + defer + and tensors + and all( + isinstance(t, torch.Tensor) + and not isinstance(t, (QuantizedTensor, QuantizedTensorStorage)) + for t in tensors + ) + ): + return None, list(tensors) + if torch.compiler.is_compiling(): + return torch.cat(tensors), None + return noop_cat(tensors), None + + @dataclasses.dataclass class _ParameterInitMeta: """ diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index bf4f9bba1d5..a241bfda380 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -11,7 +11,7 @@ import warnings from enum import Enum from abc import ABC, abstractmethod -from typing import Any, Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union from contextlib import contextmanager from types import MethodType @@ -1395,7 +1395,9 @@ def _get_weight_quantizers(self) -> List[Quantizer]: def _enable_weight_preswizzle( self, quantizer: Quantizer, - weight: torch.Tensor, + weight: Optional[torch.Tensor], + *, + weight_shape: Optional[Sequence[int]] = None, ) -> bool: """Whether to fuse scale-factor swizzling into weight quantization. @@ -1412,7 +1414,8 @@ def _enable_weight_preswizzle( if isinstance(quantizer, MXFP8Quantizer): return True if isinstance(quantizer, NVFP4Quantizer): - rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1] + shape = weight.shape if weight_shape is None else weight_shape + rows, cols = math.prod(shape[:-1]), shape[-1] arch_supported = get_device_compute_capability() >= (10, 0) if quantizer.with_rht: return arch_supported and rows % 64 == 0 and cols % 128 == 0 diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 77f008f4ecb..d14fb25efe7 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -66,7 +66,7 @@ from ._common import ( apply_normalization, check_fp8_reduce_and_update, - noop_cat, + prepare_parameter, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -78,6 +78,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo.concat import concat_shape, materialize_concat from ..dynamo import TensorOrQuantized from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -108,7 +109,7 @@ class LayerNormLinearFwdArgs: inp: torch.Tensor ln_weight: torch.Tensor ln_bias: Optional[torch.Tensor] - weight: TensorOrQuantized + weight: Optional[TensorOrQuantized] bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- @@ -180,6 +181,8 @@ class LayerNormLinearFwdArgs: # --- Misc --- cpu_offloading: bool is_grad_enabled: bool + weight_list: Optional[List[torch.Tensor]] = None + bias_list: Optional[List[torch.Tensor]] = None def any_requires_grad(self) -> bool: """Whether any differentiable input requires a gradient.""" @@ -314,8 +317,11 @@ def _layernorm_linear_forward_impl( inp = args.inp ln_weight = args.ln_weight ln_bias = args.ln_bias - weight = args.weight - bias = args.bias + weight = materialize_concat(args.weight_list) if args.weight_list else args.weight + bias = materialize_concat(args.bias_list) if args.bias_list else args.bias + if args.weight_list or args.bias_list: + args.weight, args.bias = weight, bias + args.weight_list = args.bias_list = None input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer output_quantizer = args.output_quantizer @@ -1475,9 +1481,10 @@ def forward( inp: torch.Tensor, ln_weight: torch.Tensor, ln_bias: Optional[torch.Tensor], - weight: torch.Tensor, + weight: Optional[torch.Tensor], bias: Optional[torch.Tensor], fwd_args: LayerNormLinearFwdArgs, + *parameter_parts: torch.Tensor, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: """Forward pass: compute the output and set up the autograd context. @@ -1492,6 +1499,14 @@ def forward( fwd_args.ln_bias = ln_bias fwd_args.weight = weight fwd_args.bias = bias + if parameter_parts: + weight_count = len(fwd_args.weight_list or ()) + fwd_args.weight_list = list(parameter_parts[:weight_count]) if weight_count else None + fwd_args.bias_list = list(parameter_parts[weight_count:]) or None + ctx.parameter_split_sizes = ( + tuple(t.shape[0] for t in fwd_args.weight_list or ()), + tuple(t.shape[0] for t in fwd_args.bias_list or ()), + ) ( out, ln_out_return, @@ -1544,14 +1559,24 @@ def backward( nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") - return ( - dgrad, - dgamma, - dbeta, - wgrad, - grad_bias, - None, # fwd_args - ) + if hasattr(ctx, "parameter_split_sizes"): + weight_sizes, bias_sizes = ctx.parameter_split_sizes + part_grads = [] + for grad, sizes in ((wgrad, weight_sizes), (grad_bias, bias_sizes)): + if sizes: + part_grads.extend( + grad.split(sizes) if grad is not None else (None,) * len(sizes) + ) + return ( + dgrad, + dgamma, + dbeta, + None if weight_sizes else wgrad, + None if bias_sizes else grad_bias, + None, + *part_grads, + ) + return dgrad, dgamma, dbeta, wgrad, grad_bias, None class LayerNormLinear(TransformerEngineBaseModule): @@ -2079,7 +2104,7 @@ def forward( try: # Get concatenated weight and bias tensors - weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + weight, bias, weight_list, bias_list = self._get_weight_and_bias() quantizers = ( self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) @@ -2101,7 +2126,9 @@ def forward( ) = quantizers if weight_quantizer is not None and not debug: weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( - weight_quantizer, weight_tensor + weight_quantizer, + weight, + weight_shape=concat_shape(weight_list) if weight_list else None, ) cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" @@ -2136,16 +2163,16 @@ def forward( ub_bulk_wgrad = self.ub_bulk_wgrad ub_bulk_dgrad = self.ub_bulk_dgrad - linear_bias_tensor = ( - bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None - ) + linear_bias = bias if (self.apply_bias and not self.gemm_bias_unfused_add) else None fwd_args = LayerNormLinearFwdArgs( # tensors inp=inp, ln_weight=self.layer_norm_weight, ln_bias=self.layer_norm_bias, - weight=weight_tensor, - bias=linear_bias_tensor, + weight=weight, + bias=linear_bias, + weight_list=weight_list, + bias_list=bias_list, weight_workspace=weight_workspace, # requires_grad flags input_requires_grad=inp.requires_grad, @@ -2155,9 +2182,15 @@ def forward( if self.layer_norm_bias is not None else False ), - weight_requires_grad=weight_tensor.requires_grad, + weight_requires_grad=( + any(t.requires_grad for t in weight_list) + if weight_list + else weight.requires_grad + ), bias_requires_grad=( - linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False + any(t.requires_grad for t in bias_list) + if bias_list + else linear_bias is not None and linear_bias.requires_grad ), # quantizers input_quantizer=input_quantizer, @@ -2214,13 +2247,24 @@ def forward( is_grad_enabled=is_grad_enabled, ) - if is_grad_enabled: + if is_grad_enabled and (weight_list or bias_list): out, ln_out, new_weight_workspace = _LayerNormLinear.apply( inp, self.layer_norm_weight, self.layer_norm_bias, - weight_tensor, - linear_bias_tensor, + weight, + linear_bias, + fwd_args, + *(weight_list or ()), + *(bias_list or ()), + ) + elif is_grad_enabled: + out, ln_out, new_weight_workspace = _LayerNormLinear.apply( + inp, + self.layer_norm_weight, + self.layer_norm_bias, + weight, + linear_bias, fwd_args, ) else: @@ -2229,8 +2273,8 @@ def forward( inp, self.layer_norm_weight, self.layer_norm_bias, - weight_tensor, - linear_bias_tensor, + weight, + linear_bias, fwd_args, ) @@ -2243,12 +2287,12 @@ def forward( self.end_forward() if self.gemm_bias_unfused_add: - out = out + cast_if_needed(bias_tensor, self.activation_dtype) + out = out + cast_if_needed(bias, self.activation_dtype) if self.return_bias: if self.return_layernorm_output: - return out, cast_if_needed(bias_tensor, self.activation_dtype), ln_out - return out, cast_if_needed(bias_tensor, self.activation_dtype) + return out, cast_if_needed(bias, self.activation_dtype), ln_out + return out, cast_if_needed(bias, self.activation_dtype) if self.return_layernorm_output: return out, ln_out return out @@ -2298,16 +2342,17 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): for name, q in zip(names, original_quantizers) ) - def _get_weight_and_bias_tensors(self): - # Get concatenated weight and bias tensors - unfused_weights = self._get_weight_tensors() - - weight_tensor = noop_cat(unfused_weights) + def _get_weight_and_bias(self): + weight, weight_list = prepare_parameter(self._get_weight_tensors(), defer=True) + bias, bias_list = None, None if self.use_bias: - bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names]) + bias, bias_list = prepare_parameter( + [getattr(self, name) for name in self.bias_names], + defer=self.apply_bias and not self.gemm_bias_unfused_add, + ) else: - bias_tensor = getattr(self, self.bias_names[0]) # Unused - return weight_tensor, bias_tensor + bias = getattr(self, self.bias_names[0]) + return weight, bias, weight_list, bias_list def onnx_forward( self, @@ -2332,7 +2377,11 @@ def onnx_forward( ) = self._get_quantizers(fp8_output, False, is_grad_enabled) inp_dtype = inp.dtype - weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias() + if weight_list: + weight_tensor = torch.cat(weight_list) + if bias_list: + bias_tensor = torch.cat(bias_list) ln_out, ln_out_return = onnx_layernorm( inp, self.layer_norm_weight, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 5171743ec08..3ddbcd7571e 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -34,7 +34,7 @@ from ._common import ( can_reconstruct_wgrad_input_from_original, check_fp8_reduce_and_update, - noop_cat, + prepare_parameter, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -85,6 +85,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo.concat import concat_shape, concat_to_spec, materialize_concat from ..dynamo import ( TensorSpec, TensorOrQuantized, @@ -112,7 +113,7 @@ class LinearFwdArgs: """Single-argument bag for the forward path of :class:`_Linear`.""" # --- Differentiable tensors (also passed positionally to autograd) --- - weight: TensorOrQuantized + weight: Optional[TensorOrQuantized] inp: torch.Tensor bias: Optional[torch.Tensor] @@ -180,6 +181,8 @@ class LinearFwdArgs: # --- Misc --- cpu_offloading: bool is_grad_enabled: bool + weight_list: Optional[List[torch.Tensor]] = None + bias_list: Optional[List[torch.Tensor]] = None def compile_unsupported_reason(self) -> Optional[str]: """Reason this config can't use the torch.compile custom-op path (else None).""" @@ -240,6 +243,8 @@ class LinearBwdArgs: weight_fp8: Optional[TensorOrQuantized] = None saved_weight: Optional[TensorOrQuantized] = None bias: Optional[torch.Tensor] = None + weight_list: Optional[List[torch.Tensor]] = None + bias_list: Optional[List[torch.Tensor]] = None # --- Quantizers --- input_quantizer: Optional[Quantizer] = None @@ -305,14 +310,31 @@ class LinearBwdArgs: def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: """Pull saved tensors from ``ctx`` into the fields backward consumes.""" - ( - self.inputmat, - self.weight_fp8, - self.saved_weight, - self.bias, - ) = restore_from_func_ctx( - ctx - ) # pylint: disable=unbalanced-tuple-unpacking + saved = restore_from_func_ctx(ctx) + if len(saved) == 4: + self.inputmat, self.weight_fp8, self.saved_weight, self.bias = saved + return + self.inputmat, self.weight_fp8, self.saved_weight, self.bias = saved[:4] + weight_count = ctx.fwd_tensor_list_lengths.get("weight_list", 0) + self.weight_list = saved[4 : 4 + weight_count] if weight_count else None + self.bias_list = saved[4 + weight_count :] or None + + def postprocess_gradients(self, grads): + """Split gradients outside the op so the parts can share storage.""" + wgrad, dgrad, bgrad = grads + weight_grads = None + bias_grads = None + if self.weight_list and wgrad is not None: + weight_grads = list(wgrad.split([t.shape[0] for t in self.weight_list])) + if self.bias_list and bgrad is not None: + bias_grads = list(bgrad.split([t.shape[0] for t in self.bias_list])) + return ( + None if self.weight_list else wgrad, + dgrad, + None if self.bias_list else bgrad, + weight_grads, + bias_grads, + ) def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: @@ -351,6 +373,8 @@ def _fake_workspace_valid(workspace: TensorSpec, quantizer: Optional[Quantizer]) def _linear_forward_impl( args: LinearFwdArgs, + *, + save_materialized_parameters: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: """Forward implementation for the linear layer. @@ -360,10 +384,13 @@ def _linear_forward_impl( cache). The last two are ``None`` when gradients are disabled. """ - weight = args.weight - is_dist_weight = is_distributed_weight(args.weight) + weight = materialize_concat(args.weight_list) if args.weight_list else args.weight + is_dist_weight = is_distributed_weight(weight) inp = args.inp - bias = args.bias + bias = materialize_concat(args.bias_list) if args.bias_list else args.bias + if save_materialized_parameters and (args.weight_list or args.bias_list): + args.weight, args.bias = weight, bias + args.weight_list = args.bias_list = None input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer output_quantizer = args.output_quantizer @@ -771,9 +798,9 @@ def _linear_forward_fake( "(fsdp_group is not None); use FSDP2 or MCore FSDP." ) - weight = args.weight + weight = concat_to_spec(args.weight_list) if args.weight_list else args.weight inp = args.inp - bias = args.bias + bias = concat_to_spec(args.bias_list) if args.bias_list else args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer output_quantizer = args.output_quantizer @@ -1015,7 +1042,7 @@ def _linear_setup_ctx( bwd_args.grad_output_quantizer = fwd_args.grad_output_quantizer # Differentiability summary - bwd_args.use_bias = bias is not None + bwd_args.use_bias = bias is not None or bool(fwd_args.bias_list) bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True) @@ -1098,7 +1125,16 @@ def _linear_setup_ctx( saved_weight = weight if bias_alias == "bias": saved_bias = bias - return (saved_inputmat, wt_save, saved_weight, saved_bias) + if not fwd_args.weight_list and not fwd_args.bias_list: + return saved_inputmat, wt_save, saved_weight, saved_bias + return ( + saved_inputmat, + wt_save, + saved_weight, + saved_bias, + *(fwd_args.weight_list or ()), + *(fwd_args.bias_list or ()), + ) def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: @@ -1111,10 +1147,10 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None grad_output = args.grad_output assert grad_output is not None inputmat = args.inputmat + saved_weight = materialize_concat(args.weight_list) if args.weight_list else args.saved_weight weight_fp8 = args.weight_fp8 - saved_weight = args.saved_weight is_dist_weight = is_distributed_weight(saved_weight) - bias = args.bias + bias = materialize_concat(args.bias_list) if args.bias_list else args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer grad_input_quantizer = args.grad_input_quantizer @@ -1355,6 +1391,8 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None elif bwd_args.weight_quantizer is not None: bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) + elif args.weight_list: + weight_fp8 = cast_if_needed(saved_weight, bwd_args.activation_dtype) elif ( is_dist_weight and bwd_args.fp8 @@ -1756,7 +1794,7 @@ def _linear_backward_fake( "(fsdp_group is not None); use FSDP2 or MCore FSDP." ) - weight = args.saved_weight + weight = concat_to_spec(args.weight_list) if args.weight_list else args.saved_weight out_dtype = args.activation_dtype out_features, in_features = weight.shape @@ -1811,7 +1849,7 @@ def _linear_backward_fake( # Custom op used under ``torch.compile``. _linear_op = register_custom_op_with_autograd( op_name="linear", - input_tensors_for_grad=["weight", "inp", "bias"], + input_tensors_for_grad=["weight", "inp", "bias", "weight_list", "bias_list"], fwd_arg_type=LinearFwdArgs, fwd_impl=_linear_forward_impl, fwd_fake_impl=_linear_forward_fake, @@ -1830,16 +1868,16 @@ class _Linear(torch.autograd.Function): @staticmethod def forward( ctx, - weight: torch.Tensor, + weight: Optional[torch.Tensor], inp: torch.Tensor, bias: Optional[torch.Tensor], fwd_args: LinearFwdArgs, + *parameter_parts: torch.Tensor, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Forward pass: compute linear output and set up autograd context. - ``weight``, ``inp`` and ``bias`` are positional Tensor arguments so - autograd tracks them; they are immediately re-attached to ``fwd_args`` - so every downstream helper can be invoked with a single argument. + Tensor operands, including split parameter parts, are explicit + autograd inputs and are re-attached to ``fwd_args``. ``weight_workspace`` is intentionally NOT a positional input: it is a non-differentiable cached tensor passed in via @@ -1849,12 +1887,20 @@ def forward( fwd_args.weight = weight fwd_args.inp = inp fwd_args.bias = bias + if parameter_parts: + weight_count = len(fwd_args.weight_list or ()) + fwd_args.weight_list = list(parameter_parts[:weight_count]) if weight_count else None + fwd_args.bias_list = list(parameter_parts[weight_count:]) or None + ctx.parameter_split_sizes = ( + tuple(t.shape[0] for t in fwd_args.weight_list or ()), + tuple(t.shape[0] for t in fwd_args.bias_list or ()), + ) ( out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs, - ) = _linear_forward_impl(fwd_args) + ) = _linear_forward_impl(fwd_args, save_materialized_parameters=True) if ctx is not None: bwd_args = LinearBwdArgs() tensors_to_save_from_setup = _linear_setup_ctx( @@ -1893,6 +1939,22 @@ def backward( if bwd_args.ub_name is not None: nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" result = _linear_backward_impl(bwd_args) + (None,) # fwd_args grad slot + if hasattr(ctx, "parameter_split_sizes"): + wgrad, dgrad, bgrad, _ = result + weight_sizes, bias_sizes = ctx.parameter_split_sizes + part_grads = [] + for grad, sizes in ((wgrad, weight_sizes), (bgrad, bias_sizes)): + if sizes: + part_grads.extend( + grad.split(sizes) if grad is not None else (None,) * len(sizes) + ) + result = ( + None if weight_sizes else wgrad, + dgrad, + None if bias_sizes else bgrad, + None, + *part_grads, + ) reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, # main_grad closure) so they don't outlive backward via ctx under retain_graph. @@ -1907,7 +1969,7 @@ def backward( @no_torch_dynamo() def _linear_eager( - weight_tensor: torch.Tensor, + weight_tensor: Optional[torch.Tensor], inp: torch.Tensor, bias: Optional[torch.Tensor], fwd_args: LinearFwdArgs, @@ -1915,6 +1977,15 @@ def _linear_eager( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Run ``_Linear`` eagerly, bypassing Dynamo.""" if is_grad_enabled: + if fwd_args.weight_list or fwd_args.bias_list: + return _Linear.apply( + weight_tensor, + inp, + bias, + fwd_args, + *(fwd_args.weight_list or ()), + *(fwd_args.bias_list or ()), + ) return _Linear.apply(weight_tensor, inp, bias, fwd_args) return _Linear.forward(None, weight_tensor, inp, bias, fwd_args) @@ -1952,7 +2023,9 @@ class Linear(TransformerEngineBaseModule): (preferably an OrderedDict) is provided, the keys are used as names and values as split sizes along dim 0. The resulting parameters will have names that end in ``_weight`` or ``_bias``, so trailing underscores are - stripped from any provided names. + stripped from any provided names. Under ``torch.compile``, adjacent + parts sharing storage are consumed without a concatenation copy. + Disjoint parts and returned split biases still require concatenation. device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the @@ -2415,7 +2488,9 @@ def forward( if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True - if torch.compiler.is_compiling() and _linear_op is not None: + is_compiling = torch.compiler.is_compiling() + use_compiled_op = is_compiling and _linear_op is not None + if use_compiled_op: reason = self._compile_eager_fallback_reason( inp, skip_fp8_weight_update, fp8_output, fp8_grad, is_grad_enabled, debug ) @@ -2427,7 +2502,11 @@ def forward( inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor)) try: - weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + weight, bias, weight_list, bias_list = self._get_weight_and_bias() + weight_shape = concat_shape(weight_list) if weight_list else None + weight_requires_grad = ( + any(t.requires_grad for t in weight_list) if weight_list else weight.requires_grad + ) quantizers = ( self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) @@ -2449,11 +2528,10 @@ def forward( ) = quantizers if weight_quantizer is not None and not debug: weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( - weight_quantizer, weight_tensor + weight_quantizer, weight, weight_shape=weight_shape ) - use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None - if _linear_op is None and torch.compiler.is_compiling(): + if _linear_op is None and is_compiling: warn_if_compile_disabled() if use_compiled_op: # Process groups cross the op boundary separately from quantizers. @@ -2480,9 +2558,7 @@ def forward( backward_override = None custom = is_custom(input_quantizer) or is_custom(weight_quantizer) backward_input_needs_gather = ( - weight_tensor.requires_grad - and self.parallel_mode == "column" - and self.sequence_parallel + weight_requires_grad and self.parallel_mode == "column" and self.sequence_parallel ) if debug: @@ -2500,22 +2576,24 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad - linear_bias_tensor = ( - bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None - ) + linear_bias = bias if (self.apply_bias and not self.gemm_bias_unfused_add) else None wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None fwd_args = LinearFwdArgs( # tensors - weight=weight_tensor, + weight=weight, inp=inp, - bias=linear_bias_tensor, + bias=linear_bias, + weight_list=weight_list, + bias_list=bias_list, weight_workspace=weight_workspace, # requires_grad flags input_requires_grad=inp.requires_grad, - weight_requires_grad=weight_tensor.requires_grad, + weight_requires_grad=weight_requires_grad, bias_requires_grad=( - linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False + any(t.requires_grad for t in bias_list) + if bias_list + else linear_bias is not None and linear_bias.requires_grad ), # quantizers input_quantizer=input_quantizer, @@ -2577,11 +2655,13 @@ def forward( use_compiled_op = False if use_compiled_op: - check_gemm_dims(inp, weight_tensor) + check_gemm_dims( + inp.shape, weight_shape if weight_shape is not None else weight.shape + ) out, new_weight_workspace = _linear_op(fwd_args) else: out, new_weight_workspace = _linear_eager( - weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled + weight, inp, linear_bias, fwd_args, is_grad_enabled ) if new_weight_workspace is not None and cache_name is not None: @@ -2592,10 +2672,10 @@ def forward( finally: self.end_forward() if self.gemm_bias_unfused_add: - out = out + cast_if_needed(bias_tensor, self.activation_dtype) + out = out + cast_if_needed(bias, self.activation_dtype) if self.return_bias: - return out, cast_if_needed(bias_tensor, self.activation_dtype) + return out, cast_if_needed(bias, self.activation_dtype) return out def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): @@ -2672,8 +2752,7 @@ def _compile_eager_fallback_reason( prepare_forward. Quantizer checks stay in compile_unsupported_reason.""" if debug: return "debug instrumentation (nvidia-dlfw-inspect)" - weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() - if is_distributed_weight(weight_tensor): + if any(is_distributed_weight(getattr(self, name)) for name in self.weight_names): return "a DistributedWeight (custom weight parallelism, e.g. GTP)" if isinstance(inp, (QuantizedTensor, QuantizedTensorStorage)): return "a quantized input tensor" @@ -2681,8 +2760,10 @@ def _compile_eager_fallback_reason( return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" any_requires_grad = ( inp.requires_grad - or weight_tensor.requires_grad - or (bias_tensor is not None and bias_tensor.requires_grad) + or any(getattr(self, name).requires_grad for name in self.weight_names) + or ( + self.use_bias and any(getattr(self, name).requires_grad for name in self.bias_names) + ) ) if fp8_output and is_grad_enabled and any_requires_grad: return "differentiable fp8_output=True" @@ -2722,15 +2803,18 @@ def _forward_eager_fallback( fp8_grad=fp8_grad, ) - def _get_weight_and_bias_tensors(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - # Get concatenated weight and bias tensors - unfused_weights = self._get_weight_tensors() - weight_tensor = noop_cat(unfused_weights) + def _get_weight_and_bias(self): + weights = self._get_weight_tensors() + weight, weight_list = prepare_parameter(weights, defer=True) + bias = None + bias_list = None if self.use_bias: - bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names]) - else: - bias_tensor = None - return weight_tensor, bias_tensor + biases = [getattr(self, name) for name in self.bias_names] + bias, bias_list = prepare_parameter( + biases, + defer=self.apply_bias and not self.gemm_bias_unfused_add, + ) + return weight, bias, weight_list, bias_list def onnx_forward( self, @@ -2747,7 +2831,11 @@ def onnx_forward( assert_warmed_up(self) assert not TEDebugState.debug_enabled, "Debug mode is not supported in ONNX export." - weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() + weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias() + if weight_list: + weight_tensor = torch.cat(weight_list) + if bias_list: + bias_tensor = torch.cat(bias_list) ( input_quantizer, weight_quantizer, diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index d3bf57d7ada..c2c584b27fa 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -688,14 +688,14 @@ def cast_if_needed(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return tensor.to(dtype=dtype) -def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor) -> None: +def check_gemm_dims(inp_shape: Sequence[int], weight_shape: Sequence[int]) -> None: """Check contracting dimensions at trace time for ``torch.compile``. Recipe-specific alignment requirements are checked by the runtime implementation. """ # pylint: disable=protected-access torch._check( - inp.shape[-1] == weight.shape[-1], + inp_shape[-1] == weight_shape[-1], lambda: "GEMM not possible: input last dim must equal in_features", )