From 4e0e99bb6a356cb734194aa985d837e3744b0b92 Mon Sep 17 00:00:00 2001 From: wangbluo <2538539015@qq.com> Date: Fri, 11 Oct 2024 17:31:40 +0800 Subject: [PATCH] fix the test --- .../test_layer/test_ring_attn.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_shardformer/test_layer/test_ring_attn.py b/tests/test_shardformer/test_layer/test_ring_attn.py index 1c7647a7d..38dfcb239 100644 --- a/tests/test_shardformer/test_layer/test_ring_attn.py +++ b/tests/test_shardformer/test_layer/test_ring_attn.py @@ -5,6 +5,7 @@ from flash_attn import flash_attn_qkvpacked_func, flash_attn_varlen_qkvpacked_fu from torch.testing import assert_close import colossalai +from colossalai.cluster import ProcessGroupMesh from colossalai.shardformer.layer import AttnMaskType from colossalai.shardformer.layer.attn import AttnMaskType, RingAttention from colossalai.shardformer.layer.utils import split_batch_zigzag, split_varlen_zigzag @@ -17,11 +18,15 @@ from colossalai.utils import get_current_device @parameterize("nheads", [5]) @parameterize("d", [128]) @parameterize("dtype", [torch.bfloat16, torch.float16]) -def check_ring_attn(seq_len, bs, nheads, d, dtype): +def check_ring_attn(seq_len, bs, nheads, d, dtype, sp_size, tp_size=1): torch.cuda.manual_seed(2) device = get_current_device() - sp_group = dist.group.WORLD - sp_size = dist.get_world_size() + sp_axis, tp_axis = 0, 1 + pg_mesh = ProcessGroupMesh(sp_size, tp_size) + tp_group = pg_mesh.get_group_along_axis(tp_axis) + sp_group = pg_mesh.get_group_along_axis(sp_axis) + # sp_group = dist.group.WORLD + # sp_size = dist.get_world_size() # Some outliers may seem large, but our errors are still lower than # than Megatron-LM context parallel's # (https://github.com/NVIDIA/TransformerEngine/blob/33a3d02f81c56e6f7b542c09bfa86657078d57fb/tests/pytorch/fused_attn/run_fused_attn_with_cp.py#L215) @@ -44,6 +49,7 @@ def check_ring_attn(seq_len, bs, nheads, d, dtype): AttnMaskType.CAUSAL, return_softmax=True, inner_ring_size=max(2, sp_size // 2), + tp_group=tp_group, # inner_ring_size=4 ) ring_out = ring_out.transpose(1, 2) @@ -161,12 +167,12 @@ def check_packed_seq(seqlen, bs, nheads, d, dtype): def launch_single_ring(rank, world_size, port): colossalai.launch(rank, world_size, "localhost", port) check_packed_seq() - check_ring_attn() + check_ring_attn(sp_size=world_size) def launch_double_ring(rank, world_size, port): colossalai.launch(rank, world_size, "localhost", port) - check_ring_attn() + check_ring_attn(sp_size=4, tp_size=2) @rerun_if_address_is_in_use() @@ -176,7 +182,7 @@ def test_ring_attn(world_size): @rerun_if_address_is_in_use() -@parameterize("world_size", [4]) +@parameterize("world_size", [8]) def test_double_ring(world_size): spawn(launch_double_ring, nprocs=world_size)