[Cosmos3] Add native tensor parallelism by yzhautouskay · Pull Request #15137 · huggingface/diffusers

yzhautouskay · 2026-10-27T11:47:08Z

Context

Cosmos TP/CP currently lives in the examples following #14054, where it was kept outside the core model while the library parallelism APIs were evolving. Those example helpers and their CLI behavior remain unchanged.

Diffusers now provides native model TP (#13718) and sharded checkpoint loading (#14544), so this PR integrates Cosmos with that infrastructure. Following the review direction in #14609, this PR is TP-only. Native Cosmos CP and TP+CP integration can follow separately after the generic composition work in #14725 is settled.

Cosmos native TP

Adds a TP plan for both understanding and generation attention projections and dense MLPs.
Infers rank-local attention head counts from sharded projection widths and validates that query/KV heads divide across the TP mesh.
Handles ReLU2 models with an instance plan that omits nonexistent gate projections.
Keeps the standard attention processor and backend dispatch unchanged.
Documents native TP with the existing 189-frame Super modular text-to-video workflow.

Shared TP change

modeling_utils.py now uses model._get_tp_plan() for checkpoint sharding and TP hooks, following @sayakpaul's proposed yzhautouskay@ef4beac. Cosmos models using SiLU have gate_proj layers; models using ReLU2 do not. The getter adds these layers to a copy of the TP plan only for SiLU, leaving the class-level plan unchanged. Other models retain their existing behavior because the default getter returns their _tp_plan unchanged.

Validation

53 CPU/model tests passed; the focused TP/loading suite passed 11 tests across SiLU/ReLU2 and FP32/BF16; the modular pipeline suite passed 31 tests.
Twelve old/new non-TP comparisons were bit-exact. Small-model TP4 checks covered numerical parity, both loading paths, replicated parameters and buffers.
Native and existing example TP4 produced byte-identical full Super videos with the official 189-frame, 720x1280 workflow.
Across Nano/49, Super/49 and Nano/189 on four GB200 GPUs, median latency differed by at most 1.8% and peak allocated memory by at most 0.051%. Native reserved memory was up to 1.455 GB higher.
make style, make fix-copies and git diff --check passed.

Before submitting

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.

yzhautouskay · 2026-10-27T13:40:23Z

@sayakpaul @yiyixuxu Could you please this PR to enable native TP in Cosmos3?

github-actions · 2026-10-28T07:34:00Z

Hi @yzhautouskay, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

sayakpaul

Thanks!

Some changes are a bit intrusive so I did some work in https://github.com/huggingface/diffusers/compare/cosmos3-tp-plan-resolution?expand=1 which should solve the issues in a less intrusive way. Does that work for you?

@yiyixuxu @DN6 thoughts?

yzhautouskay · 2026-11-05T09:53:48Z

@sayakpaul I updated this PR after #15147 got merged, thank you! Could we merge this one now?

sayakpaul · 2026-11-05T16:24:10Z

+## Native tensor parallelism
+
+Follow the [tensor parallelism guide](../../training/distributed_inference#tensor-parallelism) for distributed setup,
+launching, seeding, and rank-zero output handling. For Cosmos, adapt the
+[existing modular workflow](#modular-examples-for-all-existing-workflows) as follows:
+
+- Use `nvidia/Cosmos3-Super` or `nvidia/Cosmos3-Nano` with `Cosmos3OmniTransformer.from_pretrained`,
+  `subfolder="transformer"`, `dtype=torch.bfloat16`, and `parallel_config=TensorParallelConfig(tp_degree=4)`.
+- Create `Cosmos3OmniModularPipeline` from the same checkpoint, then call
+  `update_components(transformer=transformer)` before `load_components`. Place the remaining components
+  on each rank's local device and reuse the workflow's prompts, scheduler, and generation settings.
+
+Nano and Super support TP degrees 1, 2, 4, and 8; the degree must divide the query and KV head counts
+and the MLP intermediate size. For Cosmos CP or combined TP+CP, use the [existing example helpers](#context-parallelism).


We usually don't include a separate section like this in pipeline docs. Cc: @stevhliu

I just wanted to show in the context of previous "examples helpers" level of CP and TP implementation: https://github.com/yzhautouskay/diffusers/blob/f09416b7f4707d36f47aefe143d2d2d54614b8ca/docs/source/en/api/pipelines/cosmos3.md#context-parallelism

I move this doc section under it

sayakpaul · 2026-11-05T16:28:46Z

            dist.destroy_process_group()


+def _tensor_parallel_output(output):


Can we add an accompanying comment to explain where it happens?

sayakpaul · 2026-11-05T16:29:35Z


        model = model_class(**init_dict)
-        model.load_state_dict(state_dict)
+        model.load_state_dict(state_dict, assign=True)


Why do we need assign=True?

sayakpaul · 2026-11-05T16:30:00Z

        model.eval()

-        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+        inputs_on_device = send_to_device(inputs_dict, device)


Let's keep it explicit and avoid this.

sayakpaul · 2026-11-05T16:31:46Z

+        return inputs
+
+
+class TestCosmos3OmniTransformerTensorParallelConfig(Cosmos3TensorParallelTesterConfig):


sayakpaul · 2026-11-05T16:32:03Z

+    @pytest.fixture(autouse=True, params=["silu", "relu2"])
+    def activation(self, request):
+        self.hidden_act = request.param


Why do we need this fixture?

# Conflicts: # src/diffusers/models/modeling_utils.py

Add native tensor parallelism for Cosmos 3

a1b567f

github-actions Bot added documentation Improvements or additions to documentation models tests size/L PR with diff > 200 LOC labels Oct 27, 2026

sayakpaul added the no-issue-needed for PRs that do not require link to an issue label Oct 28, 2026

Merge branch 'main' into yzhautouskay/cosmos3-native-tp-only

971bf9a

sayakpaul reviewed Oct 28, 2026

View reviewed changes

sayakpaul and others added 3 commits October 28, 2026 07:51

use _get_tp_plan for overrides.

ef4beac

Refactor Cosmos 3 tensor parallel tests

77ed5c4

Simplify Cosmos 3 tensor parallelism documentation

02631f9

yzhautouskay mentioned this pull request Oct 29, 2026

Fix replicated parameter dtypes during tensor-parallel loading #15147

Merged

10 tasks

Merge upstream/main into yzhautouskay/cosmos3-native-tp-only

2e7c7c7

sayakpaul reviewed Nov 5, 2026

View reviewed changes

yzhautouskay added 2 commits November 5, 2026 08:51

Merge upstream/main into yzhautouskay/cosmos3-native-tp-only

96a7dc7

# Conflicts: # src/diffusers/models/modeling_utils.py

Fix docs

f09416b

		dist.destroy_process_group()


		def _tensor_parallel_output(output):

		return inputs


		class TestCosmos3OmniTransformerTensorParallelConfig(Cosmos3TensorParallelTesterConfig):

Conversation

yzhautouskay commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context

Cosmos native TP

Shared TP change

Validation

Before submitting

Who can review?

Uh oh!

yzhautouskay commented Oct 27, 2026

Uh oh!

github-actions Bot commented Oct 28, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yzhautouskay commented Nov 5, 2026

Uh oh!

sayakpaul Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

yzhautouskay Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yzhautouskay commented Oct 27, 2026 •

edited

Loading