[QUESTION]Create the pp groups when enable context_parallel_size > 1 and encoder_tensor_parallel_size != tensor_parallel_size

[heavyrain-lzy]

Your question
I have a quesstion about create the pp groups when enable context_parallel_size > 1 and encoder_tensor_parallel_size != tensor_parallel_size.

When enabling context_parallel, the input will be splited symmetrically in order to balance calculation. Using zip(cycle(e_ranks), d_ranks) is not right.

Megatron-LM/megatron/core/parallel_state.py

Line 602 in 4ff9e66

# Map 1 encoder tp rank to several decoder tp ranks, because

[Advait251206]

This is a subtle but important point about how parallel group construction changes when context parallelism (CP) is introduced alongside mismatched encoder/decoder tensor parallel sizes.

You are correct to question the use of:

zip(cycle(e_ranks), d_ranks)

because it implicitly assumes a many-to-one or uneven mapping pattern that breaks under CP symmetry constraints.

---

1. What changes when context_parallel_size > 1

When context parallelism is enabled:

• The sequence dimension is split across CP ranks
• Each decoder rank now processes a different slice of context
• Work distribution becomes symmetric across CP groups

So instead of:

encoder TP rank → multiple decoder TP ranks (static mapping)

You now effectively have:

encoder TP × CP shard → aligned decoder TP × CP shard

---

2. Why the original mapping breaks

This pattern:

zip(cycle(e_ranks), d_ranks)

assumes:

• encoder ranks are reused cyclically
• decoder ranks are independent list
• CP dimension does NOT exist

❌ Problem under CP:

When CP > 1:

• decoder ranks are no longer a flat list
• they are structured as:

(d_tp_rank, d_cp_rank)

So flattening them into d_ranks destroys structure.

---

3. Correct mental model

Instead of:

encoder_tp_rank → decoder_tp_rank

You now need:

encoder_tp_rank → (decoder_tp_rank, decoder_cp_rank)

or more precisely:

(encoder_tp, encoder_cp) ↔ (decoder_tp, decoder_cp)

---

4. What should happen logically

With:

encoder_tensor_parallel_size ≠ tensor_parallel_size
context_parallel_size > 1

You must ensure:

✔ Alignment condition:

Each CP shard must map independently:

for each cp group:
    map encoder_tp_group → decoder_tp_group

So mapping becomes block structured, not linear zipped.

---

5. Why zip(cycle()) is wrong here

It implicitly assumes:

• no hierarchy between TP and CP
• linear rank ordering is sufficient

But CP introduces:

a 2D process grid instead of 1D flattening

So cycle(e_ranks) causes:

• encoder ranks to be reused incorrectly across CP shards
• cross-shard mixing
• invalid communication groups

---

6. Correct approach (conceptual fix)

Instead of flattening, you should:

Step 1: reshape ranks into grid form

encoder: [TP_e × CP_e]
decoder: [TP_d × CP_d]

Step 2: group by CP index first

Then within each CP group:

• map encoder TP → decoder TP

---

7. Expected correct grouping logic

Pseudo-logic:

for cp_id in range(cp_size):
    enc_cp_group = encoder_ranks[cp_id::cp_size]
    dec_cp_group = decoder_ranks[cp_id::cp_size]

    for e_tp, d_tp in zip(enc_cp_group, dec_cp_group):
        create_group(e_tp → d_tp)

This preserves:

• CP locality
• TP alignment inside CP shard
• deterministic mapping

---

8. Key insight

The core issue is:

You cannot treat TP groups independently once CP introduces a second partition dimension.

So:

• ❌ linear zip across flattened ranks
• ✔ hierarchical grouping (CP → TP)

---

9. Why this matters in Megatron/TP/CP systems

Because:

• NCCL groups are topology-sensitive
• mis-grouping leads to:
  • deadlocks
  • silent performance degradation
  • incorrect all-reduce boundaries

---

10. Bottom line

• ✔ Your intuition is correct: zip(cycle()) is invalid under CP symmetry
• ❌ Flattened rank mapping breaks the CP × TP structure
• ✔ Correct design requires 2D grouping (CP-first, then TP mapping)
• ⚠ Encoder/decoder mismatch must be handled per-CP slice, not globally

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206