diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index d720ca7a4a..6ae070b70d 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -3,6 +3,8 @@ # See LICENSE for license information. import abc +import warnings +from typing import Optional import pytest import torch @@ -19,6 +21,8 @@ except ImportError: _opaque_available = False +from torch._dynamo.utils import counters + import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe @@ -36,7 +40,7 @@ MXFP8Quantizer, NVFP4Quantizer, ) -from utils import recipe_id +from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id from transformer_engine.pytorch.attention.dot_product_attention.backends import ( UnfusedDotProductAttention, ) @@ -540,14 +544,492 @@ def fn(q, k, v, extra): compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead") for _ in range(3): + # Compared against eager rather than only checked for finiteness: a + # replay reusing stale buffers would pass the latter. q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True) - out = compiled(q, k, v, extra) - out.sum().backward() - torch.cuda.synchronize() - assert torch.isfinite(out).all() - assert q.grad is not None - assert k.grad is not None - assert v.grad is not None + inputs = (q, k, v, extra) + replayed = _run_and_capture(compiled, inputs, {}, [q, k, v]) + eager = _run_and_capture(fn, inputs, {}, [q, k, v]) + _assert_run_matches(replayed, eager, "unfused", dtype) + + +# --------------------------------------------------------------------------- +# DotProductAttention under torch.compile +# --------------------------------------------------------------------------- + + +# Model configurations, described with the same ModelConfig the eager attention +# tests use. Which backends can run each of them is not hardcoded here -- +# get_available_attention_backends() answers that, so configurations only one +# backend supports (arbitrary masks, biases, MLA head dims, ...) are covered +# rather than avoided. +def _cfg( + model_config, + qkv_format="bshd", + packed=None, + interleave_dim=-3, + share_cu_seqlens=False, +): + """A model configuration plus what DotProductAttention is handed it as. + + `packed` is "qkv" or "kv" for the declarative packed inputs, interleaved at + `interleave_dim`, or None for separate q/k/v. + """ + return dict( + model_config=model_config, + qkv_format=qkv_format, + packed=packed, + interleave_dim=interleave_dim, + share_cu_seqlens=share_cu_seqlens, + ) + + +def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str: + """bshd + 3 @ -3 -> bs3hd; bshd + 2 @ -2 -> bsh2d; as DotProductAttention + derives it from the declaration.""" + position = len(qkv_format) + interleave_dim + 1 + return qkv_format[:position] + str(packed_dim) + qkv_format[position:] + + +_DPA_COMPILE_CONFIGS = { + "self_bshd_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal")), + "self_sbhd_no_mask": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd"), + "self_bshd_swa": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0))), + "gqa_bshd_causal": _cfg(ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal")), + "cross_bshd_no_mask": _cfg( + ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask") + ), + "self_bshd_padding_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal")), + "self_thd_padding_causal": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd" + ), + # Self attention naturally passes one cu_seqlens tensor for both q and kv, + # which flash-attn hands to two inputs of the same autograd.Function. + "self_thd_shared_cu_seqlens": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd", share_cu_seqlens=True + ), + "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv"), + "packed_qkv_bsh3d": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv", interleave_dim=-2 + ), + "packed_kv_bshd_bs2hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="kv"), + "packed_kv_thd_th2d": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), + "thd", + packed="kv", + interleave_dim=-2, + ), + # Configurations below are supported by one backend only, or take a code + # path of their own inside DotProductAttention. + "alibi_bshd_causal": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi") + ), + "post_scale_bias_bshd": _cfg( + ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss") + ), + "arbitrary_mask_bshd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary")), + "mla_bshd_causal": _cfg(ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal")), + "sink_softmax_bshd": _cfg( + ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one") + ), +} + + +def _qkv_layout(spec: dict) -> str: + qkv_format, packed = spec["qkv_format"], spec["packed"] + if packed is None: + return "_".join([qkv_format] * 3) + if packed == "qkv": + return _packed_layout(qkv_format, 3, spec["interleave_dim"]) + return f"{qkv_format}_{_packed_layout(qkv_format, 2, spec['interleave_dim'])}" + + +def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention: + config = spec["model_config"] + return te.DotProductAttention( + num_attention_heads=config.num_heads, + kv_channels=config.kv_channels, + num_gqa_groups=config.num_gqa_groups, + attention_dropout=config.dropout_p, + qkv_format=spec["qkv_format"], + attn_mask_type=config.attn_mask_type, + window_size=config.window_size, + attention_type=config.attn_type, + softmax_type=config.softmax_type, + ).to(dtype=dtype, device="cuda") + + +def _cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor: + cu = torch.zeros(seqlens.numel() + 1, dtype=torch.int32, device="cuda") + cu[1:] = torch.cumsum(seqlens, dim=0) + return cu + + +def _make_dpa_inputs(spec: dict, dtype: torch.dtype): + """Build the (args, kwargs) that `DotProductAttention.forward` is called + with, plus the list of tensors whose gradients the test compares.""" + config = spec["model_config"] + qkv_format, packed = spec["qkv_format"], spec["packed"] + b = config.batch_size + s_q, s_kv = config.max_seqlen_q, config.max_seqlen_kv + h, g = config.num_heads, config.num_gqa_groups + d_qk, d_v = config.head_dim_qk, config.head_dim_v + padded = "padding" in config.attn_mask_type + + kwargs = {} + if padded: + # Sequences shorter than the maximum, so the padding mask is not + # degenerate: with all sequences full it would be all-False, and any + # difference in how masking is compiled would be invisible. + seqlens_q = torch.randint(1, s_q, [b], dtype=torch.int32, device="cuda") + seqlens_kv = ( + seqlens_q + if config.attn_type == "self" + else (torch.randint(1, s_kv, [b], dtype=torch.int32, device="cuda")) + ) + cu_q = _cu_seqlens(seqlens_q) + cu_kv = cu_q if spec["share_cu_seqlens"] else _cu_seqlens(seqlens_kv) + kwargs.update(cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv, max_seqlen_q=s_q, max_seqlen_kv=s_kv) + t_q, t_kv = int(cu_q[-1]), int(cu_kv[-1]) + else: + t_q, t_kv = b * s_q, b * s_kv + + def _shape(s, t, heads, head_dim): + return { + "bshd": (b, s, heads, head_dim), + "sbhd": (s, b, heads, head_dim), + "thd": (t, heads, head_dim), + }[qkv_format] + + def _randn(shape): + return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + + if packed is not None: + # Declarative packed inputs: q/k/v are derived from one buffer by DPA + # itself, and the layout comes from the declaration -- the only packed + # layout that torch.compile supports. + assert h == g and d_qk == d_v, "packed inputs require uniform heads and head dims" + interleave_dim = spec["interleave_dim"] + + def _packed(seqlen, tokens, packed_dim): + leading = { + "bshd": (b, seqlen), + "sbhd": (seqlen, b), + "thd": (tokens,), + }[qkv_format] + trailing = (packed_dim, h, d_qk) if interleave_dim == -3 else (h, packed_dim, d_qk) + return _randn(leading + trailing) + + kwargs["qkv_interleave_dim"] = interleave_dim + if packed == "qkv": + qkv = _packed(s_q, t_q, 3) + kwargs["qkv_layer"] = qkv + grad_tensors, args = [qkv], () + else: + q = _randn(_shape(s_q, t_q, h, d_qk)) + kv = _packed(s_kv, t_kv, 2) + kwargs["kv_layer"] = kv + grad_tensors, args = [q, kv], (q,) + else: + q = _randn(_shape(s_q, t_q, h, d_qk)) + k = _randn(_shape(s_kv, t_kv, g, d_qk)) + v = _randn(_shape(s_kv, t_kv, g, d_v)) + grad_tensors, args = [q, k, v], (q, k, v) + + if config.attn_mask_type == "arbitrary": + kwargs["attention_mask"] = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device="cuda") + if config.attn_bias_type != "no_bias": + kwargs["core_attention_bias_type"] = config.attn_bias_type + if config.attn_bias_type == "post_scale_bias": + kwargs["core_attention_bias"] = torch.randn(1, h, s_q, s_kv, dtype=dtype, device="cuda") + + return args, kwargs, grad_tensors + + +def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) -> None: + """Skip what the backend under test cannot run, or -- for a test that + compiles it -- cannot be compiled.""" + if compiled and backend == "fused": + # FusedAttention's forward carries @no_torch_dynamo, so there is nothing + # to compile: it runs as an eager island. Drop this skip once it traces, + # and the tests below cover it as they do the others. + pytest.skip("FusedAttention is an eager island and does not compile") + available, _, _ = get_available_attention_backends( + spec["model_config"], dtype, _qkv_layout(spec) + ) + flash_supported, fused_supported, unfused_supported = available + supported = { + "flash": flash_supported, + "fused": fused_supported, + "unfused": unfused_supported, + }[backend] + if not supported: + pytest.skip(f"the {backend} backend does not support this configuration") + + +def _force_dpa_backend(monkeypatch, backend: str) -> None: + """Restrict DotProductAttention to a single backend.""" + from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import ( + _attention_backends, + ) + from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + FlashAttentionUtils, + ) + + if backend == "flash" and not FlashAttentionUtils.is_installed: + pytest.skip("flash-attn is not installed") + + for name, var in ( + ("flash", "NVTE_FLASH_ATTN"), + ("fused", "NVTE_FUSED_ATTN"), + ("unfused", "NVTE_UNFUSED_ATTN"), + ): + monkeypatch.setenv(var, "1" if name == backend else "0") + # Backend selection is cached on the attention params only, so the env vars + # above are not enough to invalidate it. + _attention_backends["backend_selection_requires_update"] = True + + +def _assert_dpa_backend(backend: str) -> None: + from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import ( + _attention_backends, + ) + + assert _attention_backends[f"use_{backend}_attention"], ( + f"expected the {backend} backend to run, selected:" + f" flash={_attention_backends['use_flash_attention']}," + f" fused={_attention_backends['use_fused_attention']}," + f" unfused={_attention_backends['use_unfused_attention']}" + ) + + +def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) -> None: + """Assert a compiled result matches the eager one. + + A backend that calls the same kernel either way has to match exactly: + compiling changes what surrounds it, not its arithmetic. The unfused + backend is instead built from PyTorch ops, which inductor fuses and + reassociates; that error scales with the softmax sums rather than with each + output element, hence an absolute tolerance taken from the tensor's scale. + """ + if backend != "unfused": + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + return + tols = dtype_tols(dtype) + torch.testing.assert_close( + actual, + expected, + rtol=tols["rtol"], + atol=max(tols["atol"], 1e-3 * expected.abs().max().item()), + ) + + +def _run_and_capture(fn, args, kwargs, grads): + """Run `fn` and backward through its output, returning the output and a copy + of every input gradient. + + The output is cloned because CUDA graphs hand back tensors owned by their + memory pool, which the next replay overwrites, and the gradients are cleared + so that the same inputs can be reused for the other run. + """ + out = fn(*args, **kwargs).clone() + out.sum().backward() + torch.cuda.synchronize() + captured = [] + for tensor in grads: + assert tensor.grad is not None + captured.append(tensor.grad.clone()) + tensor.grad = None + return out, captured + + +def _assert_run_matches(actual, expected, backend: str, dtype: torch.dtype) -> None: + """Compare an (output, gradients) pair produced by `_run_and_capture`.""" + (out, out_grads), (ref, ref_grads) = actual, expected + _assert_matches_eager(out, ref, backend, dtype) + for out_grad, ref_grad in zip(out_grads, ref_grads): + _assert_matches_eager(out_grad, ref_grad, backend, dtype) + + +def _compare_compiled_to_eager( + module, args, kwargs, grads, monkeypatch, backend: str, dtype: torch.dtype, **compile_kwargs +) -> None: + """Run the module eagerly and compiled on the same inputs, and compare both + the output and every input gradient.""" + eager = _run_and_capture(module, args, kwargs, grads) + + torch._dynamo.reset() + # Force backend selection to be re-run (and traced) inside the compiled + # region instead of being served from the cache the eager call populated. + _force_dpa_backend(monkeypatch, backend) + compiled = _run_and_capture(torch.compile(module, **compile_kwargs), args, kwargs, grads) + + _assert_dpa_backend(backend) + _assert_run_matches(compiled, eager, backend, dtype) + + +@pytest.mark.parametrize("backend", ["flash", "fused", "unfused"]) +@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys()) +def test_dpa_torch_compile(monkeypatch, backend, config): + """`DotProductAttention` under `torch.compile(fullgraph=True)` must match + eager in forward and backward, for every backend that supports the + configuration. + + `fullgraph=True` makes the test fail on any graph break, so it covers the + whole module: input unpacking, qkv layout, backend selection and the backend + itself. + """ + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS[config] + _skip_unsupported(spec, backend, dtype) + _force_dpa_backend(monkeypatch, backend) + + module = _make_dpa(spec, dtype) + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + _compare_compiled_to_eager( + module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=True + ) + + +def test_dpa_torch_compile_around_fused(monkeypatch): + """FusedAttention itself is an eager island, but everything around it is + compiled: DotProductAttention traces up to the backend call, breaks the + graph there and resumes afterwards. What crosses that break has to survive + it -- the sub-backend enum did not, and reached cuDNN as the function that + produced it.""" + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"] + _skip_unsupported(spec, "fused", dtype, compiled=False) + _force_dpa_backend(monkeypatch, "fused") + + module = _make_dpa(spec, dtype) + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + # No fullgraph: the graph break at the eager island is the point here. + _compare_compiled_to_eager(module, args, kwargs, grads, monkeypatch, "fused", dtype) + + +@pytest.mark.parametrize("backend", ["flash", "unfused"]) +def test_dpa_torch_compile_cudagraphs(monkeypatch, backend): + """`mode="reduce-overhead"`: forward and backward of DotProductAttention + are captured into CUDA graphs and replayed on subsequent iterations.""" + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"] + _force_dpa_backend(monkeypatch, backend) + + module = _make_dpa(spec, dtype) + + torch._dynamo.reset() + counters.clear() + compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead") + + for _ in range(3): + # Fresh inputs every iteration: a replay that reuses stale buffers would + # still produce finite values and non-None gradients, so only comparing + # against eager catches it. + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + replayed = _run_and_capture(compiled, args, kwargs, grads) + eager = _run_and_capture(module, args, kwargs, grads) + _assert_run_matches(replayed, eager, backend, dtype) + _assert_dpa_backend(backend) + # Without this, inductor declining to capture -- a mutated input, a CPU + # scalar -- would leave the test passing while measuring nothing. + assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("fp8_attention", [False, True], ids=["fp8_gemms_only", "fp8_attention"]) +def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention): + """FP8 attention is not supported on the compiled path and falls back to + eager. FP8 elsewhere in the model with attention in high precision -- the + common training setup -- must stay compiled.""" + dtype = torch.bfloat16 + spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"] + _force_dpa_backend(monkeypatch, "unfused") + + module = _make_dpa(spec, dtype) + args, kwargs, _ = _make_dpa_inputs(spec, dtype) + fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention) + + def fn(*args, **kwargs): + with te.autocast(enabled=True, recipe=fp8_recipe): + return module(*args, **kwargs) + + torch._dynamo.reset() + _force_dpa_backend(monkeypatch, "unfused") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + try: + torch.compile(fn)(*args, **kwargs) + except Exception: # pylint: disable=broad-except + # FP8 attention is not available on every device; what matters here + # is which path was taken, which the warning below reports. + pass + fell_back = any("Falling back to eager" in str(w.message) for w in caught) + + assert fell_back == fp8_attention + + +def _packed_views_inputs(spec, dtype, interleave_dim=-3): + """Packed q/k/v handed over as plain views, without declaring the packing.""" + config = spec["model_config"] + b, s = config.batch_size, config.max_seqlen_q + h, d = config.num_heads, config.head_dim_qk + shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d) + qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + q, k, v = [qkv.select(interleave_dim, i) for i in range(3)] + return (q, k, v), {}, [qkv] + + +def _thd_without_max_seqlen_inputs(spec, dtype): + """thd inputs that leave max_seqlen to be derived from cu_seqlens.""" + args, kwargs, grads = _make_dpa_inputs(spec, dtype) + del kwargs["max_seqlen_q"], kwargs["max_seqlen_kv"] + return args, kwargs, grads + + +_EAGER_FALLBACK_CASES = { + # name: (config name, input builder) + "packed_views_bs3hd": ( + "self_bshd_causal", + lambda spec, dtype: _packed_views_inputs(spec, dtype, -3), + ), + "packed_views_bsh3d": ( + "self_bshd_causal", + lambda spec, dtype: _packed_views_inputs(spec, dtype, -2), + ), + "thd_without_max_seqlen": ("self_thd_padding_causal", _thd_without_max_seqlen_inputs), +} + + +@pytest.mark.parametrize("backend", ["flash", "unfused"]) +@pytest.mark.parametrize("case", _EAGER_FALLBACK_CASES.keys()) +def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case): + """Calls that cannot be traced run as an eager island instead, with a + warning: recognizing packed q/k/v takes the data pointers dynamo cannot + read, and deriving max_seqlen off cu_seqlens is a device synchronization. + Both keep working, and both are avoidable -- by declaring the packing via + qkv_layer/kv_layer, or by passing max_seqlen.""" + dtype = torch.bfloat16 + config_name, make_inputs = _EAGER_FALLBACK_CASES[case] + spec = _DPA_COMPILE_CONFIGS[config_name] + _force_dpa_backend(monkeypatch, backend) + + module = _make_dpa(spec, dtype) + args, kwargs, grads = make_inputs(spec, dtype) + + eager = _run_and_capture(module, args, kwargs, grads) + + torch._dynamo.reset() + _force_dpa_backend(monkeypatch, backend) + with pytest.warns(UserWarning, match="Falling back to eager execution"): + fell_back = _run_and_capture(torch.compile(module), args, kwargs, grads) + _assert_run_matches(fell_back, eager, backend, dtype) + + # The same eager island is an error when the user asked for a full graph. + torch._dynamo.reset() + _force_dpa_backend(monkeypatch, backend) + with pytest.raises(Exception, match="torch.compiler.disable"): + torch.compile(module, fullgraph=True)(*args, **kwargs) # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6a60160f77..04bcabe7c4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -793,6 +793,18 @@ def backward( return dq, dk, dv +def _maybe_unshare_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor): + """Copy kv's cumulative sequence lengths when they are q's as well. + + Self attention passes one tensor for both, and flash-attn forwards them to + two inputs of the same autograd.Function -- which dynamo cannot trace. In + eager the duplicate is harmless, so nothing is copied there. + """ + if cu_q is cu_kv and torch.compiler.is_compiling(): + return cu_q, cu_kv.clone() + return cu_q, cu_kv + + class FlashAttention(torch.nn.Module): """Dot product attention, using HazyResearch flash-attn package: https://github.com/Dao-AILab/flash-attention @@ -1122,12 +1134,12 @@ def forward( else: func = flash_attn_with_kvcache_v3 # pylint: disable=possibly-used-before-assignment if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None): - fa_optional_forward_args_thd.append( - cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q - ) - fa_optional_forward_args_thd.append( - cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv + cu_q, cu_kv = _maybe_unshare_cu_seqlens( + cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q, + cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv, ) + fa_optional_forward_args_thd.append(cu_q) + fa_optional_forward_args_thd.append(cu_kv) fa_optional_forward_args_thd.append(max_seqlen_q) fa_optional_forward_args_thd.append(max_seqlen_kv) if use_flash_attn_4: @@ -1138,8 +1150,9 @@ def forward( if inference_params is None: fa_4_optional_forward_kwargs["deterministic"] = self.deterministic if func is flash_attn_varlen_func_v4: - fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q - fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv + cu_q, cu_kv = _maybe_unshare_cu_seqlens(cu_seqlens_q, cu_seqlens_kv) + fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q + fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv output = func( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 4cc4cab1b8..0ed01a9c9e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -196,6 +196,38 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]: + """Why this DotProductAttention call has to run outside the graph, or None. + + `call` maps `DotProductAttention.forward`'s parameter names to the arguments + this call passed, including `self`. + """ + # FP8 GEMMs with the attention itself in high precision -- the common + # training setup -- stay on the compiled path; only FP8 attention bails out. + qstate = FP8GlobalStateManager.quantization_state + fp8_recipe = qstate.fp8_recipe + if qstate.fp8_enabled and fp8_recipe is not None: + if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha: + return "FP8 attention" + + if call["self"].cp_group is not None: + return "context parallelism" + + qkv_format = call.get("qkv_format") or call["self"].qkv_format + if qkv_format == "thd" and ( + call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None + ): + # Deriving it reads the sequence lengths off cu_seqlens, which is a + # device synchronization and a data-dependent value while tracing. + return "deriving max_seqlen from cu_seqlens" + + if call.get("qkv_layer") is None and call.get("kv_layer") is None: + qkv = [call.get(name) for name in ("query_layer", "key_layer", "value_layer")] + if dpa_utils.qkv_layout_needs_detection(*qkv): + return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer" + return None + + def _unpack_packed_qkv( qkv_layer: Optional[torch.Tensor], kv_layer: Optional[torch.Tensor], @@ -715,6 +747,11 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: """ _original_recipe = self.fp8_meta.get("recipe", None) + qstate = FP8GlobalStateManager.quantization_state + if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters): + super().init_fp8_metadata(num_gemms=num_gemms) + return + # global recipe set in autocast() fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): @@ -1090,7 +1127,7 @@ def get_quantizer_roles( ] return base[:num_quantizers] - @no_torch_dynamo(recursive=False) + @no_torch_dynamo(when=_needs_eager_dpa) def forward( self, query_layer: Optional[torch.Tensor] = None, @@ -1820,18 +1857,25 @@ def forward( _attention_backends["fused_attention_backend"] = fused_attention_backend _attention_backends["use_unfused_attention"] = use_unfused_attention _attention_backends["backend_selection_requires_update"] = False + # logging.Logger methods graph-break under torch.compile, so + # selection is only logged in eager -- as in + # get_attention_backend. Note the arguments below are + # evaluated either way, so they have to stay traceable. + logger = ( + dpa_utils.no_op_logger if torch.compiler.is_compiling() else self.logger + ) if use_flash_attention: - self.logger.info( + logger.info( "Running with FlashAttention backend (version %s)", flash_attention_backend, ) elif use_fused_attention: - self.logger.info( + logger.info( "Running with FusedAttention backend (sub-backend %s)", int(fused_attention_backend), ) elif use_unfused_attention: - self.logger.info("Running with UnfusedDotProductAttention backend") + logger.info("Running with UnfusedDotProductAttention backend") else: use_flash_attention = _attention_backends["use_flash_attention"] flash_attention_backend = _attention_backends["flash_attention_backend"] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8a4c6ce286..a2d260c6b2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -346,7 +346,7 @@ def error(self, *args, **kwargs): """No-op.""" -_no_op_logger = _NoOpLogger() +no_op_logger = _NoOpLogger() @torch.compiler.assume_constant_result @@ -361,11 +361,17 @@ def _get_fused_attn_backend( *args, ): """Constant-foldable tex.get_fused_attn_backend: the result depends only on - the attention config, and the python-side enum keeps it traceable by - torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax - are taken as their string keys and resolved to the pybind enums here, so - that every argument is a python literal or a python enum.""" - return FusedAttnBackend.cast( + the attention config. Layout/bias/mask/softmax are taken as their string + keys and resolved to the pybind enums here, so that every argument is a + python literal or a python enum. + + Returns a plain int rather than a FusedAttnBackend member: dynamo + reconstructs the result of an assume_constant_result call by re-emitting the + call, which is only valid inside the frame that made it. An int survives a + graph break because it is baked into the graph as a literal, while an enum + member comes out of the reconstruction corrupted (see the cast at the call + site, which restores the enum).""" + return int( tex.get_fused_attn_backend( is_training, q_type, @@ -395,8 +401,11 @@ def get_attention_backend( Whether the `FlashAttention` backend has been selected. use_fused_attention : bool Whether the `FusedAttention` backend has been selected. - fused_attention_backend : FusedAttnBackend - If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`. + fused_attention_backend : int + If `use_fused_attention = True`, the integer value of one of `FusedAttention`'s three + sub-backends, else `None`. It is not a `FusedAttnBackend` member because that does not + survive a graph break under `torch.compile`; `FusedAttnBackend.cast` turns it into one, + and comparing it against a member works either way. use_unfused_attention : bool Whether the `UnfusedDotProductAttention` backend has been selected. available_backends : List[bool] @@ -452,7 +461,7 @@ def get_attention_backend( if torch.compiler.is_compiling(): # logging.Logger methods graph-break under torch.compile; backend # selection logs are only emitted in eager mode. - logger = _no_op_logger + logger = no_op_logger else: logger = logging.getLogger("DotProductAttention") logger.setLevel(AttentionLogging._log_level) @@ -1455,11 +1464,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cuda_graph, deterministic, ) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: + if fused_attention_backend == FusedAttnBackend.No_Backend.value: logger.debug("Disabling FusedAttention as no backend supports the provided input") use_fused_attention = False fused_attention_backend = None - elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]: + elif ( + has_score_mod and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value + ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " "F16/BF16 arbitrary-seqlen", @@ -1509,7 +1520,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend["FP8"] + fused_attention_backend == FusedAttnBackend.FP8.value and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1520,7 +1531,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] + fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value and is_training and ( device_compute_capability < (9, 0) @@ -1987,21 +1998,21 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor: tensor of shape [batch_size * max_seqlen, 1, 1] containing the indices for the valid tokens in a batch. """ + # Built with device-side ops only: reading the sequence lengths on the host + # would synchronize the device once per sequence. bs = len(cu_seqlens) - 1 seqlens = cu_seqlens[1:] - cu_seqlens[:-1] - indices = [i * max_seqlen + ii for i, j in enumerate(seqlens) for ii in range(j)] - indices = torch.Tensor(indices).unsqueeze(1).unsqueeze(1).to(dtype=torch.int64, device="cuda") - - num_nonzeros = indices.shape[0] - pad_amount = bs * max_seqlen - num_nonzeros - indices = F.pad( - input=indices, - pad=(0, 0, 0, 0, 0, pad_amount), - mode="constant", - value=float(bs * max_seqlen), + positions = torch.arange(max_seqlen, device=cu_seqlens.device) + valid = (positions.unsqueeze(0) < seqlens.unsqueeze(1)).flatten() + # Sorting the mask is stable, so the valid positions come first in order, + # and the invalid tail is replaced by the out-of-range padding index. + ordered = torch.argsort(~valid, stable=True) + indices = torch.where( + torch.arange(bs * max_seqlen, device=cu_seqlens.device) < valid.sum(), + ordered, + bs * max_seqlen, ) - - return indices + return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1) def get_full_cu_seqlens( @@ -2025,7 +2036,9 @@ def _get_cu_seqlens(batch_size, max_seqlen, device): device=device, ) - if is_in_onnx_export_mode(): + if is_in_onnx_export_mode() or torch.compiler.is_compiling(): + # The cache only saves an allocation, and its key -- is_inference_mode_enabled() + # -- is a fundamental graph break. return _get_cu_seqlens(batch_size, max_seqlen, device) is_inference = torch.is_inference_mode_enabled() @@ -2365,6 +2378,20 @@ def get_qkv_format( return qkv_format, q_format, kv_format +def qkv_layout_needs_detection(*qkv: Optional[torch.Tensor]) -> bool: + """Whether the layout of these q/k/v can only be told by inspecting memory. + + True for tensors that may be slices of a packed buffer, i.e. strided ones + that do not own their whole storage. + """ + return any( + x is not None + and not x.is_contiguous() + and x.untyped_storage().size() != x.numel() * x.element_size() + for x in qkv + ) + + def get_qkv_layout( q: torch.Tensor, k: torch.Tensor, @@ -2538,10 +2565,24 @@ def run_iteratively(q, k, v): return qkv_layout - if not is_in_onnx_export_mode(): - qkv_layout = run_iteratively(q, k, v) - else: + if is_in_onnx_export_mode(): + # Checked first: the ONNX exporter runs through dynamo, so it also sets + # is_compiling(), and it has its own handling below. qkv_layout = "not_supported" + elif torch.compiler.is_compiling(): + # run_iteratively reads data pointers and storage offsets, which dynamo + # cannot trace; unpacked q/k/v need no detection anyway. + assert not qkv_layout_needs_detection(q, k, v), ( + "q/k/v may be views into a packed buffer, whose layout cannot be detected under" + " torch.compile. Pass the packed buffer explicitly via DotProductAttention's" + " qkv_layer/kv_layer (with qkv_interleave_dim)." + ) + if is_same_q_kv_format: + qkv_layout = "_".join([qkv_format] * 3) + else: + qkv_layout = q_format + "_" + kv_format + "_" + kv_format + else: + qkv_layout = run_iteratively(q, k, v) if qkv_layout == "not_supported": # force q,k,v to be contiguous and run get_layout again q, k, v = [x.contiguous() for x in [q, k, v]] diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 046019ee58..9a33df7634 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -107,10 +107,13 @@ class FusedAttnBackend(IntEnum): ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum value-for-value, and instances of the two enums compare equal when they share the same integer value. Unlike the pybind enum, a plain-python - ``IntEnum`` is traceable by ``torch.compile``: comparisons constant-fold - cleanly and instances safely cross the ``assume_constant_result`` boundary - in ``get_attention_backend``. Lookup by name (``FusedAttnBackend["FP8"]``) - works the same way as with the dict this used to be. + ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member + constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works + the same way as with the dict this used to be. + + Members do not survive a graph break, though, so ``get_attention_backend`` + returns the sub-backend as a plain int and ``cast`` turns it back into a + member. """ No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 1b93b8254c..22e9c9ae3e 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -3,7 +3,9 @@ # See LICENSE for license information. """NVFuser functions and JIT utilities""" +import inspect import os +import warnings from functools import wraps from typing import Callable, Optional, Tuple import torch @@ -26,6 +28,8 @@ def lazy_compile(func): @wraps(func) def wrapper(*args, **kwargs): nonlocal compiled_func + if torch.compiler.is_compiling(): + return func(*args, **kwargs) if compiled_func is None: compiled_func = torch.compile(func) return compiled_func(*args, **kwargs) @@ -49,22 +53,66 @@ def wrapper(*args, **kwargs): if torch.__version__ >= "2": import torch._dynamo - def no_torch_dynamo(recursive=True): - """Decorator to disable Torch Dynamo, except during ONNX export.""" + def no_torch_dynamo(recursive=True, when=None): + """Decorator to disable Torch Dynamo, except during ONNX export. - def decorator(f): + `when` makes it conditional: it is called with the arguments of the call + keyed by parameter name -- an argument the call left out is absent, so + reading it with .get() yields its default as long as that default is + None -- and returns the reason this particular call cannot be traced, or + None to have it traced as usual. The reason is reported once per + distinct message. + """ + + def _disable(f): # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True - disabled_f = ( + return ( torch._dynamo.disable(f, recursive=recursive) if torch.__version__ >= "2.1" else torch._dynamo.disable(f) ) - @wraps(f) - def wrapper(*args, **kwargs): - if is_in_onnx_export_mode(): + def decorator(f): + disabled_f = _disable(f) + if when is None: + + @wraps(f) + def wrapper(*args, **kwargs): + if is_in_onnx_export_mode(): + return f(*args, **kwargs) + return disabled_f(*args, **kwargs) + + else: + parameters = inspect.signature(f).parameters + # Arguments are matched to names by position, which only holds + # without a *args in between. + assert not any( + p.kind is inspect.Parameter.VAR_POSITIONAL for p in parameters.values() + ), f"no_torch_dynamo(when=...) does not support *args, which {f.__name__} takes" + parameter_names = list(parameters) + + # The warning belongs inside the disabled function: warnings.warn + # graph-breaks on its own, which would mask the break that matters. + def report_and_run(reason, *args, **kwargs): + warnings.warn( + f"Falling back to eager execution under torch.compile: {reason} is" + " unsupported on the compiled path (graph-breaks under fullgraph=True).", + stacklevel=3, + ) return f(*args, **kwargs) - return disabled_f(*args, **kwargs) + + disabled_report_and_run = _disable(report_and_run) + + @wraps(f) + def wrapper(*args, **kwargs): # pylint: disable=function-redefined + if is_in_onnx_export_mode() or not torch.compiler.is_compiling(): + return f(*args, **kwargs) + call = dict(zip(parameter_names, args)) + call.update(kwargs) + reason = when(call) + if reason is None: + return f(*args, **kwargs) + return disabled_report_and_run(reason, *args, **kwargs) return wrapper @@ -72,7 +120,7 @@ def wrapper(*args, **kwargs): else: # Fallback for PyTorch < 2.0: no-op decorator - def no_torch_dynamo(recursive=True): # pylint: disable=unused-argument + def no_torch_dynamo(recursive=True, when=None): # pylint: disable=unused-argument """No-op decorator for PyTorch < 2.0.""" return lambda func: func