diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8946b5c8f9..07b0cff874 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1436,7 +1436,10 @@ def ref_fn(inp): torch.testing.assert_close(out, out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) torch.testing.assert_close( - model.weight.grad, ref_model.weight.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL + model.weight.grad, + ref_model.weight.grad, + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, ) workspace = model._fp8_workspaces.get("weight") @@ -1560,3 +1563,567 @@ def fn(inp): "Unexpected recompilation(s) across different batch sizes: " f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" ) + + +# --------------------------------------------------------------------------- +# te.GroupedLinear +# --------------------------------------------------------------------------- + +_GROUPED_M_SPLITS = [32, 16, 48, 32] # FP8-legal splits (each divisible by 8) +_GROUPED_NUM_GEMMS = len(_GROUPED_M_SPLITS) +_GROUPED_IN, _GROUPED_OUT = 64, 32 + + +def _grouped_input(dtype, device, requires_grad=False): + return torch.randn( + sum(_GROUPED_M_SPLITS), + _GROUPED_IN, + dtype=dtype, + device=device, + requires_grad=requires_grad, + ) + + +def _assert_close_grouped(fn, compiled, model, base): + """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the + forward output, the input gradient and every expert's weight / bias + gradients match.""" + num_gemms = model.num_gemms + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = fn(inp_eager) + out_eager.sum().backward() + ref_out = out_eager.detach().clone() + ref_igrad = inp_eager.grad.detach().clone() + ref_wgrads = [getattr(model, f"weight{i}").grad.detach().clone() for i in range(num_gemms)] + ref_bgrads = [getattr(model, f"bias{i}").grad.detach().clone() for i in range(num_gemms)] + + inp_compiled = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + # Clone before a later cuda-graph replay overwrites the static output buffer. + out_compiled = compiled(inp_compiled).clone() + out_compiled.sum().backward() + + torch.testing.assert_close(out_compiled, ref_out, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp_compiled.grad, ref_igrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + for i in range(num_gemms): + torch.testing.assert_close( + getattr(model, f"weight{i}").grad, + ref_wgrads[i], + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + msg=f"wgrad mismatch for expert {i}", + ) + torch.testing.assert_close( + getattr(model, f"bias{i}").grad, + ref_bgrads[i], + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + msg=f"bias grad mismatch for expert {i}", + ) + + +@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_grouped_linear_compiles(fp8_recipe, compile_mode): + """ + torch.compile(fullgraph=True) of ``te.GroupedLinear`` under every built-in + recipe (plus the bf16-only baseline), for both the default backend and + ``mode="reduce-overhead"`` (CUDA-graph trees). + """ + dtype = torch.bfloat16 + device = "cuda" + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + + def fn(inp): + if fp8_recipe is None: + return model(inp, _GROUPED_M_SPLITS) + with te.autocast(recipe=fp8_recipe): + return model(inp, _GROUPED_M_SPLITS) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup(fn, _grouped_input(dtype, device, requires_grad=True), backward=True) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_grouped_linear_compile_with_quantized_fp8_weight(compile_mode): + """torch.compile of GroupedLinear with the weights initialized as FP8 + tensors (exercises the wrapper op's input flattening inside a ``Tensor[]`` + slot group).""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + with te.quantized_model_init(enabled=True, recipe=fp8_recipe): + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, + _GROUPED_IN, + _GROUPED_OUT, + params_dtype=dtype, + device=device, + ) + assert isinstance(model.weight0, te.Float8Tensor) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, _GROUPED_M_SPLITS) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup(fn, _grouped_input(dtype, device, requires_grad=True), backward=True) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_grouped_linear_compile_is_first_microbatch(compile_mode): + """torch.compile of ``te.GroupedLinear`` across a microbatch schedule: + ``is_first_microbatch=True`` caches every expert's FP8 weight, later steps + must reuse the caches and stay numerically aligned with eager.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + ref_model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + with torch.no_grad(): + for i in range(_GROUPED_NUM_GEMMS): + getattr(ref_model, f"weight{i}").copy_(getattr(model, f"weight{i}")) + getattr(ref_model, f"bias{i}").copy_(getattr(model, f"bias{i}")) + + schedule = [True, False, False] + is_first = schedule[0] # rebound each step; closed over by the fns. + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, _GROUPED_M_SPLITS, is_first_microbatch=is_first) + + def ref_fn(inp): + with te.autocast(recipe=fp8_recipe): + return ref_model(inp, _GROUPED_M_SPLITS, is_first_microbatch=is_first) + + # Eager priming: FP8 state must exist before tracing (creating quantizers + # in-graph breaks later recompiles; upstream Dynamo bug). + is_first = None + fn(_grouped_input(dtype, device, requires_grad=True)) + is_first = schedule[0] + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup(fn, _grouped_input(dtype, device, requires_grad=True), backward=True) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + cached_workspaces = None + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for step, is_first in enumerate(schedule): + base = _grouped_input(dtype, device) + + inp_ref = base.detach().clone().requires_grad_(True) + ref_model.zero_grad(set_to_none=True) + out_ref = ref_fn(inp_ref) + out_ref.sum().backward() + + inp = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out = compiled(inp).clone() + out.sum().backward() + + torch.testing.assert_close(out, out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + for i in range(_GROUPED_NUM_GEMMS): + torch.testing.assert_close( + getattr(model, f"weight{i}").grad, + getattr(ref_model, f"weight{i}").grad, + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, + ) + + workspaces = [ + model._fp8_workspaces.get(f"weight{i}") for i in range(_GROUPED_NUM_GEMMS) + ] + assert all( + ws is not None for ws in workspaces + ), f"missing cached FP8 weight(s) after step {step}" + if step == 0: + cached_workspaces = workspaces + else: + for i in range(_GROUPED_NUM_GEMMS): + assert ( + workspaces[i] is cached_workspaces[i] + ), f"cache for expert {i} rebuilt at step {step}" + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_te_grouped_linear_m_splits_change(): + """Unchanged ``m_splits`` reuse the graph. A changed distribution makes + Dynamo mark the values dynamic; they cross the op boundary as a + ``SymInt[]`` slot, so the (recompiled, now symbolic) graph still runs the + compiled path with eager numerics -- for any later distribution too.""" + dtype = torch.bfloat16 + device = "cuda" + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + m_splits = list(_GROUPED_M_SPLITS) + + def fn(inp): + return model(inp, m_splits) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + # Two warmup calls: the second absorbs the one-time recompile from module + # attributes lazily created during call one. + for _ in range(2): + compiled(_grouped_input(dtype, device, requires_grad=True)).sum().backward() + model.zero_grad(set_to_none=True) + baseline = _dynamo_counter("stats", "unique_graphs") + + # Same splits: no recompile. + compiled(_grouped_input(dtype, device, requires_grad=True)).sum().backward() + model.zero_grad(set_to_none=True) + if baseline: + assert _dynamo_counter("stats", "unique_graphs") == baseline + + # New split distributions (same total): compiled path, same numerics. + m_splits[:2] = [16, 32] + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + m_splits[2:] = [40, 40] + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.xfail(reason="waiting for a PyTorch fix", strict=False) +def test_te_grouped_linear_compile_train_eval_switch(): + """train -> eval -> train on the same compiled ``te.GroupedLinear``, vs eager.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, _GROUPED_M_SPLITS, is_first_microbatch=True) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + def train_step(): + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + model.zero_grad(set_to_none=True) + + train_step() + + model.eval() + x = _grouped_input(dtype, device) + with torch.no_grad(): + out_eval = compiled(x) + ref_eval = fn(x) + torch.testing.assert_close(out_eval, ref_eval, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + model.train() + train_step() + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_te_grouped_linear_compile_m_splits_tensor_falls_back(): + """A device-tensor ``m_splits`` inside the compiled region cannot cross the + op boundary (its values would be unbacked); the module must warn and fall + back to eager (a graph break).""" + dtype = torch.bfloat16 + device = "cuda" + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + m_splits_tensor = torch.tensor(_GROUPED_M_SPLITS, dtype=torch.int64, device="cpu") + + def fn(inp): + return model(inp, m_splits_tensor) + + torch._dynamo.reset() + # Graph-break fallback: numerics must match eager. + compiled = torch.compile(fn) + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + + # fullgraph=True: the explicit graph break surfaces the reason. + torch._dynamo.reset() + strict = torch.compile(fn, fullgraph=True) + with pytest.raises(Exception, match="falling back to eager"): + strict(_grouped_input(dtype, device, requires_grad=True)) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_grouped_linear_compile_delayed_scaling_falls_back(): + """Delayed-scaling quantizers are not value-opaque; under torch.compile the + module must fall back to eager (a graph break) with eager numerics.""" + dtype = torch.bfloat16 + device = "cuda" + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + fp8_recipe = recipe.DelayedScaling() + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, _GROUPED_M_SPLITS) + + torch._dynamo.reset() + compiled = torch.compile(fn) # no fullgraph: the fallback graph-breaks + # Functional smoke only: delayed scaling advances amax history every call, + # so eager-vs-compiled outputs are not bit-comparable across calls. + try: + inp = _grouped_input(dtype, device, requires_grad=True) + out = compiled(inp) + out.sum().backward() + assert out.shape == (sum(_GROUPED_M_SPLITS), _GROUPED_OUT) + assert inp.grad is not None + for i in range(_GROUPED_NUM_GEMMS): + assert getattr(model, f"weight{i}").grad is not None + finally: + # Drop the delayed-scaling amax-reduction registrations, or every later + # autocast exit would run the (untraceable) fused amax reduction. + te.quantization.FP8GlobalStateManager.reset() + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_grouped_linear_compile_save_original_input(): + """torch.compile with ``save_original_input=True``: the op saves the input + by alias and re-splits/re-quantizes it in backward.""" + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, + _GROUPED_IN, + _GROUPED_OUT, + params_dtype=dtype, + device=device, + save_original_input=True, + ) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, _GROUPED_M_SPLITS) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + for _ in range(2): + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + + +_fused_grouped_cc_ok = torch.cuda.is_available() and ( + (9, 0) <= torch.cuda.get_device_capability() <= (11, 0) +) +_fused_grouped_cublas_ok = _fused_grouped_cc_ok and tex.get_cublasLt_version() >= ( + 130400 if torch.cuda.get_device_capability() < (10, 0) else 130300 +) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif( + not _fused_grouped_cublas_ok, + reason="fused grouped GEMM needs Hopper/Blackwell and a recent cuBLASLt", +) +@pytest.mark.parametrize("compile_mode", _compile_modes) +@pytest.mark.parametrize( + "fp8_recipe", + [None] + ([recipe.Float8CurrentScaling()] if fp8_available else []), + ids=lambda r: "bf16" if r is None else type(r).__name__, +) +def test_te_grouped_linear_fused_compiles(fp8_recipe, compile_mode, monkeypatch): + """torch.compile(fullgraph=True) of the fused GroupedTensor path + (``NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1``): ``m_splits`` stays a + device tensor, so changing the split distribution in place reuses the very + same graph -- no recompiles, no host sync.""" + if fp8_recipe is not None and torch.cuda.get_device_capability() < (10, 0): + if tex.get_cublasLt_version() < 130500: + pytest.skip("FP8 current-scaling fused grouped GEMM on Hopper needs cuBLASLt 13.5+") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "1") + dtype = torch.bfloat16 + device = "cuda" + model = te.GroupedLinear( + _GROUPED_NUM_GEMMS, _GROUPED_IN, _GROUPED_OUT, params_dtype=dtype, device=device + ) + m_splits_dev = torch.tensor(_GROUPED_M_SPLITS, dtype=torch.int64, device=device) + + def fn(inp): + if fp8_recipe is None: + return model(inp, m_splits_dev) + with te.autocast(recipe=fp8_recipe): + return model(inp, m_splits_dev) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup(fn, _grouped_input(dtype, device, requires_grad=True), backward=True) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + # Two warmup calls absorb the one-time recompile from lazily created + # module attributes. + for _ in range(2): + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + baseline = _dynamo_counter("stats", "unique_graphs") + + # Change the split distribution IN PLACE: same graph, no recompile. + m_splits_dev.copy_(torch.tensor([16, 32, 40, 40], dtype=torch.int64, device=device)) + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + m_splits_dev.copy_(torch.tensor([64, 8, 32, 24], dtype=torch.int64, device=device)) + _assert_close_grouped(fn, compiled, model, _grouped_input(dtype, device)) + if baseline: + assert ( + _dynamo_counter("stats", "unique_graphs") == baseline + ), "changing device-tensor m_splits values must not recompile" + + +# --------------------------------------------------------------------------- +# Custom-op list adapters (framework unit test via a toy op) +# --------------------------------------------------------------------------- + + +if _opaque_available: + from dataclasses import dataclass + from typing import List, Optional, Tuple + + from transformer_engine.pytorch.dynamo import register_custom_op + + @dataclass + class _ToyListFwdArgs: + inp: torch.Tensor + weights: List[Optional[torch.Tensor]] + scales: Tuple[float, ...] + + @dataclass + class _ToyListBwdArgs: + grad_output: Optional[torch.Tensor] = None + inp: Optional[torch.Tensor] = None + weights: List[Optional[torch.Tensor]] = None + scales: Tuple[float, ...] = () + + def _toy_list_fwd(args): + out = torch.zeros_like(args.inp) + for w, s in zip(args.weights, args.scales): + if w is not None: + out = out + args.inp * w * s + return (out, None, None) + + def _toy_list_fwd_fake(args): + out = TensorSpec( + shape=tuple(args.inp.shape), + dtype=args.inp.dtype, + requires_grad=True, + device=args.inp.device, + ) + return (out, None, None) + + def _toy_list_setup_ctx(bwd_args, fwd_args, outputs, ctx_attrs, saved): + del outputs, ctx_attrs, saved + bwd_args.inp = fwd_args.inp + bwd_args.weights = fwd_args.weights + bwd_args.scales = fwd_args.scales + return () + + def _toy_list_bwd(args): + dgrad = torch.zeros_like(args.inp) + wgrads = [] + for w, s in zip(args.weights, args.scales): + if w is None: + wgrads.append(None) + continue + dgrad = dgrad + args.grad_output * w * s + wgrads.append(args.grad_output * args.inp * s) + return (dgrad, wgrads) + + def _toy_list_bwd_fake(args): + spec = lambda t: TensorSpec(shape=tuple(t.shape), dtype=t.dtype, device=t.device) + dgrad = spec(args.inp) + wgrads = [None if w is None else spec(w) for w in args.weights] + return (dgrad, wgrads) + + _toy_list_op = register_custom_op( + op_name="toy_list_op_test", + input_tensors_for_grad=["inp", "weights"], + fwd_arg_type=_ToyListFwdArgs, + fwd_impl=_toy_list_fwd, + fwd_fake_impl=_toy_list_fwd_fake, + setup_context=_toy_list_setup_ctx, + bwd_arg_type=_ToyListBwdArgs, + bwd_impl=_toy_list_bwd, + bwd_fake_impl=_toy_list_bwd_fake, + ) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_custom_op_tensor_list_grads(): + """A toy op with a ``List[Optional[Tensor]]`` input (including a ``None`` + entry riding the sentinel) compiles fullgraph and routes one gradient per + list element.""" + assert _toy_list_op is not None + device = "cuda" + scales = (2.0, 3.0, 0.5) + + def run(fn): + torch.manual_seed(7) + inp = torch.randn(8, 4, device=device, requires_grad=True) + w0 = torch.randn(8, 4, device=device, requires_grad=True) + w2 = torch.randn(8, 4, device=device, requires_grad=True) + out = fn(inp, w0, w2) + out.sum().backward() + return out.detach().clone(), inp.grad, w0.grad, w2.grad + + def fn(inp, w0, w2): + args = _ToyListFwdArgs(inp=inp, weights=[w0, None, w2], scales=scales) + return _toy_list_op(args) + + def ref_fn(inp, w0, w2): + return inp * w0 * scales[0] + inp * w2 * scales[2] + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + out_ref, igrad_ref, w0grad_ref, w2grad_ref = run(ref_fn) + out, igrad, w0grad, w2grad = run(compiled) + torch.testing.assert_close(out, out_ref) + torch.testing.assert_close(igrad, igrad_ref) + torch.testing.assert_close(w0grad, w0grad_ref) + torch.testing.assert_close(w2grad, w2grad_ref) + + # Eager through the same op (no compile) must agree too. + out_e, igrad_e, w0grad_e, w2grad_e = run(fn) + torch.testing.assert_close(out_e, out_ref) + torch.testing.assert_close(igrad_e, igrad_ref) + torch.testing.assert_close(w0grad_e, w0grad_ref) + torch.testing.assert_close(w2grad_e, w2grad_ref) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 178892c74d..27b1c7faca 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -34,6 +34,11 @@ buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. * ``_QuantizerAdapter`` -- a quantizer, baked into the graph as a value-opaque constant. + * ``_ListAdapter`` -- generic list lift of any of the above: one slot group + for the whole list (per-element entries / concatenated payloads / bundled + metadata); gradients are list-shaped, one grad per element. + * ``_SymIntListAdapter`` -- ``List[int]`` as a ``SymInt[]`` slot, so + host-side ints Dynamo marks dynamic stay compilable. * ``_ProcessGroupAdapter`` -- a ProcessGroup, carried as its c10d registry name and re-resolved inside the op. * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, @@ -335,25 +340,33 @@ def _collect(value: Any) -> None: # --------------------------------------------------------------------------- # -def _storage_flatten( - value: Any, extra_meta: Optional[Dict[str, Any]] = None -) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: - """Split a ``QuantizedTensor`` / bare storage into ``(meta, Tensor[])``. +def _storage_flatten_dict(value: Any) -> Tuple[Dict[str, Any], List[torch.Tensor]]: + """Split a ``QuantizedTensor`` / bare storage into ``(meta_dict, Tensor[])``. The flatten context (embedding the value-opaque quantizer) plus inner names - and -- for a wrapper subclass -- the outer geometry are stashed in the bundle + and -- for a wrapper subclass -- the outer geometry are stashed in the dict so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``. - ``extra_meta`` is merged in before the bundle is built (so its ``_frozen`` - hash key stays consistent) -- used to tag the tensor-or-quantized slot ``__kind__``. """ inner_names, ctx = value.__tensor_flatten__() meta = dict(ctx) meta["_inner_names"] = list(inner_names) if isinstance(value, torch.Tensor): meta["_outer_shape"] = torch.Size(value.shape) + tensors = [getattr(value, name) for name in inner_names] + return meta, tensors + + +def _storage_flatten( + value: Any, extra_meta: Optional[Dict[str, Any]] = None +) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: + """:func:`_storage_flatten_dict` with the metadata wrapped in a bundle. + + ``extra_meta`` is merged in before the bundle is built (so its ``_frozen`` + hash key stays consistent) -- used to tag the tensor-or-quantized slot ``__kind__``. + """ + meta, tensors = _storage_flatten_dict(value) if extra_meta: meta.update(extra_meta) - tensors = [getattr(value, name) for name in inner_names] return OpaqueValueBundle(meta), tensors @@ -404,6 +417,10 @@ class _Adapter: on each call and must agree on the slot layout that ``schema_slots`` declares. """ + # True when the field is list-valued and its grad slot is a ``Tensor[]`` + # slot, so its gradient is a *list* of tensors (one per element). + is_list: bool = False + @classmethod def try_build(cls, name: str, annot: Any) -> Optional["_Adapter"]: """Decide whether this adapter type handles the field ``name`` given its @@ -584,6 +601,38 @@ def grad_slot(self) -> Optional[int]: return 0 +class _SymIntListAdapter(_Adapter): + """``List[int]`` -> one ``SymInt[]`` slot. + + Host-side ints that may become symbolic under torch.compile (e.g. MoE split + sizes marked dynamic across calls): carried as op inputs rather than baked + into a value bundle, so one symbolic graph serves every value set. The real + impl receives plain ints; the fake may see ``SymInt``s. + """ + + def __init__(self, name: str) -> None: + self.name = name + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_SymIntListAdapter"]: + annot, _ = _strip_optional(annot) + if get_origin(annot) is not list: + return None + entry_args = get_args(annot) + if len(entry_args) == 1 and entry_args[0] is int: + return cls(name) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.name, "SymInt[]")] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.name: list(getattr(owner, self.name))} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = list(args[self.name]) + + class _QuantizerAdapter(_Adapter): """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. @@ -620,6 +669,127 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.meta_slot()][self.QUANTIZER_KEY] +class _FieldProxy: + """One-field stand-in so a scalar adapter can read ``getattr(owner, name)`` + for a single list element (Dynamo-traceable, unlike ``SimpleNamespace``).""" + + def __init__(self, name: str, value: Any) -> None: + setattr(self, name, value) + + +class _ListAdapter(_Adapter): + """Generic lift of a scalar field adapter to its ``List[...]`` analog. + + A custom-op schema is fixed at registration while list lengths vary per + call, so the inner adapter's slots are aggregated instead of replicated + per element: + + * ``Tensor`` / ``Tensor?`` -> ``Tensor[]``: one entry per element + (``None`` rides the 0-element sentinel); + * ``Tensor[]`` -> ``Tensor[]``: the elements' payloads concatenated in + element order, with per-element counts stashed in the bundle slot; + * ``OpaqueValueBundle`` -> one bundle carrying the per-element bundles. + + ``to_slots`` / ``from_slots`` delegate per element to the inner adapter, so + any scalar adapter semantics (kind tagging, storage flatten, value-opaque + checks) apply unchanged. Gradients of a differentiable inner adapter are + list-shaped: one grad per element on the lifted grad slot. + """ + + is_list = True + + ITEMS_KEY = "items" + COUNTS_KEY = "counts" + + # Scalar adapters that can be lifted. An inner adapter with a ``Tensor[]`` + # slot must also declare exactly one bundle slot (it carries the counts). + _INNER_ADAPTERS: Tuple[type, ...] = ( + _TensorOrQuantizedAdapter, + _TensorAdapter, + _QuantizerAdapter, + ) + + def __init__(self, name: str, inner: _Adapter) -> None: + self.name = name + self.inner = inner + self._slots = inner.schema_slots() + self._bundle_slots = [n for n, t in self._slots if t == _OPAQUE_VALUE_BUNDLE_TYPE_NAME] + self._flat_slots = [n for n, t in self._slots if t == "Tensor[]"] + if len(self._bundle_slots) > 1 or (self._flat_slots and not self._bundle_slots): + raise TypeError( + f"cannot lift {type(inner).__name__} to a list adapter: need exactly " + "one bundle slot to carry per-element metadata" + ) + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_ListAdapter"]: + if get_origin(annot) is not list: + return None + entry_args = get_args(annot) + if len(entry_args) != 1: + return None + for adapter_cls in cls._INNER_ADAPTERS: + inner = adapter_cls.try_build(name, entry_args[0]) + if inner is not None: + return cls(name, inner) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [ + (n, "Tensor[]" if t in ("Tensor", "Tensor?", "Tensor[]") else t) for n, t in self._slots + ] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + agg: Dict[str, List[Any]] = {n: [] for n, _ in self._slots} + counts: Dict[str, List[int]] = {n: [] for n in self._flat_slots} + for value in getattr(owner, self.name): + slots = self.inner.to_slots(_FieldProxy(self.name, value)) + for slot_name, type_str in self._slots: + val = slots[slot_name] + if type_str in ("Tensor", "Tensor?"): + agg[slot_name].append(_encode_none(val)) + elif type_str == "Tensor[]": + counts[slot_name].append(len(val)) + agg[slot_name].extend(val) + else: + agg[slot_name].append(val) + for slot_name in self._bundle_slots: + agg[slot_name] = OpaqueValueBundle( + {self.ITEMS_KEY: agg[slot_name], self.COUNTS_KEY: counts} + ) + return agg + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + items: Optional[List[Any]] = None + counts: Dict[str, List[int]] = {} + if self._bundle_slots: + bundle = args[self._bundle_slots[0]] + items = bundle[self.ITEMS_KEY] + counts = bundle[self.COUNTS_KEY] + length = len(items) if items is not None else len(args[self._slots[0][0]]) + cursors = {n: 0 for n in self._flat_slots} + values: List[Any] = [] + for i in range(length): + element_args: Dict[str, Any] = {} + for slot_name, type_str in self._slots: + if type_str in ("Tensor", "Tensor?"): + element_args[slot_name] = _decode_none(args[slot_name][i]) + elif type_str == "Tensor[]": + count = counts[slot_name][i] + start = cursors[slot_name] + element_args[slot_name] = list(args[slot_name][start : start + count]) + cursors[slot_name] = start + count + else: + element_args[slot_name] = items[i] + element_kwargs: Dict[str, Any] = {} + self.inner.from_slots(element_args, element_kwargs) + values.append(element_kwargs[self.name]) + kwargs[self.name] = values + + def grad_slot(self) -> Optional[int]: + return self.inner.grad_slot() + + class _ProcessGroupAdapter(_Adapter): """``ProcessGroup`` -> its c10d registry name in one ``OpaqueValueBundle`` slot. @@ -758,6 +928,8 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: _FIELD_ADAPTERS: Tuple[type, ...] = ( _TensorOrQuantizedAdapter, _TensorAdapter, + _ListAdapter, + _SymIntListAdapter, _ProcessGroupAdapter, _QuantizerAdapter, ) @@ -803,7 +975,13 @@ def _get_adapters(cls: type) -> List[_Adapter]: def _tensor_field_names(adapters: List[_Adapter]) -> List[str]: """Names of fields carrying tensors (for building the spec view).""" - return [b.name for b in adapters if isinstance(b, (_TensorAdapter, _TensorOrQuantizedAdapter))] + tensor_adapter_types = (_TensorAdapter, _TensorOrQuantizedAdapter) + return [ + b.name + for b in adapters + if isinstance(b, tensor_adapter_types) + or (isinstance(b, _ListAdapter) and isinstance(b.inner, tensor_adapter_types)) + ] def _build_schema(adapters: List[_Adapter]) -> Tuple[str, List[str]]: @@ -850,7 +1028,12 @@ 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, tuple)): + overrides[name] = [ + (to_tensor_spec(v) if v is not None and not isinstance(v, TensorSpec) else v) + for v in value + ] + elif value is not None and not isinstance(value, TensorSpec): overrides[name] = to_tensor_spec(value) if not overrides: return obj @@ -977,8 +1160,10 @@ def _pack_fwd_result(result: Any) -> List[torch.Tensor]: def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. - Each grad occupies exactly one slot (validated against ``num_grad_inputs``); - a :class:`TensorSpec` grad is materialized into a single tensor. + Each grad occupies one slot -- or, for a list-shaped grad (list-typed + differentiable input), one slot per element -- with the top-level count + validated against ``num_grad_inputs``. A :class:`TensorSpec` grad is + materialized into a single tensor. """ grads = list(grads) if len(grads) != num_grad_inputs: @@ -986,12 +1171,18 @@ def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List f"{op_qualname} expected bwd_impl to return {num_grad_inputs} grads " f"(one per input_tensors_for_grad entry), got {len(grads)}" ) + + def _one(g: Any) -> torch.Tensor: + if isinstance(g, TensorSpec): + return _encode_none(g.create_tensor()) + return _encode_none(g) + out: List[torch.Tensor] = [] for g in grads: - if isinstance(g, TensorSpec): - out.append(_encode_none(g.create_tensor())) + if isinstance(g, (list, tuple)): + out.extend(_one(v) for v in g) else: - out.append(_encode_none(g)) + out.append(_one(g)) return out @@ -1017,23 +1208,24 @@ def _unpack_fwd_fake_result( def _resolve_grad_targets( fwd_adapters: List[_Adapter], input_tensors_for_grad: List[str], -) -> Tuple[int, List[int]]: +) -> Tuple[int, List[Tuple[int, bool]]]: """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. ``fwd_adapters`` already encode the arg dataclass's fields (they are built from it), so the type itself is not needed here. Returns ``(slot_count, grad_targets)``: the total number of input schema - slots and, for each requested input name, the schema-slot index its gradient - maps to. + slots and, for each requested input name, an ``(index, is_list)`` pair -- + the schema-slot index its gradient maps to and whether that gradient is + list-shaped (one grad per list element). """ - name_to_slot: Dict[str, int] = {} + name_to_slot: Dict[str, Tuple[int, bool]] = {} slot_offset = 0 for adapter in fwd_adapters: slots = adapter.schema_slots() grad_slot = adapter.grad_slot() if grad_slot is not None: - name_to_slot[adapter.name] = slot_offset + grad_slot + name_to_slot[adapter.name] = (slot_offset + grad_slot, adapter.is_list) slot_offset += len(slots) non_differentiable = [n for n in input_tensors_for_grad if n not in name_to_slot] @@ -1107,8 +1299,12 @@ def _register_autograd_for_op( """ def _setup_context(ctx, inputs, output): + # Tensor lists only: a SymInt[] slot is also a list but takes a plain + # ``None`` grad, not a list-shaped one. ctx.fwd_tensor_list_lengths = { - i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) + i: len(value) + for i, value in enumerate(inputs) + if isinstance(value, list) and all(isinstance(v, torch.Tensor) for v in value) } kwargs = dict(zip(fwd_arg_names, inputs)) fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_adapters) @@ -1149,31 +1345,85 @@ def _autograd_backward(ctx, *grad_outputs): out: List[Any] = [None] * slot_count for pos, length in ctx.fwd_tensor_list_lengths.items(): out[pos] = [None] * length - for pos, g in zip(grad_targets, grads): - out[pos] = g + # The bwd op's flat return packs a list-shaped grad as one slot per + # element (see ``_pack_bwd_result``); walk it with a cursor. + cursor = 0 + for pos, is_list in grad_targets: + if is_list: + length = ctx.fwd_tensor_list_lengths[pos] + out[pos] = list(grads[cursor : cursor + length]) + cursor += length + else: + out[pos] = grads[cursor] + cursor += 1 + if cursor != len(grads): + raise RuntimeError( + f"bwd op returned {len(grads)} flat grads, expected {cursor}: a " + "list-shaped grad must have exactly one entry per input list element" + ) return tuple(out) fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) -def _tensor_or_quantized_offsets(adapters: List[_Adapter]) -> List[int]: - """Start index of each ``_TensorOrQuantizedAdapter`` group in the flat args.""" - offsets: List[int] = [] +def _tensor_or_quantized_offsets(adapters: List[_Adapter]) -> List[Tuple[int, bool]]: + """``(start index, is_list)`` of each tensor-or-quantized group in the flat args.""" + offsets: List[Tuple[int, bool]] = [] pos = 0 for adapter in adapters: - if isinstance(adapter, _TensorOrQuantizedAdapter): - offsets.append(pos) + if isinstance(adapter, _TensorOrQuantizedAdapter) or ( + isinstance(adapter, _ListAdapter) + and isinstance(adapter.inner, _TensorOrQuantizedAdapter) + ): + offsets.append((pos, adapter.is_list)) pos += len(adapter.schema_slots()) return offsets +def _flatten_list_subclasses(new_args: List[Any], offset: int, subclass: type) -> None: + """Rewrite ``subclass`` entries of the list group at ``offset`` into its + storage layout: the entry's inner tensors are spliced into the concatenated + ``__tensors`` slot (in element order) and the ``__meta`` bundle is rebuilt + with the entry retagged as storage. + """ + outer = list(new_args[offset]) + if not any(isinstance(v, subclass) for v in outer): + return + bundle = new_args[offset + 2] + items = list(bundle[_ListAdapter.ITEMS_KEY]) + counts_map = {k: list(v) for k, v in bundle[_ListAdapter.COUNTS_KEY].items()} + # A lifted tensor-or-quantized group has exactly one concatenated slot. + (counts_slot, counts) = next(iter(counts_map.items())) + inner = list(new_args[offset + 1]) + for i, val in enumerate(outer): + if not isinstance(val, subclass): + continue + meta, tensors = _storage_flatten( + val, {_TensorOrQuantizedAdapter.KIND_KEY: _TensorOrQuantizedKind.STORAGE} + ) + insert_at = sum(counts[:i]) + inner[insert_at:insert_at] = tensors + outer[i] = _encode_none(None) + items[i] = meta + counts[i] = len(tensors) + new_args[offset] = outer + new_args[offset + 1] = inner + new_args[offset + 2] = OpaqueValueBundle( + {_ListAdapter.ITEMS_KEY: items, _ListAdapter.COUNTS_KEY: {counts_slot: counts}} + ) + + def _flatten_subclass_into_slots( - new_args: List[Any], slot_offsets: List[int], subclass: type + new_args: List[Any], slot_offsets: List[Tuple[int, bool]], subclass: type ) -> None: - """Rewrite each tensor-or-quantized-adapter group whose ``Tensor?`` slot holds an - instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). + """Rewrite each tensor-or-quantized group whose tensor slot holds (an) + instance(s) of ``subclass`` into the storage layout (3 slots: name / tensors + / meta); list groups are rewritten per element. """ - for offset in slot_offsets: + for offset, is_list in slot_offsets: + if is_list: + _flatten_list_subclasses(new_args, offset, subclass) + continue val = new_args[offset] if not isinstance(val, subclass): continue @@ -1186,7 +1436,7 @@ def _flatten_subclass_into_slots( def _make_slot_forwarder( - base_op: Any, slot_offsets: Sequence[int], subclasses: Sequence[type] + base_op: Any, slot_offsets: Sequence[Tuple[int, bool]], subclasses: Sequence[type] ) -> Callable[[Sequence[Any]], List[torch.Tensor]]: """Return ``call(args)`` forwarding to ``base_op``, first flattening any ``subclasses`` instance sitting in the tensor-or-quantized slot groups at @@ -1227,7 +1477,7 @@ def _register_wrapper_op( wrapper_op_name: str, schema_str: str, base_op: Any, - slot_offsets: Sequence[int] = (), + slot_offsets: Sequence[Tuple[int, bool]] = (), subclasses: Sequence[type] = (), ) -> Any: """Define the wrapper op via ``torch.library.custom_op``: forward to the base @@ -1239,7 +1489,10 @@ def _forward(*flat: Any) -> List[torch.Tensor]: return forward(flat) op_def = torch.library.custom_op( - f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", + _forward, + mutates_args=(), + schema=schema_str, ) op_def.register_fake(_forward) return op_def diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 12497d0cb0..08e12ea7a5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -4,8 +4,10 @@ """GroupedLinear API""" -from typing import Union, Optional, Callable, Tuple, List, Sequence +from dataclasses import dataclass +from typing import Any, Union, Optional, Callable, Dict, Tuple, List, Sequence from itertools import chain +import math import os import warnings import weakref @@ -42,6 +44,10 @@ requires_grad, resolve_grouped_linear_single_param_flags, get_nvtx_range_context, + warn_compile_eager_fallback, + warn_if_compile_disabled, + check_grouped_gemm_dims, + is_non_tn_fp8_gemm_supported, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -59,6 +65,14 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) +from ..cpp_extensions.gemm import get_cublas_workspace +from ..dynamo import ( + TensorSpec, + TensorOrQuantized, + register_custom_op, + is_value_opaque_quantizer, +) +from .linear import _fake_workspace_valid from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload @@ -292,7 +306,9 @@ def _split_quantize_hybrid( disable_bulk_allocation=False, ): """Grouped split+quantize for an all-hybrid, generation-validated operand.""" - from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage + from ..tensor.storage.hybrid_tensor_storage import ( + HybridQuantizedTensorStorage as HybridStorage, + ) reference = quantizers[0] rowwise_enabled = reference.rowwise_usage @@ -353,6 +369,13 @@ def _split_quantize_hybrid( ] +@torch.compiler.assume_constant_result +@functools.lru_cache(maxsize=None) +def _get_cublaslt_version() -> int: + """Cached, Dynamo-constant cuBLASLt version (the pybind call is untraceable).""" + return tex.get_cublasLt_version() + + def _split_quantize( tensor: torch.Tensor, split_sizes: List[int], @@ -399,7 +422,7 @@ def _split_quantize_and_bias( quantizers: Optional[List[Quantizer]], dtype: torch.dtype, use_bias: bool, - recipe: Recipe, + recipe_supports_native_bgrad: bool, disable_bulk_allocation: bool, ) -> Tuple[ Sequence[Union[torch.Tensor, QuantizedTensorStorage]], @@ -418,7 +441,7 @@ def _split_quantize_and_bias( and not hybrid and use_bias and not identity - and (recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8()) + and recipe_supports_native_bgrad ) if use_native_bgrad_quantize: outputs = [None] * num_splits @@ -446,6 +469,1696 @@ def _split_quantize_and_bias( return outputs, grad_biases +@dataclass(slots=True) +class GroupedLinearFwdArgs: + """Single-argument bag for the forward path of :class:`_GroupedLinear`.""" + + # --- Differentiable tensors (also passed positionally to autograd) --- + inp: torch.Tensor + weights: List[TensorOrQuantized] + biases: List[torch.Tensor] + + # --- Non-differentiable cached / user-provided tensors --- + # TensorOrQuantized entries so cached quantized workspaces can cross the op + # boundary; ``None`` entries mark cache misses. + weight_workspaces: List[TensorOrQuantized] + out: Optional[torch.Tensor] + dgrad_out: Optional[torch.Tensor] + skip_fp8_weight_update: Optional[torch.Tensor] + # Device-tensor form of the splits, used only by the fused GroupedTensor + # path (gated off the compiled path, where this is None). + m_splits_tensor: Optional[torch.Tensor] + + # --- requires_grad flags (cached so backward does not re-query) --- + input_requires_grad: bool + weights_requires_grad: bool + bias_requires_grad: bool + + # --- Quantizers (one per GEMM) --- + input_quantizers: List[Quantizer] + weight_quantizers: List[Quantizer] + output_quantizers: List[Quantizer] + grad_input_quantizers: List[Quantizer] + grad_weight_quantizers: List[Quantizer] + grad_output_quantizers: List[Quantizer] + + # --- Split geometry --- + # Host-side ints, carried through the op as a ``SymInt[]`` slot (so splits + # Dynamo marks dynamic still compile). ``None`` only transiently, when the + # caller passed a device tensor under compile (an eager-fallback reason); + # the eager wrapper fills it from the tensor. + m_splits: Optional[List[int]] + num_gemms: int + + # --- Numerical / dtype config --- + activation_dtype: torch.dtype + fp8: bool + fp8_calibration: bool + save_original_input: bool + backward_override: Optional[str] + fprop_use_split_accumulator: bool + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool + native_bgrad_recipe_ok: bool + debug: bool + + # --- Weight-workspace caching --- + is_first_microbatch: Optional[bool] + cache_weight: bool + + # --- Fused GroupedTensor path (dispatched in the autograd wrapper) --- + use_grouped_tensor_path: bool + single_grouped_param: bool + + # --- Misc --- + use_bias: bool + sequence_parallel: bool + fuse_wgrad_accumulation: bool + wgrad_store: Optional[Any] + cpu_offloading: bool + is_grad_enabled: bool + # True only when running as the torch.compile custom op: disables bulk + # allocation and packed-buffer tricks whose outputs would alias each other + # (a custom op may not return aliasing tensors). + compiled_op: bool = False + + def compile_unsupported_reason(self) -> Optional[str]: + """Reason this config can't use the torch.compile custom-op path (else None).""" + if self.debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + if is_distributed_weight(self.weights[0]): + return "a DistributedWeight (custom weight parallelism, e.g. GTP)" + if self.m_splits is None: + return "m_splits passed as a device tensor inside the compiled region" + if self.out is not None or self.dgrad_out is not None: + return "a user-provided out/dgrad_out buffer (the op would return an input alias)" + if self.use_grouped_tensor_path: + return "the fused GroupedTensor path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM)" + if self.single_grouped_param: + return "single_grouped_weight/single_grouped_bias parameter views" + if self.cpu_offloading: + return "CPU activation offloading" + if self.wgrad_store is not None: + # Non-None only when delayed wgrad compute is on (see GroupedLinear.forward). + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + if self.fp8_calibration: + # calibrate() inside the op would mutate quantizers rebuilt from + # graph constants, losing the amax updates. + return "fp8_calibration" + if any(w.requires_grad != self.weights_requires_grad for w in self.weights): + return "mixed requires_grad across the weights list" + for quantizer_list in ( + self.input_quantizers, + self.weight_quantizers, + self.output_quantizers, + self.grad_input_quantizers, + self.grad_weight_quantizers, + self.grad_output_quantizers, + ): + for quantizer in quantizer_list: + # e.g. delayed-scaling Float8Quantizer and unregistered + # custom-recipe quantizers are not value-opaque and can't cross + # the custom-op boundary. + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return "a quantizer not registered as a torch.compile value-opaque type" + return None + + +@dataclass(slots=True) +class GroupedLinearBwdArgs: + """Single-argument bag for the backward path of :class:`_GroupedLinear`.""" + + # --- Saved / restored tensors (populated at backward entry) --- + grad_output: Optional[torch.Tensor] = None + # Full (2D-viewable) input saved once when backward re-splits it itself: + # the plain (non-quantized) path and ``save_original_input``. + inputmat_full: Optional[torch.Tensor] = None + inputmats: List[TensorOrQuantized] = None + weights_fp8: List[TensorOrQuantized] = None + saved_weights: List[TensorOrQuantized] = None + biases: List[torch.Tensor] = None + dgrad_out: Optional[torch.Tensor] = None + + # --- Quantizers (one per GEMM) --- + input_quantizers: List[Quantizer] = None + weight_quantizers: List[Quantizer] = None + grad_input_quantizers: List[Quantizer] = None + grad_weight_quantizers: List[Quantizer] = None + grad_output_quantizers: List[Quantizer] = None + + # --- Split geometry --- + # ``SymInt[]`` slot, see ``GroupedLinearFwdArgs.m_splits``. + m_splits: Optional[List[int]] = None + num_gemms: int = 0 + weights_shape_1: int = 0 + + # --- Differentiability summary --- + use_bias: bool = False + requires_dgrad: bool = False + weights_requires_grad: bool = False + + # --- Numerical / dtype config --- + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + backward_override: Optional[str] = None + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD + native_bgrad_recipe_ok: bool = False + save_original_input: bool = False + debug: bool = False + + # --- Weight-grad scheduling / accumulation (eager-only, gated off compile) --- + is_first_microbatch: Optional[bool] = None + fuse_wgrad_accumulation: bool = False + wgrad_store: Optional[Any] = None + origin_weight_refs: Optional[Any] = None + origin_weights_overwrite_main_grad: bool = False + main_grad_funcs: Optional[Any] = None + + # --- FP8 reduce-and-update bookkeeping (eager wrapper only) --- + reduce_and_update_bwd_fp8_tensors: bool = False + + # --- Misc --- + cpu_offloading: bool = False + compiled_op: bool = False + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" + saved = restore_from_func_ctx(ctx) + n = self.num_gemms + self.inputmat_full = saved[0] + self.inputmats = list(saved[1 : 1 + n]) + self.weights_fp8 = list(saved[1 + n : 1 + 2 * n]) + self.saved_weights = list(saved[1 + 2 * n : 1 + 3 * n]) + self.biases = list(saved[1 + 3 * n : 1 + 4 * n]) + + +def _grouped_linear_forward_impl( + args: GroupedLinearFwdArgs, +) -> Tuple[Any, ...]: + """Forward implementation for the grouped linear layer (legacy, non-fused path). + + Returns ``(out, *new_workspaces, tensors_to_save, ctx_attrs)``. + ``new_workspaces`` are the freshly produced FP8 weight workspaces (returned + alongside ``out`` so the caller can refresh its cache). The trailing two are + ``None`` when gradients are disabled. + """ + inp = args.inp + weights = list(args.weights) + biases = list(args.biases) + num_gemms = args.num_gemms + m_splits = list(args.m_splits) + input_quantizers = args.input_quantizers + weight_quantizers = args.weight_quantizers + output_quantizers = args.output_quantizers + activation_dtype = args.activation_dtype + fp8 = args.fp8 + debug = args.debug + use_bias = args.use_bias + is_grad_enabled = args.is_grad_enabled + save_original_input = args.save_original_input + backward_override = args.backward_override + cpu_offloading = args.cpu_offloading + device = inp.device + weight_requires_grad = args.weights_requires_grad + + is_dist_weight = is_distributed_weight(weights[0]) + if is_dist_weight: + weights = materialize_weight_for_forward(weights) + + # Configure quantizers + if input_quantizers[0] is not None: + for input_quantizer in input_quantizers: + input_quantizer.set_usage( + rowwise=True, + columnwise=( + is_grad_enabled + and weight_requires_grad + and not save_original_input + and backward_override is None + ), + ) + columnwise_usage = is_grad_enabled and args.input_requires_grad + if backward_override is not None: + columnwise_usage = False + if not columnwise_usage: + columnwise_usage = ( + is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase() + ) + # No need to set the quantizer states if weight is already quantized + # for debug mode we create quantizer every iteration, thus we need to set the quantizer states + if weight_quantizers[0] is not None and ( + not isinstance(weights[0], QuantizedTensorStorage) or debug + ): + for weight_quantizer in weight_quantizers: + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif isinstance(weights[0], QuantizedTensorStorage): + # If weights are already quantized, no need to set quantizer states + weight_quantizers = [weight._quantizer for weight in weights] + if output_quantizers[0] is not None: + for output_quantizer in output_quantizers: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # Initialize input tensors + in_features = weights[0].size(-1) + if inp.size(-1) != in_features: + raise ValueError( + f"Input tensor (shape={tuple(inp.size())}) is not compatible with " + f"weight tensor (shape={tuple(weights[0].size())})" + ) + + inp_view = inp.reshape(-1, in_features) + fp8_or_debug = fp8 or debug + inputmat_full = None + if fp8_or_debug: + # Disable bulk allocation when CPU offloading is active: offloading skips small + # tensors (like scales), but bulk allocation shares storage across all tensors, + # so if scales can't be offloaded, nothing in the group can be offloaded. + # The compiled op also disables it: bulk-allocated storages alias each + # other, and the op returns them as saved tensors. + inputmats = _split_quantize( + inp_view, + m_splits, + with_quantized_output=True, + quantizers=input_quantizers, + dtype=activation_dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=cpu_offloading or args.compiled_op, + ) + else: + # Plain path: split views of one (possibly cast) buffer. Save that + # buffer once (backward re-splits it) instead of N aliasing views -- + # except under CPU offloading, which marks the individual views. + inputmat_full = cast_if_needed(inp_view, activation_dtype) + inputmats = torch.split(inputmat_full, m_splits) + if cpu_offloading: + inputmat_full = None + + if cpu_offloading: + start_offload(*inputmats) + + # Initialize weights + weights_fp8: list + new_workspaces = [None] * num_gemms + if fp8_or_debug: + weights_fp8 = [] + update_ws = args.is_first_microbatch is None or args.is_first_microbatch + for i in range(num_gemms): + weight_fp8, new_workspaces[i] = quantize_weight( + tensor=weights[i], + quantizer=weight_quantizers[i], + workspace=args.weight_workspaces[i] if args.weight_workspaces else None, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + workspace_dtype=activation_dtype, + cache=args.cache_weight, + ) + weights_fp8.append(weight_fp8) + else: + weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights] + + # Initialize biases + bias_dtype = activation_dtype + if fp8 and activation_dtype == torch.float32: + bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias + biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases + # Initialize output tensor + out = _GroupedLinear._validate_or_alloc_output( + args.out, + sum(m_splits), + weights_fp8[0].size(0), + activation_dtype, + device, + ) + + # Perform GEMM + general_grouped_gemm( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + single_output=True, + m_splits=m_splits, + bias=biases, + use_bias=use_bias, + use_split_accumulator=args.fprop_use_split_accumulator, + ) + + if args.fp8_calibration: + for i in range(num_gemms): + input_quantizers[i].calibrate(inputmats[i]) + weight_quantizers[i].calibrate(weights[i]) + + if cpu_offloading: + mark_not_offload(*weights_fp8, *weights) + + tensors_to_save = None + ctx_attrs = None + if is_grad_enabled: + # Saved-tensor layout: ``(inputmat_full, *inputmats, *weights_fp8, + # *saved_weights, *biases)`` -- 1 + 4N slots. Slots that alias a forward + # input or another op return are deduped through name-based alias tags + # (rebuilt in ``_grouped_linear_setup_ctx``): a custom op may not return + # aliasing tensors. + aliases: List[Optional[Tuple]] = [None] * (1 + 4 * num_gemms) + + # TODO: update after #1638 is merged. # pylint: disable=fixme + if weight_requires_grad: + if save_original_input: + inputmat_full = None + inputmats = [None] * num_gemms + aliases[0] = ("inp",) + else: + for inputmat in inputmats: + if isinstance(inputmat, QuantizedTensorStorage): + if backward_override is not None: + # In dequantized mode we should dequantize directly from + # fprop quantized layouts without retargeting usage. + inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + if inputmat_full is not None: + # Plain path: the views in ``inputmats`` are not saved. + inputmats = [None] * num_gemms + if inputmat_full is inp_view: + # No-op cast: inp_view is inp itself or a view of it. + inputmat_full = None + aliases[0] = ("inp",) + else: + inputmat_full = None + inputmats = [None] * num_gemms + + # Original weights are only needed by high_precision dgrad. The weakrefs + # used for fused wgrad accumulation serve a different purpose: restoring + # Python parameter attributes without keeping the parameter alive here. + save_origin_weights = backward_override == "high_precision" and args.input_requires_grad + saved_weights = [None] * num_gemms + wt_saves = list(weights_fp8) + for i in range(num_gemms): + slot = 1 + num_gemms + i + if wt_saves[i] is weights[i]: + aliases[slot] = ("weights", i) + wt_saves[i] = None + elif new_workspaces[i] is not None and wt_saves[i] is new_workspaces[i]: + aliases[slot] = ("new_weight_workspaces", i) + wt_saves[i] = None + elif ( + args.weight_workspaces + and args.weight_workspaces[i] is not None + and wt_saves[i] is args.weight_workspaces[i] + ): + aliases[slot] = ("weight_workspaces", i) + wt_saves[i] = None + if save_origin_weights: + aliases[1 + 2 * num_gemms + i] = ("weights", i) + if is_dist_weight: + # GTP: gathered workspace is transient (re-gathered in backward), don't save it. + wt_saves = [None] * num_gemms + for i in range(num_gemms): + aliases[1 + num_gemms + i] = None + aliases[1 + 2 * num_gemms + i] = ("weights", i) + + saved_biases = list(biases) + for i in range(num_gemms): + if saved_biases[i] is not None and saved_biases[i] is args.biases[i]: + aliases[1 + 3 * num_gemms + i] = ("biases", i) + saved_biases[i] = None + + tensors_to_save = ( + inputmat_full, + *inputmats, + *wt_saves, + *saved_weights, + *saved_biases, + ) + ctx_attrs = {"saved_tensor_aliases": tuple(aliases)} + + # [*, in_features] -> [*, out_features] + out = out.view(-1, *inp.shape[1:-1], out.shape[-1]) + return (out, *new_workspaces, tensors_to_save, ctx_attrs) + + +def _grouped_linear_forward_fake( + args: GroupedLinearFwdArgs, +) -> Tuple[Any, ...]: + """Shape/metadata-only twin of :func:`_grouped_linear_forward_impl` for + torch.compile. Only mirrors configs the compiled path admits (see + ``compile_unsupported_reason``): no debug, offloading, distributed weights, + calibration, or fused GroupedTensor path. + """ + inp = args.inp + weights = list(args.weights) + num_gemms = args.num_gemms + m_splits = list(args.m_splits) + input_quantizers = args.input_quantizers + weight_quantizers = args.weight_quantizers + output_quantizers = args.output_quantizers + activation_dtype = args.activation_dtype + fp8 = args.fp8 + is_grad_enabled = args.is_grad_enabled + save_original_input = args.save_original_input + backward_override = args.backward_override + weight_requires_grad = args.weights_requires_grad + in_features = weights[0].shape[-1] + out_features = weights[0].shape[0] + + # Mirror the impl's quantizer usage setup exactly (buffer layouts must agree). + if input_quantizers[0] is not None: + for input_quantizer in input_quantizers: + input_quantizer.set_usage( + rowwise=True, + columnwise=( + is_grad_enabled + and weight_requires_grad + and not save_original_input + and backward_override is None + ), + ) + columnwise_usage = is_grad_enabled and args.input_requires_grad + if backward_override is not None: + columnwise_usage = False + if not columnwise_usage: + columnwise_usage = ( + is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase() + ) + if weight_quantizers[0] is not None and not weights[0].is_quantized: + for weight_quantizer in weight_quantizers: + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif weights[0].is_quantized: + weight_quantizers = [weight.quantizer for weight in weights] + if output_quantizers[0] is not None: + for output_quantizer in output_quantizers: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # Input pipeline: quantized per-split storages (fp8) or one full cast buffer. + inputmat_full = None + inputmats: List[Optional[TensorSpec]] = [None] * num_gemms + inputmat_full_aliases_inp = False + if fp8: + inputmats = [ + TensorSpec( + shape=(m_splits[i], in_features), + dtype=activation_dtype, + quantizer=input_quantizers[i], + device=inp.device, + ) + for i in range(num_gemms) + ] + else: + inputmat_full_aliases_inp = inp.dtype == activation_dtype + inputmat_full = TensorSpec( + shape=(sum(m_splits), in_features), + dtype=activation_dtype, + device=inp.device, + ) + + # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed`` per GEMM. + new_workspaces: List[Optional[TensorSpec]] = [None] * num_gemms + weights_fp8: List[Optional[TensorSpec]] = [None] * num_gemms + weight_aliases: List[Optional[Tuple]] = [None] * num_gemms + for i in range(num_gemms): + if fp8: + if weights[i].is_quantized: + weight_aliases[i] = ("weights", i) + continue + workspace = args.weight_workspaces[i] + if workspace is not None and not _fake_workspace_valid(workspace, weight_quantizers[i]): + # quantize_weight drops a stale workspace and builds a new one. + workspace = None + if workspace is not None: + weight_aliases[i] = ("weight_workspaces", i) + continue + weightmat = TensorSpec( + shape=tuple(weights[i].shape), + dtype=activation_dtype, + quantizer=weight_quantizers[i], + device=weights[i].device, + ) + if args.cache_weight: + # Persistent cache entries are wrappers, not bare storages. + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_workspaces[i] = weightmat + weight_aliases[i] = ("new_weight_workspaces", i) + else: + weights_fp8[i] = weightmat + else: + if weights[i].dtype == activation_dtype: + weight_aliases[i] = ("weights", i) + else: + weights_fp8[i] = TensorSpec( + shape=tuple(weights[i].shape), + dtype=activation_dtype, + device=weights[i].device, + ) + + # Bias pipeline. + bias_dtype = activation_dtype + if fp8 and activation_dtype == torch.float32: + bias_dtype = torch.bfloat16 + saved_biases: List[Optional[TensorSpec]] = [None] * num_gemms + bias_aliases: List[Optional[Tuple]] = [None] * num_gemms + for i in range(num_gemms): + bias = args.biases[i] + if bias is None: + continue + if not args.use_bias or bias.dtype == bias_dtype: + bias_aliases[i] = ("biases", i) + else: + saved_biases[i] = TensorSpec( + shape=tuple(bias.shape), dtype=bias_dtype, device=bias.device + ) + + out = TensorSpec( + shape=(*tuple(inp.shape[:-1]), out_features), + dtype=activation_dtype, + quantizer=None, + requires_grad=is_grad_enabled + and (args.input_requires_grad or weight_requires_grad or args.bias_requires_grad), + device=inp.device, + ) + + tensors_to_save = None + ctx_attrs = None + if is_grad_enabled: + aliases: List[Optional[Tuple]] = [None] * (1 + 4 * num_gemms) + if weight_requires_grad: + if save_original_input: + inputmat_full = None + inputmats = [None] * num_gemms + aliases[0] = ("inp",) + else: + if fp8: + for inputmat in inputmats: + if backward_override is not None: + inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + elif inputmat_full_aliases_inp: + inputmat_full = None + aliases[0] = ("inp",) + else: + inputmat_full = None + inputmats = [None] * num_gemms + + saved_weights = [None] * num_gemms + save_origin_weights = backward_override == "high_precision" and args.input_requires_grad + for i in range(num_gemms): + aliases[1 + num_gemms + i] = weight_aliases[i] + if save_origin_weights: + aliases[1 + 2 * num_gemms + i] = ("weights", i) + aliases[1 + 3 * num_gemms + i] = bias_aliases[i] + + tensors_to_save = ( + inputmat_full, + *inputmats, + *weights_fp8, + *saved_weights, + *saved_biases, + ) + ctx_attrs = {"saved_tensor_aliases": tuple(aliases)} + + return (out, *new_workspaces, tensors_to_save, ctx_attrs) + + +def _grouped_linear_setup_ctx( + bwd_args: GroupedLinearBwdArgs, + fwd_args: GroupedLinearFwdArgs, + fwd_outputs: Tuple[Any, ...], + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate ``bwd_args`` from forward state and return the tensors to persist + (alias-tagged slots rebuilt from ``fwd_args`` / ``fwd_outputs``).""" + num_gemms = fwd_args.num_gemms + new_workspaces = list(fwd_outputs[1:]) + + weights = fwd_args.weights + weight_quantizers = fwd_args.weight_quantizers + if isinstance(weights[0], QuantizedTensorStorage) and not fwd_args.debug: + weight_quantizers = [weight._quantizer for weight in weights] + + bwd_args.input_quantizers = fwd_args.input_quantizers + bwd_args.weight_quantizers = weight_quantizers + bwd_args.grad_input_quantizers = fwd_args.grad_input_quantizers + bwd_args.grad_weight_quantizers = fwd_args.grad_weight_quantizers + bwd_args.grad_output_quantizers = fwd_args.grad_output_quantizers + + bwd_args.m_splits = fwd_args.m_splits + bwd_args.num_gemms = num_gemms + bwd_args.weights_shape_1 = weights[0].shape[1] + + bwd_args.use_bias = fwd_args.use_bias + bwd_args.requires_dgrad = fwd_args.input_requires_grad + bwd_args.weights_requires_grad = fwd_args.weights_requires_grad + + bwd_args.activation_dtype = fwd_args.activation_dtype + bwd_args.fp8 = fwd_args.fp8 + bwd_args.backward_override = fwd_args.backward_override + bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator + bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator + bwd_args.native_bgrad_recipe_ok = fwd_args.native_bgrad_recipe_ok + bwd_args.save_original_input = fwd_args.save_original_input + bwd_args.debug = fwd_args.debug + + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.fuse_wgrad_accumulation = fwd_args.fuse_wgrad_accumulation + bwd_args.wgrad_store = fwd_args.wgrad_store + bwd_args.cpu_offloading = fwd_args.cpu_offloading + bwd_args.dgrad_out = fwd_args.dgrad_out + bwd_args.compiled_op = fwd_args.compiled_op + + if fwd_args.fuse_wgrad_accumulation and fwd_args.weights_requires_grad: + # Keep weakrefs to weights to preserve attributes like main_grad + # when we need to modify the weight python objects + bwd_args.origin_weight_refs = [weakref.ref(w) for w in weights] + bwd_args.origin_weights_overwrite_main_grad = getattr( + weights[0], "overwrite_main_grad", False + ) + # MCore FSDP creates main_grad lazily before backward + if hasattr(weights[0], "__fsdp_param__"): + bwd_args.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + elif is_distributed_weight(weights[0]): + bwd_args.main_grad_funcs = [weights[i].grad_buffer for i in range(num_gemms)] + else: + bwd_args.main_grad_funcs = [lambda j=i: weights[j].main_grad for i in range(num_gemms)] + + if fwd_args.backward_override is not None: + bwd_args.fp8 = False + bwd_args.debug = False + bwd_args.grad_input_quantizers = [None] * num_gemms + bwd_args.grad_weight_quantizers = [None] * num_gemms + bwd_args.grad_output_quantizers = [None] * num_gemms + + # Rebuild alias-deduped save slots. + saved = list(tensors_to_save_from_forward) + aliases = ctx_attrs["saved_tensor_aliases"] + for slot, alias in enumerate(aliases): + if alias is None: + continue + if alias[0] == "inp": + saved[slot] = fwd_args.inp + elif alias[0] == "weights": + saved[slot] = weights[alias[1]] + elif alias[0] == "new_weight_workspaces": + saved[slot] = new_workspaces[alias[1]] + elif alias[0] == "weight_workspaces": + saved[slot] = fwd_args.weight_workspaces[alias[1]] + elif alias[0] == "biases": + saved[slot] = fwd_args.biases[alias[1]] + return tuple(saved) + + +def _grouped_linear_backward_impl( + args: GroupedLinearBwdArgs, +) -> Tuple[Optional[torch.Tensor], List[Optional[torch.Tensor]], List[Optional[torch.Tensor]]]: + """Backward implementation for the grouped linear layer. + + Caller must have populated ``args.grad_output`` and run + ``args.setup_saved_tensors(ctx)`` before invocation. Returns + ``(dgrad, wgrad_list, grad_biases)``. + """ + grad_output = args.grad_output + num_gemms = args.num_gemms + m_splits = list(args.m_splits) + inputmats = list(args.inputmats) + weights = list(args.weights_fp8) + saved_weights = list(args.saved_weights) + biases = list(args.biases) + device = grad_output.device + in_features = args.weights_shape_1 + dgrad = None + + # Plain (non-quantized) inputs are saved as one full buffer; re-split it. + if args.inputmat_full is not None and not args.save_original_input: + inputmats = list(torch.split(args.inputmat_full.reshape(-1, in_features), m_splits)) + + # Restore from weakrefs to get original weight python objects + # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) + # Only needed when fuse_wgrad_accumulation is enabled. + origin_weights = [None] * num_gemms + main_grads = [None] * num_gemms + is_dist_weight = is_distributed_weight(saved_weights[0]) + if is_dist_weight: + origin_weights = saved_weights + if args.fuse_wgrad_accumulation and args.weights_requires_grad: + main_grads = [main_grad_func() for main_grad_func in args.main_grad_funcs] + elif args.fuse_wgrad_accumulation and args.weights_requires_grad: + origin_weight_refs = args.origin_weight_refs + args.origin_weight_refs = None + origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] + assert all( + w is not None for w in origin_weights + ), "weight was removed while fuse_wgrad_accumulation=True" + main_grads = [main_grad_func() for main_grad_func in args.main_grad_funcs] + for origin_weight, main_grad in zip(origin_weights, main_grads): + if main_grad is not None: + origin_weight.main_grad = main_grad + + # Preprocess grad output + grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) + out_features = grad_output_view.shape[-1] + grad_output_reference = args.grad_output_quantizers[0] + if args.fp8 and isinstance(grad_output_reference, HybridQuantizer): + # Usage is a runtime decision, not part of generation validation. + # Apply it uniformly so dispatch can read the first parent without + # rescanning every expert. + for grad_output_quantizer in args.grad_output_quantizers: + grad_output_quantizer.set_usage( + rowwise=args.requires_dgrad, + columnwise=args.weights_requires_grad, + ) + grad_output, grad_biases = _split_quantize_and_bias( + grad_output_view, + m_splits, + fp8=args.fp8, + debug=args.debug, + quantizers=args.grad_output_quantizers, + dtype=args.activation_dtype, + use_bias=args.use_bias, + recipe_supports_native_bgrad=args.native_bgrad_recipe_ok, + disable_bulk_allocation=args.cpu_offloading, + ) + + if is_dist_weight: + accumulate_wgrad_into_param_main_grad = False + elif args.is_first_microbatch is not None: + accumulate_wgrad_into_param_main_grad = ( + args.fuse_wgrad_accumulation and not args.is_first_microbatch + ) + else: + accumulate_wgrad_into_param_main_grad = args.fuse_wgrad_accumulation + + if is_dist_weight: + weights = materialize_weight_for_backward(origin_weights) + + if args.requires_dgrad: + dgrad = _GroupedLinear._validate_or_alloc_output( + args.dgrad_out, + sum(m_splits), + in_features, + args.activation_dtype, + device, + ) + weights_for_dgrad = weights + if args.backward_override == "dequantized": + weights_for_dgrad = [ + _GroupedLinear._maybe_dequantize(weight, args.activation_dtype) + for weight in weights + ] + elif args.backward_override == "high_precision": + weights_for_dgrad = [ + _GroupedLinear._maybe_dequantize(weight, args.activation_dtype) + for weight in saved_weights + ] + # Make sure weights are available in column-wise format + # for dgrad computation. + for weight in weights_for_dgrad: + if isinstance(weight, QuantizedTensorStorage): + weight.update_usage(columnwise_usage=True) + general_grouped_gemm( + weights_for_dgrad, + grad_output, + [dgrad], + args.grad_input_quantizers, + args.activation_dtype, + single_output=True, + layout="NN", + m_splits=m_splits, + grad=True, + use_split_accumulator=args.dgrad_use_split_accumulator, + ) + + if args.weights_requires_grad: + if args.fuse_wgrad_accumulation: + wgrad_list = main_grads + elif args.compiled_op: + # Packed allocation would make the returned wgrads alias each other. + wgrad_list = [ + torch.empty( + (out_features, in_features), + dtype=args.activation_dtype, + device=device, + ) + for _ in range(num_gemms) + ] + else: + wgrad_packed = torch.empty( + num_gemms, + out_features, + in_features, + dtype=args.activation_dtype, + device=device, + ) + wgrad_list = [wgrad_packed[i] for i in range(num_gemms)] + if is_dist_weight: + # Gathered weights are no longer needed after dgrad GEMM. + del weights + + if args.save_original_input: + inp = args.inputmat_full + inp_view = inp.reshape(-1, in_features) + if args.input_quantizers[0] is not None: + for input_quantizer in args.input_quantizers: + if isinstance( + input_quantizer, + (Float8Quantizer, Float8CurrentScalingQuantizer), + ): + input_quantizer.set_usage(rowwise=True, columnwise=True) + else: + input_quantizer.set_usage(rowwise=False, columnwise=True) + inputmats = _split_quantize( + inp_view, + m_splits, + with_quantized_output=args.fp8 or args.debug, + quantizers=args.input_quantizers, + dtype=args.activation_dtype, + with_debug_quantizers=args.debug, + disable_bulk_allocation=args.cpu_offloading, + ) + elif args.backward_override == "dequantized": + inputmats = [ + _GroupedLinear._maybe_dequantize(inputmat, args.activation_dtype) + for inputmat in inputmats + ] + grouped_gemm_wgrad = functools.partial( + general_grouped_gemm, + quantization_params=args.grad_weight_quantizers, + out_dtype=args.activation_dtype, + layout="NT", + grad=True, + m_splits=m_splits, + use_bias=args.use_bias if grad_biases[0] is None else None, + bias=biases, + use_split_accumulator=args.wgrad_use_split_accumulator, + accumulate=( + accumulate_wgrad_into_param_main_grad + if not is_dist_weight and not args.origin_weights_overwrite_main_grad + else False + ), + ) + # WGRAD + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + args.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) + else: + _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) + + for i in range(num_gemms): + if grad_biases[i] is None: + grad_biases[i] = grad_biases_[i] + del grad_biases_ + + # Deallocate input tensor (in-place storage resize: not allowed on + # the compiled path, where the inputs belong to the op caller). + if not args.compiled_op: + clear_tensor_data(*inputmats) + + def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): + if args.weights_requires_grad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr(weight, "grad_added_to_main_grad"): + weight.grad_added_to_main_grad = True + if getattr(weight, "zero_out_wgrad", False): + wgrad = get_dummy_wgrad( + list(main_grad.shape), + weight.dtype, + zero=True, + ) + else: + wgrad = get_dummy_wgrad( + list(main_grad.shape), + weight.dtype, + ) + elif args.fuse_wgrad_accumulation: + wgrad = None + else: + wgrad = None + return wgrad + + if is_dist_weight: + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) + else: + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] + else: + wgrad_list = [None] * num_gemms + + if not args.use_bias or ( + args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute() and not args.fp8 + ): + grad_biases = [None] * num_gemms + + dgrad_out = None + if args.requires_dgrad: + # Input shape rederived from grad_output (out.shape == (*inp.shape[:-1], out_features)). + dgrad_out = dgrad.view(*args.grad_output.shape[:-1], in_features) + return (dgrad_out, wgrad_list, list(grad_biases)) + + +def _grouped_linear_backward_fake( + args: GroupedLinearBwdArgs, +) -> Tuple[Optional[TensorSpec], List[Optional[TensorSpec]], List[Optional[TensorSpec]]]: + """Allocation-free fake of :func:`_grouped_linear_backward_impl` on ``TensorSpec``.""" + num_gemms = args.num_gemms + grad_output = args.grad_output + in_features = args.weights_shape_1 + out_features = grad_output.shape[-1] + + # Mirror the impl's hybrid grad-output usage retarget. + grad_output_reference = args.grad_output_quantizers[0] + if args.fp8 and isinstance(grad_output_reference, HybridQuantizer): + for grad_output_quantizer in args.grad_output_quantizers: + grad_output_quantizer.set_usage( + rowwise=args.requires_dgrad, + columnwise=args.weights_requires_grad, + ) + + dgrad = None + if args.requires_dgrad: + dgrad = TensorSpec( + shape=(*tuple(grad_output.shape[:-1]), in_features), + dtype=args.activation_dtype, + device=grad_output.device, + ) + + wgrad_list: List[Optional[TensorSpec]] = [None] * num_gemms + if args.weights_requires_grad and not args.fuse_wgrad_accumulation: + wgrad_list = [ + TensorSpec( + shape=(out_features, in_features), + dtype=args.activation_dtype, + device=grad_output.device, + ) + for _ in range(num_gemms) + ] + + # FP8 backward computes bgrad while splitting grad_output whenever bias is + # used; in high precision it is fused into the wgrad GEMM, so it only + # exists when wgrad runs. + grad_biases: List[Optional[TensorSpec]] = [None] * num_gemms + if args.use_bias and (args.weights_requires_grad or args.fp8): + grad_biases = [ + TensorSpec( + shape=(out_features,), + dtype=args.activation_dtype, + device=grad_output.device, + ) + for _ in range(num_gemms) + ] + + return (dgrad, wgrad_list, grad_biases) + + +# Custom op used under ``torch.compile``. +_grouped_linear_op = register_custom_op( + op_name="grouped_linear", + input_tensors_for_grad=["inp", "weights", "biases"], + fwd_arg_type=GroupedLinearFwdArgs, + fwd_impl=_grouped_linear_forward_impl, + fwd_fake_impl=_grouped_linear_forward_fake, + setup_context=_grouped_linear_setup_ctx, + bwd_arg_type=GroupedLinearBwdArgs, + bwd_impl=_grouped_linear_backward_impl, + bwd_fake_impl=_grouped_linear_backward_fake, +) + + +# --------------------------------------------------------------------------- # +# Fused GroupedTensor path under torch.compile +# --------------------------------------------------------------------------- # + +# GroupedTensorStorage payload slots, in ``prepare_for_saving`` order. +_GX_PAYLOAD_KEYS = ( + "data", + "columnwise_data", + "scale_inv", + "columnwise_scale_inv", + "amax", + "columnwise_amax", + "scale", + "first_dims", + "last_dims", + "tensor_offsets", +) + + +@dataclass(slots=True) +class GroupedLinearFusedFwdArgs: + """Single-argument bag for the fused GroupedTensor forward path.""" + + # --- Differentiable tensors --- + inp: torch.Tensor + weights: List[TensorOrQuantized] + biases: List[torch.Tensor] + + # --- Non-differentiable tensors --- + weight_workspaces: List[TensorOrQuantized] + # Split sizes stay a device tensor on this path: no host sync. + m_splits_tensor: torch.Tensor + skip_fp8_weight_update: Optional[torch.Tensor] + + # --- requires_grad flags --- + input_requires_grad: bool + weights_requires_grad: bool + bias_requires_grad: bool + + # --- Quantizers --- + input_quantizers: List[Quantizer] + weight_quantizers: List[Quantizer] + grad_input_quantizers: List[Quantizer] + grad_weight_quantizers: List[Quantizer] + grad_output_quantizers: List[Quantizer] + + # --- Static geometry / config --- + num_gemms: int + in_features: int + out_features: int + activation_dtype: torch.dtype + fp8: bool + use_bias: bool + is_first_microbatch: Optional[bool] + cache_weight: bool + is_grad_enabled: bool + fprop_use_split_accumulator: bool + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool + compiled_op: bool = False + + def compile_unsupported_reason(self) -> Optional[str]: + """Reason this fused config can't use the compiled custom op (else None). + + The generic exclusions (debug, offloading, save_original_input, + backward_override, calibration) are already filtered out by + ``_is_grouped_tensor_path_supported`` before this path is selected. + """ + if is_distributed_weight(self.weights[0]): + return "a DistributedWeight (custom weight parallelism, e.g. GTP)" + if self.fp8 and not isinstance(self.input_quantizers[0], Float8CurrentScalingQuantizer): + # MXFP8 / NVFP4 / block-scaling grouped storages carry per-split + # host scale offsets, which cannot cross the op boundary. + return ( + "a fused-path FP8 recipe other than per-tensor current scaling " + "(grouped scale offsets are per-split host metadata)" + ) + if any(w.requires_grad != self.weights_requires_grad for w in self.weights): + return "mixed requires_grad across the weights list" + for quantizer_list in ( + self.input_quantizers, + self.weight_quantizers, + self.grad_input_quantizers, + self.grad_weight_quantizers, + self.grad_output_quantizers, + ): + for quantizer in quantizer_list: + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return "a quantizer not registered as a torch.compile value-opaque type" + return None + + +@dataclass(slots=True) +class GroupedLinearFusedBwdArgs: + """Single-argument bag for the fused GroupedTensor backward path.""" + + grad_output: Optional[torch.Tensor] = None + # Saved grouped-input payload (``_GX_PAYLOAD_KEYS`` order); the storage + # object itself is rebuilt inside the op from these plus static metadata. + gx_payload: List[torch.Tensor] = None + weights_fp8: List[TensorOrQuantized] = None + m_splits_tensor: Optional[torch.Tensor] = None + dgrad_out: Optional[torch.Tensor] = None + + input_quantizers: List[Quantizer] = None + grad_input_quantizers: List[Quantizer] = None + grad_weight_quantizers: List[Quantizer] = None + grad_output_quantizers: List[Quantizer] = None + + num_gemms: int = 0 + in_features: int = 0 + out_features: int = 0 + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + use_bias: bool = False + requires_dgrad: bool = False + weights_requires_grad: bool = False + is_first_microbatch: Optional[bool] = None + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD + # Whether the saved grouped input exists (weights require grad) and its + # storage flags (static per recipe). + gx_present: bool = False + gx_swizzled: bool = False + reduce_and_update_bwd_fp8_tensors: bool = False + compiled_op: bool = False + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" + saved = restore_from_func_ctx(ctx) + n_payload = len(_GX_PAYLOAD_KEYS) + self.gx_payload = list(saved[:n_payload]) + self.weights_fp8 = list(saved[n_payload : n_payload + self.num_gemms]) + + +def _gx_scale_inv_offsets(num_gemms: int) -> List[int]: + """Per-tensor scale offsets for the per-tensor current-scaling recipe.""" + return list(range(num_gemms + 1)) + + +def _rebuild_grouped_input(args: GroupedLinearFusedBwdArgs) -> Optional[GroupedTensorStorage]: + """Rebuild the saved grouped input storage from its flat payload.""" + if not args.gx_present: + return None + payload = dict(zip(_GX_PAYLOAD_KEYS, args.gx_payload)) + tokens = int(payload["data"].numel() // args.in_features) if not args.fp8 else None + if tokens is None: + data = payload["data"] if payload["data"] is not None else payload["columnwise_data"] + tokens = data.numel() // args.in_features + quantizer = args.input_quantizers[0] if args.fp8 else None + offsets_kwargs = {} + if args.fp8: + offsets_kwargs = { + "scale_inv_offsets": ( + _gx_scale_inv_offsets(args.num_gemms) if payload["scale_inv"] is not None else None + ), + "columnwise_scale_inv_offsets": ( + _gx_scale_inv_offsets(args.num_gemms) + if payload["columnwise_scale_inv"] is not None + else None + ), + } + return GroupedTensorStorage( + shape=(tokens, args.in_features), + dtype=args.activation_dtype, + num_tensors=args.num_gemms, + quantizer=quantizer, + data=payload["data"], + columnwise_data=payload["columnwise_data"], + scale_inv=payload["scale_inv"], + columnwise_scale_inv=payload["columnwise_scale_inv"], + amax=payload["amax"], + columnwise_amax=payload["columnwise_amax"], + scale=payload["scale"], + first_dims=payload["first_dims"], + last_dims=payload["last_dims"], + tensor_offsets=payload["tensor_offsets"], + with_gemm_swizzled_scales=args.gx_swizzled, + **offsets_kwargs, + ) + + +def _grouped_linear_fused_forward_impl(args: GroupedLinearFusedFwdArgs) -> Tuple[Any, ...]: + """Fused GroupedTensor forward: mirrors ``_forward_grouped_tensor`` with the + saved state expressed as a flat payload + alias tags.""" + inp = args.inp + num_gemms = args.num_gemms + device = inp.device + in_features = args.in_features + out_features = args.out_features + activation_dtype = args.activation_dtype + fp8 = args.fp8 + is_grad_enabled = args.is_grad_enabled + weight_requires_grad = args.weights_requires_grad + + split_sizes = args.m_splits_tensor.to(device=device) + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + + inp_view = inp.reshape(-1, in_features) + x = cast_if_needed(inp_view, activation_dtype) + if fp8: + input_quantizer = args.input_quantizers[0] + input_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and weight_requires_grad, + ) + input_quantizer.optimize_for_gemm = True + grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes) + else: + grouped_x = _GroupedLinear._make_grouped_tensor( + x, + num_gemms=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=in_features, + dtype=activation_dtype, + ) + + columnwise_usage = is_grad_enabled and args.input_requires_grad + weights_for_gemm, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( + args.weights, + args.weight_quantizers, + args.weight_workspaces, + with_quantized_compute=fp8, + columnwise_usage=columnwise_usage, + activation_dtype=activation_dtype, + is_first_microbatch=args.is_first_microbatch, + skip_fp8_weight_update=args.skip_fp8_weight_update, + cache_weight=args.cache_weight, + ) + + out = torch.empty((x.size(0), out_features), dtype=activation_dtype, device=device) + grouped_out = _GroupedLinear._make_grouped_tensor( + out, + num_gemms=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=out_features, + dtype=activation_dtype, + ) + + grouped_bias = None + if args.use_bias: + grouped_bias = _GroupedLinear._make_grouped_bias( + args.biases, + num_gemms=num_gemms, + out_features=out_features, + dtype=activation_dtype, + ) + + general_grouped_gemm_for_grouped_tensor( + weights_for_gemm, + grouped_x, + grouped_out, + layout="TN", + bias=grouped_bias, + use_split_accumulator=args.fprop_use_split_accumulator, + ) + + tensors_to_save = None + ctx_attrs = None + if is_grad_enabled: + n_payload = len(_GX_PAYLOAD_KEYS) + aliases: List[Optional[Tuple]] = [None] * (n_payload + num_gemms) + gx_payload: List[Optional[torch.Tensor]] = [None] * n_payload + gx_present = weight_requires_grad + if gx_present: + # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data + # in backward pass) + if fp8 and grouped_x.columnwise_data is not None: + grouped_x.rowwise_data = None + grouped_x.scale_inv = None + gx_payload = list(grouped_x.get_data_tensors()) + # first_dims is the split tensor itself and tensor_offsets derives + # from it: both rebuilt in backward from m_splits_tensor instead of + # being saved (they may alias the op input). + gx_payload[7] = None + gx_payload[9] = None + if not fp8 and gx_payload[0] is not None and x is inp_view: + # No-op cast: the packed data aliases the op input. + gx_payload[0] = None + aliases[0] = ("inp",) + + weights_to_save = list(weights_for_gemm) if args.input_requires_grad else [None] * num_gemms + for i in range(num_gemms): + slot = n_payload + i + if weights_to_save[i] is None: + continue + if weights_to_save[i] is args.weights[i]: + aliases[slot] = ("weights", i) + weights_to_save[i] = None + elif new_workspaces[i] is not None and weights_to_save[i] is new_workspaces[i]: + aliases[slot] = ("new_weight_workspaces", i) + weights_to_save[i] = None + elif ( + args.weight_workspaces + and args.weight_workspaces[i] is not None + and weights_to_save[i] is args.weight_workspaces[i] + ): + aliases[slot] = ("weight_workspaces", i) + weights_to_save[i] = None + + tensors_to_save = (*gx_payload, *weights_to_save) + ctx_attrs = { + "saved_tensor_aliases": tuple(aliases), + "gx_present": gx_present, + "gx_swizzled": bool(getattr(grouped_x, "_with_gemm_swizzled_scales", False)), + } + + out = out.view(-1, *inp.shape[1:-1], out.shape[-1]) + return (out, *new_workspaces, tensors_to_save, ctx_attrs) + + +def _grouped_linear_fused_forward_fake(args: GroupedLinearFusedFwdArgs) -> Tuple[Any, ...]: + """Shape/metadata-only twin of :func:`_grouped_linear_fused_forward_impl`. + + Only mirrors configs the fused compiled gate admits: bf16/fp16 or FP8 + per-tensor current scaling. + """ + inp = args.inp + num_gemms = args.num_gemms + in_features = args.in_features + out_features = args.out_features + activation_dtype = args.activation_dtype + fp8 = args.fp8 + is_grad_enabled = args.is_grad_enabled + weight_requires_grad = args.weights_requires_grad + tokens = math.prod(inp.shape[:-1]) + device = inp.device + + if fp8: + input_quantizer = args.input_quantizers[0] + input_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and weight_requires_grad, + ) + input_quantizer.optimize_for_gemm = True + + # Weight pipeline (same per-GEMM logic as the legacy fake). + columnwise_usage = is_grad_enabled and args.input_requires_grad + new_workspaces: List[Optional[TensorSpec]] = [None] * num_gemms + weights_saved: List[Optional[TensorSpec]] = [None] * num_gemms + weight_aliases: List[Optional[Tuple]] = [None] * num_gemms + for i in range(num_gemms): + if fp8: + args.weight_quantizers[i].set_usage(rowwise=True, columnwise=columnwise_usage) + if args.weights[i].is_quantized: + weight_aliases[i] = ("weights", i) + continue + workspace = args.weight_workspaces[i] + if workspace is not None and not _fake_workspace_valid( + workspace, args.weight_quantizers[i] + ): + workspace = None + if workspace is not None: + weight_aliases[i] = ("weight_workspaces", i) + continue + weightmat = TensorSpec( + shape=tuple(args.weights[i].shape), + dtype=activation_dtype, + quantizer=args.weight_quantizers[i], + device=args.weights[i].device, + ) + if args.cache_weight: + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_workspaces[i] = weightmat + weight_aliases[i] = ("new_weight_workspaces", i) + else: + weights_saved[i] = weightmat + else: + if args.weights[i].dtype == activation_dtype: + weight_aliases[i] = ("weights", i) + else: + weights_saved[i] = TensorSpec( + shape=tuple(args.weights[i].shape), + dtype=activation_dtype, + device=args.weights[i].device, + ) + + out = TensorSpec( + shape=(*tuple(inp.shape[:-1]), out_features), + dtype=activation_dtype, + requires_grad=is_grad_enabled + and (args.input_requires_grad or weight_requires_grad or args.bias_requires_grad), + device=device, + ) + + tensors_to_save = None + ctx_attrs = None + if is_grad_enabled: + n_payload = len(_GX_PAYLOAD_KEYS) + aliases: List[Optional[Tuple]] = [None] * (n_payload + num_gemms) + gx_payload: List[Optional[TensorSpec]] = [None] * n_payload + gx_present = weight_requires_grad + if gx_present: + total = tokens * in_features + + def _spec(numel, dtype): + return TensorSpec(shape=(numel,), dtype=dtype, device=device) + + if not fp8: + if inp.dtype == activation_dtype: + aliases[0] = ("inp",) + else: + gx_payload[0] = _spec(total, activation_dtype) + else: + # FP8 per-tensor current scaling; on Hopper the rowwise data is + # freed after the fprop GEMM when a columnwise copy exists. + has_columnwise = is_grad_enabled and weight_requires_grad + keep_rowwise = not has_columnwise or is_non_tn_fp8_gemm_supported() + if keep_rowwise: + gx_payload[0] = _spec(total, torch.uint8) # data + gx_payload[2] = _spec(num_gemms, torch.float32) # scale_inv + if has_columnwise and not is_non_tn_fp8_gemm_supported(): + gx_payload[1] = _spec(total, torch.uint8) # columnwise_data + gx_payload[3] = _spec(num_gemms, torch.float32) # columnwise_scale_inv + gx_payload[4] = _spec(num_gemms, torch.float32) # amax + gx_payload[6] = _spec(num_gemms, torch.float32) # scale + + for i in range(num_gemms): + if args.input_requires_grad: + aliases[n_payload + i] = weight_aliases[i] + else: + weights_saved[i] = None + + tensors_to_save = (*gx_payload, *weights_saved) + ctx_attrs = { + "saved_tensor_aliases": tuple(aliases), + "gx_present": gx_present, + "gx_swizzled": False, + } + + return (out, *new_workspaces, tensors_to_save, ctx_attrs) + + +def _grouped_linear_fused_setup_ctx( + bwd_args: GroupedLinearFusedBwdArgs, + fwd_args: GroupedLinearFusedFwdArgs, + fwd_outputs: Tuple[Any, ...], + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate the fused backward args and rebuild alias-deduped save slots.""" + num_gemms = fwd_args.num_gemms + new_workspaces = list(fwd_outputs[1:]) + + bwd_args.input_quantizers = fwd_args.input_quantizers + bwd_args.grad_input_quantizers = fwd_args.grad_input_quantizers + bwd_args.grad_weight_quantizers = fwd_args.grad_weight_quantizers + bwd_args.grad_output_quantizers = fwd_args.grad_output_quantizers + + bwd_args.m_splits_tensor = fwd_args.m_splits_tensor + + bwd_args.num_gemms = num_gemms + bwd_args.in_features = fwd_args.in_features + bwd_args.out_features = fwd_args.out_features + bwd_args.activation_dtype = fwd_args.activation_dtype + bwd_args.fp8 = fwd_args.fp8 + bwd_args.use_bias = fwd_args.use_bias + bwd_args.requires_dgrad = fwd_args.input_requires_grad + bwd_args.weights_requires_grad = fwd_args.weights_requires_grad + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator + bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator + bwd_args.gx_present = ctx_attrs["gx_present"] + bwd_args.gx_swizzled = ctx_attrs["gx_swizzled"] + bwd_args.compiled_op = fwd_args.compiled_op + + saved = list(tensors_to_save_from_forward) + for slot, alias in enumerate(ctx_attrs["saved_tensor_aliases"]): + if alias is None: + continue + if alias[0] == "inp": + saved[slot] = fwd_args.inp + elif alias[0] == "weights": + saved[slot] = fwd_args.weights[alias[1]] + elif alias[0] == "new_weight_workspaces": + saved[slot] = new_workspaces[alias[1]] + elif alias[0] == "weight_workspaces": + saved[slot] = fwd_args.weight_workspaces[alias[1]] + return tuple(saved) + + +def _grouped_linear_fused_backward_impl( + args: GroupedLinearFusedBwdArgs, +) -> Tuple[Optional[torch.Tensor], List[Optional[torch.Tensor]], List[Optional[torch.Tensor]]]: + """Fused GroupedTensor backward, mirroring ``_backward_grouped_tensor``.""" + grad_output = args.grad_output + num_gemms = args.num_gemms + device = grad_output.device + + split_sizes = args.m_splits_tensor.to(device=device) + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + + # The saved packed input may be an alias of the full ``inp`` (no-op cast): + # payload slot 0 then holds the multi-dim input; flatten it back. + if args.gx_present and args.gx_payload[0] is not None: + args.gx_payload[0] = args.gx_payload[0].reshape(-1) + grouped_x = _rebuild_grouped_input(args) + if grouped_x is not None: + grouped_x.first_dims = split_sizes + grouped_x.tensor_offsets = base_split_offsets * args.in_features + + grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) + dy_2d = cast_if_needed(grad_output_view, args.activation_dtype) + dbias_packed = None + if args.fp8: + grad_output_quantizer = args.grad_output_quantizers[0] + grad_output_quantizer.set_usage( + rowwise=args.requires_dgrad, + columnwise=args.weights_requires_grad, + ) + grad_output_quantizer.optimize_for_gemm = True + grouped_dy = tex.group_quantize(dy_2d, grad_output_quantizer, num_gemms, split_sizes) + else: + grouped_dy = _GroupedLinear._make_grouped_tensor( + dy_2d, + num_gemms=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=args.out_features, + dtype=args.activation_dtype, + ) + + grad_biases: List[Optional[torch.Tensor]] = [None] * num_gemms + if args.use_bias: + if dbias_packed is None: + dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_gemms) + grad_biases = [ + dbias_packed[i].to(dtype=args.activation_dtype, copy=args.compiled_op) + for i in range(num_gemms) + ] + + dgrad = None + if args.requires_dgrad: + for weight in args.weights_fp8: + if isinstance(weight, QuantizedTensorStorage): + weight.update_usage(columnwise_usage=True) + dgrad = torch.empty( + (dy_2d.size(0), args.in_features), dtype=args.activation_dtype, device=device + ) + grouped_dgrad = _GroupedLinear._make_grouped_tensor( + dgrad, + num_gemms=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=args.in_features, + dtype=args.activation_dtype, + ) + general_grouped_gemm_for_grouped_tensor( + list(args.weights_fp8), + grouped_dy, + grouped_dgrad, + layout="NN", + use_split_accumulator=args.dgrad_use_split_accumulator, + ) + + if args.weights_requires_grad: + if args.compiled_op: + # Packed allocation would make the returned wgrads alias each other. + wgrad_list = [ + torch.empty( + (args.out_features, args.in_features), + dtype=args.activation_dtype, + device=device, + ) + for _ in range(num_gemms) + ] + else: + wgrad_packed = torch.empty( + num_gemms, + args.out_features, + args.in_features, + dtype=args.activation_dtype, + device=device, + ) + wgrad_list = [wgrad_packed[i] for i in range(num_gemms)] + general_grouped_gemm_for_grouped_tensor( + grouped_x, + grouped_dy, + wgrad_list, + layout="NT", + use_split_accumulator=args.wgrad_use_split_accumulator, + ) + else: + wgrad_list = [None] * num_gemms + + if not args.use_bias: + grad_biases = [None] * num_gemms + + dgrad_out = None + if args.requires_dgrad: + dgrad_out = dgrad.view(*grad_output.shape[:-1], args.in_features) + return (dgrad_out, wgrad_list, grad_biases) + + +def _grouped_linear_fused_backward_fake( + args: GroupedLinearFusedBwdArgs, +) -> Tuple[Optional[TensorSpec], List[Optional[TensorSpec]], List[Optional[TensorSpec]]]: + """Allocation-free fake of :func:`_grouped_linear_fused_backward_impl`.""" + num_gemms = args.num_gemms + grad_output = args.grad_output + device = grad_output.device + + if args.fp8: + grad_output_quantizer = args.grad_output_quantizers[0] + grad_output_quantizer.set_usage( + rowwise=args.requires_dgrad, + columnwise=args.weights_requires_grad, + ) + + dgrad = None + if args.requires_dgrad: + dgrad = TensorSpec( + shape=(*tuple(grad_output.shape[:-1]), args.in_features), + dtype=args.activation_dtype, + device=device, + ) + + wgrad_list: List[Optional[TensorSpec]] = [None] * num_gemms + if args.weights_requires_grad: + wgrad_list = [ + TensorSpec( + shape=(args.out_features, args.in_features), + dtype=args.activation_dtype, + device=device, + ) + for _ in range(num_gemms) + ] + + grad_biases: List[Optional[TensorSpec]] = [None] * num_gemms + if args.use_bias: + grad_biases = [ + TensorSpec(shape=(args.out_features,), dtype=args.activation_dtype, device=device) + for _ in range(num_gemms) + ] + + return (dgrad, wgrad_list, grad_biases) + + +# Custom op for the fused GroupedTensor path under ``torch.compile``. +_grouped_linear_fused_op = register_custom_op( + op_name="grouped_linear_fused", + input_tensors_for_grad=["inp", "weights", "biases"], + fwd_arg_type=GroupedLinearFusedFwdArgs, + fwd_impl=_grouped_linear_fused_forward_impl, + fwd_fake_impl=_grouped_linear_fused_forward_fake, + setup_context=_grouped_linear_fused_setup_ctx, + bwd_arg_type=GroupedLinearFusedBwdArgs, + bwd_impl=_grouped_linear_fused_backward_impl, + bwd_fake_impl=_grouped_linear_fused_backward_fake, +) + + +@no_torch_dynamo() +def _grouped_linear_eager( + inp: torch.Tensor, + m_splits: torch.Tensor, + fwd_args: GroupedLinearFwdArgs, + weights_and_biases: Tuple[torch.Tensor, ...], + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, list]: + """Run ``_GroupedLinear`` eagerly, bypassing Dynamo.""" + if fwd_args.m_splits is None: + fwd_args.m_splits = tuple(m_splits.tolist()) + if is_grad_enabled: + return _GroupedLinear.apply(inp, fwd_args, *weights_and_biases) + return _GroupedLinear.forward(None, inp, fwd_args, *weights_and_biases) + + __all__ = ["GroupedLinear"] @@ -517,7 +2230,7 @@ def _is_grouped_tensor_path_supported( device_capability = get_device_compute_capability() if not (9, 0) <= device_capability <= (11, 0): return False - cublaslt_version = tex.get_cublasLt_version() + cublaslt_version = _get_cublaslt_version() if cublaslt_version < 130300: return False if device_capability < (10, 0) and cublaslt_version < 130400: @@ -846,364 +2559,77 @@ def _forward_grouped_tensor( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces - # pylint: disable=keyword-arg-before-vararg @staticmethod def forward( ctx, inp: torch.Tensor, - m_splits: torch.Tensor, - non_tensor_args: Tuple, - out: Optional[torch.Tensor], - dgrad_out: Optional[torch.Tensor], + fwd_args: GroupedLinearFwdArgs, *weights_and_biases, ) -> Tuple[torch.Tensor, list]: - # pylint: disable=missing-function-docstring - - # Reduce number of arguments to autograd function in order - # to reduce CPU overhead due to pytorch arg checking. - ( - use_bias, - is_first_microbatch, - fp8, - fp8_calibration, - wgrad_store, - input_quantizers, - weight_quantizers, - output_quantizers, - grad_input_quantizers, - grad_weight_quantizers, - grad_output_quantizers, - fuse_wgrad_accumulation, - cpu_offloading, - sequence_parallel, - activation_dtype, - is_grad_enabled, - weight_workspaces, - cache_weight, - skip_fp8_weight_update, - save_original_input, - delayed_scaling_input_quantizer, - unsafe_requantization_input_quantizer, - debug, - ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override - else: - backward_override = None - if backward_override == "high_precision": - save_original_input = True - elif backward_override == "dequantized": - save_original_input = False + """Forward pass: compute grouped linear output and set up autograd context. - num_gemms = len(m_splits) - weights = weights_and_biases[:num_gemms] - biases = weights_and_biases[num_gemms:] - device = inp.device - weight_requires_grad = weights[0].requires_grad - - origin_weights = weights - is_dist_weight = is_distributed_weight(weights[0]) - if is_dist_weight: - weights = materialize_weight_for_forward(weights) - - backward_needs_input = is_grad_enabled and weight_requires_grad - if backward_override is None and save_original_input and backward_needs_input: - if delayed_scaling_input_quantizer is not None: - if FP8GlobalStateManager.get_fp8_recipe().custom(): - warnings.warn( - "save_original_input is incompatible with delayed-scaling quantizers " - "(Float8Quantizer). Disabling save_original_input for this module.", - stacklevel=2, - ) - save_original_input = False - else: - raise ValueError( - "DelayedScaling recipe is not supported with save_original_input" - ) - - # Megatron-Core may enable this automatically to reuse an activation - # already retained by an upstream operation. The resolved quantizer - # generation is classified once in ``_validate_quantizer_generation``. - if save_original_input and unsafe_requantization_input_quantizer is not None: - warnings.warn( - "Ignoring save_original_input=True because the input quantizer cannot " - "safely reconstruct the backward operand from the original input " - f"({unsafe_requantization_input_quantizer}).", - stacklevel=2, - ) - save_original_input = False - - # Configure quantizers - if input_quantizers[0] is not None: - for input_quantizer in input_quantizers: - input_quantizer.set_usage( - rowwise=True, - columnwise=( - is_grad_enabled - and weight_requires_grad - and not save_original_input - and backward_override is None - ), - ) - columnwise_usage = is_grad_enabled and inp.requires_grad - if backward_override is not None: - columnwise_usage = False - if not columnwise_usage: - columnwise_usage = ( - is_fp8_activation_recompute_enabled() - and not in_fp8_activation_recompute_phase() - ) - # No need to set the quantizer states if weight is already quantized - # for debug mode we create quantizer every iteration, thus we need to set the quantizer states - if weight_quantizers[0] is not None and ( - not isinstance(weights[0], QuantizedTensorStorage) or debug - ): - for weight_quantizer in weight_quantizers: - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - elif isinstance(weights[0], QuantizedTensorStorage): - # If weights are already quantized, no need to set quantizer states - weight_quantizers = [weight._quantizer for weight in weights] - if output_quantizers[0] is not None: - for output_quantizer in output_quantizers: - output_quantizer.set_usage(rowwise=True, columnwise=False) - - # Initialize input tensors - in_features = weights[0].size(-1) - if inp.size(-1) != in_features: - raise ValueError( - f"Input tensor (shape={tuple(inp.size())}) is not compatible with " - f"weight tensor (shape={tuple(weights[0].size())})" - ) + ``inp`` and the weights / biases 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. + """ + num_gemms = fwd_args.num_gemms + fwd_args.inp = inp + fwd_args.weights = list(weights_and_biases[:num_gemms]) + fwd_args.biases = list(weights_and_biases[num_gemms:]) - if _GroupedLinear._is_grouped_tensor_path_supported( - fp8=fp8, - fp8_calibration=fp8_calibration, - debug=debug, - cpu_offloading=cpu_offloading, - backward_override=backward_override, - save_original_input=save_original_input, - activation_dtype=activation_dtype, - input_quantizers=input_quantizers, - output_quantizers=output_quantizers, - ): + if fwd_args.use_grouped_tensor_path: return _GroupedLinear._forward_grouped_tensor( ctx, inp=inp, - m_splits=m_splits, - use_bias=use_bias, - is_first_microbatch=is_first_microbatch, - fp8=fp8, - wgrad_store=wgrad_store, - input_quantizers=input_quantizers, - weight_quantizers=weight_quantizers, - grad_input_quantizers=grad_input_quantizers, - grad_weight_quantizers=grad_weight_quantizers, - grad_output_quantizers=grad_output_quantizers, - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - activation_dtype=activation_dtype, - is_grad_enabled=is_grad_enabled, - weight_workspaces=weight_workspaces, - cache_weight=cache_weight, - skip_fp8_weight_update=skip_fp8_weight_update, - weights=weights, - biases=biases, - out=out, - dgrad_out=dgrad_out, + m_splits=fwd_args.m_splits_tensor, + use_bias=fwd_args.use_bias, + is_first_microbatch=fwd_args.is_first_microbatch, + fp8=fwd_args.fp8, + wgrad_store=fwd_args.wgrad_store, + input_quantizers=fwd_args.input_quantizers, + weight_quantizers=fwd_args.weight_quantizers, + grad_input_quantizers=fwd_args.grad_input_quantizers, + grad_weight_quantizers=fwd_args.grad_weight_quantizers, + grad_output_quantizers=fwd_args.grad_output_quantizers, + fuse_wgrad_accumulation=fwd_args.fuse_wgrad_accumulation, + activation_dtype=fwd_args.activation_dtype, + is_grad_enabled=fwd_args.is_grad_enabled, + weight_workspaces=fwd_args.weight_workspaces, + cache_weight=fwd_args.cache_weight, + skip_fp8_weight_update=fwd_args.skip_fp8_weight_update, + weights=fwd_args.weights, + biases=fwd_args.biases, + out=fwd_args.out, + dgrad_out=fwd_args.dgrad_out, ) - # Convert splits to list of ints for compatibility with split functions - m_splits = m_splits.tolist() - - inp_view = inp.reshape(-1, in_features) - # Disable bulk allocation when CPU offloading is active: offloading skips small - # tensors (like scales), but bulk allocation shares storage across all tensors, - # so if scales can't be offloaded, nothing in the group can be offloaded. - inputmats = _split_quantize( - inp_view, - m_splits, - with_quantized_output=fp8 or debug, - quantizers=input_quantizers, - dtype=activation_dtype, - with_debug_quantizers=debug, - disable_bulk_allocation=cpu_offloading, - ) - - if cpu_offloading: - start_offload(*inputmats) - - # Initialize weights - weights_fp8: list - new_workspaces = [None] * num_gemms - if fp8 or debug: - weights_fp8 = [] - update_ws = is_first_microbatch is None or is_first_microbatch - for i in range(num_gemms): - weight_fp8, new_workspaces[i] = quantize_weight( - tensor=weights[i], - quantizer=weight_quantizers[i], - workspace=weight_workspaces[i] if weight_workspaces else None, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - weights_fp8.append(weight_fp8) - - else: - weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights] - - # Initialize biases - bias_dtype = activation_dtype - if fp8 and activation_dtype == torch.float32: - bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias - biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases - # Initialize output tensor - out = _GroupedLinear._validate_or_alloc_output( - out, - sum(m_splits), - weights_fp8[0].size(0), - activation_dtype, - device, - ) - - # Choose whether to use split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator - - # Perform GEMM - general_grouped_gemm( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - ) - - if fp8_calibration: - for i in range(num_gemms): - input_quantizers[i].calibrate(inputmats[i]) - weight_quantizers[i].calibrate(weights[i]) - - if cpu_offloading: - mark_not_offload(*weights_fp8, *weights) + outputs = _grouped_linear_forward_impl(fwd_args) + out = outputs[0] + new_workspaces = list(outputs[1 : 1 + num_gemms]) + tensors_to_save_from_forward, ctx_attrs = outputs[-2], outputs[-1] - if is_grad_enabled: + if ctx is not None: ctx.use_grouped_tensor_path = False - ctx.weight_quantizers = weight_quantizers - ctx.weights_shape_1 = weights[0].shape[1] - - # TODO: update after #1638 is merged. # pylint: disable=fixme - if weight_requires_grad: - if save_original_input: - inputmats = [None] * num_gemms - inputmats[0] = inp - else: - for inputmat in inputmats: - if isinstance(inputmat, QuantizedTensorStorage): - if backward_override is not None: - # In dequantized mode we should dequantize directly from - # fprop quantized layouts without retargeting usage. - inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) - else: - inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) - else: - inputmats = [None] * num_gemms - - # Original weights are only needed by high_precision dgrad. The weakrefs - # used for fused wgrad accumulation serve a different purpose: restoring - # Python parameter attributes without keeping the parameter alive here. - saved_weights = ( - weights - if backward_override == "high_precision" and inp.requires_grad - else [None] * num_gemms - ) - if is_dist_weight: - # GTP: gathered workspace is transient (re-gathered in backward), don't save it. - weights_fp8 = [None] * num_gemms - saved_weights = origin_weights - tensors_to_save, tensor_objects = prepare_for_saving( - *inputmats, - *weights_fp8, - *saved_weights, - *biases, + bwd_args = GroupedLinearBwdArgs() + tensors_to_save_from_setup = _grouped_linear_setup_ctx( + bwd_args, + fwd_args, + (out, *new_workspaces), + ctx_attrs, + tensors_to_save_from_forward, ) + tensors_to_save, tensor_objects = prepare_for_saving(*tensors_to_save_from_setup) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects - - ctx.grad_input_quantizers = grad_input_quantizers - ctx.grad_output_quantizers = grad_output_quantizers - ctx.grad_weight_quantizers = grad_weight_quantizers - - ctx.weights_requires_grad = weights[0].requires_grad - if fuse_wgrad_accumulation and ctx.weights_requires_grad: - # Keep weakrefs to weights to preserve attributes like main_grad - # when we need to modify the weight python objects - ctx.origin_weight_refs = [weakref.ref(w) for w in weights] - ctx.origin_weights_overwrite_main_grad = getattr( - weights[0], "overwrite_main_grad", False - ) - # This check is needed to ensure that main_grad is not created - # during the forward pass when using MCore FSDP as it creates - # the main_grad buffer lazily before backprop - if hasattr(weights[0], "__fsdp_param__"): - # MCore FSDP creates main_grad lazily before backward - ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] - elif is_dist_weight: - ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)] - else: - ctx.main_grad_funcs = [ - lambda j=i: weights[j].main_grad for i in range(num_gemms) - ] - ctx.device = device - ctx.output_quantizers = output_quantizers - ctx.m_splits = m_splits - ctx.num_gemms = num_gemms - ctx.activation_dtype = activation_dtype - ctx.fp8 = fp8 - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.backward_override = backward_override - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = use_bias - ctx.sequence_parallel = sequence_parallel - ctx.inp_shape = inp.shape - ctx.requires_dgrad = inp.requires_grad - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad(inp, weights[0], biases[0]): - ctx.reduce_and_update_bwd_fp8_tensors = ( - ctx.reduce_and_update_bwd_fp8_tensors - or FP8GlobalStateManager.is_first_fp8_module() + ctx.backward_objects = bwd_args + if fwd_args.fp8 and requires_grad(inp, fwd_args.weights[0], fwd_args.biases[0]): + bwd_args.reduce_and_update_bwd_fp8_tensors = ( + FP8GlobalStateManager.is_first_fp8_module() ) - ctx.wgrad_store = wgrad_store - ctx.debug = debug - ctx.save_original_input = save_original_input - ctx.input_quantizers = input_quantizers - ctx.dgrad_out = dgrad_out + if fwd_args.backward_override is not None: + bwd_args.reduce_and_update_bwd_fp8_tensors = False - # backward overrides - if backward_override is not None: - ctx.fp8 = False - ctx.debug = False - ctx.ub_overlap_ag = False - ctx.ub_overlap_rs_dgrad = False - ctx.ub_bulk_dgrad = False - ctx.ub_bulk_wgrad = False - ctx.grad_input_quantizers = [None] * num_gemms - ctx.grad_weight_quantizers = [None] * num_gemms - ctx.grad_output_quantizers = [None] * num_gemms - ctx.reduce_and_update_bwd_fp8_tensors = False - - # [*, in_features] -> [*, out_features] except first dimension changes for SP - return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + return out, new_workspaces @staticmethod def _backward_grouped_tensor( @@ -1407,262 +2833,25 @@ def backward( # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): if ctx.use_grouped_tensor_path: - return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) - - saved_tensors = restore_from_func_ctx(ctx) - N = ctx.num_gemms - inputmats = saved_tensors[:N] - weights = saved_tensors[N : 2 * N] - saved_weights = saved_tensors[2 * N : 3 * N] - biases = saved_tensors[3 * N : 4 * N] - - # Restore from weakrefs to get original weight python objects - # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) - # Only needed when fuse_wgrad_accumulation is enabled. - origin_weights = [None] * N - main_grads = [None] * N - is_dist_weight = is_distributed_weight(saved_weights[0]) - if is_dist_weight: - origin_weights = saved_weights - if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: - main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] - elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: - origin_weight_refs = ctx.origin_weight_refs - ctx.origin_weight_refs = None - origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] - assert all( - w is not None for w in origin_weights - ), "weight was removed while fuse_wgrad_accumulation=True" - main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] - for origin_weight, main_grad in zip(origin_weights, main_grads): - if main_grad is not None: - origin_weight.main_grad = main_grad - - # Preprocess grad output - grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) - grad_output_reference = ctx.grad_output_quantizers[0] - if ctx.fp8 and isinstance(grad_output_reference, HybridQuantizer): - # Usage is a runtime decision, not part of generation validation. - # Apply it uniformly so dispatch can read the first parent without - # rescanning every expert. - for grad_output_quantizer in ctx.grad_output_quantizers: - grad_output_quantizer.set_usage( - rowwise=ctx.requires_dgrad, - columnwise=ctx.weights_requires_grad, - ) - grad_output, grad_biases = _split_quantize_and_bias( - grad_output_view, - ctx.m_splits, - fp8=ctx.fp8, - debug=ctx.debug, - quantizers=ctx.grad_output_quantizers, - dtype=ctx.activation_dtype, - use_bias=ctx.use_bias, - recipe=ctx.fp8_recipe, - disable_bulk_allocation=ctx.cpu_offloading, - ) - - if is_dist_weight: - accumulate_wgrad_into_param_main_grad = False - elif ctx.is_first_microbatch is not None: - accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch - ) - else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - if is_dist_weight: - weights = materialize_weight_for_backward(origin_weights) - - if ctx.requires_dgrad: - dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8 or ctx.debug: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - dgrad_gemm_use_split_accumulator = ( - recipe.fp8_gemm_dgrad.use_split_accumulator - ) - dgrad = _GroupedLinear._validate_or_alloc_output( - ctx.dgrad_out, - sum(ctx.m_splits), - ctx.weights_shape_1, - ctx.activation_dtype, - ctx.device, - ) - weights_for_dgrad = weights - if ctx.backward_override == "dequantized": - weights_for_dgrad = [ - ( - weight.dequantize(dtype=ctx.activation_dtype) - if isinstance(weight, QuantizedTensorStorage) - else cast_if_needed(weight, ctx.activation_dtype) - ) - for weight in weights - ] - elif ctx.backward_override == "high_precision": - weights_for_dgrad = [ - ( - weight.dequantize(dtype=ctx.activation_dtype) - if isinstance(weight, QuantizedTensorStorage) - else cast_if_needed(weight, ctx.activation_dtype) - ) - for weight in saved_weights - ] - # Make sure weights are available in column-wise format - # for dgrad computation. - for weight in weights_for_dgrad: - if isinstance(weight, QuantizedTensorStorage): - weight.update_usage(columnwise_usage=True) - general_grouped_gemm( - weights_for_dgrad, - grad_output, - [dgrad], - ctx.grad_input_quantizers, - ctx.activation_dtype, - single_output=True, - layout="NN", - m_splits=ctx.m_splits, - grad=True, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - ) - - if ctx.weights_requires_grad: - wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - wgrad_gemm_use_split_accumulator = ( - recipe.fp8_gemm_wgrad.use_split_accumulator - ) - if ctx.fuse_wgrad_accumulation: - wgrad_list = main_grads - else: - wgrad_packed = torch.empty( - ctx.num_gemms, - *weights[0].size(), - dtype=ctx.activation_dtype, - device=ctx.device, - ) - wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] - if is_dist_weight: - # Gathered weights are no longer needed after dgrad GEMM. - del weights - - if ctx.save_original_input: - inp = inputmats[0] - in_features = inp.shape[-1] - inp_view = inp.reshape(-1, in_features) - if ctx.input_quantizers[0] is not None: - for input_quantizer in ctx.input_quantizers: - if isinstance( - input_quantizer, - (Float8Quantizer, Float8CurrentScalingQuantizer), - ): - input_quantizer.set_usage(rowwise=True, columnwise=True) - else: - input_quantizer.set_usage(rowwise=False, columnwise=True) - inputmats = _split_quantize( - inp_view, - ctx.m_splits, - with_quantized_output=ctx.fp8 or ctx.debug, - quantizers=ctx.input_quantizers, - dtype=ctx.activation_dtype, - with_debug_quantizers=ctx.debug, - disable_bulk_allocation=ctx.cpu_offloading, - ) - elif ctx.backward_override == "dequantized": - inputmats_dequant = [] - for inputmat in inputmats: - if isinstance(inputmat, QuantizedTensorStorage): - inputmats_dequant.append( - inputmat.dequantize(dtype=ctx.activation_dtype) - ) - else: - inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype)) - inputmats = inputmats_dequant - grouped_gemm_wgrad = functools.partial( - general_grouped_gemm, - quantization_params=ctx.grad_weight_quantizers, - out_dtype=ctx.activation_dtype, - layout="NT", - grad=True, - m_splits=ctx.m_splits, - use_bias=ctx.use_bias if grad_biases[0] is None else None, - bias=biases, - use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=( - accumulate_wgrad_into_param_main_grad - if not is_dist_weight - and not getattr(ctx, "origin_weights_overwrite_main_grad", False) - else False - ), - ) - # WGRAD - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) - else: - _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) - - for i in range(ctx.num_gemms): - if grad_biases[i] is None: - grad_biases[i] = grad_biases_[i] - del grad_biases_ - - # Deallocate input tensor - clear_tensor_data(*inputmats) - - def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): - if ctx.weights_requires_grad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( - weight, "grad_added_to_main_grad" - ): - weight.grad_added_to_main_grad = True - if getattr(weight, "zero_out_wgrad", False): - wgrad = get_dummy_wgrad( - list(main_grad.shape), - weight.dtype, - zero=True, - ) - else: - wgrad = get_dummy_wgrad( - list(main_grad.shape), - weight.dtype, - ) - elif ctx.fuse_wgrad_accumulation: - wgrad = None - else: - wgrad = None - return wgrad - - if is_dist_weight: - wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) - else: - wgrad_list = [ - handle_custom_ddp_from_mcore(weight, main_grad, wgrad) - for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) - ] - else: - wgrad_list = [None] * ctx.num_gemms - - if not ctx.use_bias or ( - ctx.wgrad_store is not None - and ctx.wgrad_store.delay_wgrad_compute() - and not ctx.fp8 - ): - grad_biases = [None] * ctx.num_gemms - - if ctx.reduce_and_update_bwd_fp8_tensors: + result = _GroupedLinear._backward_grouped_tensor(ctx, grad_output) + # Legacy return layout: (dgrad, m_splits, non_tensor_args, out, + # dgrad_out, *wgrads, *grad_biases) -> map onto (inp, fwd_args, + # *weights_and_biases). + return (result[0], None, *result[5:]) + + bwd_args: GroupedLinearBwdArgs = ctx.backward_objects + bwd_args.grad_output = grad_output + bwd_args.setup_saved_tensors(ctx) + dgrad, wgrad_list, grad_biases = _grouped_linear_backward_impl(bwd_args) + 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) so they don't outlive backward via ctx under retain_graph. + ctx.backward_objects = None + del bwd_args + + if reduce_and_update_bwd_fp8_tensors: FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - return ( - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, - None, # m_splits - None, # non_tensor_args - None, # out - None, # dgrad_out - *wgrad_list, - *grad_biases, - ) + return (dgrad, None, *wgrad_list, *grad_biases) class GroupedLinear(TransformerEngineBaseModule): @@ -2088,6 +3277,16 @@ def reset_parameters(self, defer_init=False): self.make_grouped_weights(defer_init=defer_init) elif self.single_grouped_bias: self._make_grouped_biases() + if not defer_init: + # Allocate the process-global grouped cuBLAS workspace eagerly: + # under torch.compile the first grouped GEMM can run inside + # CUDA-graph capture, and a workspace first allocated there would + # live in the graph pool. + weight = getattr(self, "weight0", None) + if weight is None: + weight = getattr(self, "weight", None) + if weight is not None and weight.device.type == "cuda": + get_cublas_workspace(weight.device.index, False, True) def set_tensor_parallel_attributes(self, defer_init=False) -> None: """Set attributes needed for TP""" @@ -2146,7 +3345,7 @@ def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None if not has_grouped_weight and has_per_gemm_weights: per_gemm_weights = [state_dict.pop(key) for key in per_gemm_weight_keys] per_gemm_weights = [ - weight.dequantize() if isinstance(weight, QuantizedTensorStorage) else weight + (weight.dequantize() if isinstance(weight, QuantizedTensorStorage) else weight) for weight in per_gemm_weights ] state_dict[grouped_weight_key] = torch.stack(per_gemm_weights, dim=0) @@ -2229,17 +3428,29 @@ def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False) return super().load_state_dict(state_dict_copy, strict=strict, assign=assign) def _load_from_state_dict( - self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, ): """Load state, including compatibility across grouped-weight checkpoint formats.""" self._remap_grouped_weight_state_dict_keys(state_dict, prefix) self._remap_grouped_bias_state_dict_keys(state_dict, prefix) super()._load_from_state_dict( - state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, ) - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -2298,11 +3509,15 @@ def forward( is_first_microbatch = False # Make sure splits are in expected format + m_splits_host: Optional[Tuple[int, ...]] = None if not isinstance(m_splits, torch.Tensor): + m_splits_host = tuple(int(s) for s in m_splits) # Convert list of ints to tensor for backward compatibility m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") elif m_splits.dtype != torch.int64: m_splits = m_splits.to(dtype=torch.int64) + if m_splits_host is None and not torch.compiler.is_compiling(): + m_splits_host = tuple(m_splits.tolist()) if m_splits.size() != (num_gemms,): raise ValueError( f"Shape of splits tensor ({tuple(m_splits.size())}) " @@ -2350,12 +3565,9 @@ def forward( for q in weight_quantizers: q.optimize_for_gemm = optimize_for_gemm - if is_grad_enabled: - linear_fn = _GroupedLinear.apply - autograd_ctx = [] - else: - linear_fn = _GroupedLinear.forward - autograd_ctx = [None] + use_compiled_op = torch.compiler.is_compiling() and _grouped_linear_op is not None + if _grouped_linear_op is None and torch.compiler.is_compiling(): + warn_if_compile_disabled() cache_weight = is_first_microbatch is not None weight_workspaces = ( @@ -2364,41 +3576,215 @@ def forward( else [None] * num_gemms ) - non_tensor_args = ( - self.apply_bias, - is_first_microbatch, - self.fp8, - self.fp8_calibration, - self.wgrad_store, - input_quantizers, - weight_quantizers, - output_quantizers, - grad_input_quantizers, - grad_weight_quantizers, - grad_output_quantizers, - self.fuse_wgrad_accumulation, - is_cpu_offload_enabled(), - self.sequence_parallel, - self.activation_dtype, - is_grad_enabled, - weight_workspaces, - cache_weight, - skip_fp8_weight_update, - self.save_original_input, - self._delayed_scaling_input_quantizer, - self._unsafe_requantization_input_quantizer, - debug, - ) - out, new_workspaces = linear_fn( - *autograd_ctx, - inp, - m_splits, - non_tensor_args, - out, - dgrad_out, - *weight_tensors, - *bias_tensors, + weight_requires_grad = weight_tensors[0].requires_grad + fprop_use_split_accumulator = _2X_ACC_FPROP + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD + native_bgrad_recipe_ok = False + if self.fp8: + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_fprop"): + fprop_use_split_accumulator = _recipe.fp8_gemm_fprop.use_split_accumulator + if hasattr(_recipe, "fp8_gemm_dgrad"): + dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator + if hasattr(_recipe, "fp8_gemm_wgrad"): + wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator + native_bgrad_recipe_ok = ( + _recipe.delayed() or _recipe.float8_current_scaling() or _recipe.mxfp8() + ) + else: + backward_override = None + + # Resolve the save_original_input runtime flips before building the + # args, so the impl and the compile-time fake see one final flag. + save_original_input = self.save_original_input + if backward_override == "high_precision": + save_original_input = True + elif backward_override == "dequantized": + save_original_input = False + backward_needs_input = is_grad_enabled and weight_requires_grad + if backward_override is None and save_original_input and backward_needs_input: + if self._delayed_scaling_input_quantizer is not None: + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError( + "DelayedScaling recipe is not supported with save_original_input" + ) + + # Megatron-Core may enable this automatically to reuse an activation + # already retained by an upstream operation. The resolved quantizer + # generation is classified once in ``_validate_quantizer_generation``. + if save_original_input and self._unsafe_requantization_input_quantizer is not None: + warnings.warn( + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({self._unsafe_requantization_input_quantizer}).", + stacklevel=2, + ) + save_original_input = False + + cpu_offloading = is_cpu_offload_enabled() + use_grouped_tensor_path = _GroupedLinear._is_grouped_tensor_path_supported( + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + debug=debug, + cpu_offloading=cpu_offloading, + backward_override=backward_override, + save_original_input=save_original_input, + activation_dtype=self.activation_dtype, + input_quantizers=input_quantizers, + output_quantizers=output_quantizers, ) + wgrad_store = ( + self.wgrad_store + if self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute() + else None + ) + + fwd_args = GroupedLinearFwdArgs( + # tensors + inp=inp, + weights=list(weight_tensors), + biases=list(bias_tensors), + weight_workspaces=weight_workspaces, + out=out, + dgrad_out=dgrad_out, + skip_fp8_weight_update=skip_fp8_weight_update, + m_splits_tensor=m_splits, + # requires_grad flags + input_requires_grad=inp.requires_grad, + weights_requires_grad=weight_requires_grad, + bias_requires_grad=(bias_tensors[0].requires_grad if self.apply_bias else False), + # quantizers + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + output_quantizers=output_quantizers, + grad_input_quantizers=grad_input_quantizers, + grad_weight_quantizers=grad_weight_quantizers, + grad_output_quantizers=grad_output_quantizers, + # split geometry + m_splits=m_splits_host, + num_gemms=num_gemms, + # numerical / dtype config + activation_dtype=self.activation_dtype, + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + save_original_input=save_original_input, + backward_override=backward_override, + fprop_use_split_accumulator=fprop_use_split_accumulator, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, + native_bgrad_recipe_ok=native_bgrad_recipe_ok, + debug=debug, + # weight-workspace caching + is_first_microbatch=is_first_microbatch, + cache_weight=cache_weight, + # fused GroupedTensor path + use_grouped_tensor_path=use_grouped_tensor_path, + single_grouped_param=self.single_grouped_weight or self.single_grouped_bias, + # misc + use_bias=self.apply_bias, + sequence_parallel=self.sequence_parallel, + fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, + wgrad_store=wgrad_store, + cpu_offloading=cpu_offloading, + is_grad_enabled=is_grad_enabled, + ) + + if use_compiled_op and use_grouped_tensor_path: + # The fused GroupedTensor path has its own custom op: m_splits + # stays a device tensor (no host sync, dropless-MoE friendly). + fused_args = GroupedLinearFusedFwdArgs( + inp=inp, + weights=list(weight_tensors), + biases=list(bias_tensors), + weight_workspaces=weight_workspaces, + m_splits_tensor=m_splits, + skip_fp8_weight_update=skip_fp8_weight_update, + input_requires_grad=inp.requires_grad, + weights_requires_grad=weight_requires_grad, + bias_requires_grad=( + bias_tensors[0].requires_grad if self.apply_bias else False + ), + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + grad_input_quantizers=grad_input_quantizers, + grad_weight_quantizers=grad_weight_quantizers, + grad_output_quantizers=grad_output_quantizers, + num_gemms=num_gemms, + in_features=self.in_features, + out_features=self.out_features, + activation_dtype=self.activation_dtype, + fp8=self.fp8, + use_bias=self.apply_bias, + is_first_microbatch=is_first_microbatch, + cache_weight=cache_weight, + is_grad_enabled=is_grad_enabled, + fprop_use_split_accumulator=fprop_use_split_accumulator, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, + ) + fused_reason = ( + fused_args.compile_unsupported_reason() + if _grouped_linear_fused_op is not None + else "custom-op registration unavailable" + ) + if fused_reason is None and (out is not None or dgrad_out is not None): + fused_reason = ( + "a user-provided out/dgrad_out buffer (the op would return an input alias)" + ) + if fused_reason is None and self.fuse_wgrad_accumulation: + fused_reason = "fuse_wgrad_accumulation (main_grad)" + if fused_reason is None and wgrad_store is not None: + fused_reason = "delayed wgrad compute (wgrad_store)" + if fused_reason is None and ( + self.single_grouped_weight or self.single_grouped_bias + ): + fused_reason = "single_grouped_weight/single_grouped_bias parameter views" + if fused_reason is not None: + torch._dynamo.graph_break( + msg=f"te.GroupedLinear (fused) falling back to eager: {fused_reason}" + ) + warn_compile_eager_fallback(fused_reason) + use_compiled_op = False + else: + fused_args.compiled_op = True + outputs = _grouped_linear_fused_op(fused_args) + out = outputs[0] + new_workspaces = list(outputs[1:]) + elif use_compiled_op: + fallback_reason = fwd_args.compile_unsupported_reason() + if fallback_reason is not None: + # Explicit break so fullgraph=True errors show the reason + # (warnings.warn below would break the graph inscrutably). + torch._dynamo.graph_break( + msg=f"te.GroupedLinear falling back to eager: {fallback_reason}" + ) + warn_compile_eager_fallback(fallback_reason) + use_compiled_op = False + + if use_compiled_op and not use_grouped_tensor_path: + fwd_args.compiled_op = True + fwd_args.m_splits_tensor = None + check_grouped_gemm_dims(inp, weight_tensors[0], fwd_args.m_splits, self.fp8) + outputs = _grouped_linear_op(fwd_args) + out = outputs[0] + new_workspaces = list(outputs[1:]) + elif not use_compiled_op: + out, new_workspaces = _grouped_linear_eager( + inp, + m_splits, + fwd_args, + (*weight_tensors, *bias_tensors), + is_grad_enabled, + ) if cache_weight: for i, ws in enumerate(new_workspaces): diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index c6e0c63ebf..27c63a5a92 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -697,6 +697,45 @@ def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None: ) +def check_grouped_gemm_dims( + inp: torch.Tensor, + weight: torch.Tensor, + m_splits: Sequence[int], + fp8: bool, +) -> None: + """Grouped analog of :func:`check_gemm_dims`: emit the grouped TN GEMM dim + constraints as ``torch._check`` guards at trace time. torch.compile path + only (``m_splits`` entries are host-side ints there); eager validation + lives in the op impl. + """ + # pylint: disable=protected-access + torch._check( + inp.shape[-1] == weight.shape[-1], + lambda: "GEMM not possible: input last dim must equal in_features", + ) + torch._check( + math.prod(inp.shape[:-1]) == sum(m_splits), + lambda: "GEMM not possible: m_splits must sum to the input's token count", + ) + if not fp8: + return + torch._check( + weight.shape[0] % 8 == 0 and weight.shape[-1] % 16 == 0, + lambda: ( + "FP8 execution requires the weight's out_features to be divisible by 8" + " and in_features to be divisible by 16" + ), + ) + torch._check( + inp.shape[-1] % 16 == 0, + lambda: "FP8 execution requires the input's last dimension to be divisible by 16", + ) + torch._check( + all(m % 8 == 0 for m in m_splits), + lambda: "FP8 execution requires every m_splits entry to be divisible by 8", + ) + + def is_bf16_compatible() -> bool: """Replaces torch.cuda.is_bf16_compatible() with an explicit check on device compute capability to enforce sm_80 or higher.