[PyTorch/Jax] Fix attention mask definition, and sliding window for decoder

README.rst

@@ -184,10 +184,35 @@ Compiling with FlashAttention-2
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Transformer Engine release v0.11.0 adds support for FlashAttention-2 in PyTorch for improved performance. 
 
-It is a known issue that FlashAttention-2 compilation is resource-intensive and requires a large amount of RAM (see `bug <https://github.com/Dao-AILab/flash-attention/issues/358>`_), which may lead to out of memory errors during the installation of Transformer Engine. Please try setting **MAX_JOBS=1** in the environment to circumvent the issue. If the errors persist, install a supported version of FlashAttention-1 (v1.0.6 to v1.0.9).
+It is a known issue that FlashAttention-2 compilation is resource-intensive and requires a large amount of RAM (see `bug <https://github.com/Dao-AILab/flash-attention/issues/358>`_), which may lead to out of memory errors during the installation of Transformer Engine. Please try setting **MAX_JOBS=1** in the environment to circumvent the issue.
 
 Note that NGC PyTorch 23.08+ containers include FlashAttention-2.
 
+Breaking Changes
+================
+
+v1.7: Padding mask definition for PyTorch
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+In an effort to unify the definition and usage of the attention mask across all three frameworks in Transformer Engine, the padding mask has changed from `True` meaning inclusion of the corresponding position in attention to exclusion of that position in our PyTorch implementation. Since v1.7, all attention mask types follow the same definition where `True` means masking out the corresponding position and `False` means including that position in attention calculation.
+
+An example of this change is,
+
+.. code-block:: bash
+
+    # for a batch of 3 sequences where `a`s, `b`s and `c`s are the useful tokens
+    # and `0`s are the padding tokens,
+    [a, a, a, 0, 0,
+     b, b, 0, 0, 0,
+     c, c, c, c, 0]
+    # the padding mask for this batch before v1.7 is,
+    [ True,  True,  True, False, False,
+      True,  True, False, False, False,
+      True,  True,  True,  True, False]
+    # and for v1.7 onwards it should be,
+    [False, False, False,  True,  True,
+     False, False,  True,  True,  True,
+     False, False, False, False,  True]
+
 FP8 Convergence
 ===============
 

tests/jax/test_fused_attn.py

