[Feat] support hsdp for Bagel (#3150)

Signed-off-by: siyuan.lei <siyuanlei37@gmail.com>
Signed-off-by: lsyyyyy <siyuanlei37@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

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** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |

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

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 `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.

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.

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

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
 

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