Skip to content

[Cosmos3] Add native tensor parallelism #15137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yzhautouskay wants to merge 11 commits into
base: main
Choose a base branch
from
Open

[Cosmos3] Add native tensor parallelism #15137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
a1b567f
Add native tensor parallelism for Cosmos 3
yzhautouskay Oct 27, 2026
971bf9a
Merge branch 'main' into yzhautouskay/cosmos3-native-tp-only
sayakpaul Oct 28, 2026
ef4beac
use _get_tp_plan for overrides.
sayakpaul Oct 28, 2026
77ed5c4
Refactor Cosmos 3 tensor parallel tests
yzhautouskay Oct 28, 2026
02631f9
Simplify Cosmos 3 tensor parallelism documentation
yzhautouskay Oct 28, 2026
2e7c7c7
Merge upstream/main into yzhautouskay/cosmos3-native-tp-only
yzhautouskay Nov 5, 2026
96a7dc7
Merge upstream/main into yzhautouskay/cosmos3-native-tp-only
yzhautouskay Nov 5, 2026
f09416b
Fix docs
yzhautouskay Nov 5, 2026
e06b3b4
Refactor tests to address review comments
yzhautouskay Nov 6, 2026
c4fc346
Bring back code snippet to the TP usage docs
yzhautouskay Nov 6, 2026
fe5488a
Refactor Cosmos 3 tensor parallel tests
yzhautouskay Nov 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 33 additions & 1 deletion docs/source/en/api/pipelines/cosmos3.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -765,7 +765,39 @@ torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py \

### Fitting large models with tensor parallelism

CP shards *activations* but replicates every weight on every rank, so it does not reduce a model's weight footprint — a model that doesn't fit on one GPU still won't fit under CP alone. To shard the **weights**, `enable_cosmos3_tensor_parallel(transformer, tp_mesh)` applies Megatron-style tensor parallelism on a second, orthogonal mesh axis:
CP shards *activations* but replicates every weight on every rank, so it does not reduce a model's weight footprint — a model that doesn't fit on one GPU still won't fit under CP alone.