@@ -66,7 +66,7 @@ def general_dot_product_attention(query: ArrayLike, key: ArrayLike, value: Array
     if mask is not None:
         if mask.ndim != logits.ndim:
             mask = jnp.expand_dims(mask, axis=-3)
-        logits = jnp.where(mask, logits, jnp.finfo(dtype).min)
+        logits = jnp.where(mask, jnp.finfo(dtype).min, logits)
 
     softmax_out = jax.nn.softmax(logits).astype(dtype)
 
@@ -90,24 +90,34 @@ def is_causal_mask(mask: AttnMaskType):
 
 def make_decoder_mask(q_tokens: ArrayLike, kv_tokens: ArrayLike) -> Array:
     """
-    Create padded causal mask
+    Create inverse padded causal mask where `True` means allowing the corresponding
+    position to participate in attention and `False` means masking out that position.
     """
     q_idxs = jnp.broadcast_to(jnp.arange(q_tokens.shape[-1], dtype=jnp.int32), q_tokens.shape)
     kv_idxs = jnp.broadcast_to(jnp.arange(kv_tokens.shape[-1], dtype=jnp.int32), kv_tokens.shape)
-    causal_mask = make_attention_mask(q_idxs, kv_idxs, jnp.greater_equal)
-    padding_mask = make_attention_mask(q_tokens > 0, kv_tokens > 0)
-    return combine_masks(causal_mask, padding_mask)
+    inv_causal_mask = make_attention_mask(q_idxs, kv_idxs, jnp.greater_equal)
+    inv_padding_mask = make_attention_mask(q_tokens > 0, kv_tokens > 0)
+    return combine_masks(inv_causal_mask, inv_padding_mask)
 
+def make_mask(q_token: ArrayLike, kv_token: ArrayLike, attn_mask_type: AttnMaskType) -> Array:
+    """
+    Create attention mask based on mask type. A `True` value in the mask means
+    masking out the corresponding position and a `False` value means allowing
+    that position to participate in attention.
+    """
+    if is_causal_mask(attn_mask_type):
+        inv_mask = make_decoder_mask(q_token, kv_token)
+    else:
+        inv_mask = make_attention_mask(q_token > 0, kv_token > 0)
+    mask = jnp.logical_not(inv_mask)
+    return mask
 
 def jax_dpa(query, key, value, bias, q_token, kv_token, dropout_rng, **kwargs):
     """
     JAX native dot product attention implementation
     """
     attn_mask_type = kwargs['attn_mask_type']
-    if is_causal_mask(attn_mask_type):
-        mask = make_decoder_mask(q_token, kv_token)
-    else:
-        mask = make_attention_mask(q_token > 0, kv_token > 0)
+    mask = make_mask(q_token, kv_token, attn_mask_type)
 
     output = general_dot_product_attention(query,
                                            key,
@@ -127,13 +137,7 @@ def customcall_fused_dpa(query, key, value, bias, q_token, kv_token, dropout_rng
     TE customcall dot product attention implementation
     """
     attn_mask_type = kwargs['attn_mask_type']
-    if is_causal_mask(attn_mask_type):
-        mask = make_decoder_mask(q_token, kv_token)
-    else:
-        mask = make_attention_mask(q_token > 0, kv_token > 0)
-
-    # mask invert
-    mask = jnp.logical_not(mask)
+    mask = make_mask(q_token, kv_token, attn_mask_type)
 
     qkv_layout = kwargs.pop('qkv_layout')
     match qkv_layout:
@@ -298,6 +302,8 @@ def test_backward(self):
         """
 
         self._setup_inputs()
+        if self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape != BiasShape.BIAS_1HSS:
+            pytest.skip("Bias gradient calculation is only supported for 1HSS bias shape.")
 
         def grad_func(func, *args, **kwargs):
             # Gradient is small, use a gradient multiplier to amplify the gradient

tests/jax/utils.py

@@ -635,7 +635,7 @@ def qkv_init(key, shape, dtype):
                 # position should only attend to those key positions that have already
                 # been generated and cached, not the remaining zero elements.
                 mask = combine_masks(
-                    mask,
+                    jnp.logical_not(mask),
                     jnp.broadcast_to(
                         jnp.arange(length) <= cur_index,
                 # (1, 1, length) represent (head dim, query length, key length)

tests/pytorch/fused_attn/test_fused_attn.py

@@ -544,27 +544,26 @@ def _run_dot_product_attention(
             attention_mask_q = torch.Tensor([]).to(dtype=torch.bool)
             for i in range(config.batch_size):
                 attention_mask_q = torch.cat([attention_mask_q,
-                    torch.Tensor([True]*seqlens_q[i] + [False]*(config.max_seqlen_q-seqlens_q[i]))
+                    torch.Tensor([False]*seqlens_q[i] + [True]*(config.max_seqlen_q-seqlens_q[i]))
                     .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
             attention_mask = attention_mask_q.to(device="cuda")
         if config.attn_type == 'cross':
             attention_mask_q = torch.Tensor([]).to(dtype=torch.bool)
             attention_mask_kv = torch.Tensor([]).to(dtype=torch.bool)
             for i in range(config.batch_size):
                 attention_mask_q = torch.cat([attention_mask_q,
-                    torch.Tensor([True]*seqlens_q[i] + [False]*(config.max_seqlen_q-seqlens_q[i]))
+                    torch.Tensor([False]*seqlens_q[i] + [True]*(config.max_seqlen_q-seqlens_q[i]))
                     .to(dtype=torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
                 attention_mask_kv = torch.cat([attention_mask_kv, torch.Tensor(
-                    [True]*seqlens_kv[i] + [False]*(config.max_seqlen_kv-seqlens_kv[i]))
+                    [False]*seqlens_kv[i] + [True]*(config.max_seqlen_kv-seqlens_kv[i]))
                     .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"))
+    window_size = None
     if swa:
         window_size, attention_mask = get_swa(config.max_seqlen_q, config.max_seqlen_kv)
     elif "causal" in config.attn_mask_type:
         window_size, attention_mask = (-1, 0), None
-    else:
-        window_size, attention_mask = None, None
 
     alibi_slopes = None
     if config.attn_bias_type == "alibi" and config.alibi_type == "custom":
@@ -858,7 +857,7 @@ def _run_transformer_layer(
         attention_mask_q = torch.Tensor([]).to(dtype=torch.bool)
         for i in range(config.batch_size):
             attention_mask_q = torch.cat([attention_mask_q,
-                torch.Tensor([True]*seqlens_q[i] + [False]*(config.max_seqlen_q-seqlens_q[i]))
+                torch.Tensor([False]*seqlens_q[i] + [True]*(config.max_seqlen_q-seqlens_q[i]))
                 .to(torch.bool).unsqueeze(0).unsqueeze(0).unsqueeze(0)], dim=0)
         attention_mask = attention_mask_q.to(device="cuda")
 
@@ -944,7 +943,7 @@ def _run_transformer_layer(
 
 model_configs_fp8_vs_f16 = {
     #  test:             b,  h, hg,   d,   sq,  skv,   p,      mask,      bias
-    "fp8_9 ": ModelConfig(2, 24, 24, 128, 2048, 2048, 0.0, "no_mask", "no_bias"),
+    "fp8_9" : ModelConfig(2, 24, 24, 128, 2048, 2048, 0.0, "no_mask", "no_bias"),
     "fp8_10": ModelConfig(2, 24, 24, 128, 2048, 2048, 0.0,  "causal", "no_bias"),
     "fp8_11": ModelConfig(2, 24, 12, 128, 2048, 2048, 0.0, "no_mask", "no_bias"),
     "fp8_12": ModelConfig(2, 24, 12, 128, 2048, 2048, 0.0,  "causal", "no_bias"),
@@ -1143,24 +1142,49 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd):
         dtype, config, False, qkv_layout)
 
     tols = dict(atol=5e-1, rtol=5e-2)
+    rmse_tol = 0.1
+    bwd_names = ['dq', 'dk', 'dv']
+    fwd_rmse = _rmse(fused_attn_fwd_fp8, fused_attn_fwd_f16)
+    fwd_range = max(fused_attn_fwd_fp8.max().item(),
+        fused_attn_fwd_f16.max().item()) - min(fused_attn_fwd_fp8.min().item(),
+        fused_attn_fwd_f16.min().item())
     if _NVTE_DEBUG:
-        print('[test_dpa_fp8_vs_f16]: ', tols)
+        print()
+        print('========== {:^25s} =========='.format('forward output'))
         print('fused_attn_fwd_fp8 min {:.6f} max {:.6f}'.format(
             fused_attn_fwd_fp8.min().item(),fused_attn_fwd_fp8.max().item()))
         print('fused_attn_fwd_f16 min {:.6f} max {:.6f}'.format(
             fused_attn_fwd_f16.min().item(), fused_attn_fwd_f16.max().item()))
-        print('fused_attn_fwd RMSE: {:.6f}'.format(
-            _rmse(fused_attn_fwd_fp8, fused_attn_fwd_f16)))
-    torch.testing.assert_close(fused_attn_fwd_fp8, fused_attn_fwd_f16, **tols)
+        print('fused_attn_fwd RMSE: {:.6f}'.format(fwd_rmse))
+        try:
+            torch.testing.assert_close(fused_attn_fwd_fp8, fused_attn_fwd_f16, **tols)
+        except Exception as e:
+            print(e)
+            print()
+    assert(fwd_rmse < rmse_tol * fwd_range
+        ), "FWD RMSE {:.5f} is over tolerance {:.5f} ({:.5f} * {:.5f})".format(
+        fwd_rmse, rmse_tol * fwd_range, rmse_tol, fwd_range)
     for i,_ in enumerate(fused_attn_bwd_f16):
+        bwd_rmse = _rmse(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i])
+        bwd_range = max(fused_attn_bwd_fp8[i].max().item(),
+            fused_attn_bwd_f16[i].max().item()) - min(fused_attn_bwd_fp8[i].min().item(),
+            fused_attn_bwd_f16[i].min().item())
         if _NVTE_DEBUG:
-            print('fused_attn_bwd_fp8 min {:.6f} max {:.6f}'.format(
+            print()
+            print('========== {:^25s} =========='.format(bwd_names[i]))
+            print('fused_attn_bwd_fp8[{}] min {:.6f} max {:.6f}'.format(i,
                 fused_attn_bwd_fp8[i].min().item(), fused_attn_bwd_fp8[i].max().item()))
-            print('fused_attn_bwd_f16 min {:.6f} max {:.6f}'.format(
+            print('fused_attn_bwd_f16[{}] min {:.6f} max {:.6f}'.format(i,
                 fused_attn_bwd_f16[i].min().item(), fused_attn_bwd_f16[i].max().item()))
-            print('fused_attn_bwd RMSE: {:.6f}'.format(
-                _rmse(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i])))
-        torch.testing.assert_close(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i], **tols)
+            print('fused_attn_bwd RMSE[{}]: {:.6f}'.format(i, bwd_rmse))
+            try:
+                torch.testing.assert_close(fused_attn_bwd_fp8[i], fused_attn_bwd_f16[i], **tols)
+            except Exception as e:
+                print(e)
+                print()
+        assert(bwd_rmse < rmse_tol * bwd_range
+            ), "BWD RMSE {:.5f} is over tolerance {:.5f} ({:.5f} * {:.5f})".format(
+            bwd_rmse, rmse_tol * bwd_range, rmse_tol, bwd_range)
 
 
 def _run_dpa_fp8_vs_f16(dtype, config, fp8_dpa, qkv_layout):
@@ -1231,7 +1255,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
             layout = layout.replace('h', 'hg')
             layout = layout.replace('t', 'tg')
         tensor_shape = [dim_to_num[j] for j in layout.split('_')]
-        tensor = 0.1 * torch.randn(tensor_shape, dtype=dtype, device="cuda")
+        tensor = torch.randn(tensor_shape, dtype=dtype, device="cuda")
         tensor_count = 1
         split_dim = 0
         for dim, l in enumerate(layout.split('_')):
@@ -1252,7 +1276,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker:
     qkv_format_kv = qkv_format_kv.replace('s', 'sq')
     out_grad_shape = [dim_to_num[i] for i in qkv_format_kv.split('_')]
     out_grad_shape_new = [*out_grad_shape[:-2], out_grad_shape[-2] * out_grad_shape[-1]]
-    out_grad = 0.1 * torch.randn(out_grad_shape_new, dtype=dtype, device="cuda")
+    out_grad = torch.randn(out_grad_shape_new, dtype=dtype, device="cuda")
 
     with fp8_autocast(enabled=fp8_dpa, fp8_recipe=fp8_recipe):
         out = dpa(inp[0], inp[1], inp[2],
@@ -1359,7 +1383,7 @@ def _run_custom_mha_fp8(dtype, config, backend):
     if backend == "FusedAttention":
         os.environ["NVTE_FUSED_ATTN"] = "1"
 
-    inp = 0.0001 * torch.randint(0, 100,
+    inp = 0.0001 * torch.randint(-100, 100,
             (config.batch_size * config.max_seqlen_q, config.num_heads * config.head_dim),
             dtype=dtype, device="cuda", requires_grad=True)
     seqlens = torch.full([config.batch_size], config.max_seqlen_q,