Skip to content

Enable Bagel FlowGRPO LoRA training with vLLM-omni rollout#1

Merged
princepride merged 1 commit into
bagel-pipelinefrom
bagel-fsdp-training
May 3, 2026
Merged

Enable Bagel FlowGRPO LoRA training with vLLM-omni rollout#1
princepride merged 1 commit into
bagel-pipelinefrom
bagel-fsdp-training

Conversation

@princepride
Copy link
Copy Markdown
Collaborator

  • Add DiffusionModelBase implementation for Bagel (bagel.py) with custom model loading, scheduler, and forward pass
  • Add build_module hook in DiffusionModelBase and DiffusersFSDPEngine to support non-standard model loading (e.g. Bagel's BagelForTraining)
  • Fix scheduler device/precision mismatch in FlowMatchSDEDiscreteScheduler (index_for_timestep nearest-neighbor, move sigmas to sample device)
  • Fix pipeline_bagel to stack trajectory tensors for proper batching
  • Add system prompt and negative prompt to OCR training data
  • Add error handling in reward_fn for vLLM reward model failures
  • Guard prompt_embeds access for models that don't produce embeddings
  • Add profiler support to diffusion trainer (config + start/stop hooks)
  • Fix prompt_ids -> prompt_token_ids rename in agent loop server call
  • Update run_bagel_flowgrpo.sh with tuned hyperparameters

Co-authored-by: Claude

What does this PR do?

Add concise overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review.

Checklist Before Starting

  • Search for similar PRs. Paste at least one query link here: ...
  • Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
    • {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
    • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
    • {type} is in feat, fix, refactor, chore, test
    • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
    • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

- Add DiffusionModelBase implementation for Bagel (bagel.py) with custom
  model loading, scheduler, and forward pass
- Add build_module hook in DiffusionModelBase and DiffusersFSDPEngine to
  support non-standard model loading (e.g. Bagel's BagelForTraining)
- Fix scheduler device/precision mismatch in FlowMatchSDEDiscreteScheduler
  (index_for_timestep nearest-neighbor, move sigmas to sample device)
- Fix pipeline_bagel to stack trajectory tensors for proper batching
- Add system prompt and negative prompt to OCR training data
- Add error handling in reward_fn for vLLM reward model failures
- Guard prompt_embeds access for models that don't produce embeddings
- Add profiler support to diffusion trainer (config + start/stop hooks)
- Fix prompt_ids -> prompt_token_ids rename in agent loop server call
- Update run_bagel_flowgrpo.sh with tuned hyperparameters

Co-authored-by: Claude
Co-authored-by: Cursor <cursoragent@cursor.com>
@princepride princepride merged commit 83ef93b into bagel-pipeline May 3, 2026
2 of 4 checks passed
@princepride princepride deleted the bagel-fsdp-training branch May 3, 2026 03:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant