[yzhautouskay]

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 loading fixes

• modeling_utils.py now uses the constructed model’s model._tp_plan instead of cls._tp_plan when calculating checkpoint shards and installing TP hooks. Cosmos ReLU2 overrides the instance plan to remove nonexistent gate_proj entries, preventing the previous module-resolution failure. Other models are unaffected because, unless they define an instance-specific plan, model._tp_plan resolves to the same class-level plan as cls._tp_plan
• model_loading_utils.py preserves the loader-selected dtype for replicated tensors. This prevents Accelerate from undoing FP32 module exceptions, such as Cosmos timestep weights during BF16 loading.

Both changes are covered by focused regression tests. Non-TP loading behavior is 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

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• 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?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

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]

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

[github-actions]

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?

[sayakpaul]

What is happening here?

[yzhautouskay]

Before the change during Cosmos BF16 TP loading, the loader correctly kept time_embedder weights in FP32, but Accelerate cast them back to BF16 using the destination meta parameter’s dtype. Other models could hit the same issue whenever those dtypes differed

So passing dtype=param.dtype makes Accelerate preserve the loader's chosen dtype

[sayakpaul]

Can we open a new PR for this before this one is merged with a simple test?

[yzhautouskay]

Created PR #15147