diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index b18abecd7e..8a1d3a1fcf 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -109,7 +109,7 @@ this for the comments and docstrings you are adding or rewriting. - **Shape-suffix tensor names.** In model code, name tensors with shape suffixes (Noam Shazeer convention: https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd), - e.g. `x_BLD`, `q_BLNH`, `out_TNH`. Capital-letter suffixes denote *logical* + e.g. `x_BLD`, `q_BLHK`, `out_THV`. Capital-letter suffixes denote *logical* tensor dimensions, not a physical sharding layout -- a name like `routed_input_RD` keeps the same suffix whether or not `R` is a local shard under EP/SP. Letters are scoped per module, not global: give each module that diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0c96ce776e..6cdf7911d5 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -76,7 +76,7 @@ When appropriate, one should consider - To add a unit test, put it in the [tests](tests/) folder and follow the existing test files. - To add a GPU integration test, add a configuration to the matching module in [torchtitan_recipes/tests](torchtitan_recipes/tests/) and a new `OverrideDefinitions` naming it in [integration_tests](tests/integration_tests/). These suites name a full Trainer configuration per run. - Updating [README](README.md) and writing a new note in the [docs](docs/) folder on installation and usage, similar to [float8.md](torchtitan/components/quantization/float8.md). -- Following the tensor shape-suffix naming convention for new model code (e.g. `x_BLD`, `q_BLNH`, `out_TNH`), with a per-module legend comment as in [attention.py](torchtitan/models/common/attention.py). Capital suffixes name logical tensor dimensions (not sharding layout) and are scoped per file. +- Following the tensor shape-suffix naming convention for new model code (e.g. `x_BLD`, `q_BLHK`, `out_THV`), with a per-module legend comment as in [attention.py](torchtitan/models/common/attention.py). Capital suffixes name logical tensor dimensions (not sharding layout) and are scoped per file. - Adding a new file with benchmark results in [benchmarks](benchmarks) folder. - Creating GitHub issues for things that cannot be addressed at the moment. - Writing a post on [PyTorch Forums](https://discuss.pytorch.org/c/distributed/torchtitan/44) and linking to it. diff --git a/tests/unit_tests/cpu/test_flex_attention.py b/tests/unit_tests/cpu/test_flex_attention.py index 0e4da9c7fc..d835892d54 100644 --- a/tests/unit_tests/cpu/test_flex_attention.py +++ b/tests/unit_tests/cpu/test_flex_attention.py @@ -29,53 +29,53 @@ def _mask(seq_len: int, batch_size: int): device="cpu", ) - def test_tnh_layout(self) -> None: + def test_thk_thv_layout(self) -> None: num_tokens, num_heads, head_dim = 8, 4, 16 - q_TNH = torch.randn(num_tokens, num_heads, head_dim) - k_TNH = torch.randn_like(q_TNH) - v_TNH = torch.randn_like(q_TNH) + q_THK = torch.randn(num_tokens, num_heads, head_dim) + k_THK = torch.randn_like(q_THK) + v_THV = torch.randn_like(q_THK) - def kernel(q_BNTH, k_BNTH, v_BNTH, **kwargs): - self.assertEqual(q_BNTH.shape, (1, num_heads, num_tokens, head_dim)) - self.assertEqual(k_BNTH.shape, q_BNTH.shape) - self.assertEqual(v_BNTH.shape, q_BNTH.shape) - lse_BNT = torch.randn(1, num_heads, num_tokens) - return q_BNTH, SimpleNamespace(lse=lse_BNT) + def kernel(q_1HTK, k_1HTK, v_1HTV, **kwargs): + self.assertEqual(q_1HTK.shape, (1, num_heads, num_tokens, head_dim)) + self.assertEqual(k_1HTK.shape, q_1HTK.shape) + self.assertEqual(v_1HTV.shape, q_1HTK.shape) + lse_1HT = torch.randn(1, num_heads, num_tokens) + return q_1HTK, SimpleNamespace(lse=lse_1HT) with patch.object(FlexAttention, "compiled_flex_attn", side_effect=kernel): - out_TNH = self.attention( - q_TNH, - k_TNH, - v_TNH, + out_THV = self.attention( + q_THK, + k_THK, + v_THV, attention_masks=self._mask(num_tokens, 1), ) - torch.testing.assert_close(out_TNH, q_TNH) + torch.testing.assert_close(out_THV, q_THK) - def test_tnh_out_transform_layout(self) -> None: + def test_thk_thv_out_transform_layout(self) -> None: num_tokens, num_heads, head_dim = 8, 4, 16 - q_TNH = torch.randn(num_tokens, num_heads, head_dim) - expected_lse_TN = torch.randn(num_tokens, num_heads) + q_THK = torch.randn(num_tokens, num_heads, head_dim) + expected_lse_TH = torch.randn(num_tokens, num_heads) - def kernel(q_BNTH, k_BNTH, v_BNTH, **kwargs): - return q_BNTH, SimpleNamespace( - lse=expected_lse_TN.transpose(0, 1).unsqueeze(0) + def kernel(q_1HTK, k_1HTK, v_1HTV, **kwargs): + return q_1HTK, SimpleNamespace( + lse=expected_lse_TH.transpose(0, 1).unsqueeze(0) ) - def out_transform(out_TNH, lse_TN): - torch.testing.assert_close(lse_TN, expected_lse_TN) - return out_TNH + def out_transform(out_THV, lse_TH): + torch.testing.assert_close(lse_TH, expected_lse_TH) + return out_THV with patch.object(FlexAttention, "compiled_flex_attn", side_effect=kernel): - out_TNH = self.attention( - q_TNH, - q_TNH, - q_TNH, + out_THV = self.attention( + q_THK, + q_THK, + q_THK, attention_masks=self._mask(num_tokens, 1), out_transform=out_transform, ) - torch.testing.assert_close(out_TNH, q_TNH) + torch.testing.assert_close(out_THV, q_THK) if __name__ == "__main__": diff --git a/tests/unit_tests/cpu/test_fused_qkv.py b/tests/unit_tests/cpu/test_fused_qkv.py index 539e79ccce..5016eb3897 100644 --- a/tests/unit_tests/cpu/test_fused_qkv.py +++ b/tests/unit_tests/cpu/test_fused_qkv.py @@ -141,22 +141,22 @@ def test_forward_outputs_are_contiguous_and_correct(self): fused = _build_fused() num_tokens = 6 x_TD = torch.randn(num_tokens, _DIM) - xq_TNH, xk_TNH, xv_TNH = fused(x_TD) + xq_THK, xk_THK, xv_THV = fused(x_TD) - self.assertEqual(xq_TNH.shape, (num_tokens, _N_HEADS, _HEAD_DIM)) - self.assertEqual(xk_TNH.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM)) - self.assertEqual(xv_TNH.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM)) - self.assertTrue(xq_TNH.is_contiguous()) - self.assertTrue(xk_TNH.is_contiguous()) - self.assertTrue(xv_TNH.is_contiguous()) + self.assertEqual(xq_THK.shape, (num_tokens, _N_HEADS, _HEAD_DIM)) + self.assertEqual(xk_THK.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM)) + self.assertEqual(xv_THV.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM)) + self.assertTrue(xq_THK.is_contiguous()) + self.assertTrue(xk_THK.is_contiguous()) + self.assertTrue(xv_THV.is_contiguous()) sd = fused.state_dict() - ref_q_TNH = (x_TD @ sd["wq.weight"].T).view(num_tokens, _N_HEADS, _HEAD_DIM) - ref_k_TNH = (x_TD @ sd["wk.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM) - ref_v_TNH = (x_TD @ sd["wv.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM) - torch.testing.assert_close(xq_TNH, ref_q_TNH) - torch.testing.assert_close(xk_TNH, ref_k_TNH) - torch.testing.assert_close(xv_TNH, ref_v_TNH) + ref_q_THK = (x_TD @ sd["wq.weight"].T).view(num_tokens, _N_HEADS, _HEAD_DIM) + ref_k_THK = (x_TD @ sd["wk.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM) + ref_v_THV = (x_TD @ sd["wv.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM) + torch.testing.assert_close(xq_THK, ref_q_THK) + torch.testing.assert_close(xk_THK, ref_k_THK) + torch.testing.assert_close(xv_THV, ref_v_THV) def test_raw_pointer_read_needs_contiguous(self): """A consumer reading the base pointer with contiguous head-major strides diff --git a/tests/unit_tests/cpu/test_model_td_layout.py b/tests/unit_tests/cpu/test_model_td_layout.py index fff290c7f5..739a0b23c8 100644 --- a/tests/unit_tests/cpu/test_model_td_layout.py +++ b/tests/unit_tests/cpu/test_model_td_layout.py @@ -20,38 +20,38 @@ class _AttentionOutput(nn.Module): - def forward(self, q_TNH, k_TNH, v_TNH, *, out_transform=None, **kwargs): - num_q_heads = q_TNH.shape[1] - num_v_heads = v_TNH.shape[1] - out_TNH = v_TNH.repeat_interleave(num_q_heads // num_v_heads, dim=1) + def forward(self, q_THK, k_THK, v_THV, *, out_transform=None, **kwargs): + num_q_heads = q_THK.shape[1] + num_v_heads = v_THV.shape[1] + out_THV = v_THV.repeat_interleave(num_q_heads // num_v_heads, dim=1) if out_transform is not None: - lse_TN = torch.zeros( - q_TNH.shape[:2], device=q_TNH.device, dtype=q_TNH.dtype + lse_TH = torch.zeros( + q_THK.shape[:2], device=q_THK.device, dtype=q_THK.dtype ) - out_TNH = out_transform(out_TNH, lse_TN) - return out_TNH + out_THV = out_transform(out_THV, lse_TH) + return out_THV class TestModelTDLayout(unittest.TestCase): - def test_sdpa_preserves_blnh_shape(self): + def test_sdpa_preserves_blhv_shape(self): attention = ScaledDotProductAttention.Config().build() - q_BLNH = torch.randn(2, 8, 4, 16) - k_BLNH = torch.randn(2, 8, 2, 16) - v_BLNH = torch.randn(2, 8, 2, 16) + q_BLHK = torch.randn(2, 8, 4, 16) + k_BLHK = torch.randn(2, 8, 2, 16) + v_BLHV = torch.randn(2, 8, 2, 16) - out_BLNH = attention(q_BLNH, k_BLNH, v_BLNH, enable_gqa=True) + out_BLHV = attention(q_BLHK, k_BLHK, v_BLHV, enable_gqa=True) - self.assertEqual(out_BLNH.shape, q_BLNH.shape) + self.assertEqual(out_BLHV.shape, q_BLHK.shape) - def test_graph_trainer_sdpa_preserves_tnh_shape(self): + def test_graph_trainer_sdpa_preserves_thv_shape(self): attention = GraphTrainerScaledDotProductAttention.Config().build() - q_TNH = torch.randn(8, 4, 16) - k_TNH = torch.randn(8, 2, 16) - v_TNH = torch.randn(8, 2, 16) + q_THK = torch.randn(8, 4, 16) + k_THK = torch.randn(8, 2, 16) + v_THV = torch.randn(8, 2, 16) - out_TNH = attention(q_TNH, k_TNH, v_TNH, enable_gqa=True) + out_THV = attention(q_THK, k_THK, v_THV, enable_gqa=True) - self.assertEqual(out_TNH.shape, q_TNH.shape) + self.assertEqual(out_THV.shape, q_THK.shape) def test_gpt_oss_attention_preserves_td_shape(self): config = gptoss_configs["debugmodel"]("standard", "varlen") diff --git a/tests/unit_tests/cpu/test_varlen_attention.py b/tests/unit_tests/cpu/test_varlen_attention.py index 360a6c4798..798c1977bc 100644 --- a/tests/unit_tests/cpu/test_varlen_attention.py +++ b/tests/unit_tests/cpu/test_varlen_attention.py @@ -5,7 +5,8 @@ # LICENSE file in the root directory of this source tree. # Shape suffix legend: -# T = packed tokens, N = attention heads, H = head dimension, D = model dim +# T = packed tokens, H = attention heads, K = query/key head dimension, +# V = value head dimension, D = model dimension import unittest from unittest.mock import patch @@ -63,11 +64,11 @@ def test_gqa_preserves_td_shape(self): positions_T = torch.tensor([0, 1, 0, 1, 2, 3]) metadata = create_varlen_metadata_for_document(positions_T) - def _identity_varlen(q_TNH, k_TNH, v_TNH, *args, **kwargs): - self.assertEqual(q_TNH.ndim, 3) - self.assertEqual(k_TNH.ndim, 3) - self.assertEqual(v_TNH.ndim, 3) - return q_TNH + def _identity_varlen(q_THK, k_THK, v_THV, *args, **kwargs): + self.assertEqual(q_THK.ndim, 3) + self.assertEqual(k_THK.ndim, 3) + self.assertEqual(v_THV.ndim, 3) + return q_THK with patch( "torchtitan.models.common.attention._varlen_attn", @@ -77,7 +78,7 @@ def _identity_varlen(q_TNH, k_TNH, v_TNH, *args, **kwargs): self.assertEqual(out_TD.shape, x_TD.shape) - def test_tnh_sharding_uses_varlen_argument_names(self): + def test_thk_thv_sharding_uses_varlen_argument_names(self): from torchtitan.models.llama3 import llama3_configs from torchtitan.models.llama3.sharding import set_llama3_sharding_config @@ -88,45 +89,45 @@ def test_tnh_sharding_uses_varlen_argument_names(self): assert sharding is not None self.assertEqual( set(sharding.in_src_shardings or {}), - {"q_TNH", "k_TNH", "v_TNH"}, + {"q_THK", "k_THK", "v_THV"}, ) - q_layout = (sharding.in_src_shardings or {})["q_TNH"] - k_dst_layout = (sharding.in_dst_shardings or {})["k_TNH"] + q_layout = (sharding.in_src_shardings or {})["q_THK"] + k_dst_layout = (sharding.in_dst_shardings or {})["k_THK"] axis_types = _per_axis_types(q_layout) self.assertEqual(axis_types[MeshAxisName.DP], spmd.S(0)) self.assertEqual(axis_types[MeshAxisName.CP], spmd.S(0)) self.assertEqual(axis_types[MeshAxisName.TP], spmd.S(1)) self.assertEqual(_per_axis_types(k_dst_layout)[MeshAxisName.CP], spmd.R) - def test_out_transform_receives_tn_lse(self): + def test_out_transform_receives_th_lse(self): num_tokens, num_heads, head_dim = 5, 2, 4 - q_TNH = torch.randn(num_tokens, num_heads, head_dim) + q_THK = torch.randn(num_tokens, num_heads, head_dim) positions_T = torch.tensor([0, 1, 0, 1, 2]) metadata = create_varlen_metadata_for_document(positions_T) inner_attention = VarlenAttention.Config().build() def _varlen_with_lse(q, k, v, *args, **kwargs): - lse_NT = torch.randn(num_heads, num_tokens) - return q, lse_NT + lse_HT = torch.randn(num_heads, num_tokens) + return q, lse_HT - def _check_shapes(out_TNH, lse_TN): - self.assertEqual(out_TNH.shape, q_TNH.shape) - self.assertEqual(lse_TN.shape, (num_tokens, num_heads)) - return out_TNH + def _check_shapes(out_THV, lse_TH): + self.assertEqual(out_THV.shape, q_THK.shape) + self.assertEqual(lse_TH.shape, (num_tokens, num_heads)) + return out_THV with patch( "torchtitan.models.common.attention._varlen_attn", side_effect=_varlen_with_lse, ): - out_TNH = inner_attention( - q_TNH, - q_TNH, - q_TNH, + out_THV = inner_attention( + q_THK, + q_THK, + q_THK, attention_masks=metadata, out_transform=_check_shapes, ) - self.assertEqual(out_TNH.shape, q_TNH.shape) + self.assertEqual(out_THV.shape, q_THK.shape) def test_llama_decoder_preserves_td_shape(self): from torchtitan.models.llama3 import llama3_configs @@ -138,8 +139,8 @@ def test_llama_decoder_preserves_td_shape(self): positions_T = torch.tensor([0, 1, 0, 1, 2, 3]) metadata = model.get_attention_masks(positions_T) - def _identity_varlen(q_TNH, k_TNH, v_TNH, *args, **kwargs): - return q_TNH + def _identity_varlen(q_THK, k_THK, v_THV, *args, **kwargs): + return q_THK with patch( "torchtitan.models.common.attention._varlen_attn", diff --git a/tests/unit_tests/cpu/test_vision_attention.py b/tests/unit_tests/cpu/test_vision_attention.py index d62fcab55f..c017956161 100644 --- a/tests/unit_tests/cpu/test_vision_attention.py +++ b/tests/unit_tests/cpu/test_vision_attention.py @@ -19,9 +19,9 @@ def __init__(self) -> None: super().__init__() self.input_shape: torch.Size | None = None - def forward(self, q_TNH, k_TNH, v_TNH, *, attention_masks): - self.input_shape = q_TNH.shape - return q_TNH + def forward(self, q_THK, k_THK, v_THV, *, attention_masks): + self.input_shape = q_THK.shape + return q_THK class TestVisionAttention(unittest.TestCase): diff --git a/tests/unit_tests/gpu/test_qk_clip.py b/tests/unit_tests/gpu/test_qk_clip.py index e28681c5b0..7bf82bff8a 100644 --- a/tests/unit_tests/gpu/test_qk_clip.py +++ b/tests/unit_tests/gpu/test_qk_clip.py @@ -37,8 +37,8 @@ class QKClipTest(unittest.TestCase): def test_attention_records_training_maxima_only(self) -> None: attention = QKClipFlexAttention.Config().build() - q_TNH = torch.randn(2, 2, 4) - max_scores_BNT = torch.tensor([[[1.0, 3.0], [4.0, 2.0]]]) + q_THK = torch.randn(2, 2, 4) + max_scores_1HT = torch.tensor([[[1.0, 3.0], [4.0, 2.0]]]) block_mask = create_block_mask( lambda _b, _h, q_idx, kv_idx: q_idx >= kv_idx, 1, @@ -48,40 +48,40 @@ def test_attention_records_training_maxima_only(self) -> None: device="cpu", _compile=False, ) - aux = SimpleNamespace(lse=None, max_scores=max_scores_BNT) + aux = SimpleNamespace(lse=None, max_scores=max_scores_1HT) attention.train() with patch( "torchtitan.models.common.attention.FlexAttention.compiled_flex_attn", - return_value=(q_TNH.transpose(0, 1).unsqueeze(0), aux), + return_value=(q_THK.transpose(0, 1).unsqueeze(0), aux), ): attention( - q_TNH, - q_TNH, - q_TNH, + q_THK, + q_THK, + q_THK, attention_masks=block_mask, ) - self.assertEqual(len(attention.max_attention_logits_N), 1) + self.assertEqual(len(attention.max_attention_logits_H), 1) torch.testing.assert_close( - attention.max_attention_logits_N[0], + attention.max_attention_logits_H[0], torch.tensor([3.0, 4.0]), ) - attention.max_attention_logits_N.clear() + attention.max_attention_logits_H.clear() attention.eval() with patch( "torchtitan.models.common.attention.FlexAttention.compiled_flex_attn", - return_value=(q_TNH.transpose(0, 1).unsqueeze(0), aux), + return_value=(q_THK.transpose(0, 1).unsqueeze(0), aux), ): attention( - q_TNH, - q_TNH, - q_TNH, + q_THK, + q_THK, + q_THK, attention_masks=block_mask, ) - self.assertFalse(attention.max_attention_logits_N) + self.assertFalse(attention.max_attention_logits_H) def test_optimizer_hook_runs_qk_clip(self) -> None: model = nn.Linear(2, 2, bias=False) @@ -176,7 +176,7 @@ def weight_module(num_rows: int) -> nn.Module: if self.rank == 0 else torch.tensor([200.0, 50.0], device=device) ) - attention.inner_attention.max_attention_logits_N.append(rank_maxima) + attention.inner_attention.max_attention_logits_H.append(rank_maxima) model = nn.Module() model.add_module("attention", attention) @@ -186,18 +186,18 @@ def weight_module(num_rows: int) -> nn.Module: ) local_scale = 0.5 if self.rank == 0 else 0.25 - q_weight_NDI = attention.wq_b.weight.to_local().view( + q_weight_HDI = attention.wq_b.weight.to_local().view( 1, attention.qk_head_dim, in_features, ) - kv_weight_NDI = attention.wkv_b.weight.to_local().view( + kv_weight_HDI = attention.wkv_b.weight.to_local().view( 1, qk_nope_head_dim + v_head_dim, in_features, ) torch.testing.assert_close( - q_weight_NDI[:, :qk_nope_head_dim], + q_weight_HDI[:, :qk_nope_head_dim], torch.full( (1, qk_nope_head_dim, in_features), local_scale**0.5, @@ -205,7 +205,7 @@ def weight_module(num_rows: int) -> nn.Module: ), ) torch.testing.assert_close( - q_weight_NDI[:, qk_nope_head_dim:], + q_weight_HDI[:, qk_nope_head_dim:], torch.full( (1, qk_rope_head_dim, in_features), local_scale, @@ -213,7 +213,7 @@ def weight_module(num_rows: int) -> nn.Module: ), ) torch.testing.assert_close( - kv_weight_NDI[:, :qk_nope_head_dim], + kv_weight_HDI[:, :qk_nope_head_dim], torch.full( (1, qk_nope_head_dim, in_features), local_scale**0.5, @@ -221,10 +221,10 @@ def weight_module(num_rows: int) -> nn.Module: ), ) torch.testing.assert_close( - kv_weight_NDI[:, qk_nope_head_dim:], + kv_weight_HDI[:, qk_nope_head_dim:], torch.ones(1, v_head_dim, in_features, device=device), ) - self.assertFalse(attention.inner_attention.max_attention_logits_N) + self.assertFalse(attention.inner_attention.max_attention_logits_H) @with_comms def test_weight_scaling_is_communication_free(self) -> None: @@ -270,7 +270,7 @@ def weight_module(num_rows: int) -> nn.Module: attention.wq_b = weight_module(num_heads * attention.qk_head_dim) attention.wkv_b = weight_module(num_heads * (qk_nope_head_dim + v_head_dim)) attention.inner_attention = QKClipFlexAttention.Config().build() - attention.inner_attention.max_attention_logits_N.append( + attention.inner_attention.max_attention_logits_H.append( torch.full((num_heads,), 400.0, device=device) ) model.add_module(f"layer_{layer_id}", attention) diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py index adba767f25..7e2d4ede00 100644 --- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py +++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py @@ -15,7 +15,7 @@ VarlenMetadata, ) -# Tensor shape suffixes: B batch, L seq len, N heads, K key head dim, +# Tensor shape suffixes: B batch, L seq len, H heads, K query/key head dim, # V value head dim. @@ -25,11 +25,11 @@ def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor: def _torch_native_gated_delta( - q_BLNK: torch.Tensor, - k_BLNK: torch.Tensor, - v_BLNV: torch.Tensor, - g_BLN: torch.Tensor, - beta_BLN: torch.Tensor, + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, + g_BLH: torch.Tensor, + beta_BLH: torch.Tensor, ) -> torch.Tensor: """Standalone math reference for the gated delta rule recurrence. @@ -37,64 +37,64 @@ def _torch_native_gated_delta( numerical baseline for the FLA kernels. Args: - q_BLNK, k_BLNK: (batch, seq, n_heads, key_head_dim) - v_BLNV: (batch, seq, n_heads, value_head_dim) - g_BLN: (batch, seq, n_heads) -- log-space decay, always negative - beta_BLN: (batch, seq, n_heads) -- update gate in (0, 1) + q_BLHK, k_BLHK: (batch, seq, num_heads, key_head_dim) + v_BLHV: (batch, seq, num_heads, value_head_dim) + g_BLH: (batch, seq, num_heads) -- log-space decay, always negative + beta_BLH: (batch, seq, num_heads) -- update gate in (0, 1) Returns: output: (batch, seq, n_heads, value_head_dim) """ - B, L, N, K = q_BLNK.shape - V = v_BLNV.shape[-1] - dtype = q_BLNK.dtype + B, L, H, K = q_BLHK.shape + V = v_BLHV.shape[-1] + dtype = q_BLHK.dtype # Upcast to float32 -- recurrence accumulates over seqlen steps - q_BLNK = _l2norm(q_BLNK.float(), dim=-1) * (K**-0.5) - k_BLNK = _l2norm(k_BLNK.float(), dim=-1) - v_BLNV, g_BLN, beta_BLN = v_BLNV.float(), g_BLN.float(), beta_BLN.float() + q_BLHK = _l2norm(q_BLHK.float(), dim=-1) * (K**-0.5) + k_BLHK = _l2norm(k_BLHK.float(), dim=-1) + v_BLHV, g_BLH, beta_BLH = v_BLHV.float(), g_BLH.float(), beta_BLH.float() - out_BLNV = torch.zeros(B, L, N, V, dtype=torch.float32, device=q_BLNK.device) - state_BNKV = torch.zeros(B, N, K, V, dtype=torch.float32, device=q_BLNK.device) + out_BLHV = torch.zeros(B, L, H, V, dtype=torch.float32, device=q_BLHK.device) + state_BHKV = torch.zeros(B, H, K, V, dtype=torch.float32, device=q_BLHK.device) for t in range(L): - q_BNK = q_BLNK[:, t] - k_BNK = k_BLNK[:, t] - v_BNV = v_BLNV[:, t] - g_BN11 = g_BLN[:, t].exp().unsqueeze(-1).unsqueeze(-1) - beta_BN1 = beta_BLN[:, t].unsqueeze(-1) + q_BHK = q_BLHK[:, t] + k_BHK = k_BLHK[:, t] + v_BHV = v_BLHV[:, t] + g_BH11 = g_BLH[:, t].exp().unsqueeze(-1).unsqueeze(-1) + beta_BH1 = beta_BLH[:, t].unsqueeze(-1) - state_BNKV = state_BNKV * g_BN11 - kv_mem_BNV = torch.einsum("bnkv,bnk->bnv", state_BNKV, k_BNK) - delta_BNV = (v_BNV - kv_mem_BNV) * beta_BN1 - state_BNKV = state_BNKV + torch.einsum("bnk,bnv->bnkv", k_BNK, delta_BNV) - out_BLNV[:, t] = torch.einsum("bnkv,bnk->bnv", state_BNKV, q_BNK) + state_BHKV = state_BHKV * g_BH11 + kv_mem_BHV = torch.einsum("bhkv,bhk->bhv", state_BHKV, k_BHK) + delta_BHV = (v_BHV - kv_mem_BHV) * beta_BH1 + state_BHKV = state_BHKV + torch.einsum("bhk,bhv->bhkv", k_BHK, delta_BHV) + out_BLHV[:, t] = torch.einsum("bhkv,bhk->bhv", state_BHKV, q_BHK) - return out_BLNV.to(dtype) + return out_BLHV.to(dtype) def _torch_native_gated_delta_varlen( - q_BLNK: torch.Tensor, - k_BLNK: torch.Tensor, - v_BLNV: torch.Tensor, - g_BLN: torch.Tensor, - beta_BLN: torch.Tensor, + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, + g_BLH: torch.Tensor, + beta_BLH: torch.Tensor, cu_seqlens_cpu: torch.Tensor, ) -> torch.Tensor: """Varlen reference: run each packed document through the batched reference.""" - out_segments_BLNV: list[torch.Tensor] = [] + out_segments_BLHV: list[torch.Tensor] = [] cu_seqlens_list = cu_seqlens_cpu.tolist() for start, end in zip(cu_seqlens_list[:-1], cu_seqlens_list[1:], strict=False): - out_segments_BLNV.append( + out_segments_BLHV.append( _torch_native_gated_delta( - q_BLNK[:, start:end], - k_BLNK[:, start:end], - v_BLNV[:, start:end], - g_BLN[:, start:end], - beta_BLN[:, start:end], + q_BLHK[:, start:end], + k_BLHK[:, start:end], + v_BLHV[:, start:end], + g_BLH[:, start:end], + beta_BLH[:, start:end], ) ) - return torch.cat(out_segments_BLNV, dim=1) + return torch.cat(out_segments_BLHV, dim=1) def _reference_causal_conv1d_varlen( @@ -108,21 +108,21 @@ def _reference_causal_conv1d_varlen( ``gdn._causal_conv1d_varlen`` for CPU runs. """ conv_kernel_size = weight.shape[-1] - out_segments_BTD: list[torch.Tensor] = [] + out_segments_1TD: list[torch.Tensor] = [] cu_seqlens_list = cu_seqlens_cpu.tolist() for start, end in zip(cu_seqlens_list[:-1], cu_seqlens_list[1:], strict=False): - x_segment_BDT = F.pad( + x_segment_1DT = F.pad( x_TD[start:end].transpose(0, 1).unsqueeze(0), [conv_kernel_size - 1, 0], ) - out_segment_BTD = F.conv1d( - x_segment_BDT, + out_segment_1TD = F.conv1d( + x_segment_1DT, weight, None, groups=weight.size(0), ).transpose(1, 2) - out_segments_BTD.append(out_segment_BTD) - return F.silu(torch.cat(out_segments_BTD, dim=1)).squeeze(0) + out_segments_1TD.append(out_segment_1TD) + return F.silu(torch.cat(out_segments_1TD, dim=1)).squeeze(0) class ReferenceGatedDeltaKernel(nn.Module): @@ -136,34 +136,34 @@ class ReferenceGatedDeltaKernel(nn.Module): def forward( self, - xq_TNK: torch.Tensor, - xk_TNK: torch.Tensor, - xv_TNV: torch.Tensor, - g_TN: torch.Tensor, - beta_TN: torch.Tensor, + xq_THK: torch.Tensor, + xk_THK: torch.Tensor, + xv_THV: torch.Tensor, + g_TH: torch.Tensor, + beta_TH: torch.Tensor, *, cu_seqlens: torch.Tensor | None = None, cu_seqlens_cpu: torch.Tensor | None = None, ) -> torch.Tensor: - if xq_TNK.shape[1] != xv_TNV.shape[1]: - assert xv_TNV.shape[1] % xq_TNK.shape[1] == 0 - repeat = xv_TNV.shape[1] // xq_TNK.shape[1] - xq_TNK = xq_TNK.repeat_interleave(repeat, dim=1) - xk_TNK = xk_TNK.repeat_interleave(repeat, dim=1) - - xq_BLNK = xq_TNK.unsqueeze(0) - xk_BLNK = xk_TNK.unsqueeze(0) - xv_BLNV = xv_TNV.unsqueeze(0) - g_BLN = g_TN.unsqueeze(0) - beta_BLN = beta_TN.unsqueeze(0) + if xq_THK.shape[1] != xv_THV.shape[1]: + assert xv_THV.shape[1] % xq_THK.shape[1] == 0 + repeat = xv_THV.shape[1] // xq_THK.shape[1] + xq_THK = xq_THK.repeat_interleave(repeat, dim=1) + xk_THK = xk_THK.repeat_interleave(repeat, dim=1) + + xq_1THK = xq_THK.unsqueeze(0) + xk_1THK = xk_THK.unsqueeze(0) + xv_1THV = xv_THV.unsqueeze(0) + g_1TH = g_TH.unsqueeze(0) + beta_1TH = beta_TH.unsqueeze(0) if cu_seqlens is None: return _torch_native_gated_delta( - xq_BLNK, xk_BLNK, xv_BLNV, g_BLN, beta_BLN + xq_1THK, xk_1THK, xv_1THV, g_1TH, beta_1TH ).squeeze(0) assert cu_seqlens_cpu is not None return _torch_native_gated_delta_varlen( - xq_BLNK, xk_BLNK, xv_BLNV, g_BLN, beta_BLN, cu_seqlens_cpu + xq_1THK, xk_1THK, xv_1THV, g_1TH, beta_1TH, cu_seqlens_cpu ).squeeze(0) @@ -346,33 +346,33 @@ def causal_conv(tensor, conv): .transpose(0, 1) ) - query_TNK = causal_conv(model.in_proj_q(x_TD), model.conv_q).reshape( + query_THK = causal_conv(model.in_proj_q(x_TD), model.conv_q).reshape( num_tokens, -1, model.key_head_dim ) - key_TNK = causal_conv(model.in_proj_k(x_TD), model.conv_k).reshape( + key_THK = causal_conv(model.in_proj_k(x_TD), model.conv_k).reshape( num_tokens, -1, model.key_head_dim ) - value_TNV = causal_conv(model.in_proj_v(x_TD), model.conv_v).reshape( + value_THV = causal_conv(model.in_proj_v(x_TD), model.conv_v).reshape( num_tokens, -1, model.value_head_dim ) - gate_TNV = model.in_proj_z(x_TD).reshape(num_tokens, -1, model.value_head_dim) - a_TN = model.in_proj_a(x_TD) - b_TN = model.in_proj_b(x_TD) - decay_TN = -torch.exp(model.A_log.float()) * F.softplus( - a_TN.float() + model.dt_bias + gate_THV = model.in_proj_z(x_TD).reshape(num_tokens, -1, model.value_head_dim) + a_TH = model.in_proj_a(x_TD) + b_TH = model.in_proj_b(x_TD) + decay_TH = -torch.exp(model.A_log.float()) * F.softplus( + a_TH.float() + model.dt_bias ) - update_gate_TN = torch.sigmoid(b_TN) - output_TNV = model.inner_gated_delta_net.kernel( - query_TNK, - key_TNK, - value_TNV, - decay_TN, - update_gate_TN, + update_gate_TH = torch.sigmoid(b_TH) + output_THV = model.inner_gated_delta_net.kernel( + query_THK, + key_THK, + value_THV, + decay_TH, + update_gate_TH, cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, ) - output_TNV = model.norm(output_TNV, gate_TNV) - return model.out_proj(output_TNV.reshape(num_tokens, -1)) + output_THV = model.norm(output_THV, gate_THV) + return model.out_proj(output_THV.reshape(num_tokens, -1)) def test_extracted_forward_matches_main(self): torch.manual_seed(42) diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py index 71a38c10f0..f9bdf1557f 100644 --- a/torchtitan/experiments/graph_trainer/common_utils.py +++ b/torchtitan/experiments/graph_trainer/common_utils.py @@ -42,18 +42,18 @@ class Config(ScaledDotProductAttention.Config): def forward( self, - q_TNH: torch.Tensor, - k_TNH: torch.Tensor, - v_TNH: torch.Tensor, + q_THK: torch.Tensor, + k_THK: torch.Tensor, + v_THV: torch.Tensor, **kwargs, ) -> torch.Tensor: - out_BLNH = super().forward( - q_TNH.unsqueeze(0), - k_TNH.unsqueeze(0), - v_TNH.unsqueeze(0), + out_1THV = super().forward( + q_THK.unsqueeze(0), + k_THK.unsqueeze(0), + v_THV.unsqueeze(0), **kwargs, ) - return out_BLNH.squeeze(0) + return out_1THV.squeeze(0) @contextmanager diff --git a/torchtitan/experiments/rl/models/attention.py b/torchtitan/experiments/rl/models/attention.py index 3db95958d9..50b6b7714c 100644 --- a/torchtitan/experiments/rl/models/attention.py +++ b/torchtitan/experiments/rl/models/attention.py @@ -334,9 +334,9 @@ def __init__(self, config: Config) -> None: def forward( self, - q_TNH: torch.Tensor, - k_TNH: torch.Tensor, - v_TNH: torch.Tensor, + q_THK: torch.Tensor, + k_THK: torch.Tensor, + v_THV: torch.Tensor, *, attention_masks: AttentionMasksType | None = None, **kwargs, @@ -344,9 +344,9 @@ def forward( """Run vLLM paged attention on local (non-DTensor) tensors. Args: - q_TNH: ``(num_tokens, num_heads, head_dim)`` - k_TNH: ``(num_tokens, num_kv_heads, head_dim)`` - v_TNH: ``(num_tokens, num_kv_heads, head_dim)`` + q_THK: ``(num_tokens, num_heads, query/key_head_dim)`` + k_THK: ``(num_tokens, num_kv_heads, query/key_head_dim)`` + v_THV: ``(num_tokens, num_kv_heads, value_head_dim)`` Returns: ``(num_tokens, num_heads, head_dim)``. @@ -357,11 +357,11 @@ def forward( "manages causal masking and the KV-cache internally." ) - out_TD = self.vllm_attn(q_TNH, k_TNH, v_TNH) + out_TD = self.vllm_attn(q_THK, k_THK, v_THV) # vLLM's flash attention backend may pad the token count (e.g. # round up to an even number), which introduces a new symbolic # shape under torch.compile. Narrow to trim this padding. - num_tokens, _, head_dim = q_TNH.shape + num_tokens, _, head_dim = q_THK.shape out_TD = out_TD.narrow(0, 0, num_tokens) return out_TD.view(num_tokens, -1, head_dim) diff --git a/torchtitan/experiments/rl/models/gdn.py b/torchtitan/experiments/rl/models/gdn.py index 6544b5ee18..47348ba6d2 100644 --- a/torchtitan/experiments/rl/models/gdn.py +++ b/torchtitan/experiments/rl/models/gdn.py @@ -429,13 +429,13 @@ def forward( query_TC: torch.Tensor, key_TC: torch.Tensor, value_TC: torch.Tensor, - a_TN: torch.Tensor, - b_TN: torch.Tensor, + a_TH: torch.Tensor, + b_TH: torch.Tensor, conv_q_weight_C1W: torch.Tensor, conv_k_weight_C1W: torch.Tensor, conv_v_weight_C1W: torch.Tensor, - A_log_N: torch.Tensor, - dt_bias_N: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_H: torch.Tensor, cu_seqlens: torch.Tensor, *, key_head_dim: int, @@ -456,17 +456,17 @@ def forward( num_tokens = mixed_qkv_TC.shape[0] # Padded rows must remain defined across vLLM graph replays. - output_TNV = mixed_qkv_TC.new_zeros( + output_THV = mixed_qkv_TC.new_zeros( num_tokens, self.local_num_v_heads, self.head_v_dim ) self._forward( mixed_qkv_TC, - a_TN, - b_TN, + a_TH, + b_TH, conv_weight_CW, None, - A_log_N, - dt_bias_N, - output_TNV, + A_log_H, + dt_bias_H, + output_THV, ) - return output_TNV + return output_THV diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py index 122cb11181..941e3f0589 100644 --- a/torchtitan/models/common/attention.py +++ b/torchtitan/models/common/attention.py @@ -6,10 +6,11 @@ # Shape suffix legend # (https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd): -# B = singleton kernel batch, T = packed tokens, D = model dimension, -# N = num heads (N is used for both query and kv heads in GQA; +# B = batch, T = packed tokens, L = sequence length, +# D = model dimension, +# H = attention heads (H is used for both query and kv heads in GQA; # the variable name xq/xk/xv disambiguates), -# H = head dimension (per-head dim) +# K = query/key head dimension, V = value head dimension. from collections.abc import Callable, Mapping from dataclasses import dataclass, field @@ -153,9 +154,9 @@ def __init__(self, config: Config) -> None: def forward( self, - q_TNH: torch.Tensor, - k_TNH: torch.Tensor, - v_TNH: torch.Tensor, + q_THK: torch.Tensor, + k_THK: torch.Tensor, + v_THV: torch.Tensor, *, attention_masks: VarlenMetadata, scale: float | None = None, @@ -195,9 +196,9 @@ def forward( varlen_attn_fn = varlen_attn if out_transform is None else varlen_attn_with_lse result = varlen_attn_fn( - q_TNH.to(torch.bfloat16), - k_TNH.to(torch.bfloat16), - v_TNH.to(torch.bfloat16), + q_THK.to(torch.bfloat16), + k_THK.to(torch.bfloat16), + v_THV.to(torch.bfloat16), cu_seq_q, cu_seq_k, max_q, @@ -207,24 +208,24 @@ def forward( **varlen_kwargs, ) - # varlen_attn returns the packed output (T, N, H), plus the LSE when an + # varlen_attn returns the packed output (T, H, V), plus the LSE when an # out_transform epilogue was requested. if out_transform is None: assert isinstance(result, torch.Tensor) - return result.to(q_TNH.dtype) + return result.to(q_THK.dtype) - out_TNH, lse_NT = result - out_TNH = out_TNH.to(q_TNH.dtype) - lse_TN = lse_NT.transpose(0, 1) - return out_transform(out_TNH, lse_TN) + out_THV, lse_HT = result + out_THV = out_THV.to(q_THK.dtype) + lse_TH = lse_HT.transpose(0, 1) + return out_transform(out_THV, lse_TH) class FlexAttention(Module): """Inner attention using ``flex_attention`` with torch.compile and CP support. - Inputs use ``[T, N, H]``. The FlexAttention kernel requires a batch - dimension, so inputs are adapted to ``[1, N, T, H]`` only at the kernel - boundary. + Query/key inputs use ``[T, H, K]`` and value inputs use ``[T, H, V]``. + The FlexAttention kernel requires a batch dimension, so inputs are adapted + to ``[1, H, T, K]`` and ``[1, H, T, V]`` only at the kernel boundary. Note: The forward function must have q, k, v as the first three arguments @@ -293,7 +294,7 @@ def compiled_flex_attn( Compiled regions are not currently compatible with SPMD typechecking, so the opaque kernel output is re-typed at the boundary instead of typechecking into Flex. Attention preserves the query's sharding (output - is (B, N, T, H) with the same kernel-batch/head/token layout as ``q``), + is (1, H, T, V) with the same kernel-batch/head/token layout as ``q``), so ``out`` takes ``q``'s full SPMD type (local type + shard-dim PartitionSpec), and ``lse`` takes the same minus the trailing (unsharded) head dim. @@ -315,7 +316,7 @@ def compiled_flex_attn( q_local = spmd.get_local_type(q) q_ps = spmd.get_partition_spec(q) spmd.assert_type(out, q_local, q_ps) - # Aux outputs are (B, N, T) = q minus the trailing head dim. + # Aux outputs are (1, H, T) = q minus the trailing head dim. aux_ps = None if q_ps is None else spmd.PartitionSpec(*q_ps[:-1]) if return_aux.lse: spmd.assert_type(aux.lse, q_local, aux_ps) @@ -325,9 +326,9 @@ def compiled_flex_attn( def forward( self, - q_TNH: torch.Tensor, - k_TNH: torch.Tensor, - v_TNH: torch.Tensor, + q_THK: torch.Tensor, + k_THK: torch.Tensor, + v_THV: torch.Tensor, *, attention_masks: BlockMask, score_mod: _score_mod_signature | None = None, @@ -343,9 +344,9 @@ def forward( attention_masks, BlockMask ), f"attention_masks must be instance of BlockMask, got {type(attention_masks)}" - q_BNTH = q_TNH.transpose(0, 1).unsqueeze(0) - k_BNTH = k_TNH.transpose(0, 1).unsqueeze(0) - v_BNTH = v_TNH.transpose(0, 1).unsqueeze(0) + q_1HTK = q_THK.transpose(0, 1).unsqueeze(0) + k_1HTK = k_THK.transpose(0, 1).unsqueeze(0) + v_1HTV = v_THV.transpose(0, 1).unsqueeze(0) aux_request = self._get_aux_request(return_lse=out_transform is not None) # 1. _compiled_flex_attn has to be a class variable, otherwise there will @@ -358,10 +359,10 @@ def forward( # an inductor sub-compile (see distributed/compile.py). A null context on # the default inductor / eager paths, so no dead metadata is emitted. with maybe_regional_inductor(FlexAttention.inductor_configs): - out_BNTH, aux = FlexAttention.compiled_flex_attn( - q_BNTH, - k_BNTH, - v_BNTH, + out_1HTV, aux = FlexAttention.compiled_flex_attn( + q_1HTK, + k_1HTK, + v_1HTV, score_mod=score_mod, block_mask=attention_masks, scale=scale, @@ -370,11 +371,11 @@ def forward( kernel_options=self.kernel_options, ) self._process_aux(aux) - out_TNH = out_BNTH.squeeze(0).transpose(0, 1) + out_THV = out_1HTV.squeeze(0).transpose(0, 1) if out_transform is None: - return out_TNH - lse_TN = aux.lse.squeeze(0).transpose(0, 1) - return out_transform(out_TNH, lse_TN) + return out_THV + lse_TH = aux.lse.squeeze(0).transpose(0, 1) + return out_transform(out_THV, lse_TH) # TODO: Verify whether SDPA support can be removed without losing performance @@ -382,8 +383,9 @@ def forward( class ScaledDotProductAttention(Module): """Inner attention using ``F.scaled_dot_product_attention`` with CP support. - ``forward()`` adapts ``(B, L, N, H)`` to the kernel's ``(B, N, L, H)`` - layout and converts the result back to ``(B, L, N, H)``. + ``forward()`` adapts Q/K from ``(B, L, H, K)`` to ``(B, H, L, K)`` and V + from ``(B, L, H, V)`` to ``(B, H, L, V)``, then converts the result back to + ``(B, L, H, V)``. Note: The forward function must have q, k, v as the first three arguments to be @@ -409,9 +411,9 @@ def __init__(self, config: Config) -> None: def forward( self, - q_BLNH: torch.Tensor, - k_BLNH: torch.Tensor, - v_BLNH: torch.Tensor, + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, *, attention_masks: AttentionMasksType | None = None, scale: float | None = None, @@ -424,21 +426,21 @@ def forward( "ScaledDotProductAttention does not support attention_masks; it " "only supports causal/non-causal attention via is_causal." ) - q_BNLH, k_BNLH, v_BNLH = ( - q_BLNH.transpose(1, 2), - k_BLNH.transpose(1, 2), - v_BLNH.transpose(1, 2), + q_BHLK, k_BHLK, v_BHLV = ( + q_BLHK.transpose(1, 2), + k_BLHK.transpose(1, 2), + v_BLHV.transpose(1, 2), ) with sdpa_kernel(self.sdpa_backends, set_priority=True): - out_BNLH = F.scaled_dot_product_attention( - q_BNLH, - k_BNLH, - v_BNLH, + out_BHLV = F.scaled_dot_product_attention( + q_BHLK, + k_BHLK, + v_BHLV, scale=scale, is_causal=is_causal, enable_gqa=enable_gqa, ) - return out_BNLH.transpose(1, 2) + return out_BHLV.transpose(1, 2) def get_causal_mask_mod() -> _mask_mod_signature: @@ -686,7 +688,7 @@ def forward( """Project input into Q, K, V tensors. Returns: - (xq, xk, xv) each with shape ``[T, N, H]``. + xq and xk have shape ``[T, H, K]``; xv has shape ``[T, H, V]``. """ raise NotImplementedError @@ -717,14 +719,14 @@ def local_qkv_head_split(x): # Drop into local region, we can't propagate S(1) -> qkv head unflatten. # TODO(pianpwk): this should be doable once spmd_types tracks sharding evenness. with spmd.local(): - x_TNH = x.view(num_tokens, -1, self.head_dim) + x = x.view(num_tokens, -1, self.head_dim) if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): spmd.assert_type( - x_TNH, + x, spmd.V, spmd.PartitionSpec(("dp", "cp"), "tp", None), ) - return x_TNH + return x xq, xk, xv = ( local_qkv_head_split(xq), @@ -952,24 +954,24 @@ def forward( attention_masks: AttentionMasksType | None, positions: torch.Tensor | None = None, ) -> torch.Tensor: - xq_TNH, xk_TNH, xv_TNH = self.qkv_linear(x_TD) + xq_THK, xk_THK, xv_THV = self.qkv_linear(x_TD) # Optional QK normalization (before RoPE, per Qwen3) if self.q_norm is not None or self.k_norm is not None: assert self.q_norm is not None and self.k_norm is not None - xq_TNH = self.q_norm(xq_TNH) - xk_TNH = self.k_norm(xk_TNH) + xq_THK = self.q_norm(xq_THK) + xk_THK = self.k_norm(xk_THK) # Apply rotary embeddings - xq_TNH, xk_TNH = self.rope(xq_TNH, xk_TNH, positions) + xq_THK, xk_THK = self.rope(xq_THK, xk_THK, positions) - out_TNH = self.inner_attention( - xq_TNH, - xk_TNH, - xv_TNH, + out_THV = self.inner_attention( + xq_THK, + xk_THK, + xv_THV, attention_masks=attention_masks, scale=self.scaling, enable_gqa=self.enable_gqa, ).contiguous() - out_TD = out_TNH.view(out_TNH.shape[0], -1) + out_TD = out_THV.view(out_THV.shape[0], -1) return self.wo(out_TD) diff --git a/torchtitan/models/common/decoder_sharding.py b/torchtitan/models/common/decoder_sharding.py index 2e43e327de..028284c1a4 100644 --- a/torchtitan/models/common/decoder_sharding.py +++ b/torchtitan/models/common/decoder_sharding.py @@ -275,7 +275,8 @@ def set_gqa_attention_sharding(attention_cfg, *, enable_sp: bool) -> None: def set_gqa_inner_attention_local_map(inner_attention_cfg) -> None: """Install a ``LocalMapConfig`` on an inner-attention config. - q/k/v use ``(T, N, H)`` layout. DP/CP shard T and TP shards N. + q/k use ``(T, H, K)`` and v uses ``(T, H, V)``. DP/CP shard T and TP + shards H. ``local_map`` converts DTensors to local tensors before the kernel runs, then wraps outputs back. @@ -298,14 +299,14 @@ def set_gqa_inner_attention_local_map(inner_attention_cfg) -> None: out_src: SpmdType = q_placements inner_attention_cfg.sharding_config = ShardingConfig( in_src_shardings={ - "q_TNH": q_placements, - "k_TNH": kv_src_placements, - "v_TNH": kv_src_placements, + "q_THK": q_placements, + "k_THK": kv_src_placements, + "v_THV": kv_src_placements, }, in_dst_shardings={ - "q_TNH": q_placements, - "k_TNH": kv_dst_placements, - "v_TNH": kv_dst_placements, + "q_THK": q_placements, + "k_THK": kv_dst_placements, + "v_THV": kv_dst_placements, }, out_src_shardings=out_src, local_map=LocalMapConfig( diff --git a/torchtitan/models/flux/sharding.py b/torchtitan/models/flux/sharding.py index 1166453508..d2a9fb4fda 100644 --- a/torchtitan/models/flux/sharding.py +++ b/torchtitan/models/flux/sharding.py @@ -45,14 +45,14 @@ def set_flux_inner_attention_local_map(inner_attention_cfg) -> None: inner_attention_cfg.sharding_config = ShardingConfig( in_src_shardings={ - "q_BLNH": q_layout, - "k_BLNH": kv_src_layout, - "v_BLNH": kv_src_layout, + "q_BLHK": q_layout, + "k_BLHK": kv_src_layout, + "v_BLHV": kv_src_layout, }, in_dst_shardings={ - "q_BLNH": q_layout, - "k_BLNH": kv_dst_layout, - "v_BLNH": kv_dst_layout, + "q_BLHK": q_layout, + "k_BLHK": kv_dst_layout, + "v_BLHV": kv_dst_layout, }, out_src_shardings=q_layout, local_map=LocalMapConfig( diff --git a/torchtitan/models/kimi_k2_7/qk_clip.py b/torchtitan/models/kimi_k2_7/qk_clip.py index 6a73f75ea5..c47a692e9f 100644 --- a/torchtitan/models/kimi_k2_7/qk_clip.py +++ b/torchtitan/models/kimi_k2_7/qk_clip.py @@ -19,6 +19,10 @@ from torchtitan.models.common.attention import FlexAttention from torchtitan.models.deepseek_v3.model import Attention +# Shape suffixes: +# T = packed tokens, H = attention heads, D = projection rows per head, +# I = input features. + class QKClipFlexAttention(FlexAttention): """FlexAttention that records the maximum score for each query head.""" @@ -29,17 +33,17 @@ class Config(FlexAttention.Config): def __init__(self, config: Config) -> None: super().__init__(config) - self.max_attention_logits_N: list[torch.Tensor] = [] + self.max_attention_logits_H: list[torch.Tensor] = [] def _get_aux_request(self, *, return_lse: bool) -> AuxRequest: return AuxRequest(lse=return_lse, max_scores=self.training) def _process_aux(self, aux: Any) -> None: if self.training: - max_scores_BNT = aux.max_scores - assert max_scores_BNT is not None + max_scores_1HT = aux.max_scores + assert max_scores_1HT is not None # Record gradient-accumulation and PP microbatches, plus AC recomputation. - self.max_attention_logits_N.append(max_scores_BNT.amax(dim=(0, 2)).detach()) + self.max_attention_logits_H.append(max_scores_1HT.amax(dim=(0, 2)).detach()) def _validate_head_sharding(weight: DTensor) -> None: @@ -56,7 +60,7 @@ def _validate_head_sharding(weight: DTensor) -> None: ) -def _replicated_scales(scales_N: torch.Tensor, weight: DTensor) -> DTensor: +def _replicated_scales(scales_H: torch.Tensor, weight: DTensor) -> DTensor: """Represent per-head scales on the same distributed mesh as ``weight``. The MAX all-reduce already leaves identical scales on every rank. @@ -65,7 +69,7 @@ def _replicated_scales(scales_N: torch.Tensor, weight: DTensor) -> DTensor: ``distribute_tensor`` would add an unnecessary broadcast. """ return DTensor.from_local( - scales_N, + scales_H, weight.device_mesh, tuple(Replicate() for _ in weight.placements), run_check=False, @@ -75,7 +79,7 @@ def _replicated_scales(scales_N: torch.Tensor, weight: DTensor) -> DTensor: @torch.no_grad() def _scale_mla_heads( weight: DTensor, - scales_N: torch.Tensor, + scales_H: torch.Tensor, *, rows_per_head: int, nope_rows_per_head: int, @@ -85,25 +89,25 @@ def _scale_mla_heads( """Scale the NoPE and remaining rows of every MLA head in place. ``weight`` is viewed as ``[num_heads, rows_per_head, in_features]``, and - ``scales_N`` contains one scale per head. The remaining rows are unchanged + ``scales_H`` contains one scale per head. The remaining rows are unchanged when ``remaining_scale_exponent`` is ``None``. """ - num_heads = scales_N.numel() + num_heads = scales_H.numel() if weight.ndim != 2 or weight.shape[0] != num_heads * rows_per_head: raise ValueError("QK clip scales do not match the MLA weight shape.") _validate_head_sharding(weight) - scales_N11 = _replicated_scales(scales_N, weight).view(-1, 1, 1) - heads_NDI = weight.view(num_heads, rows_per_head, weight.shape[1]) - heads_NDI[:, :nope_rows_per_head].mul_(scales_N11.pow(nope_scale_exponent)) + scales_H11 = _replicated_scales(scales_H, weight).view(-1, 1, 1) + heads_HDI = weight.view(num_heads, rows_per_head, weight.shape[1]) + heads_HDI[:, :nope_rows_per_head].mul_(scales_H11.pow(nope_scale_exponent)) if remaining_scale_exponent is not None: - heads_NDI[:, nope_rows_per_head:].mul_(scales_N11.pow(remaining_scale_exponent)) + heads_HDI[:, nope_rows_per_head:].mul_(scales_H11.pow(remaining_scale_exponent)) @torch.no_grad() def _clip_mla_weights( attention: Attention, - scales_N: torch.Tensor, + scales_H: torch.Tensor, *, alpha: float, ) -> None: @@ -111,7 +115,7 @@ def _clip_mla_weights( # Query: NoPE rows take ``scale ** alpha``, RoPE rows take the full scale. _scale_mla_heads( cast(DTensor, q_projection.weight), - scales_N, + scales_H, rows_per_head=attention.qk_head_dim, nope_rows_per_head=attention.qk_nope_head_dim, nope_scale_exponent=alpha, @@ -121,7 +125,7 @@ def _clip_mla_weights( # V rows stay unchanged. _scale_mla_heads( cast(DTensor, attention.wkv_b.weight), - scales_N, + scales_H, rows_per_head=attention.qk_nope_head_dim + attention.v_head_dim, nope_rows_per_head=attention.qk_nope_head_dim, nope_scale_exponent=1.0 - alpha, @@ -150,27 +154,27 @@ def qk_clip( inner_attentions = [layer.inner_attention for layer in attention_layers] # Each entry holds one layer's local maximum logit per query head. - layer_max_logits_N = [ - torch.stack(inner_attention.max_attention_logits_N).amax(dim=0) + layer_max_logits_H = [ + torch.stack(inner_attention.max_attention_logits_H).amax(dim=0) for inner_attention in inner_attentions ] - num_heads_per_layer = [logits.numel() for logits in layer_max_logits_N] - max_logits_N = torch.cat(layer_max_logits_N) + num_heads_per_layer = [logits.numel() for logits in layer_max_logits_H] + max_logits_H = torch.cat(layer_max_logits_H) if reduction_mesh.size() > 1: dist.all_reduce( - max_logits_N, + max_logits_H, op=dist.ReduceOp.MAX, group=reduction_mesh.get_group(), ) - scales_N = threshold / max_logits_N.clamp_min(threshold) + scales_H = threshold / max_logits_H.clamp_min(threshold) - for attention, layer_scales_N in zip( + for attention, layer_scales_H in zip( attention_layers, - scales_N.split(num_heads_per_layer), + scales_H.split(num_heads_per_layer), strict=True, ): - _clip_mla_weights(attention, layer_scales_N, alpha=alpha) - attention.inner_attention.max_attention_logits_N.clear() + _clip_mla_weights(attention, layer_scales_H, alpha=alpha) + attention.inner_attention.max_attention_logits_H.clear() def register_qk_clip_hook( diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index ac53adc104..ebc7618bcc 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -20,6 +20,11 @@ from torchtitan.models.common.nn_modules import Conv1d from torchtitan.protocols.module import Module +# Shape suffixes: +# T = packed tokens, D = model dimension, C = projection channels, +# H = attention heads, K = query/key head dimension, V = value head dimension, +# W = convolution kernel width. + class KimiRMSNormGated(Module): """Per-head RMSNorm followed by a sigmoid output gate.""" @@ -34,15 +39,15 @@ def __init__(self, config: Config): self.eps = config.eps self.weight = nn.Parameter(torch.empty(config.dim)) - def forward(self, x_TNV: torch.Tensor, gate_TNV: torch.Tensor) -> torch.Tensor: - input_dtype = x_TNV.dtype - normalized_TNV = F.rms_norm( - x_TNV.float(), - (x_TNV.shape[-1],), + def forward(self, x_THV: torch.Tensor, gate_THV: torch.Tensor) -> torch.Tensor: + input_dtype = x_THV.dtype + normalized_THV = F.rms_norm( + x_THV.float(), + (x_THV.shape[-1],), self.weight.float(), self.eps, ) - return (normalized_TNV * gate_TNV.float().sigmoid()).to(input_dtype) + return (normalized_THV * gate_THV.float().sigmoid()).to(input_dtype) class KDAKernel(Module): @@ -65,43 +70,43 @@ def __init__(self, config: Config): def forward( self, - q_BTNK: torch.Tensor, - k_BTNK: torch.Tensor, - v_BTNV: torch.Tensor, - raw_gate_BTNK: torch.Tensor, - raw_beta_BTN: torch.Tensor, - A_log_N: torch.Tensor, - dt_bias_NK: torch.Tensor, + q_1THK: torch.Tensor, + k_1THK: torch.Tensor, + v_1THV: torch.Tensor, + raw_gate_1THK: torch.Tensor, + raw_beta_1TH: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_HK: torch.Tensor, *, cu_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: - if not q_BTNK.is_cuda: + if not q_1THK.is_cuda: raise RuntimeError("Attention Gym KDA requires CUDA tensors.") - capability = torch.cuda.get_device_capability(q_BTNK.device) + capability = torch.cuda.get_device_capability(q_1THK.device) if capability not in {(10, 0), (10, 3)}: raise RuntimeError( "Attention Gym KDA requires Blackwell SM100/SM103; " f"got CUDA capability {capability}." ) - gate_BTNK = bound_gate( - raw_gate_BTNK, + gate_1THK = bound_gate( + raw_gate_1THK, # TODO: The long-term solution is to specify mixed precision per FQN # instead of per layer. https://github.com/pytorch/pytorch/issues/156784 - A_log_N.float(), - dt_bias_NK.float(), + A_log_H.float(), + dt_bias_HK.float(), lower_bound=self.lower_bound, impl="fused", ) - output_BTNV, _ = chunk_kda( - l2norm(q_BTNK), - l2norm(k_BTNK), - v_BTNV, - gate_BTNK, - raw_beta_BTN.float().sigmoid(), + output_1THV, _ = chunk_kda( + l2norm(q_1THK), + l2norm(k_1THK), + v_1THV, + gate_1THK, + raw_beta_1TH.float().sigmoid(), cu_seqlens=cu_seqlens, ) - return output_BTNV + return output_1THV class InnerKDA(Module): @@ -128,18 +133,18 @@ def forward( query_TC: torch.Tensor, key_TC: torch.Tensor, value_TC: torch.Tensor, - raw_gate_TNK: torch.Tensor, - raw_beta_TN: torch.Tensor, + raw_gate_THK: torch.Tensor, + raw_beta_TH: torch.Tensor, conv_q_weight_C1W: torch.Tensor, conv_k_weight_C1W: torch.Tensor, conv_v_weight_C1W: torch.Tensor, - A_log_N: torch.Tensor, - dt_bias_NK: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_HK: torch.Tensor, cu_seqlens: torch.Tensor | None, ) -> torch.Tensor: - raw_gate_BTNK = raw_gate_TNK.unsqueeze(0) - raw_beta_BTN = raw_beta_TN.unsqueeze(0) - mixed_qkv_BTC = torch.cat( + raw_gate_1THK = raw_gate_THK.unsqueeze(0) + raw_beta_1TH = raw_beta_TH.unsqueeze(0) + mixed_qkv_1TC = torch.cat( (query_TC, key_TC, value_TC), dim=-1, ).unsqueeze(0) @@ -147,30 +152,30 @@ def forward( (conv_q_weight_C1W, conv_k_weight_C1W, conv_v_weight_C1W), dim=0, ) - conv_output_BTC = causal_conv1d( - mixed_qkv_BTC, + conv_output_1TC = causal_conv1d( + mixed_qkv_1TC, conv_weight_C1W[:, 0], activation="silu", cu_seqlens=cu_seqlens, ) - assert isinstance(conv_output_BTC, torch.Tensor) + assert isinstance(conv_output_1TC, torch.Tensor) - q_BTC, k_BTC, v_BTC = conv_output_BTC.chunk(3, dim=-1) - q_BTNK, k_BTNK, v_BTNV = ( + q_1TC, k_1TC, v_1TC = conv_output_1TC.chunk(3, dim=-1) + q_1THK, k_1THK, v_1THV = ( tensor.unflatten(-1, (-1, self.head_dim)) - for tensor in (q_BTC, k_BTC, v_BTC) + for tensor in (q_1TC, k_1TC, v_1TC) ) - output_BTNV = self.kernel( - q_BTNK, - k_BTNK, - v_BTNV, - raw_gate_BTNK, - raw_beta_BTN, - A_log_N, - dt_bias_NK, + output_1THV = self.kernel( + q_1THK, + k_1THK, + v_1THV, + raw_gate_1THK, + raw_beta_1TH, + A_log_H, + dt_bias_HK, cu_seqlens=cu_seqlens, ) - return output_BTNV.squeeze(0) + return output_1THV.squeeze(0) class KDA(Module): @@ -251,16 +256,16 @@ def forward( f"got {type(attention_masks).__name__}." ) num_tokens = x_TD.shape[0] - raw_gate_TNK = self.forget_b(self.forget_a(x_TD)).reshape( + raw_gate_THK = self.forget_b(self.forget_a(x_TD)).reshape( num_tokens, self.num_heads, self.head_dim ) - raw_beta_TN = self.beta(x_TD).reshape(num_tokens, self.num_heads) - out_TNV = self.inner_kda( + raw_beta_TH = self.beta(x_TD).reshape(num_tokens, self.num_heads) + out_THV = self.inner_kda( self.q_proj(x_TD), self.k_proj(x_TD), self.v_proj(x_TD), - raw_gate_TNK, - raw_beta_TN, + raw_gate_THK, + raw_beta_TH, self.q_conv.weight, self.k_conv.weight, self.v_conv.weight, @@ -269,5 +274,5 @@ def forward( cu_seqlens, ) - output_gate_TNV = self.output_gate(x_TD).view_as(out_TNV) - return self.output_proj(self.output_norm(out_TNV, output_gate_TNV).flatten(-2)) + output_gate_THV = self.output_gate(x_TD).view_as(out_THV) + return self.output_proj(self.output_norm(out_THV, output_gate_THV).flatten(-2)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index a052387ed1..8dce8c44df 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -35,8 +35,8 @@ from .vision_encoder import KimiK3VisionEncoder # Shape suffixes: -# T = packed tokens, D = model dimension, H = heads, -# K = key head dimension, V = value head dimension, +# T = packed tokens, D = model dimension, C = projection channels, H = heads, +# K = query/key head dimension, V = value head dimension, # N = attention-residual entries. diff --git a/torchtitan/models/qwen3_5/gdn.py b/torchtitan/models/qwen3_5/gdn.py index 5831629242..aa774351dd 100644 --- a/torchtitan/models/qwen3_5/gdn.py +++ b/torchtitan/models/qwen3_5/gdn.py @@ -26,6 +26,11 @@ from torchtitan.models.common.attention import VarlenMetadata from torchtitan.protocols.module import Module +# Shape suffixes: +# T = packed tokens, D = model dimension, C = projection channels, +# H = attention heads, K = query/key head dimension, V = value head dimension, +# W = convolution kernel width. + GatedDeltaBackend = Literal["fla_chunked", "fla_fused_recurrent"] spmd.register_local_autograd_function(ChunkGatedDeltaRuleFunction) @@ -61,7 +66,7 @@ def _causal_conv1d_varlen( from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d - out_BTD, _ = _fla_causal_conv1d( + out_1TD, _ = _fla_causal_conv1d( x=x_TD.unsqueeze(0), weight=weight.squeeze(1), bias=None, @@ -70,7 +75,7 @@ def _causal_conv1d_varlen( cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, ) - return out_BTD.squeeze(0) + return out_1TD.squeeze(0) class RMSNormGated(Module): @@ -254,27 +259,27 @@ def __init__(self, config: Config): def forward( self, - xq_TNK: torch.Tensor, - xk_TNK: torch.Tensor, - xv_TNV: torch.Tensor, - g_TN: torch.Tensor, - beta_TN: torch.Tensor, + xq_THK: torch.Tensor, + xk_THK: torch.Tensor, + xv_THV: torch.Tensor, + g_TH: torch.Tensor, + beta_TH: torch.Tensor, *, cu_seqlens: torch.Tensor | None = None, cu_seqlens_cpu: torch.Tensor | None = None, ) -> torch.Tensor: # Expand Q/K heads to match V when n_value_heads > n_key_heads - if xq_TNK.shape[1] != xv_TNV.shape[1]: - assert xv_TNV.shape[1] % xq_TNK.shape[1] == 0 - repeat = xv_TNV.shape[1] // xq_TNK.shape[1] - xq_TNK = xq_TNK.repeat_interleave(repeat, dim=1) - xk_TNK = xk_TNK.repeat_interleave(repeat, dim=1) - - xq_BTNK = xq_TNK.unsqueeze(0) - xk_BTNK = xk_TNK.unsqueeze(0) - xv_BTNV = xv_TNV.unsqueeze(0) - g_BTN = g_TN.unsqueeze(0) - beta_BTN = beta_TN.unsqueeze(0) + if xq_THK.shape[1] != xv_THV.shape[1]: + assert xv_THV.shape[1] % xq_THK.shape[1] == 0 + repeat = xv_THV.shape[1] // xq_THK.shape[1] + xq_THK = xq_THK.repeat_interleave(repeat, dim=1) + xk_THK = xk_THK.repeat_interleave(repeat, dim=1) + + xq_1THK = xq_THK.unsqueeze(0) + xk_1THK = xk_THK.unsqueeze(0) + xv_1THV = xv_THV.unsqueeze(0) + g_1TH = g_TH.unsqueeze(0) + beta_1TH = beta_TH.unsqueeze(0) if is_in_batch_invariant_mode() and cu_seqlens is not None: if cu_seqlens_cpu is None: @@ -282,11 +287,11 @@ def forward( "Batch-invariant Gated DeltaNet requires CPU cu_seqlens." ) return _recurrent_gdn_fwd( - xq_BTNK, - xk_BTNK, - xv_BTNV, - g_BTN, - beta_BTN, + xq_1THK, + xk_1THK, + xv_1THV, + g_1TH, + beta_1TH, cu_seqlens, cu_seqlens_cpu, ).squeeze(0) @@ -297,22 +302,22 @@ def forward( "Qwen3.5 FLA varlen DeltaNet requires a CPU cu_seqlens tensor." ) result = _fla_chunk_gated_delta_rule( - xq_BTNK, - xk_BTNK, - xv_BTNV, - g_BTN, - beta_BTN, + xq_1THK, + xk_1THK, + xv_1THV, + g_1TH, + beta_1TH, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, ) elif self.backend == "fla_fused_recurrent": result = _fla_fused_recurrent_gated_delta_rule( - xq_BTNK, - xk_BTNK, - xv_BTNV, - g_BTN, - beta=beta_BTN, + xq_1THK, + xk_1THK, + xv_1THV, + g_1TH, + beta=beta_1TH, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, ) @@ -347,13 +352,13 @@ def forward( query_TC: torch.Tensor, key_TC: torch.Tensor, value_TC: torch.Tensor, - a_TN: torch.Tensor, - b_TN: torch.Tensor, + a_TH: torch.Tensor, + b_TH: torch.Tensor, conv_q_weight_C1W: torch.Tensor, conv_k_weight_C1W: torch.Tensor, conv_v_weight_C1W: torch.Tensor, - A_log_N: torch.Tensor, - dt_bias_N: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_H: torch.Tensor, cu_seqlens: torch.Tensor, *, key_head_dim: int, @@ -401,23 +406,23 @@ def causal_conv( .transpose(0, 1) ) - xq_TNK = causal_conv(query_TC, conv_q_weight_C1W).reshape( + xq_THK = causal_conv(query_TC, conv_q_weight_C1W).reshape( num_tokens, -1, key_head_dim ) - xk_TNK = causal_conv(key_TC, conv_k_weight_C1W).reshape( + xk_THK = causal_conv(key_TC, conv_k_weight_C1W).reshape( num_tokens, -1, key_head_dim ) - xv_TNV = causal_conv(value_TC, conv_v_weight_C1W).reshape( + xv_THV = causal_conv(value_TC, conv_v_weight_C1W).reshape( num_tokens, -1, value_head_dim ) - g_TN = -torch.exp(A_log_N.float()) * F.softplus(a_TN.float() + dt_bias_N) - beta_TN = torch.sigmoid(b_TN) + g_TH = -torch.exp(A_log_H.float()) * F.softplus(a_TH.float() + dt_bias_H) + beta_TH = torch.sigmoid(b_TH) return self.kernel( - xq_TNK, - xk_TNK, - xv_TNV, - g_TN, - beta_TN, + xq_THK, + xk_THK, + xv_THV, + g_TH, + beta_TH, cu_seqlens=cu_seqlens if cu_seqlens_host is not None else None, cu_seqlens_cpu=cu_seqlens_cpu, ) @@ -514,15 +519,15 @@ def forward( key_TC = self.in_proj_k(x_TD) value_TC = self.in_proj_v(x_TD) gate_TC = self.in_proj_z(x_TD) - a_TN = self.in_proj_a(x_TD) - b_TN = self.in_proj_b(x_TD) + a_TH = self.in_proj_a(x_TD) + b_TH = self.in_proj_b(x_TD) - output_TNV = self.inner_gated_delta_net( + output_THV = self.inner_gated_delta_net( query_TC, key_TC, value_TC, - a_TN, - b_TN, + a_TH, + b_TH, self.conv_q.weight, self.conv_k.weight, self.conv_v.weight, @@ -533,7 +538,7 @@ def forward( value_head_dim=self.value_head_dim, cu_seqlens_host=cu_seqlens_host, ) - gate_TNV = gate_TC.view(num_tokens, -1, self.value_head_dim) - output_TNV = self.norm(output_TNV, gate_TNV) - out_TD = output_TNV.reshape(num_tokens, -1) + gate_THV = gate_TC.view(num_tokens, -1, self.value_head_dim) + output_THV = self.norm(output_THV, gate_THV) + out_TD = output_THV.reshape(num_tokens, -1) return self.out_proj(out_TD) diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 069f3d2e62..7de7c2a84d 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -50,6 +50,12 @@ from .sharding import annotate_deltanet_cu_seqlens, set_qwen35_sharding_config from .vision_encoder import Qwen35VisionEncoder +# Shape suffixes: +# T = packed tokens, D = model dimension, C = projection channels, +# H = attention heads, +# K = query/key head dimension, V = value head dimension, +# R = rotary dimension, P = non-rotary dimension. + Qwen35AttentionMaskDict = dict[str, BlockMask | VarlenMetadata | None] @@ -138,41 +144,41 @@ def forward( num_tokens = x_TD.shape[0] # wq is 2x wider: produces query + gate - xq_gate_TN2H = self.wq(x_TD).view(num_tokens, -1, self.head_dim * 2) - xq_TNH, gate_TNH = xq_gate_TN2H.chunk(2, dim=-1) - xk_TNH = self.wk(x_TD).view(num_tokens, -1, self.head_dim) - xv_TNH = self.wv(x_TD).view(num_tokens, -1, self.head_dim) + xq_gate_THC = self.wq(x_TD).view(num_tokens, -1, self.head_dim * 2) + xq_THK, gate_THV = xq_gate_THC.chunk(2, dim=-1) + xk_THK = self.wk(x_TD).view(num_tokens, -1, self.head_dim) + xv_THV = self.wv(x_TD).view(num_tokens, -1, self.head_dim) # QK norm (before RoPE) - xq_TNH = self.q_norm(xq_TNH) - xk_TNH = self.k_norm(xk_TNH) + xq_THK = self.q_norm(xq_THK) + xk_THK = self.k_norm(xk_THK) # Partial RoPE: only first rotary_dim elements get positional encoding assert self.rotary_dim <= self.head_dim - xq_TNR, xq_TNP = ( - xq_TNH[..., : self.rotary_dim], - xq_TNH[..., self.rotary_dim :], + xq_THR, xq_THP = ( + xq_THK[..., : self.rotary_dim], + xq_THK[..., self.rotary_dim :], ) - xk_TNR, xk_TNP = ( - xk_TNH[..., : self.rotary_dim], - xk_TNH[..., self.rotary_dim :], + xk_THR, xk_THP = ( + xk_THK[..., : self.rotary_dim], + xk_THK[..., self.rotary_dim :], ) - xq_TNR, xk_TNR = self.rope(xq_TNR, xk_TNR, positions) - xq_TNH = torch.cat([xq_TNR, xq_TNP], dim=-1) - xk_TNH = torch.cat([xk_TNR, xk_TNP], dim=-1) - - out_TNH = self.inner_attention( - xq_TNH, - xk_TNH, - xv_TNH, + xq_THR, xk_THR = self.rope(xq_THR, xk_THR, positions) + xq_THK = torch.cat([xq_THR, xq_THP], dim=-1) + xk_THK = torch.cat([xk_THR, xk_THP], dim=-1) + + out_THV = self.inner_attention( + xq_THK, + xk_THK, + xv_THV, attention_masks=attention_masks, scale=self.scaling, enable_gqa=self.enable_gqa, ).contiguous() # Output gating - out_TNH = out_TNH * torch.sigmoid(gate_TNH) - out_TD = out_TNH.view(num_tokens, -1) + out_THV = out_THV * torch.sigmoid(gate_THV) + out_TD = out_THV.view(num_tokens, -1) return self.wo(out_TD) diff --git a/torchtitan/models/qwen3_5/sharding.py b/torchtitan/models/qwen3_5/sharding.py index 8b4c2a1453..bce9c70221 100644 --- a/torchtitan/models/qwen3_5/sharding.py +++ b/torchtitan/models/qwen3_5/sharding.py @@ -330,7 +330,7 @@ def _set_deltanet_sharding( deltanet_cfg.out_proj.sharding_config = rowwise_config(output_sp=enable_sp) # The projections are 2D [T, C], while the norm and recurrence output are - # 3D [T, N, H]. Both shard the feature/head axis on TP. + # 3D [T, H, V]. Both shard the feature/head axis on TP. projected_placement = dense_activation_placement(tp=spmd.S(1), cp=spmd.S(0)) head_placement = attention_activation_placement() parameter_placement = dense_param_placement(tp=spmd.S(0)) @@ -365,26 +365,26 @@ def _set_deltanet_sharding( "query_TC": projected_placement, "key_TC": projected_placement, "value_TC": projected_placement, - "a_TN": projected_placement, - "b_TN": projected_placement, + "a_TH": projected_placement, + "b_TH": projected_placement, "conv_q_weight_C1W": parameter_placement, "conv_k_weight_C1W": parameter_placement, "conv_v_weight_C1W": parameter_placement, - "A_log_N": parameter_placement, - "dt_bias_N": parameter_placement, + "A_log_H": parameter_placement, + "dt_bias_H": parameter_placement, "cu_seqlens": cu_seqlens_placement, }, in_dst_shardings={ "query_TC": projected_placement, "key_TC": projected_placement, "value_TC": projected_placement, - "a_TN": projected_placement, - "b_TN": projected_placement, + "a_TH": projected_placement, + "b_TH": projected_placement, "conv_q_weight_C1W": parameter_placement, "conv_k_weight_C1W": parameter_placement, "conv_v_weight_C1W": parameter_placement, - "A_log_N": parameter_placement, - "dt_bias_N": parameter_placement, + "A_log_H": parameter_placement, + "dt_bias_H": parameter_placement, "cu_seqlens": cu_seqlens_placement, }, out_src_shardings=head_placement, diff --git a/torchtitan/overrides/helion_rope.py b/torchtitan/overrides/helion_rope.py index baa96b30f9..13c0ca5c2f 100644 --- a/torchtitan/overrides/helion_rope.py +++ b/torchtitan/overrides/helion_rope.py @@ -803,13 +803,13 @@ def _complex_eligible( if _HELION_IMPORT_ERROR is None: - def _helion_cossin_rope_fwd_tnh(xq, xk, cache, pos): + def _helion_cossin_rope_fwd_thk(xq, xk, cache, pos): xq_out, xk_out = _helion_cossin_rope_fwd( xq.unsqueeze(0), xk.unsqueeze(0), cache, pos.unsqueeze(0) ) return xq_out.squeeze(0), xk_out.squeeze(0) - def _helion_complex_rope_fwd_tnh(xq, xk, cache, pos): + def _helion_complex_rope_fwd_thk(xq, xk, cache, pos): xq_out, xk_out = _helion_complex_rope_fwd( xq.unsqueeze(0), xk.unsqueeze(0), cache, pos.unsqueeze(0) ) @@ -860,11 +860,11 @@ def _apply_helion_cossin_rope( ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xq_TNH + ), # xq_THK ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xk_TNH + ), # xk_THK {"dp": spmd.R, "cp": spmd.R, "tp": spmd.R}, # rope_cache_MD ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.R}, @@ -875,13 +875,13 @@ def _apply_helion_cossin_rope( ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xq_out_TNH + ), # xq_out_THK ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xk_out_TNH + ), # xk_out_THK ), - )(_helion_cossin_rope_fwd_tnh)(xq, xk, cache, pos) + )(_helion_cossin_rope_fwd_thk)(xq, xk, cache, pos) return _from_local(xq_out, query), _from_local(xk_out, key) def _apply_helion_complex_rope( @@ -920,11 +920,11 @@ def _apply_helion_complex_rope( ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xq_TNH + ), # xq_THK ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xk_TNH + ), # xk_THK {"dp": spmd.R, "cp": spmd.R, "tp": spmd.R}, # rope_cache_real ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.R}, @@ -935,13 +935,13 @@ def _apply_helion_complex_rope( ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xq_out_TNH + ), # xq_out_THK ( {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}, spmd.PartitionSpec(("dp", "cp"), "tp", None), - ), # xk_out_TNH + ), # xk_out_THK ), - )(_helion_complex_rope_fwd_tnh)(xq, xk, cache_real, pos) + )(_helion_complex_rope_fwd_thk)(xq, xk, cache_real, pos) return _from_local(xq_out, query), _from_local(xk_out, key) else: