# BAGEL-7B-MoT FlowGRPO training

Last updated: 07/06/2026

[BAGEL-7B-MoT](https://github.com/ByteDance-Seed/BAGEL) is a
Mixture-of-Transformers model supporting both image understanding and
generation.  Unlike Qwen-Image, BAGEL is a **non-diffusers** model — it
cannot be loaded by diffusers and uses its own weight-loading path via
``NonDiffusersModelBase``.  See
[How to Integrate a Non-Diffusers Model for FlowGRPO Training](../../../docs/contributing/integrating_a_non_diffusers_model.md)
for the integration architecture.

## Prerequisites

- Install VeRL-Omni (see [installation guide](../../../docs/start/install.md)).

- 4 GPUs or 8 NPUs. Run commands from the repository root.

- Download the checkpoint:

  ```bash
  huggingface-cli download ByteDance-Seed/BAGEL-7B-MoT --local-dir ~/models/ByteDance-Seed/BAGEL-7B-MoT
  ```

## OCR training

We use an OCR (optical character recognition) dataset that provides
ground-truth text for evaluating image-generation quality.  Prompts are
stored in standard chat-message format for the agent loop (see
``bagel_ocr.py``).

### Prepare the dataset

Preprocess the raw OCR data into parquet:

```bash
export WORKSPACE=${WORKSPACE:-$HOME}

python3 examples/flowgrpo_trainer/data_process/bagel_ocr.py \
  --model_path ~/models/ByteDance-Seed/BAGEL-7B-MoT \
  --input_dir ~/data/ocr \
  --output_dir $WORKSPACE/data/ocr/bagel
```

This produces ``$WORKSPACE/data/ocr/bagel/train.parquet`` and
``test.parquet``.

### Run training

For GPU:
```bash
bash examples/flowgrpo_trainer/bagel/run_bagel_ocr_lora.sh
```

For NPU:  
```bash
bash examples/flowgrpo_trainer/bagel/run_bagel_ocr_lora_npu.sh
```

The launch script uses a [Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct)
reward model with vLLM rollout (TP=4) and the ``genrm_ocr.py`` custom reward
function.

## PickScore training

PickScore evaluates image-text alignment using a
[CLIP-based model](https://huggingface.co/yuvalkirstain/PickScore_v1).  The
reward function lives entirely in ``verl_omni/utils/reward_score/pickscore_reward.py``
— there is **no** separate vLLM reward model deployment, so the GPU is shared
between the actor and the reward computation.

### Prepare the dataset

The raw PickScore dataset (``train.txt`` / ``test.txt``) should be downloaded
from the [flow_grpo repository](https://github.com/yifan123/flow_grpo/tree/main/dataset/pickscore).

Preprocess for BAGEL:

```bash
python3 examples/flowgrpo_trainer/data_process/bagel_pickscore.py \
  --model_path ~/models/ByteDance-Seed/BAGEL-7B-MoT \
  --input_dir ~/data/pickscore \
  --output_dir $WORKSPACE/data/pickscore/bagel
```

This produces ``$WORKSPACE/data/pickscore/bagel/train.parquet`` and
``test.parquet``.

### Run LoRA training

```bash
bash examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora.sh
```

Key configuration differences from OCR:
- No ``reward.reward_model.*`` flags — PickScore runs as a custom reward
  function on the rollout GPU.
- Higher ``noise_level`` (``1.3`` vs ``0.7``) and SDE window
  (``sde_window_size=2``, ``range=[0,7]``) to provide sufficient exploration
  for text-alignment learning.

### Run full-weight (non-LoRA) training

A full-weight training variant is available that trains the entire
generation pathway (``moe_gen`` parameters) while keeping the understanding
pathway frozen:

```bash
bash examples/flowgrpo_trainer/bagel/run_bagel_pickscore.sh
```

Key differences from the LoRA variant:

| Aspect | LoRA | Full-weight |
|---|---|---|
| Script | ``examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora.sh`` | ``examples/flowgrpo_trainer/bagel/run_bagel_pickscore.sh`` |
| Strategy | default | ``fsdp2`` (required for mixed ``requires_grad``) |
| Trainable params | Low-rank adapters on ``*_moe_gen`` | All ``moe_gen`` parameters (``requires_grad`` set by ``configure_trainable_params``) |
| ``lora_rank`` / ``lora_alpha`` | 64 / 128 | N/A |
| ``sde_window_size`` | 2 | 3 (more exploration for full-weight) |

**Why FSDP2 is required.** FSDP1 does not natively support mixed
``requires_grad`` within a single wrapped module — some parameters frozen,
others trainable.  FSDP2 handles this correctly and also reshards layer
parameters after forward, reducing peak memory during gradient
checkpointing.  The understanding pathway (``moe_und``) is not a LoRA
wrapper replacement but simply has ``requires_grad=False`` set by the
``configure_trainable_params`` hook.

## Key differences from Qwen-Image

| Aspect | Qwen-Image | BAGEL-7B-MoT |
|---|---|---|
| Model loading | diffusers | Custom ``from_pretrained`` via ``NonDiffusersModelBase`` |
| Architecture | Auto-detected | Explicit: ``+actor_rollout_ref.model.architecture=OmniBagelForConditionalGeneration`` |
| Deploy config | Not needed | ``bagel_deploy_config.yaml`` (single-stage topology) |
| LoRA targets | ``*_proj`` layers | ``*_proj`` + ``*_moe_gen`` (MoT dual-pathway) |
| FSDP prefixes | ``transformer_blocks.`` | ``layers.`` |
| CFG | Standard true CFG | 3-branch (gen / text-uncond / img-uncond) with global renormalisation |
| Timestep convention | ``t / 1000`` | Raw sigma with SD3-style shift of 3.0 |

## Further reading

- [How to Integrate a Non-Diffusers Model for FlowGRPO Training](../../../docs/contributing/integrating_a_non_diffusers_model.md) — full integration guide using BAGEL as the worked example
- [vLLM-Omni BAGEL docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/)