Skip to content

Allow tensor parallelism next to replicated data parallelism - #4193

Draft
qgallouedec wants to merge 1 commit into
mainfrom
zero-sync-dp-tp
Draft

Allow tensor parallelism next to replicated data parallelism#4193
qgallouedec wants to merge 1 commit into
mainfrom
zero-sync-dp-tp

Conversation

@qgallouedec

Copy link
Copy Markdown
Member

dp_replicate_size > 1 with dp_shard_size == 1 was rejected whenever tp_size > 1, so a tensor-parallel model could not be held as one replica among several: the only 2-D layouts on offer were the ones that shard the parameters, FSDP + TP and HSDP + TP.

Replicating is a valid layout, and for colocated RL it is the necessary one, since generation reads the parameters in place and cannot do that while they are sharded. Measured on Qwen3-14B with GRPO over one copy of the weights: 4067 trained tokens/s on 8 H100 at tp=4 with two replicas against 2822 at tp=8 with one, 8220 on 16 GPUs with four replicas, and 14120 on 32 with eight, replicas agreeing on the weights bit for bit after training at every layout including across four nodes.

Context parallelism keeps the restriction, and the message no longer claims it applies to tp.

What does this PR do?

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

dp_replicate_size > 1 with dp_shard_size == 1 was rejected whenever tp_size > 1, so a tensor-parallel
model could not be held as one replica among several: the only 2-D layouts on offer were the ones
that shard the parameters, FSDP + TP and HSDP + TP.

Replicating is a valid layout, and for colocated RL it is the necessary one, since generation reads
the parameters in place and cannot do that while they are sharded. Measured on Qwen3-14B with GRPO
over one copy of the weights: 4067 trained tokens/s on 8 H100 at tp=4 with two replicas against 2822
at tp=8 with one, 8220 on 16 GPUs with four replicas, and 14120 on 32 with eight, replicas agreeing
on the weights bit for bit after training at every layout including across four nodes.

Context parallelism keeps the restriction, and the message no longer claims it applies to tp.
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants