diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index 69506c928..42675c7c0 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/vllm_omni/diffusion/distributed/hsdp.py @@ -88,6 +88,7 @@ def _create_hsdp_mesh( def apply_hsdp_to_model( model: nn.Module, hsdp_config: HSDPInferenceConfig, + target_device: torch.device | None = None, ) -> nn.Module: """ Apply HSDP sharding to a model that already has weights loaded. @@ -98,6 +99,10 @@ def apply_hsdp_to_model( Args: model: Model instance with weights already loaded hsdp_config: HSDP configuration with HSDP mesh dimensions + target_device: Worker's execution device. When the model declares + _hsdp_ignored_modules, those modules are excluded from FSDP's + mesh-driven device placement, so the caller must specify where to + put them. Optional only when there are no ignored modules. Returns: HSDP-wrapped model ready for inference @@ -158,6 +163,40 @@ def apply_hsdp_to_model( if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0: 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. + # See _hsdp_ignored_modules on each model class. These params keep their + # original dtype/storage and are not subject to MixedPrecisionPolicy on the root + # FSDP wrap. Useful for small auxiliary modules (e.g., timestep embedders) that + # need to stay in a higher-precision dtype than the bulk of the model. + ignored_module_names = getattr(model, "_hsdp_ignored_modules", []) or [] + ignored_params: set[nn.Parameter] = set() + # FSDP wraps move sharded params to mesh devices automatically, but + # ignored modules stay wherever load_weights left them (CPU under HSDP). + # The caller (diffusers_loader / similar) is responsible for telling us + # which device the worker should run on, matching the convention used + # for VAEs / encoders / resident_modules. + if ignored_module_names and target_device is None: + raise ValueError( + f"Model {type(model).__name__} declares _hsdp_ignored_modules=" + f"{ignored_module_names} but apply_hsdp_to_model was called " + "without target_device. The caller must pass target_device so " + "the ignored modules can be placed on the worker's execution device." + ) + for mod_name in ignored_module_names: + sub_mod = getattr(model, mod_name, None) + if sub_mod is None: + logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name) + continue + sub_mod.to(target_device) + ignored_params.update(sub_mod.parameters()) + if ignored_params: + logger.info( + "HSDP excluding %d parameter tensors from sharding (modules: %s, moved to %s)", + len(ignored_params), + ignored_module_names, + target_device, + ) + # Apply HSDP sharding, this will automatically handle weight distribution shard_model( model, @@ -165,6 +204,7 @@ def apply_hsdp_to_model( mp_policy=mp_policy, mesh=device_mesh, hsdp_shard_conditions=hsdp_shard_conditions, + ignored_params=ignored_params if ignored_params else None, ) for param in model.parameters(): @@ -181,8 +221,17 @@ def shard_model( mp_policy: MixedPrecisionPolicy | None = None, mesh: DeviceMesh | None = None, hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]], + ignored_params: set[nn.Parameter] | None = None, ) -> None: - """Apply HSDP sharding to model modules based on shard conditions.""" + """Apply HSDP sharding to model modules based on shard conditions. + + ignored_params (if provided) are excluded from the root fully_shard + wrap, so they are not collected into the root flat-parameter, are not + subject to MixedPrecisionPolicy, and retain their original dtype. + Per-submodule shard wraps do not receive ignored_params because the + ignored modules are expected to live at the root level, not inside any + block matched by hsdp_shard_conditions. + """ hsdp_kwargs: dict[str, Any] = { "reshard_after_forward": reshard_after_forward, "mesh": mesh, @@ -198,5 +247,12 @@ def shard_model( if num_sharded == 0: raise ValueError("No modules were sharded") - fully_shard(model, **hsdp_kwargs) - logger.info("Sharded %d modules + root", num_sharded) + root_kwargs = dict(hsdp_kwargs) + if ignored_params: + root_kwargs["ignored_params"] = ignored_params + fully_shard(model, **root_kwargs) + logger.info( + "Sharded %d modules + root (ignored_params=%d)", + num_sharded, + len(ignored_params) if 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 e5ef5e319..b2339843e 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -661,7 +661,7 @@ class DiffusersPipelineLoader: # Apply HSDP sharding to all transformers for name, trans in transformers_to_shard: logger.debug("Applying HSDP to %s", name) - apply_hsdp_to_model(trans, hsdp_config) + apply_hsdp_to_model(trans, hsdp_config, target_device=target_device) # # HSDP only shards transformer modules. All other runtime modules must # # be placed on the execution device explicitly after sharding. diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 163158a39..590866846 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -976,6 +976,12 @@ class Cosmos3VFMTransformer(nn.Module): _hsdp_shard_conditions = [_is_transformer_block] + # 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), + # the dtype upcast would be silently reverted, causing dtype mismatch in forward. + _hsdp_ignored_modules = ["time_embedder"] + _sp_plan = { "gen_sp_prepare": { 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),