Add THD format support for Context Parallel

qa/L1_pytorch_context_parallel_test/test.sh

@@ -6,5 +6,5 @@ set -e
 
 : ${TE_PATH:=/opt/transformerengine}
 
-pip install pytest==6.2.5 onnxruntime==1.13.1
+pip install pytest==7.2.0 onnxruntime==1.13.1
 pytest -v -s $TE_PATH/tests/pytorch/fused_attn/test_fused_attn_with_cp.py

tests/pytorch/fused_attn/run_fused_attn_with_cp.py

@@ -6,6 +6,7 @@
 import torch
 import torch.distributed as dist
 from transformer_engine.pytorch.attention import DotProductAttention
+import transformer_engine_extensions as tex
 from test_fused_attn_with_cp import model_configs_flash_attn, model_configs_fused_attn
 
 dtypes={'fp16' : torch.float16, 'bf16' : torch.bfloat16}
@@ -58,12 +59,27 @@ def run_dpa_with_cp(dtype='bf16', model=None, qkv_format='bshd', kernel_backend=
         q_input_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim)
         kv_input_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim)
         attn_output_shape = (config.batch_size, config.max_seqlen_q, config.num_heads*config.head_dim)
+        cu_seqlens_q = None
+        cu_seqlens_kv = None
     elif qkv_format == "sbhd":
         q_input_shape = (config.max_seqlen_q, config.batch_size, config.num_heads, config.head_dim)
         kv_input_shape = (config.max_seqlen_kv, config.batch_size, config.num_gqa_groups, config.head_dim)
         attn_output_shape = (config.max_seqlen_q, config.batch_size, config.num_heads*config.head_dim)
+        cu_seqlens_q = None
+        cu_seqlens_kv = None
+    elif qkv_format == "thd":
+        seqlens_q = torch.randint(world_size * 2, config.max_seqlen_q + 1, [config.batch_size]).to(torch.int32)
+        seqlens_q = seqlens_q - seqlens_q % (world_size * 2)
+        cu_seqlens_q = torch.cat([torch.zeros([1], dtype=torch.int32), seqlens_q.cumsum(0)])
+        cu_seqlens_kv = cu_seqlens_q
+        q_input_shape = (cu_seqlens_q[-1], config.num_heads, config.head_dim)
+        kv_input_shape = (cu_seqlens_kv[-1], config.num_gqa_groups, config.head_dim)
+        attn_output_shape = (cu_seqlens_q[-1], config.num_heads*config.head_dim)
+        cu_seqlens_q = cu_seqlens_q.to(torch.int32).cuda()
+        cu_seqlens_kv = cu_seqlens_kv.to(torch.int32).cuda()
     else:
         assert False, f"{qkv_format} is an unsupported qkv_format!"
+
     q = torch.randn(q_input_shape, dtype=dtypes[dtype]).cuda()
     k = torch.randn(kv_input_shape, dtype=dtypes[dtype]).cuda()
     v = torch.randn(kv_input_shape, dtype=dtypes[dtype]).cuda()
@@ -79,6 +95,9 @@ def run_dpa_with_cp(dtype='bf16', model=None, qkv_format='bshd', kernel_backend=
     # make sure all GPU ranks have same inputs
     for x in [q, k, v, dout] + ([] if bias is None else [bias]):
         dist.broadcast(x, 0, group=cp_comm_group)
+    if qkv_format == "thd":
+        for x in [cu_seqlens_q, cu_seqlens_kv]:
+            dist.broadcast(x, 0, group=cp_comm_group)
 
     # run core_attn without CP
     for x in [q, k, v]:
@@ -87,28 +106,48 @@ def run_dpa_with_cp(dtype='bf16', model=None, qkv_format='bshd', kernel_backend=
         q, k, v,
         core_attention_bias_type=config.attn_bias_type,
         core_attention_bias=bias,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
     )
     out.backward(dout)
 
     # run core_attn wit CP
-    q_, k_, v_, dout_, *rest = [x.clone().detach() for x in [q, k, v, dout] + ([] if bias is None else [bias])]
-    bias_ = rest[0] if len(rest) else None
-    seq_dim = qkv_format.index('s')
-    q_, k_, v_, dout_ = [x.view(*x.shape[:seq_dim], 2*world_size, x.shape[seq_dim]//(2*world_size), *x.shape[(seq_dim+1):]) \
-        for x in [q_, k_, v_, dout_]]
-    seq_idx = torch.tensor([rank, 2*world_size-rank-1], device=q_.device)
-    q_, k_, v_, dout_ = [x.index_select(seq_dim, seq_idx) for x in [q_, k_, v_, dout_]]
-    q_, k_, v_, dout_ = [x.view(*x.shape[:seq_dim], -1, *x.shape[(seq_dim+2):]) for x in [q_, k_, v_, dout_]]
+    if qkv_format == "bshd" or qkv_format == "sbhd":
+        q_, k_, v_, dout_, *rest = [x.clone().detach() for x in [q, k, v, dout] + ([] if bias is None else [bias])]
+        bias_ = rest[0] if len(rest) else None
+        seq_dim = qkv_format.index('s')
+        q_, k_, v_, dout_ = [x.view(*x.shape[:seq_dim], 2*world_size, x.shape[seq_dim]//(2*world_size), *x.shape[(seq_dim+1):]) \
+            for x in [q_, k_, v_, dout_]]
+        seq_idx = torch.tensor([rank, 2*world_size-rank-1], device=q_.device)
+        q_, k_, v_, dout_ = [x.index_select(seq_dim, seq_idx) for x in [q_, k_, v_, dout_]]
+        q_, k_, v_, dout_ = [x.view(*x.shape[:seq_dim], -1, *x.shape[(seq_dim+2):]) for x in [q_, k_, v_, dout_]]
+    elif qkv_format == "thd":
+        q_, k_, v_, dout_ = [x.clone().detach() for x in [q, k, v, dout]]
+        seq_idx_q  = tex.thd_get_partitioned_indices(cu_seqlens_q, q_.size(0), world_size, rank)
+        seq_idx_kv = tex.thd_get_partitioned_indices(cu_seqlens_kv, k_.size(0), world_size, rank)
+        q_, dout_ = [x.index_select(0, seq_idx_q) for x in [q_, dout_]]
+        k_, v_ = [x.index_select(0, seq_idx_kv) for x in [k_, v_]]
+        cu_seqlens_q = cu_seqlens_q // world_size
+        cu_seqlens_kv = cu_seqlens_kv // world_size
+        bias_ = None
+    else:
+        assert False, f"{qkv_format} is an unsupported qkv_format!"
     q_, k_, v_ = [x.requires_grad_() for x in [q_, k_, v_]]
     if bias_ is not None:
         bias_ = bias_.view(*bias_.shape[:-2], 2*world_size, bias_.shape[-2]//(2*world_size), bias_.shape[-1])
         bias_ = bias_.index_select(2, seq_idx)
         bias_ = bias_.view(*bias_.shape[:2], -1, bias_.shape[-1])
     core_attn.set_context_parallel_group(cp_comm_group, cp_comm_ranks, torch.cuda.Stream())
+    max_seqlen_q  = config.max_seqlen_q
+    max_seqlen_kv = config.max_seqlen_kv
     out_ = core_attn(
         q_, k_, v_,
         core_attention_bias_type=config.attn_bias_type,
         core_attention_bias=bias_,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_kv=max_seqlen_kv,
     )
     out_.backward(dout_)
 
@@ -120,11 +159,20 @@ def run_dpa_with_cp(dtype='bf16', model=None, qkv_format='bshd', kernel_backend=
     tols = dict(atol=5e-3, rtol=5e-3)
     if dtype == 'bf16':
         tols = dict(atol=2.5e-2, rtol=2.5e-2)
-    dq, dk, dv, out = [x.view(*x.shape[:seq_dim], 2*world_size, x.shape[seq_dim]//(2*world_size), *x.shape[(seq_dim+1):]) \
-        for x in [q.grad, k.grad, v.grad, out]]
-    dq, dk, dv, out = [x.index_select(seq_dim, seq_idx) for x in [dq, dk, dv, out]]
-    dq_, dk_, dv_, out_ = [x.view(*x.shape[:seq_dim], 2, x.shape[seq_dim]//2, *x.shape[(seq_dim+1):]) \
-        for x in [q_.grad, k_.grad, v_.grad, out_]]
+
+    if qkv_format == "bshd" or qkv_format == "sbhd":
+        dq, dk, dv, out = [x.view(*x.shape[:seq_dim], 2*world_size, x.shape[seq_dim]//(2*world_size), *x.shape[(seq_dim+1):]) \
+            for x in [q.grad, k.grad, v.grad, out]]
+        dq, dk, dv, out = [x.index_select(seq_dim, seq_idx) for x in [dq, dk, dv, out]]
+        dq_, dk_, dv_, out_ = [x.view(*x.shape[:seq_dim], 2, x.shape[seq_dim]//2, *x.shape[(seq_dim+1):]) \
+            for x in [q_.grad, k_.grad, v_.grad, out_]]
+    elif qkv_format == "thd":
+        dq, out = [x.index_select(0, seq_idx_q).contiguous().view(-1) for x in [q.grad, out]]
+        dk, dv = [x.index_select(0, seq_idx_kv).contiguous().view(-1) for x in [k.grad, v.grad]]
+        dq_, dk_, dv_, out_ = [x.view(-1) for x in [q_.grad, k_.grad, v_.grad, out_]]
+    else:
+        assert False, f"{qkv_format} is an unsupported qkv_format!"
+
     if qkv_format == "bshd":
         torch.testing.assert_close(out_[:, 0], out[:, 0], **tols)
         torch.testing.assert_close(dq_[:, 0], dq[:, 0], **tols)
@@ -143,6 +191,11 @@ def run_dpa_with_cp(dtype='bf16', model=None, qkv_format='bshd', kernel_backend=
         torch.testing.assert_close(dq_[1], dq[1], **tols)
         torch.testing.assert_close(dk_[1], dk[1], **tols)
         torch.testing.assert_close(dv_[1], dv[1], **tols)
+    elif qkv_format == "thd":
+        torch.testing.assert_close(out_, out, **tols)
+        torch.testing.assert_close(dq_, dq, **tols)
+        torch.testing.assert_close(dk_, dk, **tols)
+        torch.testing.assert_close(dv_, dv, **tols)
     else:
         assert False, f"{qkv_format} is an unsupported qkv_format!"
 

tests/pytorch/fused_attn/test_fused_attn_with_cp.py

@@ -33,7 +33,7 @@ def get_bash_arguments(**kwargs):
 @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.")
 @pytest.mark.parametrize("dtype", ['bf16', 'fp16'])
 @pytest.mark.parametrize("model", model_configs_flash_attn.keys())
-@pytest.mark.parametrize("qkv_format", ['bshd', 'sbhd'])
+@pytest.mark.parametrize("qkv_format", ['bshd', 'sbhd', 'thd'])
 def test_cp_with_flash_attention(dtype, model, qkv_format):
     subprocess.run(
         get_bash_arguments(