[feature] Add HSDP pre-sharded option (#7948)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

docs/configuration/hsdp_pre_sharded_loading.md

@@ -0,0 +1,22 @@
+# Pre-sharded HSDP loading
+
+Use `--use-hsdp --hsdp-shard-size 8 --hsdp-weight-load-strategy pre_sharded`
+to shard transformer parameters before loading checkpoint values. Set the shard
+size to the available GPU count (or configure replication explicitly).
+The default strategy, `full`, retains the existing load-then-shard behavior.
+The same `hsdp_weight_load_strategy` option is available through Python and stage
+engine configuration.
+
+The pre-sharded strategy reads rank-local slices from ordinary Hugging Face
+safetensors checkpoints into FSDP-owned storage. No converted checkpoint is
+required. Cosmos 3 and Cosmos 3 Edge transformer blocks release their initialized
+parameter storage to meta as they are constructed; nonpersistent buffers retain
+their values.
+Encoders and VAEs follow their existing loading paths.
+
+This strategy requires the default diffusion pipeline loader, complete dedicated
+transformer checkpoint sources, and supported checkpoint key mappings. Tensor
+layouts must match the runtime layout; tensor transforms are unsupported.
+Checkpoint tensor dtypes must match runtime tensor dtypes at preflight; dtype conversion is unsupported.
+Quantization and LoRA are unsupported. Every worker needs checkpoint access.
+`num_weight_load_threads` controls the safetensors reader's threads per rank.

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]))

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     HostWeightPlanResult,
     TensorBinding,
+    build_checkpoint_binding_plan,
 )
 from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
@@ -2089,3 +2090,206 @@ def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
             )
         assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
         assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))
+
+
+def test_pre_sharded_hsdp_rejects_quantization_before_loading():
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=object(),
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    with pytest.raises(ValueError, match="does not support quantization"):
+        loader._load_model_with_hsdp(torch.device("cpu"))
+
+
+def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=None,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    discovered = PipelineModules(
+        dits=[model.transformer],
+        dit_names=["transformer"],
+        vaes=[],
+        encoders=[],
+        encoder_names=[],
+        resident_modules=[],
+        resident_names=[],
+    )
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run"))  # type: ignore[method-assign]
+    mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered)
+    pre_sharded_load = mocker.patch.object(
+        loader,
+        "_load_model_with_pre_sharded_hsdp",
+        return_value=model,
+    )
+
+    assert loader._load_model_with_hsdp(torch.device("cpu")) is model
+    pre_sharded_load.assert_called_once()
+
+
+def test_pre_sharded_hsdp_rejects_tensor_transform():
+    def transform(tensor):
+        return tensor.transpose(-2, -1)
+
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+                transform=transform,
+            )
+        }
+    )
+
+    with pytest.raises(ValueError, match="tensor transforms are unsupported"):
+        DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_accepts_direct_safetensors_binding():
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+            )
+        }
+    )
+
+    DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_uses_fsdp_destination_directly():
+    target = torch.empty(8, 4, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+
+    checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": target},
+    )
+    checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"]
+
+    assert checkpoint_target is target
+    assert checkpoint_target.shape == (8, 4, 3)
+    assert checkpoint_target.device.type == "cpu"
+    assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr()
+    assert local_bytes == target.numel() * target.element_size()
+
+
+def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly():
+    runtime_target = torch.empty(2, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+    checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": runtime_target},
+    )
+    checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor)
+
+    assert torch.equal(runtime_target, checkpoint_tensor)
+
+
+def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations():
+    binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors")
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.first": binding,
+            "transformer.second": binding,
+        }
+    )
+
+    with pytest.raises(ValueError, match="multiple runtime tensors"):
+        DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)},
+        )
+
+
+def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots():
+    root = nn.Module()
+    root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+    root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+
+    names = {module: name or "root" for name, module in root.named_modules()}
+    modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+        root,
+        nested_hsdp_roots={root.left},
+    )
+
+    assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"]
+
+
+@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"])
+def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch):
+    import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod
+
+    monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None)
+
+    class Pipeline(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+
+        @staticmethod
+        def remap_checkpoint_key(name):
+            return "transformer.weight" if name == "transformer.renamed" else name
+
+    checkpoint = tmp_path / "model.safetensors"
+    save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint))
+    source = SimpleNamespace(
+        model_or_path=str(tmp_path),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    model = Pipeline()
+    result = build_checkpoint_binding_plan(
+        model,
+        dit_modules=(("transformer", model.transformer),),
+        sources=(source,),
+        model_path=None,
+        tensor_parallel_size=1,
+        online_quantization=False,
+    )
+
+    if checkpoint_key == "renamed":
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed"
+    else:
+        assert result.plan is None
+        assert "no checkpoint binding" in result.fallback_reason

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -868,6 +868,93 @@ def test_qwen_gen_mlp_remains_gated() -> None:
     assert hasattr(model.gen_layers[0].mlp, "gate_proj")
 
 
+@pytest.mark.parametrize("edge", [False, True])
+@pytest.mark.parametrize(
+    "use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")]
+)
+def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch):
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    released = []
+
+    def release_and_check(block):
+        parameters = dict(block.named_parameters())
+        metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()}
+        release_module_parameters_to_meta(block)
+        for name, parameter in block.named_parameters():
+            assert parameter is parameters[name]
+            assert parameter.device.type == "meta"
+            assert parameter.__dict__ == metadata[name]
+        released.append(block)
+
+    monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check)
+    monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check)
+    model_cls = (
+        transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer
+    )
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+            hsdp_weight_load_strategy=strategy,
+        )
+    )
+    blocks = [*model.language_model.layers, *model.gen_layers]
+    expected_meta = use_hsdp and strategy == "pre_sharded"
+    assert len(released) == (len(blocks) if expected_meta else 0)
+    assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters())
+    assert model.time_embedder.linear_1.weight.device.type != "meta"
+    assert all(buffer.device.type != "meta" for buffer in model.buffers())
+
+
+@pytest.mark.parametrize("edge", [False, True])
+def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch):
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer
+
+    model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=True),
+            hsdp_weight_load_strategy="pre_sharded",
+        )
+    )
+    groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)]
+    plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups)
+    processed = []
+    # Exercise model selection with real vLLM linears, without invoking the
+    # CPU-specific weight packing hook on meta tensors in CPU-only CI.
+    monkeypatch.setattr(
+        UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer)
+    )
+    DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")]
+    assert expected
+    assert processed == expected
+
+    model._hsdp_pre_sharded_meta_post_load = False
+    with pytest.raises(ValueError, match="has not declared"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    model._hsdp_pre_sharded_meta_post_load = True
+
+    class UnsupportedMethod(UnquantizedLinearMethod):
+        pass
+
+    expected[-1].quant_method = UnsupportedMethod()
+    processed.clear()
+    with pytest.raises(ValueError, match="UnsupportedMethod"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    assert not processed
+
+
 def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -11,6 +11,7 @@ from pydantic import ValidationError
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
+from vllm_omni.config.stage_config import build_stage_runtime_overrides
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -945,3 +946,33 @@ def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
     assert launched["model"] == "generic-diffusion"
     assert launched["stage_id"] == 0
     assert runtime.stage_pools[0].clients == [client]
+
+
+def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None:
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0]
+
+    assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("prefix", ["", "stage_0_"])
+def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None:
+    overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"})
+
+    assert overrides["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("strategy", ["full", "pre_sharded"])
+def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None:
+    config = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"hsdp_weight_load_strategy": strategy}
+    )
+
+    assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy
+
+
+@pytest.mark.parametrize("strategy", ["memory_limited", "invalid"])
+def test_removed_hsdp_loading_strategies_are_rejected(strategy):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with pytest.raises(ValueError, match="hsdp_weight_load_strategy"):
+        OmniDiffusionConfig(hsdp_weight_load_strategy=strategy)

