[Cosmos3] Add native tensor parallelism

docs/source/en/api/pipelines/cosmos3.md

@@ -725,6 +725,21 @@ 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
+
+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
 
 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.

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.
 >

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(".")

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, cls._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, cls._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(

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}.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
     # 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,25 @@ 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:
+            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,

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,18 +345,16 @@ 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_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.
@@ -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"])