Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/distributed/test_cosmos3_pre_sharded.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare full and pre-sharded Cosmos loading with real two-rank FSDP."""
+
+from contextlib import ExitStack
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.gpu]
+
+
+class _SDPAAttention(nn.Module):
+    """Use deterministic native attention while testing loading and FSDP."""
+
+    def __init__(self, *, causal, softmax_scale, **kwargs):
+        super().__init__()
+        self.causal = causal
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value):
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+def _load_cosmos_worker(rank, rendezvous, checkpoint_dir, edge):
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+
+    from safetensors.torch import save_file
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+    from vllm.config.load import LoadConfig
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+    from vllm.utils.torch_utils import set_default_torch_dtype
+
+    from vllm_omni.diffusion.distributed import hsdp
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.platforms import current_omni_platform
+
+    device = torch.device("cuda", rank)
+    current_omni_platform.set_device(device)
+    dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=120))
+    try:
+        with ExitStack() as stack:
+            stack.enter_context(set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))))
+            stack.enter_context(set_default_torch_dtype(torch.bfloat16))
+            # HSDP spans both workers, while every vLLM linear uses TP=1.
+            for module in (linear, parameter):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_rank", return_value=0))
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+            for module in (transformer_cosmos3, transformer_cosmos3_edge):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+                stack.enter_context(patch.object(module, "FrameworkAttention", _SDPAAttention))
+            stack.enter_context(patch.object(transformer_cosmos3, "_get_ulysses_state", return_value=(1, 0, None)))
+            stack.enter_context(
+                patch.object(hsdp, "get_world_group", return_value=SimpleNamespace(world_size=2, rank_in_group=rank))
+            )
+
+            model_cls = (
+                transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer
+                if edge
+                else transformer_cosmos3.Cosmos3VFMTransformer
+            )
+            tf_config = dict(
+                hidden_size=16,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=8,
+                intermediate_size=32,
+                vocab_size=32,
+                latent_patch_size=2,
+                latent_channel=48,
+                temporal_compression_factor=4,
+                rope_scaling={"mrope_section": [2, 1, 1]},
+            )
+            if edge:
+                tf_config.update(
+                    backbone_type=transformer_cosmos3_edge.COSMOS3_EDGE_BACKBONE_TYPE,
+                    qk_norm_for_text=False,
+                    use_und_k_norm_for_gen=True,
+                )
+            config = SimpleNamespace(
+                tf_model_config=tf_config,
+                dtype=torch.bfloat16,
+                quantization_config=None,
+                lora_path=None,
+                num_weight_load_threads=1,
+                hsdp_weight_load_strategy="full",
+                parallel_config=SimpleNamespace(
+                    use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=2, tensor_parallel_size=1
+                ),
+            )
+
+            class Pipeline(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.transformer = model_cls(config)
+                    self.transformer.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+                    self.vae = nn.Linear(2, 2)
+                    self.weights_sources = [
+                        DiffusersPipelineLoader.ComponentSource(checkpoint_dir, None, None, "transformer.", False)
+                    ]
+
+                def load_weights(self, weights):
+                    state = dict(weights)
+                    self.transformer.load_state_dict(
+                        {name.removeprefix("transformer."): t for name, t in state.items()}
+                    )
+                    self.transformer.post_load_weights()
+                    self.transformer.eval()
+                    self.transformer.validate_loaded_weights(set(state))
+                    return set(state)
+
+            full = Pipeline()
+            if rank == 0:
+                torch.manual_seed(42)
+                with torch.no_grad():
+                    for name, tensor in full.transformer.named_parameters():
+                        tensor.copy_(torch.randn_like(tensor) * 0.02)
+                        if "norm" in name:
+                            tensor.fill_(1)
+                save_file(full.transformer.state_dict(), str(Path(checkpoint_dir) / "model.safetensors"))
+            dist.barrier()
+
+            def load(model):
+                loader = DiffusersPipelineLoader(LoadConfig(), config)
+                with patch.object(loader, "_init_from_load_format", return_value=model):
+                    return loader._load_model_with_hsdp(device)
+
+            full = load(full)
+            config.hsdp_weight_load_strategy = "pre_sharded"
+            pre_sharded = load(Pipeline())
+            full_parameters = dict(full.transformer.named_parameters())
+            from torch.distributed.tensor import DTensor
+
+            for name, actual in pre_sharded.transformer.named_parameters():
+                expected = full_parameters[name]
+                if name.startswith("time_embedder."):
+                    assert not isinstance(actual, DTensor)
+                    assert actual.dtype == torch.float32
+                else:
+                    assert isinstance(actual, DTensor)
+                    actual = actual.full_tensor()
+                    expected = expected.full_tensor()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            assert pre_sharded.vae.weight.device == device
+            torch.testing.assert_close(pre_sharded.transformer.scratch, full.transformer.scratch)
+            assert all(t.device.type != "meta" for t in pre_sharded.transformer.buffers())
+
+            inputs = dict(
+                hidden_states=torch.linspace(-1, 1, 48 * 4, device=device).reshape(1, 48, 1, 2, 2),
+                timestep=torch.tensor([1.0], device=device),
+                text_ids=torch.tensor([[1, 2]], device=device),
+                text_mask=torch.ones(1, 2, dtype=torch.long, device=device),
+                video_shape=(1, 2, 2),
+                fps=24.0,
+            )
+            with torch.no_grad():
+                # Also exercise the cached UND path on the second denoising step.
+                for step in (1.0, 0.5):
+                    inputs["timestep"].fill_(step)
+                    expected = full.transformer(**inputs)
+                    actual = pre_sharded.transformer(**inputs)
+                    assert torch.isfinite(actual).all()
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("edge", [False, True], ids=["cosmos3", "cosmos3_edge"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 2 or not dist.is_nccl_available(),
+    reason="requires two CUDA GPUs and NCCL",
+)
+def test_cosmos3_pre_sharded_matches_full_hsdp(tmp_path, edge):
+    torch.multiprocessing.spawn(
+        _load_cosmos_worker,
+        args=(f"file://{tmp_path / 'rendezvous'}", str(tmp_path), edge),
+        nprocs=2,
+        join=True,
+    )

tests/diffusion/distributed/test_hsdp.py

@@ -3,11 +3,13 @@
 """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
 import gc
