diff --git a/tests/pytorch/test_fused_attn.py b/tests/pytorch/test_fused_attn.py index 1b08eeeb10..596ba27c85 100644 --- a/tests/pytorch/test_fused_attn.py +++ b/tests/pytorch/test_fused_attn.py @@ -152,10 +152,16 @@ def _is_flash_attention_2_available() -> bool: @functools.cache def _is_flash_attention_2_1() -> bool: - """Check if flash-attn 2.0+ is available""" + """Check if flash-attn 2.1+ is available""" Version = packaging.version.Version return Version(version("flash-attn")) >= Version("2.1") +@functools.cache +def _is_flash_attention_2_3() -> bool: + """Check if flash-attn 2.3+ is available""" + Version = packaging.version.Version + return Version(version("flash-attn")) >= Version("2.3") + def _is_flash_attention_supported(config: ModelConfig) -> bool: """Check if FlashAttention supports a model configuration""" if get_device_compute_capability() < (8, 0): @@ -192,6 +198,17 @@ def _is_unfused_attention_supported(config: ModelConfig) -> bool: param_types.append(torch.bfloat16) param_types_lean = [torch.bfloat16] +def get_swa(seq_q, seq_kv, w=None): + """Generate a random sliding window size (left, right) if w is None, + and create its equivalent attention mask in [seq_q, seq_kv] shape""" + if w is None: + w = torch.randint(0, seq_kv, [2], dtype=torch.int32, device="cuda") + m = torch.ones(seq_q, seq_kv, dtype=torch.bool, device="cuda") + mu = torch.triu(m, diagonal=seq_kv-seq_q-w[0]) + ml = torch.tril(mu, diagonal=seq_kv-seq_q+w[1]) + ml = ~ ml + return w, ml + @pytest.mark.skipif(_cudnn_version() < (8,9,1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model_configs", [model_configs_base]) @@ -199,7 +216,8 @@ def _is_unfused_attention_supported(config: ModelConfig) -> bool: @pytest.mark.parametrize("ckpt_attn", [False]) @pytest.mark.parametrize("workspace_opt", [True, False]) @pytest.mark.parametrize("qkv_layout", [None]) -def test_dot_product_attention(dtype, model_configs, model, ckpt_attn, workspace_opt, qkv_layout): +@pytest.mark.parametrize("swa", [False]) +def test_dot_product_attention(dtype, model_configs, model, ckpt_attn, workspace_opt, qkv_layout, swa): """Test DotProductAttention module""" # Get configs @@ -224,36 +242,43 @@ def test_dot_product_attention(dtype, model_configs, model, ckpt_attn, workspace fused_attn_supported, fused_attn_backend = _is_fused_attention_supported( config, dtype, qkv_layout=qkv_layout, ) + if swa: + fused_attn_supported = False flash_attn_supported = _is_flash_attention_supported(config) if (len(fused_attn_backend) + flash_attn_supported + unfused_attn_supported) < 2: pytest.skip("Less than two backends to compare.") # UnfusedDotProductAttention backend if unfused_attn_supported: + if swa: + attn_mask_type = config.attn_mask_type + config.attn_mask_type = "arbitrary" unfused_attn_fwd, unfused_attn_bwd = _run_dot_product_attention( - dtype, config, "UnfusedDotProductAttention", ckpt_attn, qkv_layout, workspace_opt, + dtype, config, "UnfusedDotProductAttention", ckpt_attn, qkv_layout, workspace_opt, swa, ) + if swa: + config.attn_mask_type = attn_mask_type # FusedAttention backend if fused_attn_supported: if len(fused_attn_backend) == 1: fused_attn_fwd, fused_attn_bwd = _run_dot_product_attention( - dtype, config, "FusedAttention", ckpt_attn, qkv_layout, workspace_opt, + dtype, config, "FusedAttention", ckpt_attn, qkv_layout, workspace_opt, swa, ) if len(fused_attn_backend) == 2: os.environ["NVTE_FUSED_ATTN_BACKEND"] = "0" fused_attn_fwd, fused_attn_bwd = _run_dot_product_attention( - dtype, config, "FusedAttention", ckpt_attn, qkv_layout, workspace_opt, + dtype, config, "FusedAttention", ckpt_attn, qkv_layout, workspace_opt, swa, ) os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" fused_attn_fwd_1, fused_attn_bwd_1 = _run_dot_product_attention( - dtype, config, "FusedAttention", ckpt_attn, qkv_layout, workspace_opt, + dtype, config, "FusedAttention", ckpt_attn, qkv_layout, workspace_opt, swa, ) # FlashAttention backend if flash_attn_supported: flash_attn_fwd, flash_attn_bwd = _run_dot_product_attention( - dtype, config, "FlashAttention", ckpt_attn, qkv_layout, workspace_opt, + dtype, config, "FlashAttention", ckpt_attn, qkv_layout, workspace_opt, swa, ) if unfused_attn_supported and fused_attn_supported: @@ -279,7 +304,7 @@ def test_dot_product_attention(dtype, model_configs, model, ckpt_attn, workspace @pytest.mark.parametrize("model", ["base_1_1", "base_2_1"]) def test_dpa_checkpoint(dtype, model_configs, model): """Test DotProductAttention module with checkpointing""" - test_dot_product_attention(dtype, model_configs, model, True, True, None) + test_dot_product_attention(dtype, model_configs, model, True, True, None, False) model_configs_mask = { # test: b, h, hg, d, sq, skv, p, mask, bias @@ -303,7 +328,7 @@ def test_dpa_checkpoint(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_mask.keys()) def test_dpa_mask(dtype, model_configs, model): """Test DotProductAttention module with different mask types""" - test_dot_product_attention(dtype, model_configs, model, False, True, None) + test_dot_product_attention(dtype, model_configs, model, False, True, None, False) model_configs_bias = { # test: b, h, hg, d, sq, skv, p, mask, bias @@ -339,7 +364,22 @@ def test_dpa_mask(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_bias.keys()) def test_dpa_bias(dtype, model_configs, model): """Test DotProductAttention module with different bias types""" - test_dot_product_attention(dtype, model_configs, model, False, True, None) + test_dot_product_attention(dtype, model_configs, model, False, True, None, False) + +model_configs_swa = { + # test: b, h, hg, d, sq, skv, p, mask, bias + "swa_1_0": ModelConfig(4, 16, 16, 64, 128, 128, 0.0, "no_mask", "no_bias"), + "swa_1_1": ModelConfig(2, 16, 16, 64, 128, 256, 0.0, "no_mask", "no_bias"), + "swa_1_2": ModelConfig(4, 24, 24, 128, 2048, 2048, 0.0, "no_mask", "no_bias"), + "swa_1_3": ModelConfig(2, 24, 24, 128, 2048, 4096, 0.0, "no_mask", "no_bias"), +} +@pytest.mark.skipif(not _is_flash_attention_2_3(), reason="Flash-attn 2.3+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_swa]) +@pytest.mark.parametrize("model", model_configs_swa.keys()) +def test_dpa_sliding_window(dtype, model_configs, model): + """Test DotProductAttention module with sliding window attention""" + test_dot_product_attention(dtype, model_configs, model, False, True, None, True) qkv_layouts = [ 'sb3hd', 'sbh3d', 'sbhd_sb2hd', 'sbhd_sbh2d', 'sbhd_sbhd_sbhd', @@ -367,7 +407,7 @@ def test_dpa_bias(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", qkv_layouts) def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with different QKV layouts""" - test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout) + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False) def _run_dot_product_attention( dtype: torch.dtype, @@ -376,6 +416,7 @@ def _run_dot_product_attention( ckpt_attn: bool, qkv_layout: str, workspace_opt: bool, + swa: bool, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Run DotProductAttention module with one forward pass and one backward pass""" @@ -433,6 +474,10 @@ def _run_dot_product_attention( .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0) attention_mask = ( attention_mask_q.to(device="cuda"), attention_mask_kv.to(device="cuda")) + if swa: + window_size, attention_mask = get_swa(config.max_seqlen_q, config.max_seqlen_kv) + else: + window_size, attention_mask = None, None # Create input tensors dim_to_num = { @@ -515,6 +560,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: # Run a forward and backward pass out = block(inp[0], inp[1], inp[2], + window_size=window_size, attention_mask=attention_mask, qkv_format=qkv_format, cu_seqlens_q=cu_seqlens_q, diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py index f5b01cabc1..24a90af39c 100644 --- a/transformer_engine/pytorch/attention.py +++ b/transformer_engine/pytorch/attention.py @@ -57,6 +57,7 @@ _flash_attn_version_required = packaging.version.Version("1.0.6") _flash_attn_2_available = _flash_attn_version >= packaging.version.Version("2") _flash_attn_2_1_plus = _flash_attn_version >= packaging.version.Version("2.1") +_flash_attn_2_3_plus = _flash_attn_version >= packaging.version.Version("2.3") if _flash_attn_2_available: from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_forward_func # pylint: disable=no-name-in-module @@ -1248,6 +1249,24 @@ def run_iteratively(q, k, v): return qkv_layout, q, k, v +def check_set_window_size( + attn_mask_type: str, + window_size: Tuple[int, int] = None, + ): + """Check if sliding window size is compliant with mask type and if not, + assert or set it to the appropriate size + """ + if "causal" in attn_mask_type: + if window_size is None: + window_size = (-1, 0) + else: + assert ( + window_size[1] == 0 + ), "window_size[1] should be 0 when self_attn_mask_type includes 'causal'!" + else: + if window_size is None: + window_size = (-1, -1) + return window_size class FlashAttention(torch.nn.Module): """Dot product attention, using HazyResearch flash-attn package: @@ -1286,12 +1305,15 @@ def forward( cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_kv: Optional[torch.Tensor] = None, attn_mask_type: str = "causal", + window_size: Optional[Tuple[int, int]] = None, cp_group: Optional[dist_group_type] = None, cp_global_ranks: List[int] = None, cp_stream: torch.cuda.Stream = None, ) -> torch.Tensor: """flash-attn fprop""" + window_size = check_set_window_size(attn_mask_type, window_size) + assert ( query_layer.dtype in [torch.float16, torch.bfloat16] and key_layer.dtype in [torch.float16, torch.bfloat16] @@ -1402,6 +1424,9 @@ def forward( max_seqlen_kv = seqlens_kv.max().item() if context_parallel: + assert ( + window_size in ((-1, -1), (-1, 0)) + ), "Sliding window attention is not supported with context parallelism." with self.attention_dropout_ctx(): output = flash_attn_forward_func_with_cp( query_layer, key_layer, value_layer, @@ -1417,6 +1442,8 @@ def forward( fa_optional_forward_kwargs = {} if not _flash_attn_2_available: fa_optional_forward_kwargs["deterministic"] = self.deterministic + if _flash_attn_2_3_plus: + fa_optional_forward_kwargs["window_size"] = window_size output = flash_attn_forward_func( query_layer, key_layer, value_layer, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, @@ -1875,6 +1902,12 @@ class DotProductAttention(torch.nn.Module): :attr:`cu_seqlens_kv` in the shape of [batch_size + 1] or :attr:`attention_mask` in the shape [batch_size, 1, 1, max_seq_len]. For the "`arbitrary`" mask, users need to provide a mask that is broadcastable to the shape of softmax input. + window_size: Optional[Tuple[int, int]], default = `None` + sliding window size for local attention, where query at position i attends to keys + in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding + window and causal mask specifically. Similar to :attr:`attn_mask_type`, it can + be overridden by :attr:`window_size` in `forward` as well. attention_type: str, default = `self` type of attention, either "`self`" and "`cross`". layer_number: int, default = `None` @@ -1918,6 +1951,7 @@ def __init__( attention_dropout: float = 0.0, qkv_format: str = "sbhd", attn_mask_type: str = "causal", + window_size: Optional[Tuple[int, int]] = None, sequence_parallel: bool = False, tp_size: int = 1, get_rng_state_tracker: Optional[Callable] = None, @@ -1935,6 +1969,8 @@ def __init__( if attn_mask_type == "causal_padding": attn_mask_type = "padding_causal" self.attn_mask_type = attn_mask_type + self.window_size = window_size + self.window_size = check_set_window_size(attn_mask_type, self.window_size) self.tp_size = tp_size if tp_group is None else get_distributed_world_size(tp_group) self.tp_group = tp_group self.get_rng_state_tracker = get_rng_state_tracker @@ -1969,8 +2005,8 @@ def __init__( if _flash_attn_2_available and self.deterministic: self.use_flash_attention = False warnings.warn( - "Disabling usage of FlashAttention since version 2 does not support deterministic" - "execution. In order to use FA with deterministic behavior, please install" + "Disabling usage of FlashAttention since version 2 does not support deterministic " + "execution. In order to use FA with deterministic behavior, please install " "FlashAttention version 1." ) @@ -2065,6 +2101,7 @@ def forward( cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_kv: Optional[torch.Tensor] = None, attn_mask_type: Optional[str] = None, + window_size: Optional[Tuple[int, int]] = None, checkpoint_core_attention: bool = False, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, @@ -2138,6 +2175,8 @@ def forward( attn_mask_type: {`no_mask`, `padding`, `causal`, `padding,causal`, `causal,padding`, `arbitrary`}, default = `None`. Type of attention mask passed into softmax operation. 'padding,causal' and 'causal,padding' are equivalent. + window_size: Optional[Tuple[int, int]], default = `None` + sliding window size for local attention. checkpoint_core_attention : bool, default = `False` If true, forward activations for attention are recomputed during the backward pass in order to save memory that would @@ -2159,6 +2198,8 @@ def forward( assert (key_layer.shape == value_layer.shape ), "Keys and values must have the same shape!" + if attn_mask_type is not None: + window_size = check_set_window_size(attn_mask_type, window_size) if attn_mask_type is None: attn_mask_type = self.attn_mask_type else: @@ -2169,6 +2210,9 @@ def forward( assert (attn_mask_type in AttnMaskTypes ), f"Attention mask type {attn_mask_type} is not supported!" + if window_size is None: + window_size = self.window_size + if qkv_format is None: qkv_format = self.qkv_format @@ -2220,6 +2264,7 @@ def forward( # is: FlashAttention > FusedAttention (cuDNN) > UnfusedDotProductAttention. use_flash_attention = self.use_flash_attention use_fused_attention = self.use_fused_attention + use_unfused_attention = True # The following section filters out some backends based on # certain asserts before executing the forward pass. @@ -2249,9 +2294,11 @@ def forward( and self.device_compute_capability not in ((8, 0), (9, 0)))): use_flash_attention = False + # Filter: MQA/GQA. if not _flash_attn_2_available and self.num_gqa_groups != self.num_attention_heads: use_flash_attention = False + # Filter: cross attention + causal mask. if (_flash_attn_2_1_plus and "causal" in attn_mask_type and max_seqlen_q != max_seqlen_kv): @@ -2262,16 +2309,26 @@ def forward( ) use_flash_attention = False + # Filter: bias. if core_attention_bias_type != "no_bias" or core_attention_bias is not None: use_flash_attention = False + # Filter: sliding window attention. + # UnfusedDotProductAttention can support SWA via arbitrary attention mask. + if window_size not in ((-1, -1), (-1, 0)): + use_fused_attention = False + context_parallel = (self.cp_group is not None + and get_distributed_world_size(self.cp_group) != 1) + if (not _flash_attn_2_3_plus) or context_parallel: + use_flash_attention = False + # Filter: ONNX export. if is_in_onnx_export_mode(): use_flash_attention = False use_fused_attention = False # Filter: Attention mask type. - # attn_mask_type(s) | supported backends + # attn_mask_type(s) | supported backends # ------------------------------------------------ # no_mask | All # padding | UnfusedDotProductAttention, FlashAttention, FusedAttention @@ -2282,6 +2339,8 @@ def forward( if attn_mask_type == "arbitrary": use_flash_attention = False use_fused_attention = False + if "causal" in attn_mask_type and max_seqlen_q != max_seqlen_kv: + use_unfused_attention = False if use_fused_attention: fused_attention_backend = tex.get_fused_attn_backend( @@ -2303,6 +2362,24 @@ def forward( use_fused_attention = (use_fused_attention and is_backend_avail) + # Filter: determinism. + # backend | deterministic + # --------------------------------------------------------- + # flash-attn v1 | yes + # flash-attn v2 | no + # FusedAttnBackend["F16_max512_seqlen"] | yes + # FusedAttnBackend["F16_arbitrary_seqlen"] | workspace optimization path: yes; otherwise: no + # UnfusedDotProductAttention | yes + # + # Note that FusedAttnBackend["F16_arbitrary_seqlen"] only has workspace optimization path + # on sm90 architectures. + # + if (use_fused_attention + and fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] + and self.deterministic + and self.device_compute_capability != (9, 0)): + use_fused_attention = False + # Select FusedAttention on sm90 and FlashAttention on others for performance if (use_flash_attention and use_fused_attention @@ -2321,6 +2398,7 @@ def forward( cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, attn_mask_type=attn_mask_type, + window_size=window_size, cp_group=self.cp_group, cp_global_ranks=self.cp_global_ranks, cp_stream=self.cp_stream) @@ -2360,29 +2438,32 @@ def forward( if _NVTE_DEBUG: print("[DotProductAttention]: using unfused DPA") - if checkpoint_core_attention: - return self._checkpointed_attention_forward( - self.unfused_attention, - query_layer, - key_layer, - value_layer, - qkv_layout = qkv_layout, - cu_seqlens_q = cu_seqlens_q, - cu_seqlens_kv = cu_seqlens_kv, - attn_mask_type = attn_mask_type, - attention_mask = attention_mask, - core_attention_bias_type = core_attention_bias_type, - core_attention_bias = core_attention_bias) - return self.unfused_attention(query_layer, - key_layer, - value_layer, - qkv_layout = qkv_layout, - cu_seqlens_q = cu_seqlens_q, - cu_seqlens_kv = cu_seqlens_kv, - attn_mask_type = attn_mask_type, - attention_mask = attention_mask, - core_attention_bias_type = core_attention_bias_type, - core_attention_bias = core_attention_bias) + if use_unfused_attention: + if checkpoint_core_attention: + return self._checkpointed_attention_forward( + self.unfused_attention, + query_layer, + key_layer, + value_layer, + qkv_layout = qkv_layout, + cu_seqlens_q = cu_seqlens_q, + cu_seqlens_kv = cu_seqlens_kv, + attn_mask_type = attn_mask_type, + attention_mask = attention_mask, + core_attention_bias_type = core_attention_bias_type, + core_attention_bias = core_attention_bias) + return self.unfused_attention(query_layer, + key_layer, + value_layer, + qkv_layout = qkv_layout, + cu_seqlens_q = cu_seqlens_q, + cu_seqlens_kv = cu_seqlens_kv, + attn_mask_type = attn_mask_type, + attention_mask = attention_mask, + core_attention_bias_type = core_attention_bias_type, + core_attention_bias = core_attention_bias) + + raise Exception("No dot product attention support for the provided inputs!") class MultiheadAttention(torch.nn.Module): @@ -2427,6 +2508,12 @@ class MultiheadAttention(torch.nn.Module): arg is useful for dynamically changing mask types, e.g. a different mask for training and inference. The init arg is useful for cases involving compilation/tracing, e.g. ONNX export. + window_size: Optional[Tuple[int, int]], default = `None` + sliding window size for local attention, where query at position i attends to keys + in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding + window and causal mask specifically. Similar to :attr:`attn_mask_type`, it can + be overridden by :attr:`window_size` in `forward` as well. num_gqa_groups : int, default = `None` number of GQA groups in the transformer layer. Grouped Query Attention is described in @@ -2518,6 +2605,7 @@ def __init__( output_layer_init_method: Optional[Callable] = None, layer_number: Optional[int] = None, attn_mask_type: str = "causal", + window_size: Optional[Tuple[int, int]] = None, tp_group: Optional[dist_group_type] = None, tp_size: int = 1, num_gqa_groups: Optional[int] = None, @@ -2546,6 +2634,8 @@ def __init__( super().__init__() self.attn_mask_type = attn_mask_type + self.window_size = window_size + self.window_size = check_set_window_size(attn_mask_type, self.window_size) self.layer_number = layer_number self.input_layernorm = input_layernorm self.attention_type = attention_type @@ -2759,6 +2849,7 @@ def forward( attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, encoder_output: Optional[torch.Tensor] = None, attn_mask_type: Optional[str] = None, + window_size: Optional[Tuple[int, int]] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[InferenceParams] = None, @@ -2789,6 +2880,8 @@ def forward( attn_mask_type: {'no_mask', 'padding', 'causal', 'padding_causal', 'arbitrary'}, default = `None` type of attention mask passed into softmax operation. + window_size: Optional[Tuple[int, int]], default = `None` + sliding window size for local attention. encoder_output : Optional[torch.Tensor], default = `None` Output of the encoder block to be fed into the decoder block if using `layer_type="decoder"`. @@ -2823,8 +2916,12 @@ def forward( """ # hidden_states: [sq, b, h] + if attn_mask_type is not None: + window_size = check_set_window_size(attn_mask_type, window_size) if attn_mask_type is None: attn_mask_type = self.attn_mask_type + if window_size is None: + window_size = self.window_size if "padding" in attn_mask_type and attention_mask is not None: for i,_ in enumerate(attention_mask): @@ -3037,6 +3134,7 @@ def forward( cu_seqlens_kv=None, attention_mask=attention_mask, attn_mask_type=attn_mask_type, + window_size=window_size, checkpoint_core_attention=checkpoint_core_attention, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 0ee28da34c..2647e78a7b 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -12,7 +12,11 @@ import transformer_engine_extensions as tex from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm -from transformer_engine.pytorch.attention import InferenceParams, MultiheadAttention +from transformer_engine.pytorch.attention import ( + InferenceParams, + MultiheadAttention, + check_set_window_size, +) from transformer_engine.pytorch.jit import ( set_jit_fusion_options, warmup_jit_bias_dropout_add_all_dtypes, @@ -134,6 +138,12 @@ class TransformerLayer(torch.nn.Module): arg is useful for dynamically changing mask types, e.g. a different mask for training and inference. The init arg is useful for cases involving compilation/tracing, e.g. ONNX export. + window_size: Optional[Tuple[int, int]], default = `None` + sliding window size for local attention, where query at position i attends to keys + in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding + window and causal mask specifically. Similar to :attr:`self_attn_mask_type`, it can + be overridden by :attr:`window_size` in `forward` as well. zero_centered_gamma : bool, default = 'False' if set to 'True', gamma parameter in LayerNorm is initialized to 0 and the LayerNorm formula changes to @@ -220,6 +230,7 @@ def __init__( layer_number: Optional[int] = None, kv_channels: Optional[int] = None, self_attn_mask_type: str = "causal", + window_size: Optional[Tuple[int, int]] = None, tp_group: Optional[dist_group_type] = None, tp_size: int = 1, params_dtype: Optional[torch.dtype] = None, @@ -251,6 +262,8 @@ def __init__( ), "Userbuffer communication backend not available." self.self_attn_mask_type = self_attn_mask_type + self.window_size = window_size + self.window_size = check_set_window_size(self_attn_mask_type, self.window_size) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype ub_tp_comm_overlap = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_OVERLAP", "1"))) ub_bulk_wgrad = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_BULK_WGRAD", "1"))) @@ -491,6 +504,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, self_attn_mask_type: Optional[str] = None, + window_size: Optional[Tuple[int, int]] = None, encoder_output: Optional[torch.Tensor] = None, enc_dec_attn_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, is_first_microbatch: Optional[bool] = None, @@ -521,6 +535,8 @@ def forward( self_attn_mask_type: {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'}, default = `causal` Type of attention mask passed into softmax operation. + window_size: Optional[Tuple[int, int]], default = `None` + sliding window size for local attention. encoder_output : Optional[torch.Tensor], default = `None` Output of the encoder block to be fed into the decoder block if using `layer_type="decoder"`. @@ -562,8 +578,12 @@ def forward( to efficienly calculate and store the context during inference. """ + if self_attn_mask_type is not None: + window_size = check_set_window_size(self_attn_mask_type, window_size) if self_attn_mask_type is None: self_attn_mask_type = self.self_attn_mask_type + if window_size is None: + window_size = self.window_size assert ( self_attn_mask_type in AttnMaskTypes @@ -594,6 +614,7 @@ def forward( hidden_states, attention_mask=attention_mask, attn_mask_type=self_attn_mask_type, + window_size=window_size, inference_params=inference_params, is_first_microbatch=is_first_microbatch, checkpoint_core_attention=checkpoint_core_attention, @@ -619,6 +640,7 @@ def forward( inter_attention_outputs = self.inter_attention( hidden_states, attention_mask=enc_dec_attn_mask, + window_size=window_size, encoder_output=encoder_output, is_first_microbatch=is_first_microbatch, checkpoint_core_attention=checkpoint_core_attention,