diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 69c81737d..702ea4e87 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -107,7 +107,7 @@ The following tables show which models support each feature: | Model | ⚑TeaCache | ⚑Cache-DiT | πŸ”€SP (Ulysses & Ring) | πŸ”€CFG-Parallel | πŸ”€Tensor-Parallel | πŸ”€HSDP | πŸ’ΎCPU Offload (Layerwise) | πŸ’ΎVAE-Patch-Parallel | πŸ’ΎQuantization | πŸ”„Step Execution | |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| -| **Bagel** | βœ… | βœ… | βœ… | βœ… | βœ… | ❌ | βœ… | ❌ | ❌ | ❌ | +| **Bagel** | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | ❌ | ❌ | ❌ | | **FLUX.1-dev** | βœ… | βœ… | ❌ | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | | **FLUX.1-schnell** | ❌ | βœ… | ❌ | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | | **FLUX.2-klein** | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md index 9955fd90d..9c2e31ff3 100644 --- a/examples/offline_inference/bagel/README.md +++ b/examples/offline_inference/bagel/README.md @@ -195,6 +195,25 @@ stages: devices: "0,1" ``` +#### Hybrid Sharded Data Parallel (HSDP) + +For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. + +1. **Enable HSDP**: Set `use_hsdp: true`. +2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`). +3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups. +4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`). + +Example configuration for HSDP across 4 GPUs: +```yaml + - stage_id: 1 + devices: "0,1,2,3" + parallel_config: + use_hsdp: true + hsdp_shard_size: 4 + hsdp_replicate_size: 1 +``` + Then pass the custom deploy YAML: ```bash diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md index 763927222..ea8abe43f 100644 --- a/examples/online_serving/bagel/README.md +++ b/examples/online_serving/bagel/README.md @@ -68,6 +68,25 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size Or set `tensor_parallel_size` per stage in a custom deploy YAML. +#### Hybrid Sharded Data Parallel (HSDP) + +For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. + +1. **Enable HSDP**: Set `use_hsdp: true`. +2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`). +3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups. +4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`). + +Example configuration for HSDP across 4 GPUs: +```yaml + - stage_id: 1 + devices: "0,1,2,3" + parallel_config: + use_hsdp: true + hsdp_shard_size: 4 + hsdp_replicate_size: 1 +``` + ### Multi-Node Deployment Deploy each stage on a **separate node** for better resource utilization. Replace `` with the actual IP address of your orchestrator node. diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index ecefc1c6b..1374a4752 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -150,6 +150,19 @@ def _get_diffusion_feature_cases(model: str): id="single_card_layerwise_offload", marks=SINGLE_CARD_FEATURE_MARKS, ), + # Hybrid Sharded Data Parallel (4 GPUs) + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--use-hsdp", + "--hsdp-shard-size", + "4", + ], + ), + id="parallel_hsdp_4", + marks=PARALLEL_FEATURE_MARKS, + ), ] @@ -172,6 +185,7 @@ def test_bagel( - Ulysses-SP (degree=2) - Ring-Attention (degree=2) - Layerwise Offloading + - Hybrid Sharded Data Parallel (size=4) Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py, which checks output dimensions and basic correctness. diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 4c6fb0709..eb89e7e81 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model if TYPE_CHECKING: @@ -278,7 +279,7 @@ class DiffusersPipelineLoader: with set_default_torch_dtype(od_config.dtype): if od_config.parallel_config.use_hsdp: model = self._load_model_with_hsdp( - od_config, load_format=load_format, custom_pipeline_name=custom_pipeline_name + od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name ) else: with target_device: @@ -529,6 +530,7 @@ class DiffusersPipelineLoader: def _load_model_with_hsdp( self, od_config: OmniDiffusionConfig, + target_device: torch.device, load_format: str = "default", custom_pipeline_name: str | None = None, ) -> nn.Module: @@ -578,4 +580,19 @@ class DiffusersPipelineLoader: for name, trans in transformers_to_shard: logger.debug("Applying HSDP to %s", name) apply_hsdp_to_model(trans, hsdp_config) + + # # HSDP only shards transformer modules. All other runtime modules must + # # be placed on the execution device explicitly after sharding. + discovered_modules = ModuleDiscovery.discover(model) + 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) + return model diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index c01782790..d33b26f60 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -755,6 +755,12 @@ class Qwen2MoTDecoderLayer(nn.Module): class Qwen2MoTModel(Qwen2PreTrainedModel): _layerwise_offload_blocks_attrs = ["layers"] + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return "layers" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + def __init__( self, config, @@ -792,10 +798,10 @@ class Qwen2MoTModel(Qwen2PreTrainedModel): def forward( self, - packed_query_sequence: torch.Tensor, - query_lens: torch.Tensor, - packed_query_position_ids: torch.Tensor, - packed_query_indexes: torch.Tensor, + packed_query_sequence: torch.Tensor | None = None, + query_lens: torch.Tensor | None = None, + packed_query_position_ids: torch.Tensor | None = None, + packed_query_indexes: torch.Tensor | None = None, past_key_values: NaiveCache | None = None, key_values_lens: torch.Tensor | None = None, packed_key_value_indexes: torch.Tensor | None = None, @@ -804,7 +810,20 @@ class Qwen2MoTModel(Qwen2PreTrainedModel): mode="und", packed_vae_token_indexes=None, packed_text_indexes=None, + packed_text_ids: torch.Tensor | None = None, + return_embeddings_only: bool = False, ) -> BaseNavitOutputWithPast: + if packed_query_sequence is None: + if packed_text_ids is None: + raise ValueError("Either packed_query_sequence or packed_text_ids must be provided.") + packed_query_sequence = self.embed_tokens(packed_text_ids) + + if return_embeddings_only: + return BaseNavitOutputWithPast( + packed_query_sequence=packed_query_sequence, + past_key_values=past_key_values, + ) + # create position embeddings to be shared across the decoder layers cos, sin = self.rotary_emb(packed_query_sequence, packed_query_position_ids.unsqueeze(0)) cos = cos.squeeze(0) @@ -896,10 +915,10 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel): def forward( self, - packed_query_sequence: torch.Tensor, - query_lens: torch.Tensor, - packed_query_position_ids: torch.Tensor, - packed_query_indexes: torch.Tensor, + packed_query_sequence: torch.Tensor | None = None, + query_lens: torch.Tensor | None = None, + packed_query_position_ids: torch.Tensor | None = None, + packed_query_indexes: torch.Tensor | None = None, past_key_values: NaiveCache | None = None, key_values_lens: torch.Tensor | None = None, packed_key_value_indexes: torch.Tensor | None = None, @@ -908,6 +927,8 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel): mode="und", packed_vae_token_indexes=None, packed_text_indexes=None, + packed_text_ids: torch.Tensor | None = None, + return_embeddings_only: bool = False, ) -> BaseNavitOutputWithPast: outputs = self.model( packed_query_sequence=packed_query_sequence, @@ -922,6 +943,8 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel): mode=mode, packed_vae_token_indexes=packed_vae_token_indexes, packed_text_indexes=packed_text_indexes, + packed_text_ids=packed_text_ids, + return_embeddings_only=return_embeddings_only, ) return outputs @@ -1242,14 +1265,12 @@ class Bagel(nn.Module): packed_key_value_indexes: torch.LongTensor, key_values_lens: torch.IntTensor, ): - packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) - extra_inputs = {} if self.use_moe: extra_inputs = {"mode": "und"} output = self.language_model.forward( - packed_query_sequence=packed_text_embedding, + packed_text_ids=packed_text_ids, query_lens=text_token_lens, packed_query_position_ids=packed_text_position_ids, packed_query_indexes=packed_text_indexes, @@ -1355,10 +1376,6 @@ class Bagel(nn.Module): key_values_lens: torch.IntTensor, packed_key_value_indexes: torch.Tensor, ): - packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) - packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) - packed_sequence[packed_text_indexes] = packed_text_embedding - padded_latent = vae_model.encode(padded_images) p = self.latent_patch_size @@ -1371,9 +1388,6 @@ class Bagel(nn.Module): packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids) packed_timestep_embeds = self.time_embedder(packed_timesteps) packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + packed_pos_embed - if packed_latent.dtype != packed_sequence.dtype: - packed_latent = packed_latent.to(packed_sequence.dtype) - packed_sequence[packed_vae_token_indexes] = packed_latent extra_inputs = {} if self.use_moe: @@ -1384,7 +1398,7 @@ class Bagel(nn.Module): } output = self.language_model.forward( - packed_query_sequence=packed_sequence, + packed_text_ids=packed_text_ids, query_lens=packed_seqlens, packed_query_position_ids=packed_position_ids, packed_query_indexes=packed_indexes, @@ -1477,7 +1491,10 @@ class Bagel(nn.Module): packed_key_value_indexes: torch.LongTensor, key_values_lens: torch.IntTensor, ): - packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) + packed_text_embedding = self.language_model.forward( + packed_text_ids=packed_text_ids, + return_embeddings_only=True, + ).packed_query_sequence packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) packed_sequence[packed_text_indexes] = packed_text_embedding @@ -1675,7 +1692,6 @@ class Bagel(nn.Module): curr_tokens = packed_start_tokens while step < max_length: generated_sequence.append(curr_tokens) - packed_text_embedding = self.language_model.model.embed_tokens(curr_tokens) query_lens = torch.ones_like(curr_tokens) packed_query_indexes = torch.cumsum(key_values_lens, dim=0) + torch.arange( 0, @@ -1690,7 +1706,7 @@ class Bagel(nn.Module): packed_key_value_indexes = torch.cat(uppacked, dim=0) output = self.language_model( - packed_query_sequence=packed_text_embedding, + packed_text_ids=curr_tokens, query_lens=query_lens, packed_query_position_ids=packed_query_position_ids, packed_query_indexes=packed_query_indexes, @@ -2283,7 +2299,10 @@ class Bagel(nn.Module): packed_position_ids, ) - packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) + packed_text_embedding = self.language_model.forward( + packed_text_ids=packed_text_ids, + return_embeddings_only=True, + ).packed_query_sequence packed_sequence = packed_text_embedding.new_zeros((int(local_seqlens.sum()), self.hidden_size)) packed_sequence[local_text_indexes] = packed_text_embedding @@ -2338,7 +2357,10 @@ class Bagel(nn.Module): return self._gather_vae_for_sp(local_v_t) # Original non-SP path - packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) + packed_text_embedding = self.language_model.forward( + packed_text_ids=packed_text_ids, + return_embeddings_only=True, + ).packed_query_sequence packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) packed_sequence[packed_text_indexes] = packed_text_embedding @@ -2393,7 +2415,10 @@ class Bagel(nn.Module): cfg_batched: dict | None = None, ): # Build query sequence (identical for all CFG branches) - packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) + packed_text_embedding = self.language_model.forward( + packed_text_ids=packed_text_ids, + return_embeddings_only=True, + ).packed_query_sequence packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) packed_sequence[packed_text_indexes] = packed_text_embedding diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index d08a2cdc8..d8be715d7 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -251,6 +251,7 @@ class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerM self.language_model = Qwen2MoTForCausalLM( llm_config, parallel_config=parallel_config, quant_config=quant_config, prefix="bagel.language_model" ) + self.transformer = self.language_model.model ae_params: AutoEncoderParams = default_ae_params() self.vae = AutoEncoder(ae_params) @@ -284,11 +285,15 @@ class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerM ) ] - # When quantization is enabled, vLLM linear layers live on meta - # device until the weight loader materializes them. Calling - # .to(device) would fail on those meta tensors, so we skip it - # entirely and let the weight loader handle device placement. - if quant_config is None and not od_config.enable_layerwise_offload: + # Defer device placement to the weight-loading/offload path in three cases: + # 1. Quantization: When quantization is enabled, vLLM linear layers live on meta + # device until the weight loader materializes them. Calling .to(device) would fail on those meta tensors, + # so we skip it entirely and let the weight loader handle device placement. + # 2. Layerwise offload: modules should be initialized on CPU first, then + # selectively materialized/moved by the offloader. + # 3. HSDP: weights should be loaded on CPU first and sharded afterwards, + # rather than eagerly placing the full model on one GPU. + if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp): self.to(self.device) self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler