diff --git a/examples/flowgrpo_trainer/bagel/README.md b/examples/flowgrpo_trainer/bagel/README.md index d50b1880c..57c78040c 100644 --- a/examples/flowgrpo_trainer/bagel/README.md +++ b/examples/flowgrpo_trainer/bagel/README.md @@ -127,6 +127,26 @@ 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. +## V1 (sync) training + +A V1-trainer counterpart of the PickScore LoRA recipe is available, using +TransferQueue + ReplayBuffer via `verl_omni.trainer.main_diffusion_v1`: + +```bash +bash examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora_v1.sh +``` + +Same dataset preparation and config as the PickScore LoRA recipe above; the +only differences are the entrypoint and: + +```text +trainer.use_v1=true +trainer.v1.trainer_mode=sync +``` + +See the [Diffusion V1 guide](../../../docs/start/diffusion_v1.md) for the +general V1 trainer architecture (documented there using the SD3.5 recipe). + ## Key differences from Qwen-Image | Aspect | Qwen-Image | BAGEL-7B-MoT | diff --git a/examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora_v1.sh b/examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora_v1.sh new file mode 100755 index 000000000..63c5de011 --- /dev/null +++ b/examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora_v1.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +# Bagel LoRA RL (V1 trainer: TransferQueue + ReplayBuffer + sync mode), vllm_omni +# rollout (FlowGRPO) with PickScore reward. +# +# V1 counterpart of run_bagel_pickscore_lora.sh via main_diffusion_v1 +# (trainer.v1.trainer_mode=sync). +# +# Prerequisite: preprocess the PickScore dataset for BAGEL: +# python examples/flowgrpo_trainer/data_process/bagel_pickscore.py \ +# --model_path ~/models/ByteDance-Seed/BAGEL-7B-MoT \ +# --input_dir ~/data/pickscore \ +# --output_dir ~/data/pickscore/bagel +# +# Raw dataset (train.txt / test.txt) from: +# https://github.com/yifan123/flow_grpo/tree/main/dataset/pickscore +set -x + +# Set WORKSPACE to any writable directory; defaults to $HOME +WORKSPACE=${WORKSPACE:-$HOME} + +pickscore_train_path=$WORKSPACE/data/pickscore/bagel/train.parquet +pickscore_test_path=$WORKSPACE/data/pickscore/bagel/test.parquet + +BAGEL_DEPLOY_CONFIG=${BAGEL_DEPLOY_CONFIG:-"$(dirname "$0")/bagel_deploy_config.yaml"} + +model_name=~/models/ByteDance-Seed/BAGEL-7B-MoT +reward_function_path=verl_omni/utils/reward_score/pickscore_reward.py + +NUM_GPUS_ACTOR_ROLLOUT_REWARD=4 +ROLLOUT_TP=1 + +ENGINE=vllm_omni + +# enable reward model on 0'th gpu, it is a temporary workaround +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + data.train_files=$pickscore_train_path \ + data.val_files=$pickscore_test_path \ + data.train_batch_size=48 \ + data.max_prompt_length=256 \ + data.trust_remote_code=True \ + algorithm.global_std=False \ + actor_rollout_ref.model.path=$model_name \ + actor_rollout_ref.model.tokenizer_path=$model_name \ + +actor_rollout_ref.model.architecture=OmniBagelForConditionalGeneration \ + actor_rollout_ref.model.trust_remote_code=True \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + actor_rollout_ref.model.lora_dtype=float32 \ + actor_rollout_ref.model.target_modules="['q_proj_moe_gen','k_proj_moe_gen','v_proj_moe_gen','o_proj_moe_gen','mlp_moe_gen.gate_proj','mlp_moe_gen.up_proj','mlp_moe_gen.down_proj']" \ + actor_rollout_ref.model.fsdp_layer_prefixes="['layers.']" \ + actor_rollout_ref.actor.optim.lr=1e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.0001 \ + actor_rollout_ref.actor.ppo_mini_batch_size=24 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=12 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=12 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP \ + actor_rollout_ref.rollout.name=$ENGINE \ + actor_rollout_ref.rollout.n=16 \ + actor_rollout_ref.rollout.agent.num_workers=$((NUM_GPUS_ACTOR_ROLLOUT_REWARD / ROLLOUT_TP)) \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.pipeline.num_inference_steps=15 \ + actor_rollout_ref.rollout.pipeline.max_sequence_length=256 \ + actor_rollout_ref.rollout.algo.noise_level=1.3 \ + actor_rollout_ref.rollout.algo.sde_type="sde" \ + actor_rollout_ref.rollout.algo.sde_window_size=2 \ + actor_rollout_ref.rollout.algo.sde_window_range="[0,7]" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=50 \ + actor_rollout_ref.rollout.val_kwargs.algo.noise_level=0.0 \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.deploy_config=$BAGEL_DEPLOY_CONFIG \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=12 \ + reward.num_workers=1 \ + reward.custom_reward_function.path=$reward_function_path \ + reward.custom_reward_function.name=compute_score_pickscore \ + trainer.logger='["console", "wandb"]' \ + trainer.project_name=flow_grpo \ + trainer.experiment_name=bagel_pickscore_lora_v1 \ + trainer.log_val_generations=8 \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=$NUM_GPUS_ACTOR_ROLLOUT_REWARD \ + trainer.nnodes=1 \ + trainer.save_freq=30 \ + trainer.test_freq=30 \ + trainer.total_epochs=15 \ + trainer.total_training_steps=300 \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync "$@" diff --git a/tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh b/tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh new file mode 100755 index 000000000..cc22cd1a3 --- /dev/null +++ b/tests/special_e2e/run_flowgrpo_bagel_pickscore_v1.sh @@ -0,0 +1,149 @@ +#!/usr/bin/env bash +# FlowGRPO BAGEL PickScore LoRA e2e smoke test, V1 trainer (sync mode). +# V1 counterpart of run_flowgrpo_bagel_pickscore.sh via main_diffusion_v1. +# +# Requires: vllm-omni +# Builds offline if missing: +# ~/models/tiny-random/BAGEL-MoT +# ~/models/tiny-random/PickScore +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "${REPO_ROOT}" + +NUM_GPUS=${NUM_GPUS:-4} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/tiny-random/BAGEL-MoT} +PICKSCORE_PATH=${PICKSCORE_PATH:-${HOME}/models/tiny-random/PickScore} +DATA_DIR=${DATA_DIR:-${HOME}/data/dummy_bagel_pickscore} +dummy_train_path=${TRAIN_FILES:-${DATA_DIR}/train.parquet} +dummy_test_path=${VAL_FILES:-${DATA_DIR}/test.parquet} +TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-2} +BAGEL_DEPLOY_CONFIG=${BAGEL_DEPLOY_CONFIG:-"${REPO_ROOT}/examples/flowgrpo_trainer/bagel/bagel_deploy_config.yaml"} + +ENGINE=vllm_omni +max_prompt_length=64 + +# Smoke: prefer FA3 when available; fall back like the Qwen-Image e2e. +ATTN_BACKEND=_flash_3_varlen_hub +ROLLOUT_ATTN_BACKEND=FLASH_ATTN +if ! python3 -c 'from verl_omni.utils.diffusion_attention import fa3_available; raise SystemExit(0 if fa3_available() else 1)' >/dev/null 2>&1; then + ATTN_BACKEND=native + ROLLOUT_ATTN_BACKEND=TORCH_SDPA +fi + +_LEAK_FILE=$(mktemp) +_LEAK_PID="" +cleanup_leak_monitor() { + [[ -n "${_LEAK_PID}" ]] && kill "${_LEAK_PID}" 2>/dev/null || true + rm -f "${_LEAK_FILE}" +} +trap cleanup_leak_monitor EXIT + +start_leak_monitor() { + : > "${_LEAK_FILE}" + while true; do + if nvidia-smi -i 0 2>&1 | grep -q "vLLMOmniHttpServer"; then + echo "LEAK" >> "${_LEAK_FILE}" + fi + sleep 1 + done & + _LEAK_PID=$! +} + +check_leak_monitor() { + kill "${_LEAK_PID}" 2>/dev/null || true + _LEAK_PID="" + if grep -q "LEAK" "${_LEAK_FILE}" 2>/dev/null; then + echo "" + echo "FAIL: unexpected vLLMOmniHttpServer process(es) detected on GPU-0 —" + ray stop --force 2>/dev/null || true + exit 1 + fi +} + +n_resp_per_prompt=2 +micro_bsz_per_gpu=1 +micro_bsz=$((micro_bsz_per_gpu * NUM_GPUS)) +mini_bsz=${micro_bsz} +train_batch_size=$((mini_bsz * n_resp_per_prompt)) + +python3 tests/special_e2e/build_bagel_tiny_random.py --output-dir "${MODEL_PATH}" +python3 tests/special_e2e/build_pickscore_tiny_random.py --output-dir "${PICKSCORE_PATH}" +python3 tests/special_e2e/create_dummy_bagel_pickscore_data.py \ + --local_save_dir "${DATA_DIR}" \ + --model_path "${MODEL_PATH}" \ + --train_size "${train_batch_size}" \ + --val_size 4 \ + --max_prompt_length "${max_prompt_length}" + +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 +export PICKSCORE_PATH + +start_leak_monitor +python3 -m verl_omni.trainer.main_diffusion_v1 \ + data.train_files=${dummy_train_path} \ + data.val_files=${dummy_test_path} \ + data.train_batch_size=${train_batch_size} \ + data.max_prompt_length=${max_prompt_length} \ + data.trust_remote_code=True \ + algorithm.global_std=False \ + actor_rollout_ref.model.algorithm=flow_grpo \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.model.tokenizer_path=${MODEL_PATH} \ + +actor_rollout_ref.model.architecture=OmniBagelForConditionalGeneration \ + actor_rollout_ref.model.trust_remote_code=True \ + actor_rollout_ref.model.attn_backend=${ATTN_BACKEND} \ + actor_rollout_ref.model.lora_rank=8 \ + actor_rollout_ref.model.lora_alpha=16 \ + actor_rollout_ref.model.lora_dtype=float32 \ + actor_rollout_ref.model.target_modules="['q_proj_moe_gen','k_proj_moe_gen','v_proj_moe_gen','o_proj_moe_gen','mlp_moe_gen.gate_proj','mlp_moe_gen.up_proj','mlp_moe_gen.down_proj']" \ + actor_rollout_ref.model.fsdp_layer_prefixes="['layers.']" \ + actor_rollout_ref.actor.optim.lr=1e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.0001 \ + actor_rollout_ref.actor.ppo_mini_batch_size=${mini_bsz} \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${micro_bsz_per_gpu} \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${micro_bsz_per_gpu} \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=${ENGINE} \ + actor_rollout_ref.rollout.rollout_attn_backend=${ROLLOUT_ATTN_BACKEND} \ + actor_rollout_ref.rollout.n=${n_resp_per_prompt} \ + actor_rollout_ref.rollout.agent.num_workers=1 \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.pipeline.num_inference_steps=4 \ + actor_rollout_ref.rollout.pipeline.height=256 \ + actor_rollout_ref.rollout.pipeline.width=256 \ + actor_rollout_ref.rollout.pipeline.max_sequence_length=${max_prompt_length} \ + actor_rollout_ref.rollout.enforce_eager=True \ + actor_rollout_ref.rollout.algo.noise_level=1.0 \ + actor_rollout_ref.rollout.algo.sde_type="sde" \ + actor_rollout_ref.rollout.algo.sde_window_size=2 \ + actor_rollout_ref.rollout.algo.sde_window_range="[0,4]" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=4 \ + actor_rollout_ref.rollout.val_kwargs.algo.noise_level=0.0 \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.deploy_config=${BAGEL_DEPLOY_CONFIG} \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${micro_bsz_per_gpu} \ + reward.num_workers=1 \ + reward.custom_reward_function.path=tests/special_e2e/bagel_pickscore_reward.py \ + reward.custom_reward_function.name=compute_score_pickscore \ + trainer.logger=console \ + trainer.project_name=verl-test \ + trainer.experiment_name=flowgrpo-bagel-pickscore-v1-e2e \ + trainer.log_val_generations=0 \ + trainer.n_gpus_per_node=${NUM_GPUS} \ + trainer.nnodes=1 \ + trainer.val_before_train=False \ + trainer.test_freq=-1 \ + trainer.save_freq=-1 \ + trainer.resume_mode=disable \ + trainer.total_training_steps=${TOTAL_TRAIN_STEPS} \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync \ + "$@" +check_leak_monitor + +echo "FlowGRPO BAGEL PickScore LoRA V1 (sync) e2e test passed (training completed successfully)." diff --git a/verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py b/verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py index e88e39f6e..b49d024fd 100644 --- a/verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py +++ b/verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py @@ -106,17 +106,17 @@ def _prompt_token_ids_to_batch( micro_batch: TensorDict, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: - """Pad BAGEL-native ``prompt_token_ids`` from the data pipeline. + """Pad BAGEL-native ``prompts`` token ids from the data pipeline. Args: - micro_batch: Batch containing pre-tokenized ``prompt_token_ids``, + micro_batch: Batch containing pre-tokenized ``prompts``, one sequence per sample, already in BAGEL format. device: Target device for the padded tensors. Returns: ``(text_token_ids, text_attention_mask)`` with shape ``(B, max_len)``. """ - prompt_token_ids = micro_batch["prompt_token_ids"] + prompt_token_ids = micro_batch["prompts"] if isinstance(prompt_token_ids, NonTensorStack): prompt_token_ids = [ tu.unwrap_non_tensor_data(prompt_token_ids[i]) for i in range(micro_batch.batch_size[0])