diff --git a/docs/configuration/hsdp_pre_sharded_loading.md b/docs/configuration/hsdp_pre_sharded_loading.md new file mode 100644 index 000000000..cb8507b3d --- /dev/null +++ b/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. diff --git a/tests/diffusion/distributed/test_cosmos3_pre_sharded.py b/tests/diffusion/distributed/test_cosmos3_pre_sharded.py new file mode 100644 index 000000000..57af9c72b --- /dev/null +++ b/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, + ) diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 3c3805a82..b1911e97c 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/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])) diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 8a2ee0900..60c0cd7d6 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 58fe4b10b..56c33b9e0 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/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 diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1d106932d..496d19320 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 526f4c9e5..d77e44194 100644 --- a/vllm_omni/config/omni_config.py +++ b/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 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0c46e9c22..5f959050a 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5e54637c2..0811d2634 100644 --- a/vllm_omni/diffusion/data.py +++ b/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, diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index 82c1d1986..7e20d8471 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/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, ) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7cbff273c..c74808efb 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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 diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index f0ea8024d..3bb9f0c5a 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/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", ] diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 7241082e8..b060ea420 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/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: diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py index c9a7d856e..fcc19123b 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py +++ b/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( diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index f55996a76..05a6a0a31 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/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, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b9ed35de0..458a6237d 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ac1c3cda1..9c3aeb5b5 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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,