[Doc] Fix TP and attention-TP group layout in initialize_model_parallel docstring (#34862)

Co-authored-by: NanoByte0513 <167996578+NanoByte0513@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
qiaonan02
2026-08-19 08:15:03 +08:00
committed by GitHub
co-authored by NanoByte0513 Xiaoyu Zhang
parent c14312a664
commit 64e404263e
@@ -2372,12 +2372,16 @@ def initialize_model_parallel(
Let's say we use 2 GPUs for attention context parallelism (attn_cp_size=2) and 4 GPUs for
attention tensor parallelism (attn_tp_size=4). As for MoE part, we use 2 GPUs for moe data
parallelism (moe_dp_size=2) and 4 GPUs for moe expert parallelism (moe_ep_size=4). The present
parallelism (moe_dp_size=2) and 4 GPUs for moe expert parallelism (moe_ep_size=4). Note that
this implies tensor_model_parallel_size=8 (attn_tp_size = tp_size // attn_cp_size //
attn_dp_size), so all 8 GPUs form a single tensor model-parallel group. The present
function will create the following groups:
2 tensor model-parallel groups:
[g0, g1, g2, g3], [g4, g5, g6, g7]
1 tensor model-parallel group:
[g0, g1, g2, g3, g4, g5, g6, g7]
4 attention context-parallel groups:
[g0, g4], [g1, g5], [g2, g6], [g3, g7]
2 attention tensor-parallel groups:
[g0, g1, g2, g3], [g4, g5, g6, g7]
2 moe expert-parallel groups:
[g0, g1, g2, g3], [g4, g5, g6, g7]
4 moe data-parallel groups: