[Feat] support layerwise offload for Bagel (#2734)

Signed-off-by: Nick Cao <ncao@redhat.com>
Signed-off-by: siyuan.lei <siyuanlei37@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -195,6 +195,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
 | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
 | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally

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

tests/e2e/online_serving/test_bagel_expansion.py

@@ -9,6 +9,7 @@ Coverage:
 - Tensor-Parallel
 - Ulysses-SP
 - Ring-Attention
+- Layerwise Offloading
 
 assert_diffusion_response validates successful generation and the expected
 512x512 resolution.
@@ -31,7 +32,7 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 def _get_diffusion_feature_cases(model: str):
     """Return L4 diffusion feature cases for Bagel.
     TeaCache, Cache-DiT, CFG-Parallel, Tensor-Parallel,
-    Ulysses-SP, Ring-Attention.
+    Ulysses-SP, Ring-Attention, Layerwise Offloading.
     """
 
     return [
@@ -111,6 +112,15 @@ def _get_diffusion_feature_cases(model: str):
             id="sp_ring_2",
             marks=PARALLEL_FEATURE_MARKS,
         ),
+        # Layerwise Offloading (single-card)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--enable-layerwise-offload"],
+            ),
+            id="single_card_layerwise_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
     ]
 
 
@@ -132,6 +142,7 @@ def test_bagel(
     - Tensor-Parallel (size=2)
     - Ulysses-SP (degree=2)
     - Ring-Attention (degree=2)
+    - Layerwise Offloading
 
     Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
     which checks output dimensions and basic correctness.

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -738,6 +738,8 @@ class Qwen2MoTDecoderLayer(nn.Module):
 
 
 class Qwen2MoTModel(Qwen2PreTrainedModel):
+    _layerwise_offload_blocks_attrs = ["layers"]
+
     def __init__(
         self,
         config,

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -288,7 +288,7 @@ class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerM
         # 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:
+        if quant_config is None and not od_config.enable_layerwise_offload:
             self.to(self.device)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler