diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 6c7ea7b5db4d..2cb27bd1cd70 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/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. 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. > 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..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, 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( diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 0295e1c0567b..06fd0967aea1 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}.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, diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index de19a7edd212..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,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"]) 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