From a1b567f4eb7cb7dce13905b459add2380afc110b Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Tue, 27 Oct 2026 04:41:41 -0700
Subject: [PATCH 1/4] Add native tensor parallelism for Cosmos 3

---
 docs/source/en/api/pipelines/cosmos3.md       |  88 ++++++++++
 src/diffusers/models/model_loading_utils.py   |   2 +-
 src/diffusers/models/modeling_utils.py        |   4 +-
 .../transformers/transformer_cosmos3.py       |  34 +++-
 tests/models/test_model_loading_utils.py      |  71 ++++++++
 .../test_cosmos3_tensor_parallel.py           | 164 ++++++++++++++++++
 6 files changed, 353 insertions(+), 10 deletions(-)
 create mode 100644 tests/models/test_model_loading_utils.py
 create mode 100644 tests/models/transformers/test_cosmos3_tensor_parallel.py

diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index 6c7ea7b5db4d..988d71ba3bba 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -725,6 +725,94 @@ pipe.transformer.compile_repeated_blocks(fullgraph=True)
 SeaCache also supports the Cosmos 3 Ulysses context-parallel and DTensor-based tensor-parallel helpers documented
 below. Cache decisions are synchronized across ranks. Full-model compilation and other model-sharding strategies are not claimed.
 
+## Native tensor parallelism
+
+`Cosmos3OmniTransformer` supports native Diffusers tensor parallelism (TP). Passing a
+[`TensorParallelConfig`] to `from_pretrained` reads each rank's weight slices directly from the
+checkpoint. This reduces loading memory compared with loading the full transformer on each GPU
+before calling `enable_parallelism`.
+
+The TP degree must divide both the query and KV head counts and the MLP intermediate size.
+The released Nano and Super checkpoints support degrees 1, 2, 4, and 8. Both attention streams
+and their MLP projections are sharded; embeddings, normalization, timestep embeddings, and
+modality input/output projections remain replicated.
+
+The following example adds native TP to the Super [text-to-video](#text-to-video) workflow
+with [`Cosmos3OmniModularPipeline`], using the same prompt files, scheduler, and generation settings.
+Save it as `cosmos3_native_tp.py` and launch it with
+`torchrun --standalone --nproc-per-node=4 cosmos3_native_tp.py`:
+
+```python
+import json
+import os
+
+import torch
+import torch.distributed as dist
+
+from diffusers import Cosmos3OmniModularPipeline, Cosmos3OmniTransformer
+from diffusers.models._modeling_parallel import TensorParallelConfig
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import export_to_video
+
+local_rank = int(os.environ["LOCAL_RANK"])
+torch.cuda.set_device(local_rank)
+dist.init_process_group("nccl")
+
+json_prompt = json.load(open("assets/example_t2v_prompt.json"))
+negative_prompt = json.load(open("assets/negative_prompt.json"))
+
+tp_degree = 4
+model_id = "nvidia/Cosmos3-Super"
+transformer = Cosmos3OmniTransformer.from_pretrained(
+    model_id,
+    subfolder="transformer",
+    dtype=torch.bfloat16,
+    parallel_config=TensorParallelConfig(tp_degree=tp_degree),
+)
+pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id)
+pipe.update_components(transformer=transformer)
+pipe.load_components(dtype=torch.bfloat16)
+pipe.enable_safety_checker()
+pipe.to(torch.device("cuda", local_rank))
+pipe.update_components(
+    scheduler=UniPCMultistepScheduler.from_config(
+        pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
+    )
+)
+
+videos = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    num_frames=189,
+    height=720,
+    width=1280,
+    num_inference_steps=35,
+    guidance_scale=6.0,
+    fps=24.0,
+    generator=torch.Generator().manual_seed(42),
+    output="videos",
+)
+if dist.get_rank() == 0:
+    export_to_video(videos, "cosmos3_modular_t2v.mp4", fps=24, macro_block_size=1)
+dist.destroy_process_group()
+```
+
+Set `model_id="nvidia/Cosmos3-Nano"` to use Nano with the same code. The transformer is loaded
+separately from its `transformer` subfolder so its weights can be sharded during loading.
+`update_components` installs it in the modular pipeline; `load_components` resolves and loads
+the remaining components without replacing that transformer.
+
+This example uses a TP degree of 4. The TP degree is the size of the TP group, which need not
+span the entire distributed world. For a TP group within a larger job, pass its one-dimensional
+`DeviceMesh` to `TensorParallelConfig(mesh=tp_mesh)`.
+
+Use the same prompt and seed on every rank in the TP group. Native TP uses the standard Cosmos attention
+processor with local head counts inferred from projection widths. The existing example TP/CP
+helpers and their CLI commands remain available; do not apply the example TP helper to a
+transformer that is already sharded with native TP. Native TP+CP composition is not supported
+here. The generic [TP restrictions](../../training/distributed_inference#tensor-parallelism),
+including quantization, offloading, LoRA, and saving a sharded model, also apply.
+
 ## Context parallelism
 
 For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens.
diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py
index b6d9d7997d81..4fd60b73cf19 100644
--- a/src/diffusers/models/model_loading_utils.py
+++ b/src/diffusers/models/model_loading_utils.py
@@ -453,7 +453,7 @@ def _load_shard_file_tp(
                     param = param.to(dtype)
 
             if spec is None:
-                set_module_tensor_to_device(model, key, device, value=param)
+                set_module_tensor_to_device(model, key, device, value=param, dtype=param.dtype)
                 continue
 
             path, _, param_name = key.rpartition(".")
diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
index 469906613a5e..7dcbb5448e63 100644
--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -1388,7 +1388,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
 
             parallel_config = model._resolve_parallel_config(parallel_config)
             tp_config = parallel_config.tensor_parallel_config
-            tp_shard_specs = resolve_tp_shard_specs(model, cls._tp_plan, tp_config._mesh.size())
+            tp_shard_specs = resolve_tp_shard_specs(model, model._tp_plan, tp_config._mesh.size())
             # Each rank opens every shard file but only reads its own slices, so threading the files buys
             # nothing and would have several threads calling `register_parameter` on the same modules.
             if is_parallel_loading_enabled:
@@ -1546,7 +1546,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
                 # was recorded by `_resolve_parallel_config` before loading.
                 from ..hooks.tensor_parallel import apply_tensor_parallel
 
-                apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
+                apply_tensor_parallel(model, tp_config, model._tp_plan, weights_already_sharded=True)
             else:
                 model.enable_parallelism(config=parallel_config)
 
diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py
index 0295e1c0567b..a928c05ad58c 100644
--- a/src/diffusers/models/transformers/transformer_cosmos3.py
+++ b/src/diffusers/models/transformers/transformer_cosmos3.py
@@ -21,6 +21,7 @@
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import PeftAdapterMixin
 from ...utils import BaseOutput
+from .._modeling_parallel import TensorParallelConfig
 from ..attention import AttentionMixin, AttentionModuleMixin
 from ..attention_dispatch import dispatch_attention_fn
 from ..cache_utils import CacheMixin
@@ -64,13 +65,13 @@ def __call__(
         gen_seq: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Per-pathway projections
-        q_und = attn.to_q(und_seq).view(-1, attn.num_attention_heads, attn.head_dim)
-        k_und = attn.to_k(und_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
-        v_und = attn.to_v(und_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
-        q_gen = attn.add_q_proj(gen_seq).view(-1, attn.num_attention_heads, attn.head_dim)
-        k_gen = attn.add_k_proj(gen_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
-        v_gen = attn.add_v_proj(gen_seq).view(-1, attn.num_key_value_heads, attn.head_dim)
+        # Infer local head counts from the projection widths after tensor parallel sharding.
+        q_und = attn.to_q(und_seq).unflatten(-1, (-1, attn.head_dim))
+        k_und = attn.to_k(und_seq).unflatten(-1, (-1, attn.head_dim))
+        v_und = attn.to_v(und_seq).unflatten(-1, (-1, attn.head_dim))
+        q_gen = attn.add_q_proj(gen_seq).unflatten(-1, (-1, attn.head_dim))
+        k_gen = attn.add_k_proj(gen_seq).unflatten(-1, (-1, attn.head_dim))
+        v_gen = attn.add_v_proj(gen_seq).unflatten(-1, (-1, attn.head_dim))
 
         q_und = attn.norm_q(q_und)
         k_und = attn.norm_k(k_und)
@@ -377,6 +378,15 @@ class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attentio
     _repeated_blocks = ["Cosmos3VLTextMoTDecoderLayer"]
     _skip_layerwise_casting_patterns = ["embed_tokens", "time_embedder", "norm"]
     _keep_in_fp32_modules = ["time_embedder"]
+    _tp_plan = {
+        **{
+            f"layers.*.self_attn.{name}": "colwise"
+            for name in ("to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj")
+        },
+        **{f"layers.*.self_attn.{name}": "rowwise" for name in ("to_out", "to_add_out")},
+        **{f"layers.*.{mlp}.{name}": "colwise" for mlp in ("mlp", "mlp_moe_gen") for name in ("gate_proj", "up_proj")},
+        **{f"layers.*.{mlp}.down_proj": "rowwise" for mlp in ("mlp", "mlp_moe_gen")},
+    }
     # Optional context-parallelism seams. They default to ``None`` (no-op) so the
     # model itself carries no CP logic. `forward` applies `_cp_shard_fn` to the
     # per-pathway hidden states + rotary embeddings before the decoder layers, and
@@ -392,6 +402,14 @@ class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attentio
     # config.json — the value is a load-time runtime hint, not part of the model architecture.
     ignore_for_config = ["dtype"]
 
+    def _resolve_parallel_config(self, config):
+        tp = config if isinstance(config, TensorParallelConfig) else getattr(config, "tensor_parallel_config", None)
+        if tp is not None:
+            degree = tp.mesh.size() if tp.mesh is not None else tp.tp_degree
+            if self.config.num_attention_heads % degree or self.config.num_key_value_heads % degree:
+                raise ValueError("Cosmos3 TP degree must divide both query and KV head counts.")
+        return super()._resolve_parallel_config(config)
+
     @register_to_config
     def __init__(
         self,
@@ -438,6 +456,8 @@ def __init__(
         # Text-model layers live directly on the transformer (flat layout). The published
         # checkpoint must be re-keyed with the leading `model.` prefix stripped — see
         # scripts/build_flat_layout_repo.py for the rewrite.
+        if hidden_act == "relu2":
+            self._tp_plan = {name: style for name, style in self._tp_plan.items() if not name.endswith("gate_proj")}
         self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
         self.layers = nn.ModuleList(
             [
diff --git a/tests/models/test_model_loading_utils.py b/tests/models/test_model_loading_utils.py
new file mode 100644
index 000000000000..8b47a0241199
--- /dev/null
+++ b/tests/models/test_model_loading_utils.py
@@ -0,0 +1,71 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+import torch.distributed as dist
+from safetensors.torch import save_file
+from torch.distributed.device_mesh import init_device_mesh
+
+from diffusers.models._modeling_parallel import TensorParallelConfig
+from diffusers.models.model_loading_utils import _load_shard_file_tp
+
+
+@pytest.fixture
+def cpu_tp_config(tmp_path):
+    if not dist.is_available() or not dist.is_gloo_available():
+        pytest.skip("This loader test requires Gloo.")
+    dist.init_process_group("gloo", init_method=(tmp_path / "rendezvous").as_uri(), rank=0, world_size=1)
+    try:
+        mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("tp",))
+        config = TensorParallelConfig(mesh=mesh)
+        config.setup(rank=0, world_size=1, device=torch.device("cpu"), mesh=mesh)
+        yield config
+    finally:
+        dist.destroy_process_group()
+
+
+class TestTensorParallelReplicatedLoading:
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_preserves_requested_and_fp32_dtypes(self, tmp_path, cpu_tp_config, dtype):
+        with torch.device("meta"):
+            model = torch.nn.ModuleDict(
+                {
+                    "time_embedder": torch.nn.Linear(4, 4, dtype=torch.bfloat16),
+                    "projection": torch.nn.Linear(4, 4, dtype=torch.bfloat16),
+                }
+            )
+        model.register_buffer("indexes", torch.empty(4, device="meta", dtype=torch.int64))
+        generator = torch.Generator().manual_seed(0)
+        weights = {
+            name: torch.randn(parameter.shape, generator=generator) for name, parameter in model.named_parameters()
+        }
+        weights["indexes"] = torch.arange(4)
+        checkpoint = tmp_path / "weights.safetensors"
+        save_file(weights, checkpoint)
+        _load_shard_file_tp(
+            str(checkpoint),
+            model,
+            model.state_dict(),
+            {},
+            cpu_tp_config,
+            dtype=dtype,
+            keep_in_fp32_modules=["time_embedder"],
+            unexpected_keys=[],
+        )
+        for name, parameter in model.named_parameters():
+            expected_dtype = torch.float32 if name.startswith("time_embedder.") else dtype
+            assert parameter.dtype == expected_dtype
+            torch.testing.assert_close(parameter, weights[name].to(expected_dtype), atol=0, rtol=0)
+        torch.testing.assert_close(model.indexes, weights["indexes"], atol=0, rtol=0)
diff --git a/tests/models/transformers/test_cosmos3_tensor_parallel.py b/tests/models/transformers/test_cosmos3_tensor_parallel.py
new file mode 100644
index 000000000000..9ae1e20caf81
--- /dev/null
+++ b/tests/models/transformers/test_cosmos3_tensor_parallel.py
@@ -0,0 +1,164 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from datetime import timedelta
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.distributed.tensor import DTensor
+
+from diffusers import Cosmos3OmniTransformer
+from diffusers.hooks.tensor_parallel import resolve_tp_shard_specs
+from diffusers.models._modeling_parallel import TensorParallelConfig
+from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME
+
+from ...testing_utils import is_tensor_parallel, require_torch_multi_accelerator, torch_device
+from .test_models_transformer_cosmos3 import Cosmos3OmniTransformerTesterConfig
+
+
+class Cosmos3TensorParallelTesterConfig(Cosmos3OmniTransformerTesterConfig):
+    def get_init_dict(self):
+        return {
+            **super().get_init_dict(),
+            "head_dim": 8,
+            "hidden_size": 64,
+            "intermediate_size": 128,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 4,
+            "rope_axes_dim": [2, 1, 1],
+            "attention_bias": True,
+        }
+
+    def get_dummy_inputs(self, und_len=5, gen_len=7, dtype=torch.float32):
+        inputs = super().get_dummy_inputs(width=gen_len)
+        device = inputs["input_ids"].device
+        indexes = torch.arange(und_len, und_len + gen_len, device=device)
+        inputs.update(
+            input_ids=torch.arange(und_len, device=device) % self.get_init_dict()["vocab_size"],
+            text_indexes=torch.arange(und_len, device=device),
+            position_ids=torch.arange(und_len + gen_len, device=device).expand(3, -1),
+            und_len=und_len,
+            sequence_length=und_len + gen_len,
+            vision_tokens=[tokens.to(dtype=dtype) for tokens in inputs["vision_tokens"]],
+            vision_sequence_indexes=indexes,
+            vision_mse_loss_indexes=indexes,
+        )
+        return inputs
+
+    @pytest.fixture(params=["silu", "relu2"])
+    def model_config(self, request):
+        activation = request.param
+        return {
+            **self.get_init_dict(),
+            "hidden_act": activation,
+            "qk_norm_for_text": activation == "silu",
+            "use_und_k_norm_for_gen": activation == "relu2",
+        }
+
+
+def _cosmos_tp_worker(rank, world_size, rendezvous, checkpoint, dtype):
+    torch.cuda.set_device(rank)
+    device = torch.device("cuda", rank)
+    dist.init_process_group(
+        "nccl", init_method=rendezvous, rank=rank, world_size=world_size, timeout=timedelta(minutes=3)
+    )
+    try:
+        torch.manual_seed(2026)
+        reference = Cosmos3OmniTransformer.from_pretrained(checkpoint, dtype=dtype).to(device).eval()
+        native = copy.deepcopy(reference)
+        specs = resolve_tp_shard_specs(native, native._tp_plan, world_size)
+        # Planned weights come from rank zero even when other ranks start with different values.
+        if rank != 0:
+            with torch.no_grad():
+                for name, parameter in native.named_parameters():
+                    if name in specs:
+                        parameter.add_(1)
+        native.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
+        loaded = Cosmos3OmniTransformer.from_pretrained(
+            checkpoint, dtype=dtype, parallel_config=TensorParallelConfig(tp_degree=world_size)
+        ).eval()
+
+        assert loaded.rotary_emb.inv_freq.dtype == torch.float32
+        torch.testing.assert_close(loaded.rotary_emb.inv_freq.to(device), reference.rotary_emb.inv_freq)
+        for name, parameter in loaded.named_parameters():
+            other = native.get_parameter(name)
+            assert isinstance(parameter, DTensor) == (name in specs)
+            assert parameter.device == device
+            if isinstance(parameter, DTensor):
+                assert parameter.placements == other.placements
+                parameter, other = parameter.to_local(), other.to_local()
+            torch.testing.assert_close(parameter, other, atol=0, rtol=0)
+            if name.startswith("time_embedder."):
+                assert parameter.dtype == torch.float32
+
+        tolerance = 0.025 if dtype == torch.bfloat16 else 3e-5
+        with torch.no_grad(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+            for u, g in ((5, 7), (8, 12), (7, 5)):
+                inputs = Cosmos3TensorParallelTesterConfig().get_dummy_inputs(und_len=u, gen_len=g, dtype=dtype)
+                expected = reference(**inputs).sample[0]
+                actual = native(**inputs).sample[0]
+                torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance)
+                torch.testing.assert_close(loaded(**inputs).sample[0], actual, atol=0, rtol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+class TestCosmos3TensorParallelConfig(Cosmos3TensorParallelTesterConfig):
+    def test_plan_covers_both_streams(self, model_config):
+        model = self.model_class(**model_config)
+        specs = resolve_tp_shard_specs(model, model._tp_plan, 2)
+        for name, parameter in model.named_parameters():
+            if name.startswith("layers.") and parameter.ndim == 2:
+                assert name in specs
+        assert "embed_tokens.weight" not in specs
+        assert any("gate_proj" in key for key in specs) == (model_config["hidden_act"] == "silu")
+        assert any("gate_proj" in key for key in Cosmos3OmniTransformer._tp_plan)
+
+    @pytest.mark.parametrize("heads,kv_heads,degree", [(8, 1, 2), (6, 4, 4), (8, 4, 3)])
+    def test_invalid_head_split(self, heads, kv_heads, degree):
+        config = self.get_init_dict()
+        config.update(num_attention_heads=heads, num_key_value_heads=kv_heads)
+        model = Cosmos3OmniTransformer(**config)
+        with pytest.raises(ValueError, match="both query and KV"):
+            model._resolve_parallel_config(TensorParallelConfig(tp_degree=degree))
+        assert model._parallel_config is None
+
+
+@is_tensor_parallel
+@require_torch_multi_accelerator
+class TestCosmos3TensorParallelLoading(Cosmos3TensorParallelTesterConfig):
+    """Cosmos' nested modality inputs and list outputs require a worker beyond the flat-tensor TP mixin."""
+
+    @pytest.fixture
+    def checkpoint(self, tmp_path, model_config):
+        torch.manual_seed(2026)
+        model = self.model_class(**model_config)
+        checkpoint = tmp_path / "checkpoint"
+        model.save_pretrained(checkpoint, max_shard_size="100KB")
+        return checkpoint
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_checkpoint_and_loaded_model(self, tmp_path, checkpoint, dtype):
+        if torch_device != "cuda":
+            pytest.skip("This worker uses NCCL.")
+        assert (checkpoint / SAFE_WEIGHTS_INDEX_NAME).is_file()
+        mp.spawn(
+            _cosmos_tp_worker,
+            args=(2, (tmp_path / "rendezvous").as_uri(), str(checkpoint), dtype),
+            nprocs=2,
+            join=True,
+        )

From ef4beace3ab530a2a34922b352e4d88289bcb939 Mon Sep 17 00:00:00 2001
From: Sayak Paul <spsayakpaul@gmail.com>
Date: Wed, 28 Oct 2026 10:44:46 +0000
Subject: [PATCH 2/4] use _get_tp_plan for overrides.

---
 src/diffusers/models/modeling_utils.py            | 15 ++++++++++-----
 .../models/transformers/transformer_cosmos3.py    | 15 ++++++++++++---
 tests/models/testing_utils/parallelism.py         |  2 +-
 .../transformers/test_cosmos3_tensor_parallel.py  |  7 ++++---
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
index 7dcbb5448e63..b0211c453eb2 100644
--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -1388,7 +1388,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
 
             parallel_config = model._resolve_parallel_config(parallel_config)
             tp_config = parallel_config.tensor_parallel_config
-            tp_shard_specs = resolve_tp_shard_specs(model, model._tp_plan, tp_config._mesh.size())
+            tp_plan = model._get_tp_plan()
+            tp_shard_specs = resolve_tp_shard_specs(model, tp_plan, tp_config._mesh.size())
             # Each rank opens every shard file but only reads its own slices, so threading the files buys
             # nothing and would have several threads calling `register_parameter` on the same modules.
             if is_parallel_loading_enabled:
@@ -1546,7 +1547,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
                 # was recorded by `_resolve_parallel_config` before loading.
                 from ..hooks.tensor_parallel import apply_tensor_parallel
 
-                apply_tensor_parallel(model, tp_config, model._tp_plan, weights_already_sharded=True)
+                apply_tensor_parallel(model, tp_config, tp_plan, weights_already_sharded=True)
             else:
                 model.enable_parallelism(config=parallel_config)
 
@@ -1739,6 +1740,9 @@ def _resolve_parallel_config(
         self._parallel_config = config
         return config
 
+    def _get_tp_plan(self):
+        return self._tp_plan
+
     def enable_parallelism(
         self,
         *,
@@ -1767,16 +1771,17 @@ def enable_parallelism(
         if tp_config is not None:
             from ..hooks.tensor_parallel import _check_tp_supported, _tp_degree, resolve_tp_shard_specs
 
+            tp_plan = self._get_tp_plan()
             # Before `_resolve_parallel_config`, which records the config on the model: a model that fails these
             # checks is left untouched.
             _check_tp_supported(
                 self.__class__.__name__,
-                self._tp_plan,
+                tp_plan,
                 getattr(self.config, "num_attention_heads", None),
                 tp_config,
                 model=self,
             )
-            resolve_tp_shard_specs(self, self._tp_plan, _tp_degree(tp_config))
+            resolve_tp_shard_specs(self, tp_plan, _tp_degree(tp_config))
 
         config = self._resolve_parallel_config(config)
 
@@ -1832,7 +1837,7 @@ def enable_parallelism(
         if config.tensor_parallel_config is not None:
             from ..hooks.tensor_parallel import apply_tensor_parallel
 
-            apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
+            apply_tensor_parallel(self, config.tensor_parallel_config, tp_plan)
 
     @classmethod
     def _load_pretrained_model(
diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py
index a928c05ad58c..06fd0967aea1 100644
--- a/src/diffusers/models/transformers/transformer_cosmos3.py
+++ b/src/diffusers/models/transformers/transformer_cosmos3.py
@@ -384,7 +384,7 @@ class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attentio
             for name in ("to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj")
         },
         **{f"layers.*.self_attn.{name}": "rowwise" for name in ("to_out", "to_add_out")},
-        **{f"layers.*.{mlp}.{name}": "colwise" for mlp in ("mlp", "mlp_moe_gen") for name in ("gate_proj", "up_proj")},
+        **{f"layers.*.{mlp}.up_proj": "colwise" for mlp in ("mlp", "mlp_moe_gen")},
         **{f"layers.*.{mlp}.down_proj": "rowwise" for mlp in ("mlp", "mlp_moe_gen")},
     }
     # Optional context-parallelism seams. They default to ``None`` (no-op) so the
@@ -402,6 +402,17 @@ class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attentio
     # config.json — the value is a load-time runtime hint, not part of the model architecture.
     ignore_for_config = ["dtype"]
 
+    def _get_tp_plan(self):
+        tp_plan = super()._get_tp_plan().copy()
+        if self.config.hidden_act == "silu":
+            tp_plan.update(
+                {
+                    "layers.*.mlp.gate_proj": "colwise",
+                    "layers.*.mlp_moe_gen.gate_proj": "colwise",
+                }
+            )
+        return tp_plan
+
     def _resolve_parallel_config(self, config):
         tp = config if isinstance(config, TensorParallelConfig) else getattr(config, "tensor_parallel_config", None)
         if tp is not None:
@@ -456,8 +467,6 @@ def __init__(
         # Text-model layers live directly on the transformer (flat layout). The published
         # checkpoint must be re-keyed with the leading `model.` prefix stripped — see
         # scripts/build_flat_layout_repo.py for the rewrite.
-        if hidden_act == "relu2":
-            self._tp_plan = {name: style for name, style in self._tp_plan.items() if not name.endswith("gate_proj")}
         self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
         self.layers = nn.ModuleList(
             [
diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py
index de19a7edd212..8bc6c0485386 100644
--- a/tests/models/testing_utils/parallelism.py
+++ b/tests/models/testing_utils/parallelism.py
@@ -347,7 +347,7 @@ def _tensor_parallel_from_pretrained_worker(
             output = output.full_tensor()
 
         if rank == 0:
-            specs = resolve_tp_shard_specs(model, model_class._tp_plan, world_size)
+            specs = resolve_tp_shard_specs(model, model._get_tp_plan(), world_size)
             state_dict = model.state_dict()
 
             # The numerics check alone wouldn't catch a planned parameter left unsharded.
diff --git a/tests/models/transformers/test_cosmos3_tensor_parallel.py b/tests/models/transformers/test_cosmos3_tensor_parallel.py
index 9ae1e20caf81..a0dcbfa313a5 100644
--- a/tests/models/transformers/test_cosmos3_tensor_parallel.py
+++ b/tests/models/transformers/test_cosmos3_tensor_parallel.py
@@ -80,7 +80,7 @@ def _cosmos_tp_worker(rank, world_size, rendezvous, checkpoint, dtype):
         torch.manual_seed(2026)
         reference = Cosmos3OmniTransformer.from_pretrained(checkpoint, dtype=dtype).to(device).eval()
         native = copy.deepcopy(reference)
-        specs = resolve_tp_shard_specs(native, native._tp_plan, world_size)
+        specs = resolve_tp_shard_specs(native, native._get_tp_plan(), world_size)
         # Planned weights come from rank zero even when other ranks start with different values.
         if rank != 0:
             with torch.no_grad():
@@ -119,14 +119,15 @@ def _cosmos_tp_worker(rank, world_size, rendezvous, checkpoint, dtype):
 
 class TestCosmos3TensorParallelConfig(Cosmos3TensorParallelTesterConfig):
     def test_plan_covers_both_streams(self, model_config):
+        class_plan = self.model_class._tp_plan.copy()
         model = self.model_class(**model_config)
-        specs = resolve_tp_shard_specs(model, model._tp_plan, 2)
+        specs = resolve_tp_shard_specs(model, model._get_tp_plan(), 2)
         for name, parameter in model.named_parameters():
             if name.startswith("layers.") and parameter.ndim == 2:
                 assert name in specs
         assert "embed_tokens.weight" not in specs
         assert any("gate_proj" in key for key in specs) == (model_config["hidden_act"] == "silu")
-        assert any("gate_proj" in key for key in Cosmos3OmniTransformer._tp_plan)
+        assert self.model_class._tp_plan == class_plan
 
     @pytest.mark.parametrize("heads,kv_heads,degree", [(8, 1, 2), (6, 4, 4), (8, 4, 3)])
     def test_invalid_head_split(self, heads, kv_heads, degree):

From 77ed5c4d2666fc49e2ec230347ae0083f2ab75b1 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Wed, 28 Oct 2026 07:26:13 -0700
Subject: [PATCH 3/4] Refactor Cosmos 3 tensor parallel tests

---
 tests/models/test_model_loading_utils.py      |  71 --------
 tests/models/testing_utils/parallelism.py     |  34 ++--
 .../test_cosmos3_tensor_parallel.py           | 165 ------------------
 .../test_models_transformer_cosmos3.py        | 120 +++++++++++++
 4 files changed, 141 insertions(+), 249 deletions(-)
 delete mode 100644 tests/models/test_model_loading_utils.py
 delete mode 100644 tests/models/transformers/test_cosmos3_tensor_parallel.py

diff --git a/tests/models/test_model_loading_utils.py b/tests/models/test_model_loading_utils.py
deleted file mode 100644
index 8b47a0241199..000000000000
--- a/tests/models/test_model_loading_utils.py
+++ /dev/null
@@ -1,71 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import pytest
-import torch
-import torch.distributed as dist
-from safetensors.torch import save_file
-from torch.distributed.device_mesh import init_device_mesh
-
-from diffusers.models._modeling_parallel import TensorParallelConfig
-from diffusers.models.model_loading_utils import _load_shard_file_tp
-
-
-@pytest.fixture
-def cpu_tp_config(tmp_path):
-    if not dist.is_available() or not dist.is_gloo_available():
-        pytest.skip("This loader test requires Gloo.")
-    dist.init_process_group("gloo", init_method=(tmp_path / "rendezvous").as_uri(), rank=0, world_size=1)
-    try:
-        mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("tp",))
-        config = TensorParallelConfig(mesh=mesh)
-        config.setup(rank=0, world_size=1, device=torch.device("cpu"), mesh=mesh)
-        yield config
-    finally:
-        dist.destroy_process_group()
-
-
-class TestTensorParallelReplicatedLoading:
-    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-    def test_preserves_requested_and_fp32_dtypes(self, tmp_path, cpu_tp_config, dtype):
-        with torch.device("meta"):
-            model = torch.nn.ModuleDict(
-                {
-                    "time_embedder": torch.nn.Linear(4, 4, dtype=torch.bfloat16),
-                    "projection": torch.nn.Linear(4, 4, dtype=torch.bfloat16),
-                }
-            )
-        model.register_buffer("indexes", torch.empty(4, device="meta", dtype=torch.int64))
-        generator = torch.Generator().manual_seed(0)
-        weights = {
-            name: torch.randn(parameter.shape, generator=generator) for name, parameter in model.named_parameters()
-        }
-        weights["indexes"] = torch.arange(4)
-        checkpoint = tmp_path / "weights.safetensors"
-        save_file(weights, checkpoint)
-        _load_shard_file_tp(
-            str(checkpoint),
-            model,
-            model.state_dict(),
-            {},
-            cpu_tp_config,
-            dtype=dtype,
-            keep_in_fp32_modules=["time_embedder"],
-            unexpected_keys=[],
-        )
-        for name, parameter in model.named_parameters():
-            expected_dtype = torch.float32 if name.startswith("time_embedder.") else dtype
-            assert parameter.dtype == expected_dtype
-            torch.testing.assert_close(parameter, weights[name].to(expected_dtype), atol=0, rtol=0)
-        torch.testing.assert_close(model.indexes, weights["indexes"], atol=0, rtol=0)
diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py
index 8bc6c0485386..029085683920 100644
--- a/tests/models/testing_utils/parallelism.py
+++ b/tests/models/testing_utils/parallelism.py
@@ -20,6 +20,7 @@
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
+from accelerate.utils import send_to_device
 
 from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig
 from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
@@ -248,6 +249,14 @@ def _custom_mesh_worker(
             dist.destroy_process_group()
 
 
+def _tensor_parallel_output(output):
+    from torch.distributed.tensor import DTensor
+
+    if isinstance(output, (list, tuple)):
+        output = torch.cat([sample.flatten() for sample in output])
+    return output.full_tensor() if isinstance(output, DTensor) else output
+
+
 def _tensor_parallel_worker(
     rank,
     world_size,
@@ -285,17 +294,17 @@ def _tensor_parallel_worker(
             device = torch.device(f"{device_type}:{rank}")
 
         model = model_class(**init_dict)
-        model.load_state_dict(state_dict)
+        model.load_state_dict(state_dict, assign=True)
         model.to(device)
         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)
 
         # Shard the model across all ranks; the device mesh is built from `tp_degree` on the active accelerator.
         model.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
 
         with torch.no_grad():
-            output = model(**inputs_on_device, return_dict=False)[0]
+            output = _tensor_parallel_output(model(**inputs_on_device, return_dict=False)[0])
 
         if rank == 0:
             return_dict["status"] = "success"
@@ -313,7 +322,7 @@ def _tensor_parallel_worker(
 
 
 def _tensor_parallel_from_pretrained_worker(
-    rank, world_size, master_port, model_class, checkpoint_dir, inputs_dict, return_dict
+    rank, world_size, master_port, model_class, checkpoint_dir, inputs_dict, return_dict, dtype=None
 ):
     """Worker for `from_pretrained(..., parallel_config=...)`, i.e. sharding while reading the checkpoint.
 
@@ -336,15 +345,13 @@ def _tensor_parallel_from_pretrained_worker(
         from diffusers.hooks.tensor_parallel import resolve_tp_shard_specs
 
         model = model_class.from_pretrained(
-            checkpoint_dir, parallel_config=TensorParallelConfig(tp_degree=world_size)
+            checkpoint_dir, dtype=dtype, parallel_config=TensorParallelConfig(tp_degree=world_size)
         ).eval()
 
         device = torch.device(f"{torch_device}:{rank}")
-        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)
         with torch.no_grad():
-            output = model(**inputs_on_device, return_dict=False)[0]
-        if isinstance(output, DTensor):
-            output = output.full_tensor()
+            output = _tensor_parallel_output(model(**inputs_on_device, return_dict=False)[0])
 
         if rank == 0:
             specs = resolve_tp_shard_specs(model, model._get_tp_plan(), world_size)
@@ -369,6 +376,7 @@ def _tensor_parallel_from_pretrained_worker(
 
             unplanned = sorted(k for k, v in state_dict.items() if isinstance(v, DTensor) and k not in specs)
             assert not unplanned, f"Parameters not covered by `_tp_plan` were loaded as DTensors: {unplanned}"
+            return_dict["parameter_dtypes"] = {name: parameter.dtype for name, parameter in model.named_parameters()}
 
             return_dict["status"] = "success"
             return_dict["output"] = output.float().cpu().tolist()
@@ -404,10 +412,10 @@ def test_tensor_parallel_inference(self, batch_size: int = 1):
         model = self.model_class(**init_dict).eval().to(torch_device)
         state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
         with torch.no_grad():
-            ref_output = model(**inputs_dict, return_dict=False)[0].float().cpu()
+            ref_output = _tensor_parallel_output(model(**inputs_dict, return_dict=False)[0]).float().cpu()
 
         # Move all tensors to CPU for multiprocessing
-        inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+        inputs_dict = send_to_device(inputs_dict, "cpu")
 
         master_port = _find_free_port()
         manager = mp.Manager()
@@ -450,7 +458,7 @@ def _tp_checkpoint_and_reference(self, tmp_path, world_size, sharded):
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**init_dict).eval().to(torch_device)
         with torch.no_grad():
-            reference = model(**inputs_dict, return_dict=False)[0].float().cpu()
+            reference = _tensor_parallel_output(model(**inputs_dict, return_dict=False)[0]).float().cpu()
 
         checkpoint_dir = str(tmp_path / "checkpoint")
         max_shard_size = int(compute_module_persistent_sizes(model)[""] * 0.75) if sharded else "5GB"
@@ -463,7 +471,7 @@ def _tp_checkpoint_and_reference(self, tmp_path, world_size, sharded):
             actual_num_shards = len([file for file in os.listdir(checkpoint_dir) if file.endswith(".safetensors")])
             assert actual_num_shards == expected_num_shards > 1
 
-        inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+        inputs_dict = send_to_device(inputs_dict, "cpu")
         return checkpoint_dir, inputs_dict, reference
 
     @pytest.mark.parametrize("sharded", [False, True], ids=["unsharded", "sharded"])
diff --git a/tests/models/transformers/test_cosmos3_tensor_parallel.py b/tests/models/transformers/test_cosmos3_tensor_parallel.py
deleted file mode 100644
index a0dcbfa313a5..000000000000
--- a/tests/models/transformers/test_cosmos3_tensor_parallel.py
+++ /dev/null
@@ -1,165 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import copy
-from datetime import timedelta
-
-import pytest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from torch.distributed.tensor import DTensor
-
-from diffusers import Cosmos3OmniTransformer
-from diffusers.hooks.tensor_parallel import resolve_tp_shard_specs
-from diffusers.models._modeling_parallel import TensorParallelConfig
-from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME
-
-from ...testing_utils import is_tensor_parallel, require_torch_multi_accelerator, torch_device
-from .test_models_transformer_cosmos3 import Cosmos3OmniTransformerTesterConfig
-
-
-class Cosmos3TensorParallelTesterConfig(Cosmos3OmniTransformerTesterConfig):
-    def get_init_dict(self):
-        return {
-            **super().get_init_dict(),
-            "head_dim": 8,
-            "hidden_size": 64,
-            "intermediate_size": 128,
-            "num_attention_heads": 8,
-            "num_key_value_heads": 4,
-            "rope_axes_dim": [2, 1, 1],
-            "attention_bias": True,
-        }
-
-    def get_dummy_inputs(self, und_len=5, gen_len=7, dtype=torch.float32):
-        inputs = super().get_dummy_inputs(width=gen_len)
-        device = inputs["input_ids"].device
-        indexes = torch.arange(und_len, und_len + gen_len, device=device)
-        inputs.update(
-            input_ids=torch.arange(und_len, device=device) % self.get_init_dict()["vocab_size"],
-            text_indexes=torch.arange(und_len, device=device),
-            position_ids=torch.arange(und_len + gen_len, device=device).expand(3, -1),
-            und_len=und_len,
-            sequence_length=und_len + gen_len,
-            vision_tokens=[tokens.to(dtype=dtype) for tokens in inputs["vision_tokens"]],
-            vision_sequence_indexes=indexes,
-            vision_mse_loss_indexes=indexes,
-        )
-        return inputs
-
-    @pytest.fixture(params=["silu", "relu2"])
-    def model_config(self, request):
-        activation = request.param
-        return {
-            **self.get_init_dict(),
-            "hidden_act": activation,
-            "qk_norm_for_text": activation == "silu",
-            "use_und_k_norm_for_gen": activation == "relu2",
-        }
-
-
-def _cosmos_tp_worker(rank, world_size, rendezvous, checkpoint, dtype):
-    torch.cuda.set_device(rank)
-    device = torch.device("cuda", rank)
-    dist.init_process_group(
-        "nccl", init_method=rendezvous, rank=rank, world_size=world_size, timeout=timedelta(minutes=3)
-    )
-    try:
-        torch.manual_seed(2026)
-        reference = Cosmos3OmniTransformer.from_pretrained(checkpoint, dtype=dtype).to(device).eval()
-        native = copy.deepcopy(reference)
-        specs = resolve_tp_shard_specs(native, native._get_tp_plan(), world_size)
-        # Planned weights come from rank zero even when other ranks start with different values.
-        if rank != 0:
-            with torch.no_grad():
-                for name, parameter in native.named_parameters():
-                    if name in specs:
-                        parameter.add_(1)
-        native.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
-        loaded = Cosmos3OmniTransformer.from_pretrained(
-            checkpoint, dtype=dtype, parallel_config=TensorParallelConfig(tp_degree=world_size)
-        ).eval()
-
-        assert loaded.rotary_emb.inv_freq.dtype == torch.float32
-        torch.testing.assert_close(loaded.rotary_emb.inv_freq.to(device), reference.rotary_emb.inv_freq)
-        for name, parameter in loaded.named_parameters():
-            other = native.get_parameter(name)
-            assert isinstance(parameter, DTensor) == (name in specs)
-            assert parameter.device == device
-            if isinstance(parameter, DTensor):
-                assert parameter.placements == other.placements
-                parameter, other = parameter.to_local(), other.to_local()
-            torch.testing.assert_close(parameter, other, atol=0, rtol=0)
-            if name.startswith("time_embedder."):
-                assert parameter.dtype == torch.float32
-
-        tolerance = 0.025 if dtype == torch.bfloat16 else 3e-5
-        with torch.no_grad(), torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
-            for u, g in ((5, 7), (8, 12), (7, 5)):
-                inputs = Cosmos3TensorParallelTesterConfig().get_dummy_inputs(und_len=u, gen_len=g, dtype=dtype)
-                expected = reference(**inputs).sample[0]
-                actual = native(**inputs).sample[0]
-                torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance)
-                torch.testing.assert_close(loaded(**inputs).sample[0], actual, atol=0, rtol=0)
-    finally:
-        dist.destroy_process_group()
-
-
-class TestCosmos3TensorParallelConfig(Cosmos3TensorParallelTesterConfig):
-    def test_plan_covers_both_streams(self, model_config):
-        class_plan = self.model_class._tp_plan.copy()
-        model = self.model_class(**model_config)
-        specs = resolve_tp_shard_specs(model, model._get_tp_plan(), 2)
-        for name, parameter in model.named_parameters():
-            if name.startswith("layers.") and parameter.ndim == 2:
-                assert name in specs
-        assert "embed_tokens.weight" not in specs
-        assert any("gate_proj" in key for key in specs) == (model_config["hidden_act"] == "silu")
-        assert self.model_class._tp_plan == class_plan
-
-    @pytest.mark.parametrize("heads,kv_heads,degree", [(8, 1, 2), (6, 4, 4), (8, 4, 3)])
-    def test_invalid_head_split(self, heads, kv_heads, degree):
-        config = self.get_init_dict()
-        config.update(num_attention_heads=heads, num_key_value_heads=kv_heads)
-        model = Cosmos3OmniTransformer(**config)
-        with pytest.raises(ValueError, match="both query and KV"):
-            model._resolve_parallel_config(TensorParallelConfig(tp_degree=degree))
-        assert model._parallel_config is None
-
-
-@is_tensor_parallel
-@require_torch_multi_accelerator
-class TestCosmos3TensorParallelLoading(Cosmos3TensorParallelTesterConfig):
-    """Cosmos' nested modality inputs and list outputs require a worker beyond the flat-tensor TP mixin."""
-
-    @pytest.fixture
-    def checkpoint(self, tmp_path, model_config):
-        torch.manual_seed(2026)
-        model = self.model_class(**model_config)
-        checkpoint = tmp_path / "checkpoint"
-        model.save_pretrained(checkpoint, max_shard_size="100KB")
-        return checkpoint
-
-    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-    def test_checkpoint_and_loaded_model(self, tmp_path, checkpoint, dtype):
-        if torch_device != "cuda":
-            pytest.skip("This worker uses NCCL.")
-        assert (checkpoint / SAFE_WEIGHTS_INDEX_NAME).is_file()
-        mp.spawn(
-            _cosmos_tp_worker,
-            args=(2, (tmp_path / "rendezvous").as_uri(), str(checkpoint), dtype),
-            nprocs=2,
-            join=True,
-        )
diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py
index c897d5bfdc69..2588f21c12ae 100644
--- a/tests/models/transformers/test_models_transformer_cosmos3.py
+++ b/tests/models/transformers/test_models_transformer_cosmos3.py
@@ -15,11 +15,15 @@
 
 import pytest
 import torch
+import torch.multiprocessing as mp
+from accelerate.utils import send_to_device
 
 from diffusers import Cosmos3OmniTransformer, SeaCacheConfig
 from diffusers.hooks import sea_cache as sea_cache_module
 from diffusers.hooks._helpers import TransformerBlockRegistry
 from diffusers.hooks.sea_cache import _SEA_CACHE_ROOT_HOOK
+from diffusers.hooks.tensor_parallel import resolve_tp_shard_specs
+from diffusers.models._modeling_parallel import TensorParallelConfig
 from diffusers.models.cache_utils import CacheMixin
 from diffusers.models.transformers.transformer_cosmos3 import (
     Cosmos3NemotronRMSNorm,
@@ -27,6 +31,7 @@
     Cosmos3PackedMoTAttention,
     Cosmos3VLTextMoTDecoderLayer,
 )
+from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
@@ -36,9 +41,17 @@
     MemoryTesterMixin,
     ModelTesterMixin,
     SeaCacheTesterMixin,
+    TensorParallelTesterMixin,
     TorchCompileTesterMixin,
     TrainingTesterMixin,
 )
+from ..testing_utils.common import cast_inputs_to_dtype
+from ..testing_utils.parallelism import (
+    _find_free_port,
+    _tensor_parallel_from_pretrained_worker,
+    _tensor_parallel_output,
+    _tensor_parallel_worker,
+)
 
 
 enable_full_determinism()
@@ -710,3 +723,110 @@ def test_mixed_precision_training(self):
 
 class TestCosmos3OmniTransformerAttention(Cosmos3OmniTransformerTesterConfig, AttentionTesterMixin):
     pass
+
+
+class Cosmos3TensorParallelTesterConfig(Cosmos3OmniTransformerTesterConfig):
+    hidden_act = "relu2"
+
+    def get_init_dict(self):
+        return {
+            **super().get_init_dict(),
+            "head_dim": 8,
+            "hidden_size": 64,
+            "intermediate_size": 128,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 4,
+            "rope_axes_dim": [2, 1, 1],
+            "attention_bias": True,
+            "hidden_act": self.hidden_act,
+            "qk_norm_for_text": self.hidden_act == "silu",
+            "use_und_k_norm_for_gen": self.hidden_act == "relu2",
+        }
+
+    def get_dummy_inputs(self, batch_size=1):
+        inputs = super().get_dummy_inputs(width=7)
+        device = inputs["input_ids"].device
+        und_len, gen_len = 5 * batch_size, 7 * batch_size
+        indexes = torch.arange(und_len, und_len + gen_len, device=device)
+        inputs.update(
+            input_ids=torch.arange(und_len, device=device) % self.get_init_dict()["vocab_size"],
+            text_indexes=torch.arange(und_len, device=device),
+            position_ids=torch.arange(und_len + gen_len, device=device).expand(3, -1),
+            und_len=und_len,
+            sequence_length=und_len + gen_len,
+            vision_tokens=[inputs["vision_tokens"][0] + index for index in range(batch_size)],
+            vision_token_shapes=inputs["vision_token_shapes"] * batch_size,
+            vision_sequence_indexes=indexes,
+            vision_mse_loss_indexes=indexes,
+            vision_timesteps=torch.ones(gen_len, device=device),
+            vision_noisy_frame_indexes=inputs["vision_noisy_frame_indexes"] * batch_size,
+        )
+        return inputs
+
+
+class TestCosmos3OmniTransformerTensorParallelConfig(Cosmos3TensorParallelTesterConfig):
+    @pytest.mark.parametrize("hidden_act", ["silu", "relu2"])
+    def test_plan_covers_both_streams(self, hidden_act):
+        class_plan = self.model_class._tp_plan.copy()
+        self.hidden_act = hidden_act
+        model = self.model_class(**self.get_init_dict())
+        plan = model._get_tp_plan()
+        specs = resolve_tp_shard_specs(model, plan, 2)
+        for name, parameter in model.named_parameters():
+            if name.startswith("layers.") and parameter.ndim == 2:
+                assert name in specs
+        assert "embed_tokens.weight" not in specs
+        assert any("gate_proj" in key for key in specs) == (hidden_act == "silu")
+        plan.clear()
+        assert model._get_tp_plan()
+        assert "_tp_plan" not in model.__dict__
+        assert self.model_class._tp_plan == class_plan
+
+    @pytest.mark.parametrize("heads,kv_heads,degree", [(8, 1, 2), (6, 4, 4), (8, 4, 3)])
+    def test_invalid_head_split(self, heads, kv_heads, degree):
+        config = self.get_init_dict()
+        config.update(num_attention_heads=heads, num_key_value_heads=kv_heads)
+        model = self.model_class(**config)
+        with pytest.raises(ValueError, match="both query and KV"):
+            model._resolve_parallel_config(TensorParallelConfig(tp_degree=degree))
+        assert model._parallel_config is None
+
+
+class TestCosmos3OmniTransformerTensorParallel(Cosmos3TensorParallelTesterConfig, TensorParallelTesterMixin):
+    @pytest.fixture(autouse=True, params=["silu", "relu2"])
+    def activation(self, request):
+        self.hidden_act = request.param
+
+    @pytest.mark.parametrize("sharded", [False, True], ids=["unsharded", "sharded"])
+    def test_tensor_parallel_bfloat16(self, tmp_path, sharded):
+        init_dict = self.get_init_dict()
+        model = self.model_class(**init_dict).eval()
+        model.save_pretrained(tmp_path, max_shard_size="100KB" if sharded else "5GB")
+        assert (tmp_path / SAFE_WEIGHTS_INDEX_NAME).is_file() == sharded
+        reference = self.model_class.from_pretrained(tmp_path, dtype=torch.bfloat16).to(torch_device).eval()
+        expected_dtypes = {name: parameter.dtype for name, parameter in reference.named_parameters()}
+        assert {dtype for name, dtype in expected_dtypes.items() if name.startswith("time_embedder.")} == {
+            torch.float32
+        }
+        assert torch.bfloat16 in expected_dtypes.values()
+        inputs = cast_inputs_to_dtype(self.get_dummy_inputs(), torch.float32, torch.bfloat16)
+        with torch.no_grad():
+            expected = _tensor_parallel_output(reference(**inputs, return_dict=False)[0]).float().cpu()
+        inputs = send_to_device(inputs, "cpu")
+        state_dict = {name: tensor.cpu() for name, tensor in reference.state_dict().items()}
+
+        with mp.Manager() as manager:
+            for loading in (False, True):
+                result = manager.dict()
+                args = (2, _find_free_port(), self.model_class)
+                if loading:
+                    worker = _tensor_parallel_from_pretrained_worker
+                    args += (str(tmp_path), inputs, result, torch.bfloat16)
+                else:
+                    worker = _tensor_parallel_worker
+                    args += (init_dict, inputs, result, state_dict)
+                mp.spawn(worker, args=args, nprocs=2, join=True)
+                assert result.get("status") == "success", result.get("error", "Unknown error")
+                torch.testing.assert_close(expected, torch.tensor(result["output"]), atol=0.025, rtol=0.025)
+                if loading:
+                    assert result["parameter_dtypes"] == expected_dtypes

From 02631f9f41ac9d19800b73cda7afc05be86787b8 Mon Sep 17 00:00:00 2001
From: Yuliya Zhautouskaya <yzhautouskay@nvidia.com>
Date: Wed, 28 Oct 2026 07:59:31 -0700
Subject: [PATCH 4/4] Simplify Cosmos 3 tensor parallelism documentation

---
 docs/source/en/api/pipelines/cosmos3.md       | 97 +++----------------
 .../en/training/distributed_inference.md      |  3 +
 2 files changed, 15 insertions(+), 85 deletions(-)

diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md
index 988d71ba3bba..2cb27bd1cd70 100644
--- a/docs/source/en/api/pipelines/cosmos3.md
+++ b/docs/source/en/api/pipelines/cosmos3.md
@@ -727,91 +727,18 @@ below. Cache decisions are synchronized across ranks. Full-model compilation and
 
 ## Native tensor parallelism
 
-`Cosmos3OmniTransformer` supports native Diffusers tensor parallelism (TP). Passing a
-[`TensorParallelConfig`] to `from_pretrained` reads each rank's weight slices directly from the
-checkpoint. This reduces loading memory compared with loading the full transformer on each GPU
-before calling `enable_parallelism`.
-
-The TP degree must divide both the query and KV head counts and the MLP intermediate size.
-The released Nano and Super checkpoints support degrees 1, 2, 4, and 8. Both attention streams
-and their MLP projections are sharded; embeddings, normalization, timestep embeddings, and
-modality input/output projections remain replicated.
-
-The following example adds native TP to the Super [text-to-video](#text-to-video) workflow
-with [`Cosmos3OmniModularPipeline`], using the same prompt files, scheduler, and generation settings.
-Save it as `cosmos3_native_tp.py` and launch it with
-`torchrun --standalone --nproc-per-node=4 cosmos3_native_tp.py`:
-
-```python
-import json
-import os
-
-import torch
-import torch.distributed as dist
-
-from diffusers import Cosmos3OmniModularPipeline, Cosmos3OmniTransformer
-from diffusers.models._modeling_parallel import TensorParallelConfig
-from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
-from diffusers.utils import export_to_video
-
-local_rank = int(os.environ["LOCAL_RANK"])
-torch.cuda.set_device(local_rank)
-dist.init_process_group("nccl")
-
-json_prompt = json.load(open("assets/example_t2v_prompt.json"))
-negative_prompt = json.load(open("assets/negative_prompt.json"))
-
-tp_degree = 4
-model_id = "nvidia/Cosmos3-Super"
-transformer = Cosmos3OmniTransformer.from_pretrained(
-    model_id,
-    subfolder="transformer",
-    dtype=torch.bfloat16,
-    parallel_config=TensorParallelConfig(tp_degree=tp_degree),
-)
-pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id)
-pipe.update_components(transformer=transformer)
-pipe.load_components(dtype=torch.bfloat16)
-pipe.enable_safety_checker()
-pipe.to(torch.device("cuda", local_rank))
-pipe.update_components(
-    scheduler=UniPCMultistepScheduler.from_config(
-        pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
-    )
-)
-
-videos = pipe(
-    prompt=json.dumps(json_prompt),
-    negative_prompt=json.dumps(negative_prompt),
-    num_frames=189,
-    height=720,
-    width=1280,
-    num_inference_steps=35,
-    guidance_scale=6.0,
-    fps=24.0,
-    generator=torch.Generator().manual_seed(42),
-    output="videos",
-)
-if dist.get_rank() == 0:
-    export_to_video(videos, "cosmos3_modular_t2v.mp4", fps=24, macro_block_size=1)
-dist.destroy_process_group()
-```
-
-Set `model_id="nvidia/Cosmos3-Nano"` to use Nano with the same code. The transformer is loaded
-separately from its `transformer` subfolder so its weights can be sharded during loading.
-`update_components` installs it in the modular pipeline; `load_components` resolves and loads
-the remaining components without replacing that transformer.
-
-This example uses a TP degree of 4. The TP degree is the size of the TP group, which need not
-span the entire distributed world. For a TP group within a larger job, pass its one-dimensional
-`DeviceMesh` to `TensorParallelConfig(mesh=tp_mesh)`.
-
-Use the same prompt and seed on every rank in the TP group. Native TP uses the standard Cosmos attention
-processor with local head counts inferred from projection widths. The existing example TP/CP
-helpers and their CLI commands remain available; do not apply the example TP helper to a
-transformer that is already sharded with native TP. Native TP+CP composition is not supported
-here. The generic [TP restrictions](../../training/distributed_inference#tensor-parallelism),
-including quantization, offloading, LoRA, and saving a sharded model, also apply.
+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).
 
 ## Context parallelism
 
diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md
index 818488b4cfce..dcffd14b2c88 100644
--- a/docs/source/en/training/distributed_inference.md
+++ b/docs/source/en/training/distributed_inference.md
@@ -490,6 +490,9 @@ torchrun --nproc-per-node 4 tensor_parallel_flux.py
 
 `tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.
 
+For a TP group within a larger distributed job, pass its one-dimensional `DeviceMesh` to
+`TensorParallelConfig(mesh=tp_mesh)` instead.
+
 > [!CAUTION]
 > Loading with a tensor-parallel `parallel_config` isn't supported yet with `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, or non-safetensors weights; each raises rather than quietly falling back to loading the full checkpoint.
 >