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..2733ffc535a --- /dev/null +++ b/tests/unit_tests/cpu/test_flex_mask_type.py @@ -0,0 +1,119 @@ +# 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 +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: + 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_masks(flex_kwargs: dict) -> dict: + attn_config = BaseAttention.Config( + n_heads=1, + inner_attention=FlexInnerAttention.Config(block_size=BLOCK_SIZE, **flex_kwargs), + ) + decoder = Decoder.__new__(Decoder) + return { + "mask": decoder._create_flex_attention_mask_for_document( + PACKED_POSITIONS, attn_config + ) + } + + +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) + + return masks + + +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__": + 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.