From 755a2d9d9628e48b78f6ae39d7fce1b983501cfd Mon Sep 17 00:00:00 2001 From: lsyyyyy Date: Sat, 25 Apr 2026 15:44:56 +0800 Subject: [PATCH] [Feat] support layerwise offload for Bagel (#2734) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Nick Cao Signed-off-by: siyuan.lei Signed-off-by: 汪志鹏 Co-authored-by: Nick Cao Co-authored-by: Claude Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> Co-authored-by: 汪志鹏 --- docs/user_guide/diffusion/cpu_offload_diffusion.md | 1 + docs/user_guide/diffusion_features.md | 2 +- tests/e2e/online_serving/test_bagel_expansion.py | 13 ++++++++++++- .../diffusion/models/bagel/bagel_transformer.py | 2 ++ vllm_omni/diffusion/models/bagel/pipeline_bagel.py | 2 +- 5 files changed, 17 insertions(+), 3 deletions(-) diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index b85c7e408..1d3f1811a 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 11ed1330c..be1602788 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/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index 21142439b..21fdc314c 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/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. diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index d1254f845..cff09db3c 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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, diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 0290e9502..d08a2cdc8 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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