From d57622af7e1ef2cfb802e661c49d2756212ea7a4 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 5 Oct 2026 05:01:50 +0000 Subject: [PATCH 1/2] [attention] Add mask_type to FlexInnerAttention --- tests/unit_tests/cpu/test_flex_mask_type.py | 91 +++++++++++++++++++++ torchtitan/models/common/__init__.py | 2 + torchtitan/models/common/attention.py | 19 ++++- torchtitan/models/common/decoder.py | 11 +-- torchtitan/models/gpt_oss/model.py | 12 +-- torchtitan/models/muse_glimmer/model.py | 13 +-- 6 files changed, 123 insertions(+), 25 deletions(-) create mode 100644 tests/unit_tests/cpu/test_flex_mask_type.py diff --git a/tests/unit_tests/cpu/test_flex_mask_type.py b/tests/unit_tests/cpu/test_flex_mask_type.py new file mode 100644 index 00000000000..020f31415cf --- /dev/null +++ b/tests/unit_tests/cpu/test_flex_mask_type.py @@ -0,0 +1,91 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from torch.nn.attention.flex_attention import and_masks +from torchtitan.models.common.attention import ( + BaseAttention, + FlexInnerAttention, + get_causal_mask_mod, + get_document_mask_mod, + get_flex_causal_mask_mods, +) +from torchtitan.models.common.decoder import Decoder + +PACKED_POSITIONS = torch.tensor( + [0, 1, 2, 0, 1, 0, 1, 2, 0, 1, 0, 1, 2, 3, 0, 1], dtype=torch.int32 +) +SINGLE_DOC_POSITIONS = torch.arange(16, dtype=torch.int32) + + +def _dense(mask_mod, seq_len: int) -> torch.Tensor: + b, h = torch.tensor(0), torch.tensor(0) + return torch.tensor( + [ + [ + bool(mask_mod(b, h, torch.tensor(q), torch.tensor(kv))) + for kv in range(seq_len) + ] + for q in range(seq_len) + ] + ) + + +def _decoder_block_mask(positions: torch.Tensor, mask_type: str) -> torch.Tensor: + attn_config = BaseAttention.Config( + n_heads=1, + inner_attention=FlexInnerAttention.Config(block_size=4, mask_type=mask_type), + ) + decoder = Decoder.__new__(Decoder) + mask = decoder._create_flex_attention_mask_for_document(positions, attn_config) + return mask.to_dense() + + +class TestFlexMaskType(unittest.TestCase): + def test_default_is_block_causal(self): + self.assertEqual(FlexInnerAttention.Config().mask_type, "block_causal") + + def test_mask_mods_match_mask_type(self): + seq_len = PACKED_POSITIONS.shape[0] + causal = _dense(get_causal_mask_mod(), seq_len) + same_doc = _dense(get_document_mask_mod(PACKED_POSITIONS), seq_len) + + block_causal_mods = get_flex_causal_mask_mods( + PACKED_POSITIONS, FlexInnerAttention.Config(mask_type="block_causal") + ) + causal_mods = get_flex_causal_mask_mods( + PACKED_POSITIONS, FlexInnerAttention.Config(mask_type="causal") + ) + + self.assertEqual(len(causal_mods), 1) + torch.testing.assert_close( + _dense(and_masks(*block_causal_mods), seq_len), causal & same_doc + ) + torch.testing.assert_close(_dense(and_masks(*causal_mods), seq_len), causal) + + def test_decoder_mask_follows_mask_type(self): + block_causal = _decoder_block_mask(PACKED_POSITIONS, "block_causal") + causal = _decoder_block_mask(PACKED_POSITIONS, "causal") + num_blocks = causal.shape[-1] + + torch.testing.assert_close( + causal[0, 0].bool(), + torch.ones(num_blocks, num_blocks, dtype=torch.bool).tril(), + ) + # Packed documents mask out blocks that lie entirely in earlier documents. + self.assertLess(block_causal.sum().item(), causal.sum().item()) + + def test_mask_types_agree_for_single_document(self): + torch.testing.assert_close( + _decoder_block_mask(SINGLE_DOC_POSITIONS, "block_causal"), + _decoder_block_mask(SINGLE_DOC_POSITIONS, "causal"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/models/common/__init__.py b/torchtitan/models/common/__init__.py index 88a1f3229f7..f82e2b3fd8f 100644 --- a/torchtitan/models/common/__init__.py +++ b/torchtitan/models/common/__init__.py @@ -22,6 +22,7 @@ get_document_mask_mod, get_efficient_causal_mask_mod_for_packed_document, get_fixed_block_mask_mod, + get_flex_causal_mask_mods, get_sliding_window_mask_mod, GQAttention, InnerAttention, @@ -68,6 +69,7 @@ "get_document_mask_mod", "get_efficient_causal_mask_mod_for_packed_document", "get_fixed_block_mask_mod", + "get_flex_causal_mask_mods", "get_sliding_window_mask_mod", "GQAttention", "GroupNorm", diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py index 9ccb20e0fdc..a088a10d861 100644 --- a/torchtitan/models/common/attention.py +++ b/torchtitan/models/common/attention.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from typing import Any, ClassVar, NamedTuple +from typing import Any, ClassVar, Literal, NamedTuple import spmd_types as spmd import torch @@ -63,6 +63,7 @@ "get_document_mask_mod", "get_efficient_causal_mask_mod_for_packed_document", "get_fixed_block_mask_mod", + "get_flex_causal_mask_mods", "get_sliding_window_mask_mod", "local_head_split", ] @@ -258,6 +259,11 @@ class FlexInnerAttention(InnerAttention): class Config(InnerAttention.Config): block_size: int | tuple[int, int] = _DEFAULT_SPARSE_BLOCK_SIZE kernel_options: dict = field(default_factory=dict) + mask_type: Literal["block_causal", "causal"] = "block_causal" + """Causal mask for the flex path. ``"block_causal"``: causal within each + packed document (positions reset to 0 at document starts). ``"causal"``: + plain causal over the whole sequence, which does not depend on + ``positions`` and is only correct when each sequence holds one document.""" inductor_configs: ClassVar[dict[str, bool]] = { "wrap_inductor_compiled_regions": True, @@ -499,6 +505,17 @@ def document_mask( return document_mask +def get_flex_causal_mask_mods( + positions: torch.Tensor, + config: FlexInnerAttention.Config, +) -> list[_mask_mod_signature]: + """Return the causal mask_mods selected by ``config.mask_type``.""" + mask_mods = [get_causal_mask_mod()] + if config.mask_type == "block_causal": + mask_mods.append(get_efficient_causal_mask_mod_for_packed_document(positions)) + return mask_mods + + def get_efficient_causal_mask_mod_for_packed_document( positions: torch.Tensor, ) -> _mask_mod_signature: diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index bb95ea8b1a7..2c4f5de2bd6 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -23,8 +23,7 @@ create_attention_mask, create_varlen_metadata_for_document, FlexInnerAttention, - get_causal_mask_mod, - get_efficient_causal_mask_mod_for_packed_document, + get_flex_causal_mask_mods, VarlenInnerAttention, ) from torchtitan.models.common.decoder_sharding import decoder_input_sharding @@ -292,14 +291,12 @@ def _create_flex_attention_mask_for_document( positions: torch.Tensor, attn_config: BaseAttention.Config, ) -> BlockMask: - """Build the standard causal + packed-document flex-attention mask.""" + """Build the causal flex-attention mask selected by ``mask_type``.""" + assert isinstance(attn_config.inner_attention, FlexInnerAttention.Config) return self._create_flex_attention_mask( positions, attn_config, - [ - get_causal_mask_mod(), - get_efficient_causal_mask_mod_for_packed_document(positions), - ], + get_flex_causal_mask_mods(positions, attn_config.inner_attention), ) def preprocess_inputs( diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py index 0c2c112ec6b..6d182f2ae47 100644 --- a/torchtitan/models/gpt_oss/model.py +++ b/torchtitan/models/gpt_oss/model.py @@ -23,8 +23,7 @@ BaseAttention, create_varlen_metadata_for_document, FlexInnerAttention, - get_causal_mask_mod, - get_efficient_causal_mask_mod_for_packed_document, + get_flex_causal_mask_mods, get_sliding_window_mask_mod, QKVLinear, VarlenInnerAttention, @@ -292,12 +291,9 @@ def get_attention_masks( max_context_length=max_context_length, ) elif isinstance(inner_attn, FlexInnerAttention.Config): - base_mask_mods = [ - get_causal_mask_mod(), - get_efficient_causal_mask_mod_for_packed_document(positions), - ] - # Full-attention (causal + document) mask, used by layers without a - # sliding window. + base_mask_mods = get_flex_causal_mask_mods(positions, inner_attn) + # Full-attention (causal) mask, used by layers without a sliding + # window. masks: dict[str, BlockMask] = { "basic_mask": self._create_flex_attention_mask( positions, attn_cfg, base_mask_mods diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py index 361f49f76d7..de224385f04 100644 --- a/torchtitan/models/muse_glimmer/model.py +++ b/torchtitan/models/muse_glimmer/model.py @@ -31,8 +31,7 @@ create_attention_mask, create_varlen_metadata_for_document, FlexInnerAttention, - get_causal_mask_mod, - get_efficient_causal_mask_mod_for_packed_document, + get_flex_causal_mask_mods, get_sliding_window_mask_mod, GQAttention, VarlenInnerAttention, @@ -627,14 +626,10 @@ def get_attention_masks( f"sliding-window masks, got {type(inner_attn).__name__}" ) - # Language models always use block-causal (per-document) masking: the - # dataloaders emit per-document positions, and the efficient packed-doc - # mask ANDed with the causal mask yields same-document causal attention. + # ``mask_type="block_causal"`` (default) yields same-document causal + # attention from the per-document positions the dataloaders emit. seq_len = positions.shape[0] - base_mods = [ - get_causal_mask_mod(), - get_efficient_causal_mask_mod_for_packed_document(positions), - ] + base_mods = get_flex_causal_mask_mods(positions, inner_attn) # Match the base Decoder mask-building so the configured block size and # batch-invariance handling are honored. From 50eaf45c7d14ece191632a8e24c2984c72f5ff92 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 5 Oct 2026 05:10:38 +0000 Subject: [PATCH 2/2] [attention] Test mask_type in GPT-OSS and Muse Glimmer mask builders --- tests/unit_tests/cpu/test_flex_mask_type.py | 108 ++++++++++++-------- 1 file changed, 68 insertions(+), 40 deletions(-) diff --git a/tests/unit_tests/cpu/test_flex_mask_type.py b/tests/unit_tests/cpu/test_flex_mask_type.py index 020f31415cf..2733ffc535a 100644 --- a/tests/unit_tests/cpu/test_flex_mask_type.py +++ b/tests/unit_tests/cpu/test_flex_mask_type.py @@ -16,11 +16,18 @@ get_flex_causal_mask_mods, ) from torchtitan.models.common.decoder import Decoder +from torchtitan.models.gpt_oss import build_model_config as build_gpt_oss_config +from torchtitan.models.gpt_oss.model import GptOssModel +from torchtitan.models.muse_glimmer import ( + build_model_config as build_muse_glimmer_config, +) +from torchtitan.models.muse_glimmer.model import MuseGlimmerModel PACKED_POSITIONS = torch.tensor( [0, 1, 2, 0, 1, 0, 1, 2, 0, 1, 0, 1, 2, 3, 0, 1], dtype=torch.int32 ) SINGLE_DOC_POSITIONS = torch.arange(16, dtype=torch.int32) +BLOCK_SIZE = 4 def _dense(mask_mod, seq_len: int) -> torch.Tensor: @@ -36,55 +43,76 @@ def _dense(mask_mod, seq_len: int) -> torch.Tensor: ) -def _decoder_block_mask(positions: torch.Tensor, mask_type: str) -> torch.Tensor: +def _decoder_masks(flex_kwargs: dict) -> dict: attn_config = BaseAttention.Config( n_heads=1, - inner_attention=FlexInnerAttention.Config(block_size=4, mask_type=mask_type), + inner_attention=FlexInnerAttention.Config(block_size=BLOCK_SIZE, **flex_kwargs), ) decoder = Decoder.__new__(Decoder) - mask = decoder._create_flex_attention_mask_for_document(positions, attn_config) - return mask.to_dense() - - -class TestFlexMaskType(unittest.TestCase): - def test_default_is_block_causal(self): - self.assertEqual(FlexInnerAttention.Config().mask_type, "block_causal") - - def test_mask_mods_match_mask_type(self): - seq_len = PACKED_POSITIONS.shape[0] - causal = _dense(get_causal_mask_mod(), seq_len) - same_doc = _dense(get_document_mask_mod(PACKED_POSITIONS), seq_len) - - block_causal_mods = get_flex_causal_mask_mods( - PACKED_POSITIONS, FlexInnerAttention.Config(mask_type="block_causal") - ) - causal_mods = get_flex_causal_mask_mods( - PACKED_POSITIONS, FlexInnerAttention.Config(mask_type="causal") + return { + "mask": decoder._create_flex_attention_mask_for_document( + PACKED_POSITIONS, attn_config ) + } - self.assertEqual(len(causal_mods), 1) - torch.testing.assert_close( - _dense(and_masks(*block_causal_mods), seq_len), causal & same_doc - ) - torch.testing.assert_close(_dense(and_masks(*causal_mods), seq_len), causal) - def test_decoder_mask_follows_mask_type(self): - block_causal = _decoder_block_mask(PACKED_POSITIONS, "block_causal") - causal = _decoder_block_mask(PACKED_POSITIONS, "causal") - num_blocks = causal.shape[-1] +def _model_masks(build_config, model_cls): + def masks(flex_kwargs: dict) -> dict: + config = build_config("debugmodel", seq_len=16, attn_backend="flex") + for layer in config.layers: + inner = layer.attention.inner_attention + assert isinstance(inner, FlexInnerAttention.Config) + inner.block_size = BLOCK_SIZE + for key, value in flex_kwargs.items(): + setattr(inner, key, value) + model = model_cls.__new__(model_cls) + model.config = config + return model.get_attention_masks(PACKED_POSITIONS) - torch.testing.assert_close( - causal[0, 0].bool(), - torch.ones(num_blocks, num_blocks, dtype=torch.bool).tril(), - ) - # Packed documents mask out blocks that lie entirely in earlier documents. - self.assertLess(block_causal.sum().item(), causal.sum().item()) + return masks - def test_mask_types_agree_for_single_document(self): - torch.testing.assert_close( - _decoder_block_mask(SINGLE_DOC_POSITIONS, "block_causal"), - _decoder_block_mask(SINGLE_DOC_POSITIONS, "causal"), - ) + +class TestFlexMaskType(unittest.TestCase): + def test_mask_mods_match_mask_type(self): + for positions in (PACKED_POSITIONS, SINGLE_DOC_POSITIONS): + seq_len = positions.shape[0] + causal = _dense(get_causal_mask_mod(), seq_len) + same_doc = _dense(get_document_mask_mod(positions), seq_len) + for mask_type, expected in ( + ("block_causal", causal & same_doc), + ("causal", causal), + ): + with self.subTest( + single_doc=positions is SINGLE_DOC_POSITIONS, mask_type=mask_type + ): + mask_mods = get_flex_causal_mask_mods( + positions, FlexInnerAttention.Config(mask_type=mask_type) + ) + torch.testing.assert_close( + _dense(and_masks(*mask_mods), seq_len), expected + ) + + def test_mask_builders_follow_mask_type(self): + num_blocks = PACKED_POSITIONS.shape[0] // BLOCK_SIZE + causal_blocks = torch.ones(num_blocks, num_blocks, dtype=torch.bool).tril() + builders = { + "decoder": _decoder_masks, + "gpt_oss": _model_masks(build_gpt_oss_config, GptOssModel), + "muse_glimmer": _model_masks(build_muse_glimmer_config, MuseGlimmerModel), + } + for name, build_masks in builders.items(): + default_masks = build_masks({}) + causal_masks = build_masks({"mask_type": "causal"}) + for key, causal_mask in causal_masks.items(): + with self.subTest(builder=name, mask=key): + torch.testing.assert_close( + causal_mask.to_dense()[0, 0].bool(), causal_blocks + ) + # Default (block_causal) drops blocks in earlier documents. + self.assertLess( + default_masks[key].to_dense().sum().item(), + causal_mask.to_dense().sum().item(), + ) if __name__ == "__main__":