+from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor import DeviceMesh, DTensor
 from vllm.utils.network_utils import get_file_store_init_method
 
@@ -18,6 +20,7 @@ from vllm_omni.diffusion.distributed.hsdp import (
     _unshardable_parameters,
     shard_model,
 )
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
@@ -338,3 +341,92 @@ class TestHSDPShardConditions:
                 matched.append(name)
         assert "blocks.0" in matched
         assert "blocks.1" in matched
+
+
+def test_pre_sharded_to_empty_materializes_fsdp_internal_storage(cpu_process_group):
+    with torch.device("meta"):
+        model = nn.Module()
+        model.transformer = nn.Sequential(nn.Linear(2, 2, bias=False))
+    model.transformer.register_buffer("scratch", torch.tensor([3.0]), persistent=False)
+    shard_model(
+        model.transformer,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "0"],
+    )
+
+    load_plan = SimpleNamespace(
+        roots=[SimpleNamespace(name="transformer", module=model.transformer)],
+        groups=[SimpleNamespace(module=model.transformer[0])],
+        binding_names=frozenset({"transformer.0.weight"}),
+    )
+    sharded_weight = model.transformer[0].weight
+
+    DiffusersPipelineLoader._materialize_pre_sharded_hsdp_state(model, load_plan, torch.device("cpu"))
+
+    fsdp_param = fully_shard.state(model.transformer[0])._fsdp_param_group.fsdp_params[0]
+    assert model.transformer[0].weight is sharded_weight
+    assert model.transformer[0].weight.to_local().device.type == "cpu"
+    assert fsdp_param._sharded_param_data.device.type == "cpu"
+    assert (
+        fsdp_param._sharded_param_data.untyped_storage().data_ptr()
+        == model.transformer[0].weight.to_local().untyped_storage().data_ptr()
+    )
+    assert torch.equal(model.transformer.scratch, torch.tensor([3.0]))
+
+
+def test_pre_sharded_load_restores_checkpoint_values(cpu_process_group, tmp_path, monkeypatch):
+    from safetensors.torch import save_file
+    from vllm.config.load import LoadConfig
+
+    from vllm_omni.diffusion.distributed import hsdp as hsdp_module
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    class Transformer(nn.Module):
+        _hsdp_shard_conditions = [lambda name, module: name == "block"]
+        _hsdp_ignored_modules = ["time_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.block = nn.Linear(3, 4, bias=False)
+            self.time_embedder = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16)
+            self.root_weight = nn.Parameter(torch.zeros(2))
+            self.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+            release_module_parameters_to_meta(self.block)
+
+        def post_load_weights(self):
+            self.time_embedder.to(torch.float32)
+
+    model = nn.Module()
+    model.transformer = Transformer()
+    model.vae = nn.Linear(2, 2)
+    tensors = {
+        "block.weight": torch.arange(12, dtype=torch.float32).reshape(4, 3),
+        "root_weight": torch.tensor([2.0, 3.0]),
+        "time_embedder.weight": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
+    }
+    save_file(tensors, str(tmp_path / "model.safetensors"))
+    model.weights_sources = [DiffusersPipelineLoader.ComponentSource(str(tmp_path), None, None, "transformer.", False)]
+    config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization_config=None,
+        lora_path=None,
+        hsdp_weight_load_strategy="pre_sharded",
+        num_weight_load_threads=1,
+        parallel_config=SimpleNamespace(use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=1),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: model)
+    monkeypatch.setattr(hsdp_module, "get_world_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0))
+    monkeypatch.setattr(hsdp_module, "current_omni_platform", SimpleNamespace(device_type="cpu"))
+    loaded = loader._load_model_with_hsdp(torch.device("cpu"))
+    assert loaded is model
+    for name, expected in tensors.items():
+        actual = dict(model.transformer.named_parameters())[name]
+        if name.startswith("time_embedder."):
+            assert not isinstance(actual, DTensor)
+            assert actual.dtype == torch.float32
+            torch.testing.assert_close(actual, expected.float())
+        else:
+            assert isinstance(actual, DTensor)
+            torch.testing.assert_close(actual.to_local(), expected)
+    torch.testing.assert_close(model.transformer.scratch, torch.tensor([7.0]))