For TP-only inference, follow the [tensor parallelism guide](../../training/distributed_inference#tensor-parallelism)
for distributed setup, launching, seeding, and rank-zero output handling. After initializing the process group
and selecting each rank's CUDA device, load the modular pipeline as follows:

```python
import torch
from diffusers import Cosmos3OmniModularPipeline, Cosmos3OmniTransformer, TensorParallelConfig

model_id = "nvidia/Cosmos3-Super"
device = torch.device("cuda", torch.cuda.current_device())
transformer = Cosmos3OmniTransformer.from_pretrained(
model_id,
subfolder="transformer",
dtype=torch.bfloat16,
parallel_config=TensorParallelConfig(tp_degree=4),
)
pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id)
pipe.update_components(transformer=transformer)
pipe.load_components(dtype=torch.bfloat16)
pipe.enable_safety_checker()
pipe.to(device)
```

Continue with the [existing modular workflow](#modular-examples-for-all-existing-workflows)'s scheduler,
prompts, and generation settings. To use Nano, set `model_id = "nvidia/Cosmos3-Nano"`.

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.

To combine TP with CP, the existing example helper
`enable_cosmos3_tensor_parallel(transformer, tp_mesh)` shards the weights on a second, orthogonal mesh axis:

- The attention and MLP projections are column/row sharded across the TP group (`to_q/to_k/to_v` + `add_q/k/v` and the MLPs' `gate/up` are column-parallel; `to_out/to_add_out` and the MLPs' `down` are row-parallel with an all-reduce). Each rank ends up owning `query_heads / tp` query heads and `kv_heads / tp` KV heads.
- TP composes with CP on a 2-D `(tp, cp)` device mesh: TP splits heads/weights persistently, CP shards the sequence on top. The constraints are `tp` divides the KV heads (8), and `tp * cp` divides the query heads (32 for `Nano`, 64 for `Super`).
Expand Down
15 changes: 10 additions & 5 deletions src/diffusers/models/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1391,7 +1391,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()

@sayakpaul sayakpaul Nov 6, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we isolate the _get_tp_plan() stuff into a separate PR please? That way, the reviewing scope remains specific.

@yzhautouskay yzhautouskay Nov 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Created #15197

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:
Expand Down Expand Up @@ -1560,7 +1561,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
module = model.get_submodule(module_path) if module_path else model
module._buffers[buffer_name] = buffer.to(tp_device)

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)

Expand Down Expand Up @@ -1753,6 +1754,9 @@ def _resolve_parallel_config(
self._parallel_config = config
return config

def _get_tp_plan(self):
return self._tp_plan

def enable_parallelism(
self,
*,
Expand Down Expand Up @@ -1781,16 +1785,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))
logger.warning(
"Sharding an in-memory model costs full checkpoint memory on every rank. Prefer "
"`from_pretrained(..., parallel_config=...)`, which reads only each rank's shard."
Expand Down Expand Up @@ -1850,7 +1855,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(
Expand Down
43 changes: 36 additions & 7 deletions src/diffusers/models/transformers/transformer_cosmos3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down
41 changes: 35 additions & 6 deletions tests/models/testing_utils/parallelism.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,7 @@ def _tensor_parallel_worker(
return_dict,
state_dict,
device_type=torch_device,
assign=False,
):
"""Worker function for tensor parallel inference testing.

Expand All @@ -267,6 +268,8 @@ def _tensor_parallel_worker(
to the unsharded model up to floating-point reduction order).
"""
try:
from torch.distributed.tensor import DTensor

os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = str(master_port)
os.environ["RANK"] = str(rank)
Expand All @@ -285,17 +288,25 @@ 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=assign)
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()}
for k, v in inputs_on_device.items():
if isinstance(v, list):
inputs_on_device[k] = [x.to(device) if isinstance(x, torch.Tensor) else x for x in v]

# 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]
if isinstance(output, (list, tuple)):
# Flatten list outputs to compare predictions with different shapes.
output = torch.cat([sample.flatten() for sample in output])
if isinstance(output, DTensor):
output = output.full_tensor()

if rank == 0:
return_dict["status"] = "success"
Expand All @@ -313,7 +324,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.

Expand All @@ -336,18 +347,23 @@ 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()}
for k, v in inputs_on_device.items():
if isinstance(v, list):
inputs_on_device[k] = [x.to(device) if isinstance(x, torch.Tensor) else x for x in v]
with torch.no_grad():
output = model(**inputs_on_device, return_dict=False)[0]
if isinstance(output, (list, tuple)):
output = torch.cat([sample.flatten() for sample in output])
if isinstance(output, DTensor):
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.
Expand All @@ -369,6 +385,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()
Expand Down Expand Up @@ -404,10 +421,16 @@ 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 = model(**inputs_dict, return_dict=False)[0]
if isinstance(ref_output, (list, tuple)):
ref_output = torch.cat([sample.flatten() for sample in ref_output])
ref_output = ref_output.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()}
for k, v in inputs_dict.items():
if isinstance(v, list):
inputs_dict[k] = [x.cpu() if isinstance(x, torch.Tensor) else x for x in v]

master_port = _find_free_port()
manager = mp.Manager()
Expand Down Expand Up @@ -450,7 +473,10 @@ 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 = model(**inputs_dict, return_dict=False)[0]
if isinstance(reference, (list, tuple)):
reference = torch.cat([sample.flatten() for sample in reference])
reference = reference.float().cpu()

checkpoint_dir = str(tmp_path / "checkpoint")
max_shard_size = int(compute_module_persistent_sizes(model)[""] * 0.75) if sharded else "5GB"
Expand All @@ -464,6 +490,9 @@ def _tp_checkpoint_and_reference(self, tmp_path, world_size, sharded):
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()}
for k, v in inputs_dict.items():
if isinstance(v, list):
inputs_dict[k] = [x.cpu() if isinstance(x, torch.Tensor) else x for x in v]
return checkpoint_dir, inputs_dict, reference

@pytest.mark.parametrize("sharded", [False, True], ids=["unsharded", "sharded"])
Expand Down
Loading
Loading