anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Config Refactor]: Remove bagel yaml (#2936)

Browse Source 
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
This commit is contained in:
汪志鹏andSYLAR authored and GitHub committed 2026-04-24 21:21:30 +08:00
1 parent 22ffb85063
commit 37e5954acf
31 files changed
+1067 -1420
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/examples/offline_inference/bagel.md
+208 -185
View File
@@ -2,46 +2,61 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
## Set up
## Setup
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
## Run examples
## Architecture
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
Get into the bagel folder
| Topology | Stages | Description |
| :------- | :----- | :---------- |
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
## Quick Start
```bash
cd examples/offline_inference/bagel
# Default two-stage mode (auto-detected)
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
# Single-stage mode
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
### Modality Control
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
## Modality Control
#### Text to Image (text2img)
Control the mode using the `--modality` argument:
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
| Modality | Input | Output | Description |
| :------- | :---- | :----- | :---------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation (language model mode) |
Generate images from text prompts:
### Text to Image (text2img)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
--prompts "A cute cat" \
--steps 50
```
#### Image to Image (img2img)
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
- **Stages Used**: Stage 1 (DiT) only
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
Transform images based on text prompts:
### Image to Image (img2img)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
@@ -50,13 +65,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompts "Let the woman wear a blue dress"
```
#### Image to Text (img2text)
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
Generate text descriptions from images:
### Image to Text (img2text)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
@@ -65,196 +74,210 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompts "Describe this image in detail"
```
#### Text to Text (text2text)
- **Pipeline**: Text → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: No visual components involved, operates as pure language model
Pure text generation:
### Text to Text (text2text)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?"
# You can load prompts from a text file (one prompt per line):
# Load prompts from a text file (one prompt per line):
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--txt-prompts /path/to/prompts.txt
```
### Inference Steps
## Think Mode
Control the number of inference steps for image generation:
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
```bash
# You can adjust steps to 100 to improve image quality
# Think + text2img: plan before generating
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--steps 50 \
--prompts "A cute cat"
```
--prompts "A futuristic city with flying cars" \
--think \
--max-think-tokens 1000
### Key arguments
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
#### 📌 Command Line Arguments (end2end.py)
| Argument | Type | Default | Description |
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts directly |
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
| `--steps` | int | `50` | Number of inference steps |
| `--stage-configs-path` | string | `None` | Custom stage config file path |
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--enable-stats` | flag | `False` | Enable statistics logging |
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
| `--batch-timeout` | int | `5` | Batch timeout |
| `--init-timeout` | int | `300` | Initialization timeout |
------
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
**Stage 0 - Thinker (LLM Stage)**
| Parameter | Value | Description |
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `tensor_parallel_size` | `1` | Tensor parallel size |
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
------
**Stage 1 - DiT (Diffusion Stage)**
| Parameter | Value | Description |
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
| `engine_input_source` | `[0]` | Input source from Stage 0 |
------
#### Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)).
1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
Example configuration for TP=2 on GPUs 0 and 1:
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
------
#### 🔗 Runtime Configuration
| Parameter | Value | Description |
| :-------------------- | :------ | :------------------------------- |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## Using Mooncake Connector
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
### Prerequisites
Install the Mooncake transfer engine:
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
### Step 1: Start the Mooncake Master
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
### Step 2: Run Offline Inference with Mooncake
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
```bash
cd examples/offline_inference/bagel
# Text to Image with Mooncake
# Think + img2img: reason about the edit
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--modality img2img \
--image-path /path/to/image.jpg \
--prompts "Make it look like a watercolor painting" \
--think
# Image to Text with Mooncake
# Think + img2text: reason before describing
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--prompts "What is happening in this image?" \
--think
# Text to Text with Mooncake
# Think + text2text: chain-of-thought reasoning
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--prompts "Solve: 23 * 47" \
--think
```
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
Think mode parameters:
------
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--think` | `False` | Enable thinking mode |
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
## Classifier-Free Guidance (CFG)
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A photorealistic portrait" \
--cfg-text-scale 6.0 \
--cfg-img-scale 2.0 \
--negative-prompt "blurry, low quality, distorted" \
--cfg-interval 0.4 1.0 \
--cfg-renorm-type global \
--cfg-renorm-min 0.0
```
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
## Deployment Topologies
### Two-Stage (Default)
The default topology auto-detected from the model. No extra flags needed.
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
```
The pipeline is defined in [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
### Single-Stage
Pass the single-stage deploy config via `--deploy-config`:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
See [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
### Tensor Parallelism (TP)
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
```yaml
# Example: TP=2 on GPUs 0,1 for the Thinker stage
stages:
- stage_id: 0
tensor_parallel_size: 2
devices: "0,1"
```
Then pass the custom deploy YAML:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config /path/to/custom_bagel.yaml
```
### FP8 Quantization
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--quantization fp8
```
## Command Line Reference
### Core Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts |
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
| `--output` | string | `.` | Output directory for saved images |
| `--steps` | int | `50` | Number of diffusion inference steps |
| `--seed` | int | `None` | Random seed for reproducibility |
### Think Mode Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
### CFG Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
### Engine Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
| `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` |
| `--worker-backend` | choice | `process` | `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
| `--log-stats` | flag | `False` | Enable statistics logging |
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
## FAQ
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Two-stage VRAM usage:**
| Stage | VRAM |
| :---- | :--- |
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
| Stage 1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
## Example materials