vllm_omni/config/omni_config.py

@@ -817,6 +817,7 @@ class _DiffusionConfigProjection:
     prompt_embed_cache_size: int = Field(default=32, ge=1)
     enable_session_state_manager: bool = False
     diffusion_load_format: str = "default"
+    hsdp_weight_load_strategy: str = "full"
     diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_pipeline_cls: Any = None

vllm_omni/config/stage_config.py

@@ -492,6 +492,7 @@ class StageDeployConfig:
     auxiliary_text_encoder: str | None = None
 
     # Runtime optimizations used by diffusion loading/execution.
+    hsdp_weight_load_strategy: str | None = None
     enable_multithread_weight_load: bool | None = None
     enable_broadcast_weight_load: bool | None = None
     num_weight_load_threads: int | None = None

vllm_omni/diffusion/data.py

@@ -983,6 +983,9 @@ class OmniDiffusionConfig:
     enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = 4
 
+    # Shard meta parameters before reading rank-local HF safetensors slices.
+    hsdp_weight_load_strategy: str = "full"
+
     # Enable sleep mode
     enable_sleep_mode: bool = False
 
@@ -1196,6 +1199,10 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if self.hsdp_weight_load_strategy not in {"full", "pre_sharded"}:
+            raise ValueError(
+                f"hsdp_weight_load_strategy must be 'full' or 'pre_sharded', got {self.hsdp_weight_load_strategy!r}"
+            )
         from vllm_omni.diffusion.offloader.config import (
             OffloadStrategy,
             materialize_legacy_offload_flags,

vllm_omni/diffusion/distributed/hsdp.py

@@ -46,6 +46,14 @@ class HSDPInferenceConfig:
     reshard_after_forward: bool = True
 
 
+@dataclass
+class HSDPShardContext:
+    """Reusable state for incrementally sharding one HSDP model."""
+
+    hsdp_kwargs: dict[str, Any]
+    ignored_params: set[nn.Parameter] | None = None
+
+
 def _create_hsdp_mesh(
     device_type: str,
     replicate_size: int,
@@ -111,6 +119,38 @@ def apply_hsdp_to_model(
     Returns:
         HSDP-wrapped model ready for inference
     """
+    context = prepare_hsdp_shard_context(
+        model,
+        hsdp_config,
+        target_device=target_device,
+    )
+
+    hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
+    if not hsdp_shard_conditions:
+        raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
+
+    # Apply HSDP sharding, this will automatically handle weight distribution
+    shard_model(
+        model,
+        hsdp_shard_conditions=hsdp_shard_conditions,
+        context=context,
+    )
+
+    logger.info("HSDP applied to model: %s", type(model).__name__)
+    return model
+
+
+def prepare_hsdp_shard_context(
+    model: nn.Module,
+    hsdp_config: HSDPInferenceConfig,
+    target_device: torch.device | None = None,
+) -> HSDPShardContext:
+    """Prepare mesh, precision, and ignored-parameter state for HSDP.
+
+    The returned context can be reused to shard completed child modules while
+    checkpoint weights are still streaming, then to shard the root once all
+    children have been loaded.
+    """
     if not hsdp_config.enabled:
         raise ValueError("HSDP is not enabled in config")
 
@@ -160,7 +200,7 @@ def apply_hsdp_to_model(
     )
 
     hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
-    if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0:
+    if not hsdp_shard_conditions:
         raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
 
     # Collect parameters of any modules the model wants excluded from FSDP sharding.
@@ -188,7 +228,11 @@ def apply_hsdp_to_model(
         except AttributeError:
             logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name)
             continue
-        sub_mod.to(target_device)
+        has_meta_tensor = any(tensor.device.type == "meta" for tensor in (*sub_mod.parameters(), *sub_mod.buffers()))
+        if has_meta_tensor:
+            sub_mod.to_empty(device=target_device)
+        else:
+            sub_mod.to(target_device)
         ignored_params.update(sub_mod.parameters())
     if ignored_params:
         logger.info(
@@ -224,21 +268,36 @@ def apply_hsdp_to_model(
             target_device,
         )
 
-    # Apply HSDP sharding, this will automatically handle weight distribution
-    shard_model(
-        model,
-        reshard_after_forward=hsdp_config.reshard_after_forward,
-        mp_policy=mp_policy,
-        mesh=device_mesh,
-        hsdp_shard_conditions=hsdp_shard_conditions,
-        ignored_params=ignored_params if ignored_params else None,
+    return HSDPShardContext(
+        hsdp_kwargs={
+            "reshard_after_forward": hsdp_config.reshard_after_forward,
+            "mesh": device_mesh,
+            "mp_policy": mp_policy,
+        },
+        ignored_params=ignored_params or None,
     )
 
-    for param in model.parameters():
+
+def shard_hsdp_module(module: nn.Module, context: HSDPShardContext) -> None:
+    """Shard one completed child module with a prepared HSDP context."""
+    module_kwargs = dict(context.hsdp_kwargs)
+    if context.ignored_params:
+        module_ignored_params = context.ignored_params.intersection(module.parameters())
+        if module_ignored_params:
+            module_kwargs["ignored_params"] = module_ignored_params
+    fully_shard(module, **module_kwargs)
+    for param in module.parameters():
         param.requires_grad = False
 
-    logger.info("HSDP applied to model: %s", type(model).__name__)
-    return model
+
+def finalize_hsdp_root(model: nn.Module, context: HSDPShardContext) -> None:
+    """Shard the model root after all selected child modules are sharded."""
+    root_kwargs = dict(context.hsdp_kwargs)
+    if context.ignored_params:
+        root_kwargs["ignored_params"] = context.ignored_params
+    fully_shard(model, **root_kwargs)
+    for param in model.parameters():
+        param.requires_grad = False
 
 
 def shard_model(
@@ -249,6 +308,7 @@ def shard_model(
     mesh: DeviceMesh | None = None,
     hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]],
     ignored_params: set[nn.Parameter] | None = None,
+    context: HSDPShardContext | None = None,
 ) -> None:
     """Apply HSDP sharding to model modules based on shard conditions.
 
@@ -259,32 +319,30 @@ def shard_model(
     This is required for packed integer parameters inside sharded transformer
     blocks; the root wrap receives the full set for all remaining parameters.
     """
-    hsdp_kwargs: dict[str, Any] = {
-        "reshard_after_forward": reshard_after_forward,
-        "mesh": mesh,
-        "mp_policy": mp_policy,
-    }
+    if context is None:
+        context = HSDPShardContext(
+            hsdp_kwargs={
+                "reshard_after_forward": reshard_after_forward,
+                "mesh": mesh,
+                "mp_policy": mp_policy,
+            },
+            ignored_params=ignored_params,
+        )
+    elif any(value is not None for value in (mp_policy, mesh, ignored_params)) or not reshard_after_forward:
+        raise ValueError("Pass either context or individual HSDP shard options, not both")
 
     num_sharded = 0
     for name, module in reversed(list(model.named_modules())):
         if any(cond(name, module) for cond in hsdp_shard_conditions):
-            module_kwargs = dict(hsdp_kwargs)
-            if ignored_params:
-                module_ignored_params = ignored_params.intersection(module.parameters())
-                if module_ignored_params:
-                    module_kwargs["ignored_params"] = module_ignored_params
-            fully_shard(module, **module_kwargs)
+            shard_hsdp_module(module, context)
             num_sharded += 1
 
     if num_sharded == 0:
         raise ValueError("No modules were sharded")
 
-    root_kwargs = dict(hsdp_kwargs)
-    if ignored_params:
-        root_kwargs["ignored_params"] = ignored_params
-    fully_shard(model, **root_kwargs)
+    finalize_hsdp_root(model, context)
     logger.info(
         "Sharded %d modules + root (ignored_params=%d)",
         num_sharded,
-        len(ignored_params) if ignored_params else 0,
+        len(context.ignored_params) if context.ignored_params else 0,
     )

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -16,6 +16,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
@@ -30,7 +31,14 @@ from vllm.utils.torch_utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    HSDPShardContext,
+    apply_hsdp_to_model,
+    finalize_hsdp_root,
+    prepare_hsdp_shard_context,
+    shard_hsdp_module,
+)
 from vllm_omni.diffusion.lora.manager import LoRABackend
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
@@ -38,6 +46,8 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
 from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
+    TensorBinding,
+    build_checkpoint_binding_plan,
     build_checkpoint_mmap_plan,
     has_online_quantization,
 )
@@ -49,7 +59,7 @@ from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload,
 )
-from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, PipelineModules
 from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan
 from vllm_omni.diffusion.registry import initialize_model
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -138,6 +148,36 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ
     )
 
 
+def _is_distributed_rank_zero() -> bool:
+    return (
+        not torch.distributed.is_available()
+        or not torch.distributed.is_initialized()
+        or torch.distributed.get_rank() == 0
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class _HSDPLoadRoot:
+    name: str
+    module: nn.Module
+    context: HSDPShardContext
+
+
+@dataclasses.dataclass(frozen=True)
+class _HSDPModuleLoadGroup:
+    name: str
+    module: nn.Module
+    context: HSDPShardContext
+
+
+@dataclasses.dataclass(frozen=True)
+class _PreShardedHSDPLoadPlan:
+    roots: list[_HSDPLoadRoot]
+    groups: list[_HSDPModuleLoadGroup]
+    bindings: dict[str, TensorBinding]
+    binding_names: frozenset[str]
+
+
 class DiffusersPipelineLoader(HWRLoaderMixin):
     """Model loader that can load diffusers pipeline components from disk."""
 
@@ -1222,6 +1262,451 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     raise ValueError(f"Unknown load_format: {load_format}")
         return model
 
+    def _load_model_with_pre_sharded_hsdp(
+        self,
+        model: nn.Module,
+        discovered_modules: PipelineModules,
+        hsdp_config: HSDPInferenceConfig,
+        target_device: torch.device,
+    ) -> nn.Module:
+        """Shard meta DiT parameters before reading their rank-local checkpoint slices."""
+        from torch.distributed.checkpoint import HuggingFaceStorageReader
+        from torch.distributed.checkpoint import load as dcp_load
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+
+        load_plan = self._prepare_pre_sharded_hsdp_load(
+            model,
+            discovered_modules,
+            hsdp_config,
+            target_device,
+        )
+        self._check_no_unplanned_meta_tensors(model, load_plan.binding_names)
+        self._validate_pre_sharded_hsdp_bindings(load_plan)
+
+        self._process_pre_sharded_hsdp_groups_on_meta(load_plan)
+
+        for group in load_plan.groups:
+            logger.debug("Applying meta-first HSDP shard to %s", group.name)
+            shard_hsdp_module(group.module, group.context)
+        for root in load_plan.roots:
+            logger.debug("Finalizing meta-first HSDP root %s", root.name)
+            finalize_hsdp_root(root.module, root.context)
+
+        self._materialize_pre_sharded_hsdp_state(model, load_plan, target_device)
+
+        targets = self._get_pre_sharded_hsdp_targets(model, load_plan)
+        checkpoint_states, local_payload_bytes = self._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            targets,
+        )
+
+        thread_count = int(getattr(self.od_config, "num_weight_load_threads", 4))
+        if _is_distributed_rank_zero():
+            directory_label = "directory" if len(checkpoint_states) == 1 else "directories"
+            logger.info(
+                "Pre-sharded HSDP loading %.2f GiB of rank-local tensor slices directly into "
+                "FSDP-owned storage from %d HF safetensors %s with %d reader thread(s) per rank",
+                local_payload_bytes / 1024**3,
+                len(checkpoint_states),
+                directory_label,
+                thread_count,
+            )
+
+        for checkpoint_dir in sorted(checkpoint_states, key=str):
+            reader = HuggingFaceStorageReader(str(checkpoint_dir), thread_count=thread_count)
+            # Each worker has the full checkpoint available and each DTensor
+            # already describes its rank-local destination slice. Avoid DCP's
+            # coordinator collectives, which can leave rank 0 with less GPU
+            # memory headroom than the other HSDP workers.
+            dcp_load(checkpoint_states[checkpoint_dir], storage_reader=reader, no_dist=True)
+
+        remaining_meta = self._pre_sharded_meta_target_names(model, load_plan)
+        if remaining_meta:
+            raise ValueError(
+                "Pre-sharded HSDP checkpoint loading left tensors on meta; "
+                f"first missing tensors: {remaining_meta[:16]}"
+            )
+
+        self._finalize_pre_sharded_hsdp_loaded_weights(load_plan)
+
+        self.counter_after_loading_weights = time.perf_counter()
+        logger.info_once(
+            "Loading weights took %.2f seconds",
+            self.counter_after_loading_weights - self.counter_before_loading_weights,
+        )
+        self._move_non_hsdp_modules(discovered_modules, target_device)
+        return model
+
+    @staticmethod
+    def _validate_pre_sharded_hsdp_bindings(load_plan: _PreShardedHSDPLoadPlan) -> None:
+        unsupported = [name for name, binding in load_plan.bindings.items() if binding.transform is not None]
+        if unsupported:
+            raise ValueError(
+                "Pre-sharded HSDP requires checkpoint tensors in runtime layout; "
+                f"tensor transforms are unsupported: {unsupported[:16]}"
+            )
+        non_safetensors = sorted(
+            {
+                binding.file_path
+                for binding in load_plan.bindings.values()
+                if not binding.file_path.endswith(".safetensors")
+            }
+        )
+        if non_safetensors:
+            raise ValueError(
+                f"Pre-sharded HSDP requires existing HF safetensors files; unsupported files: {non_safetensors[:8]}"
+            )
+
+    @staticmethod
+    def _get_pre_sharded_hsdp_targets(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> dict[str, torch.Tensor]:
+        targets: dict[str, torch.Tensor] = {
+            **dict(model.named_buffers()),
+            **dict(model.named_parameters()),
+        }
+        missing_targets = sorted(load_plan.binding_names - targets.keys())
+        if missing_targets:
+            raise ValueError(
+                f"Pre-sharded HSDP changed checkpoint target names; first missing targets: {missing_targets[:16]}"
+            )
+        return targets
+
+    @staticmethod
+    def _materialize_pre_sharded_hsdp_state(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+        target_device: torch.device,
+    ) -> None:
+        """Materialize the installed FSDP rank-local state before checkpoint I/O.
+
+        ``fully_shard`` supports meta initialization through ``Module.to_empty``:
+        its registered ``_apply`` hooks update the internal sharded storage to
+        point at the newly allocated rank-local tensors. Modules are materialized
+        in the same HSDP-group order and child-first traversal used by ordinary
+        ``fully_shard`` initialization instead of the forward traversal used by
+        a recursive ``to_empty`` call.
+        PyTorch's swap-on-conversion mode preserves Parameter identity so that
+        post-load kernel objects created before sharding cannot retain stale
+        references to the original meta Parameters.
+
+        A complete checkpoint plan covers every parameter and persistent buffer
+        in each DiT.  Non-persistent buffers are not checkpoint entries, so keep
+        their initialized values across ``to_empty``.
+        """
+        unplanned_parameters: list[str] = []
+        saved_buffers: dict[str, torch.Tensor] = {}
+        for root in load_plan.roots:
+            prefix = f"{root.name}."
+            for local_name, _parameter in root.module.named_parameters():
+                full_name = prefix + local_name
+                if full_name not in load_plan.binding_names:
+                    unplanned_parameters.append(full_name)
+            for local_name, buffer in root.module.named_buffers():
+                full_name = prefix + local_name
+                if full_name in load_plan.binding_names:
+                    continue
+                if buffer.device.type == "meta":
+                    raise ValueError(f"Pre-sharded HSDP cannot preserve an unplanned meta buffer: {full_name}")
+                saved_buffers[full_name] = buffer.detach().cpu().clone()
+
+        if unplanned_parameters:
+            raise ValueError(
+                "Pre-sharded HSDP requires checkpoint coverage for every DiT parameter; "
+                f"first unplanned parameters: {unplanned_parameters[:16]}"
+            )
+
+        previous_swap_mode = torch.__future__.get_swap_module_params_on_conversion()
+        torch.__future__.set_swap_module_params_on_conversion(True)
+        try:
+            sharded_group_roots = {group.module for group in load_plan.groups}
+            materialized_modules: set[nn.Module] = set()
+            for group in load_plan.groups:
+                for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+                    group.module,
+                    nested_hsdp_roots=sharded_group_roots - {group.module},
+                ):
+                    if module not in materialized_modules:
+                        module.to_empty(device=target_device, recurse=False)
+                        materialized_modules.add(module)
+            for root in load_plan.roots:
+                for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+                    root.module,
+                    nested_hsdp_roots=sharded_group_roots,
+                ):
+                    if module not in materialized_modules:
+                        module.to_empty(device=target_device, recurse=False)
+                        materialized_modules.add(module)
+        finally:
+            torch.__future__.set_swap_module_params_on_conversion(previous_swap_mode)
+
+        if saved_buffers:
+            materialized_buffers = dict(model.named_buffers())
+            with torch.no_grad():
+                for name, value in saved_buffers.items():
+                    target = materialized_buffers[name]
+                    target.copy_(value.to(device=target.device, dtype=target.dtype))
+
+    @staticmethod
+    def _hsdp_managed_modules_post_order(
+        root: nn.Module,
+        *,
+        nested_hsdp_roots: set[nn.Module],
+    ) -> list[nn.Module]:
+        """Mirror FSDP's child-first managed-module traversal.
+
+        Nested FSDP roots own their own parameters and are materialized in their
+        corresponding load group, so the enclosing traversal must not descend
+        into them.
+        """
+        modules: list[nn.Module] = []
+        visited: set[nn.Module] = set()
+
+        def visit(module: nn.Module) -> None:
+            if module in visited:
+                return
+            visited.add(module)
+            for child in module.children():
+                if child not in nested_hsdp_roots:
+                    visit(child)
+            modules.append(module)
+
+        visit(root)
+        return modules
+
+    @staticmethod
+    def _prepare_pre_sharded_checkpoint_states(
+        load_plan: _PreShardedHSDPLoadPlan,
+        targets: dict[str, torch.Tensor],
+    ) -> tuple[dict[Path, dict[str, torch.Tensor]], int]:
+        from torch.distributed.tensor import DTensor
+
+        checkpoint_states: dict[Path, dict[str, torch.Tensor]] = {}
+        checkpoint_owners: dict[tuple[Path, str], str] = {}
+        local_payload_bytes = 0
+        for name, binding in load_plan.bindings.items():
+            target = targets[name]
+
+            checkpoint_dir = Path(binding.file_path).parent
+            owner_key = (checkpoint_dir, binding.checkpoint_key)
+            if previous_owner := checkpoint_owners.get(owner_key):
+                raise ValueError(
+                    "Pre-sharded HSDP cannot bind one checkpoint tensor to multiple runtime tensors: "
+                    f"{binding.checkpoint_key!r} maps to both {previous_owner!r} and {name!r}"
+                )
+            checkpoint_owners[owner_key] = name
+
+            if target.device.type == "meta":
+                raise ValueError(f"Pre-sharded HSDP target {name!r} was not materialized")
+
+            local_target = target.to_local() if isinstance(target, DTensor) else target
+            local_payload_bytes += local_target.numel() * local_target.element_size()
+            checkpoint_states.setdefault(checkpoint_dir, {})[binding.checkpoint_key] = target
+        return checkpoint_states, local_payload_bytes
+
+    @staticmethod
+    def _pre_sharded_meta_target_names(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> list[str]:
+        from torch.distributed.tensor import DTensor
+
+        targets = {
+            **dict(model.named_buffers()),
+            **dict(model.named_parameters()),
+        }
+        return [
+            name
+            for name in load_plan.binding_names
+            if (target := targets[name]).device.type == "meta"
+            or (isinstance(target, DTensor) and target.to_local().device.type == "meta")
+        ]
+
+    @staticmethod
+    def _process_pre_sharded_hsdp_groups_on_meta(load_plan: _PreShardedHSDPLoadPlan) -> None:
+        modules_to_process: list[tuple[nn.Module, QuantizeMethodBase]] = []
+        for group in load_plan.groups:
+            for _, module in group.module.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
+                    continue
+                modules_to_process.append((module, quant_method))
+
+        if not modules_to_process:
+            return
+        for root in load_plan.roots:
+            if not getattr(root.module, "_hsdp_pre_sharded_meta_post_load", False):
+                raise ValueError(
+                    f"Model {type(root.module).__name__} has not declared its post-load processing safe "
+                    "for pre-sharded meta initialization"
+                )
+        for module, quant_method in modules_to_process:
+            if type(quant_method) is not UnquantizedLinearMethod:
+                raise ValueError(
+                    "Pre-sharded HSDP does not support the selected post-load processing; "
+                    f"{type(module).__name__} uses {type(quant_method).__name__}"
+                )
+        for module, quant_method in modules_to_process:
+            quant_method.process_weights_after_loading(module)
+
+    def _prepare_pre_sharded_hsdp_load(
+        self,
+        model: nn.Module,
+        discovered_modules: PipelineModules,
+        hsdp_config: HSDPInferenceConfig,
+        target_device: torch.device,
+    ) -> _PreShardedHSDPLoadPlan:
+        outer_dit_names, outer_dits = discovered_modules.outermost_dits()
+        if not outer_dits:
+            raise ValueError("No DiT modules discovered for HSDP sharding")
+
+        sources = self._get_weight_sources(model)
+        plan_result = build_checkpoint_binding_plan(
+            model,
+            dit_modules=tuple(zip(outer_dit_names, outer_dits)),
+            sources=sources,
+            model_path=str(getattr(self.od_config, "model", "")) or None,
+            tensor_parallel_size=int(getattr(self.parallel_config, "tensor_parallel_size", 1)),
+            online_quantization=False,
+        )
+        if plan_result.plan is None:
+            raise ValueError(f"Pre-sharded HSDP checkpoint loading is incompatible: {plan_result.fallback_reason}")
+        plan = plan_result.plan
+        unsupported_sources = [source.prefix for source in sources if source.prefix not in plan.planned_source_prefixes]
+        if unsupported_sources:
+            raise ValueError(
+                "Pre-sharded HSDP checkpoint loading requires dedicated DiT sources; "
+                f"unsupported source prefixes: {unsupported_sources}"
+            )
+
+        roots: list[_HSDPLoadRoot] = []
+        groups: list[_HSDPModuleLoadGroup] = []
+        for outer_name, outer_dit in zip(outer_dit_names, outer_dits):
+            conditions = getattr(outer_dit, "_hsdp_shard_conditions", None)
+            if not conditions:
+                raise ValueError(f"Model {type(outer_dit).__name__} has no _hsdp_shard_conditions defined")
+
+            context = prepare_hsdp_shard_context(
+                outer_dit,
+                hsdp_config,
+                target_device=target_device,
+            )
+            roots.append(_HSDPLoadRoot(outer_name, outer_dit, context))
+            for local_name, module in reversed(list(outer_dit.named_modules())):
+                if not any(condition(local_name, module) for condition in conditions):
+                    continue
+                if not local_name:
+                    raise ValueError("_hsdp_shard_conditions must not select the DiT root")
+                full_name = f"{outer_name}.{local_name}" if local_name else outer_name
+                groups.append(_HSDPModuleLoadGroup(full_name, module, context))
+
+        parameters = dict(model.named_parameters())
+        targets: dict[str, torch.Tensor] = {**dict(model.named_buffers()), **parameters}
+        preserved_params = {param for root in roots for param in root.context.ignored_params or ()}
+        binding_names = frozenset(plan.bindings)
+        self._release_hsdp_checkpoint_targets_to_meta(
+            model,
+            targets,
+            parameters,
+            binding_names,
+            preserved_params,
+        )
+
+        return _PreShardedHSDPLoadPlan(
+            roots=roots,
+            groups=groups,
+            bindings=plan.bindings,
+            binding_names=binding_names,
+        )
+
+    @staticmethod
+    def _check_no_unplanned_meta_tensors(model: nn.Module, planned_names: frozenset[str]) -> None:
+        unplanned_meta = [
+            name
+            for name, tensor in (*model.named_parameters(), *model.named_buffers())
+            if tensor.device.type == "meta" and name not in planned_names
+        ]
+        if unplanned_meta:
+            raise ValueError(
+                "Pre-sharded HSDP initialized tensors on meta that are not covered by the "
+                f"checkpoint plan. First unplanned tensors: {unplanned_meta[:16]}"
+            )
+
+    @staticmethod
+    def _release_hsdp_checkpoint_targets_to_meta(
+        model: nn.Module,
+        targets: dict[str, torch.Tensor],
+        parameters: dict[str, nn.Parameter],
+        binding_names: frozenset[str],
+        preserved_params: set[nn.Parameter],
+    ) -> None:
+        released: dict[str, tuple[int, float]] = {}
+        for name in binding_names:
+            target = targets[name]
+            if target in preserved_params or target.device.type == "meta":
+                continue
+            logical_gib = target.numel() * target.element_size() / 1024**3
+            device_type = target.device.type
+            count, gib = released.get(device_type, (0, 0.0))
+            released[device_type] = (count + 1, gib + logical_gib)
+            if name in parameters:
+                replacement = DiffusersPipelineLoader._make_parameter_like(
+                    target,
+                    torch.empty_like(target, device="meta"),
+                )
+                torch.utils.swap_tensors(target, replacement)
+            else:
+                replacement = torch.empty_like(target, device="meta")
+                torch.utils.swap_tensors(target, replacement)
+        if released:
+            logger.info("Released HSDP checkpoint-covered tensors to meta before pre-sharded load: %s", released)
+
+    @staticmethod
+    def _make_parameter_like(template: nn.Parameter, tensor: torch.Tensor) -> nn.Parameter:
+        """Create a Parameter on new storage while preserving loader metadata."""
+        if type(template) is nn.Parameter:
+            replacement = nn.Parameter(tensor, requires_grad=template.requires_grad)
+            replacement.__dict__.update(getattr(template, "__dict__", {}))
+            return replacement
+
+        replacement = template.__class__.__new__(template.__class__, tensor)
+        replacement.__dict__.update(getattr(template, "__dict__", {}))
+        replacement.requires_grad_(template.requires_grad)
+        return replacement
+
+    @staticmethod
+    def _finalize_pre_sharded_hsdp_loaded_weights(
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> None:
+        """Run model post-load hooks for tensors materialized by pre-sharded loading."""
+        for root in load_plan.roots:
+            post_load = getattr(root.module, "post_load_weights", None)
+            if callable(post_load):
+                post_load()
+            root.module.eval()
+            validate = getattr(root.module, "validate_loaded_weights", None)
+            if callable(validate):
+                prefix = f"{root.name}."
+                validate({name.removeprefix(prefix) for name in load_plan.binding_names if name.startswith(prefix)})
+
+    @staticmethod
+    def _move_non_hsdp_modules(
+        discovered_modules: PipelineModules,
+        target_device: torch.device,
+    ) -> None:
+        modules_to_move: list[nn.Module] = []
+        if discovered_modules.vaes is not None:
+            modules_to_move.extend(discovered_modules.vaes)
+        if discovered_modules.encoders is not None:
+            modules_to_move.extend(discovered_modules.encoders)
+        if discovered_modules.resident_modules is not None:
+            modules_to_move.extend(discovered_modules.resident_modules)
+        for module in modules_to_move:
+            module.to(target_device)
+
     def _load_model_with_hsdp(
         self,
         target_device: torch.device,
@@ -1251,6 +1736,29 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         # mapping (QKV fusion, etc.).
         if load_format == "diffusers":
             raise ValueError("HSDP is not supported with the diffusers adapter load format")
+        strategy = getattr(self.od_config, "hsdp_weight_load_strategy", "full")
+        if strategy == "pre_sharded":
+            if load_format != "default":
+                raise ValueError(
+                    f"hsdp_weight_load_strategy={strategy!r} currently supports only diffusion_load_format='default'"
+                )
+            if self.quant_config is not None:
+                raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support quantization yet")
+            if getattr(self.od_config, "lora_path", None):
+                raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support LoRA yet")
+        elif strategy != "full":
+            raise ValueError(f"Unknown hsdp_weight_load_strategy: {strategy!r}")
+
+        if strategy == "pre_sharded":
+            model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
+            discovered_modules = ModuleDiscovery.discover(model)
+            return self._load_model_with_pre_sharded_hsdp(
+                model,
+                discovered_modules,
+                hsdp_config,
+                target_device,
+            )
+
         # Same host-fallback bound as the ordinary path: with a quant config the
         # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so
         # over-wide NPU-unquantizable fallback weights would otherwise round-trip

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Loader-owned host-weight plans shared with diffusion offload backends."""
+"""Loader-owned host-weight plans shared by direct-checkpoint consumers."""
 
 from __future__ import annotations
 
@@ -205,7 +205,7 @@ def _build_source_map(
     if not model_to_ckpt:
         raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources")
     logger.info(
-        "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight",
+        "Indexed %d runtime tensor names from %d checkpoint keys for checkpoint preflight",
         len(model_to_ckpt),
         indexed_keys,
     )
@@ -313,28 +313,28 @@ def _validate_source_metadata(
                     )
 
 
-def build_checkpoint_mmap_plan(
+def build_checkpoint_binding_plan(
     pipeline: nn.Module,
     *,
     dit_modules: Sequence[tuple[str, nn.Module]],
     sources: Sequence[object],
     model_path: str | None,
     tensor_parallel_size: int,
-    use_hsdp: bool,
     online_quantization: bool,
+    use_pipeline_key_remap: bool = True,
     has_distilled_lora: bool = False,
 ) -> HostWeightPlanResult:
-    """Build a complete direct-checkpoint plan or return a fallback reason."""
+    """Build a complete binding plan for the pre-sharded HSDP consumer."""
     if tensor_parallel_size != 1:
         return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
-    if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
     if online_quantization:
         return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
     if has_distilled_lora:
         return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader")
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
+    if not callable(remap_fn) and use_pipeline_key_remap:
+        remap_fn = getattr(pipeline, "remap_checkpoint_key", None)
     if not callable(remap_fn):
         remap_fn = None
     adapter = get_direct_mmap_adapter(pipeline)
@@ -387,10 +387,37 @@ def build_checkpoint_mmap_plan(
     )
 
 
+def build_checkpoint_mmap_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    sources: Sequence[object],
+    model_path: str | None,
+    tensor_parallel_size: int,
+    use_hsdp: bool,
+    online_quantization: bool,
+    has_distilled_lora: bool = False,
+) -> HostWeightPlanResult:
+    """Build the DLO mmap plan without changing its HSDP eligibility."""
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+    return build_checkpoint_binding_plan(
+        pipeline,
+        dit_modules=dit_modules,
+        sources=sources,
+        model_path=model_path,
+        tensor_parallel_size=tensor_parallel_size,
+        online_quantization=online_quantization,
+        use_pipeline_key_remap=False,
+        has_distilled_lora=has_distilled_lora,
+    )
+
+
 __all__ = [
     "HostWeightPlan",
     "HostWeightPlanResult",
     "TensorBinding",
+    "build_checkpoint_binding_plan",
     "build_checkpoint_mmap_plan",
     "has_online_quantization",
 ]

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload_strategy,
@@ -1084,6 +1085,7 @@ class Cosmos3LanguageModel(nn.Module):
         rope_theta: float,
         mrope_section: list[int],
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -1093,21 +1095,21 @@ class Cosmos3LanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3UndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3UndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         # TODO: Not used right now, will be used in the future for prompt upsampler.
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
@@ -1228,6 +1230,10 @@ class Cosmos3VFMTransformer(nn.Module):
 
     _hsdp_shard_conditions = [_is_transformer_block]
 
+    # Standard unquantized vLLM linears need no value-dependent post-load
+    # processing on accelerators. Edge inherits the same loading contract.
+    _hsdp_pre_sharded_meta_post_load = True
+
     # Modules whose parameters must NOT be FSDP-sharded at the root level.
     # time_embedder is cast to fp32 by post_load_weights for precision; if it
     # were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16),
@@ -1346,6 +1352,10 @@ class Cosmos3VFMTransformer(nn.Module):
         language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs(
             getattr(od_config, "quantization_config", None)
         )
+        release_completed_blocks_to_meta = bool(
+            getattr(getattr(od_config, "parallel_config", None), "use_hsdp", False)
+            and getattr(od_config, "hsdp_weight_load_strategy", "full") == "pre_sharded"
+        )
         mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config)
         if mixed_precision_config is None:
             if mixed_precision_source == "additional_config_disabled":
@@ -1372,6 +1382,7 @@ class Cosmos3VFMTransformer(nn.Module):
             rope_theta=self.rope_theta,
             mrope_section=self.mrope_section,
             quant_config=language_model_quant_config,
+            release_completed_blocks_to_meta=release_completed_blocks_to_meta,
             prefix="language_model",
             **self._language_model_kwargs(),
         )
@@ -1399,24 +1410,24 @@ class Cosmos3VFMTransformer(nn.Module):
             self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
             self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
 
-        self.gen_layers = nn.ModuleList(
-            [
-                Cosmos3GenDecoderLayer(
-                    layer_idx=i,
-                    hidden_size=self.hidden_size,
-                    intermediate_size=self.intermediate_size,
-                    num_attention_heads=self.num_attention_heads,
-                    num_key_value_heads=self.num_key_value_heads,
-                    head_dim=self.head_dim,
-                    rms_norm_eps=self.rms_norm_eps,
-                    quant_config=gen_layers_quant_config,
-                    mlp_cls=self._gen_mlp_cls,
-                    qk_norm=self.qk_norm_for_diffusion,
-                    prefix=f"gen_layers.{i}",
-                )
-                for i in range(self.num_hidden_layers)
-            ]
-        )
+        self.gen_layers = nn.ModuleList()
+        for i in range(self.num_hidden_layers):
+            layer = Cosmos3GenDecoderLayer(
+                layer_idx=i,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                num_attention_heads=self.num_attention_heads,
+                num_key_value_heads=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                rms_norm_eps=self.rms_norm_eps,
+                quant_config=gen_layers_quant_config,
+                mlp_cls=self._gen_mlp_cls,
+                qk_norm=self.qk_norm_for_diffusion,
+                prefix=f"gen_layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.gen_layers.append(layer)
 
         self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None
         if mixed_precision_config is not None:

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelL
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 
 from .transformer_cosmos3 import (
     Cosmos3VFMTransformer,
@@ -258,6 +259,7 @@ class Cosmos3EdgeLanguageModel(nn.Module):
         mrope_section: list[int],
         use_und_k_norm_for_gen: bool,
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -267,22 +269,22 @@ class Cosmos3EdgeLanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3EdgeUndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    use_und_k_norm_for_gen=use_und_k_norm_for_gen,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3EdgeUndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                use_und_k_norm_for_gen=use_und_k_norm_for_gen,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(

vllm_omni/diffusion/models/utils.py

@@ -168,6 +168,28 @@ def init_parameters(
         init_parameters(child, dtype, device)
 
 
+def release_module_parameters_to_meta(module: nn.Module) -> None:
+    """Release a fully constructed module's parameter storage to meta.
+
+    The module must finish construction first so initialization code and vLLM
+    weight-loader metadata follow the ordinary CPU path. ``swap_tensors`` keeps
+    each registered Parameter object's identity while dropping only its storage.
+    Persistent buffers intentionally remain materialized.
+    """
+    for param in module.parameters():
+        if param.device.type == "meta":
+            continue
+
+        meta_tensor = torch.empty_like(param, device="meta")
+        if type(param) is nn.Parameter:
+            meta_param = nn.Parameter(meta_tensor, requires_grad=param.requires_grad)
+        else:
+            meta_param = param.__class__.__new__(param.__class__, meta_tensor)
+            meta_param.requires_grad_(param.requires_grad)
+        meta_param.__dict__.update(getattr(param, "__dict__", {}))
+        torch.utils.swap_tensors(param, meta_param)
+
+
 def init_parameters_preserving_vllm_linear(
     module: nn.Module,
     dtype: torch.dtype | None,

vllm_omni/engine/arg_utils.py

@@ -558,6 +558,7 @@ class OrchestratorArgs:
     # === Diffusion model config ===
     num_gpus: int | None = None
     model_class_name: str | None = None
+    hsdp_weight_load_strategy: str | None = None
     diffusion_load_format: str | None = None
     lora_path: list[str] | None = None
     lora_backend: str | None = None

vllm_omni/entrypoints/cli/serve.py

@@ -560,6 +560,15 @@ class OmniServeCommand(CLISubcommand):
                 "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
             ),
         )
+        omni_config_group.add_argument(
+            "--hsdp-weight-load-strategy",
+            choices=["full", "pre_sharded"],
+            default=None,
+            help=(
+                "HSDP checkpoint loading strategy: 'full' loads then shards; "
+                "'pre_sharded' shards meta parameters first and reads rank-local HF safetensors slices."
+            ),
+        )
         omni_config_group.add_argument(
             "--lora-path",
             type=str,