[QUESTION] why all rank0 of tp group make datasets？

[ltm920716]

Your question
hi,
I am very confused about the logic bellow. Here, datasets are built for each rank 0 of the TP group. If it is a rank 0 in the TP group of the middle layer, is it necessary to build datasets?
If the tp group is [[0,1,2,3], [4,5,6,7]], pp-group is [[0,4], [1,5], [2,6], [3,7]], world-size=8, that is global rank0 and rank4 make datasets both，that is right？

help please， thanks！
(

Megatron-LM/megatron/training/training.py

Line 1690 in 55622ff

if is_distributed or mpu.get_tensor_model_parallel_rank() == 0:

)

[Advait251206]

This is a very good observation, and the confusion comes from how Megatron-LM mixes data parallel responsibilities with tensor model parallel ranks.

Let’s break it down clearly.

---

🧠 Key line in question

if is_distributed or mpu.get_tensor_model_parallel_rank() == 0:
    build_dataset()

This means:

“Only TP rank 0 (within each TP group) builds the dataset.”

---

🔷 Your setup

You have:

• TP groups:

  [0,1,2,3], [4,5,6,7]
  

• PP groups:

  [0,4], [1,5], [2,6], [3,7]
  

So:

• TP rank 0 in group 1 → global ranks 0 and 4
• Each TP group has its own rank 0

---

❓ Your question

Why do both global rank 0 and rank 4 build datasets? Is that correct?

✔️ Yes — it is expected in current Megatron design.

But the important nuance is:

They are NOT duplicating dataset per GPU — they are duplicating per TP group.

---

🧩 Why dataset is built per TP rank 0

Megatron uses:

🔹 Tensor Parallelism (TP)

• Splits model computation across GPUs
• BUT dataset input pipeline is still needed per TP group leader

So:

• TP rank 0 acts as "data loader owner" for that TP group

---

🔹 Why NOT only global rank 0?

Because in TP + PP hybrid:

• Global rank 0 is only one process
• But each pipeline stage still needs input batches
• TP groups are independent execution shards

So Megatron ensures:

Each TP group gets a local dataset reader (rank 0 of that group)

---

🧠 Important insight

Even though it looks like duplication:

❌ Not actually full dataset duplication across all GPUs

✅ Only dataset initialization responsibility is duplicated

What happens in reality:

• Only TP rank 0 processes load metadata / dataset index
• Data is still sharded via sampler / distributed sampler logic
• Other TP ranks do not independently re-construct dataset

---

🔁 Why PP + TP combination causes confusion

In your mapping:

TP groups:
  [0,1,2,3] → rank 0 loads dataset
  [4,5,6,7] → rank 4 loads dataset

Yes, both 0 and 4 do dataset setup — but they belong to:

• different TP groups
• same PP stage structure but different tensor partitions

So this is expected duplication of loader ownership, not data duplication

---

⚠️ Why Megatron does this design

Because:

1. Avoid cross-process file system contention

Only one rank per TP group touches dataset files.

2. Simplify dataloader logic

No need for global coordination across PP × TP mesh.

3. Ensure fault isolation

Each TP group can reconstruct dataset independently if restarted.

---

🧭 Mental model

Think of it like:

Each TP group = independent "mini data-parallel world"

Inside each world:

• rank 0 = dataset manager
• others = workers

---

📌 Final answer

• Yes, both global rank 0 and rank 4 build dataset in your example
• This is by design
• It is per TP-group responsibility, not global duplication
• It does NOT mean full dataset is redundantly loaded across all GPUs

---

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