diff --git a/examples/diffusionnft_trainer/README.md b/examples/diffusionnft_trainer/README.md index 9afe08073..45fb19ab1 100644 --- a/examples/diffusionnft_trainer/README.md +++ b/examples/diffusionnft_trainer/README.md @@ -77,6 +77,9 @@ Launch the example from the repository root: bash examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora.sh ``` +For CUDA V1 sync (TransferQueue + ReplayBuffer), use +`examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_v1.sh`. + ### NPU For Huawei Ascend NPUs, use the NPU-optimized script: @@ -85,6 +88,9 @@ For Huawei Ascend NPUs, use the NPU-optimized script: bash examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_npu.sh ``` +For NPU V1 sync, use +`examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_npu_v1.sh`. + This script uses a 16-NPU global distribution strategy with: - `actor_rollout_ref.model.attn_backend='_native_npu'` - `actor_rollout_ref.rollout.tensor_model_parallel_size=2` diff --git a/examples/diffusionnft_trainer/minimax_h3/README.md b/examples/diffusionnft_trainer/minimax_h3/README.md index 26ed1665d..614227c48 100644 --- a/examples/diffusionnft_trainer/minimax_h3/README.md +++ b/examples/diffusionnft_trainer/minimax_h3/README.md @@ -192,6 +192,9 @@ export DATA_DIR=/path/to/h3_t2va_data bash examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_t2va_lora.sh ``` +For CUDA V1 sync (TransferQueue + ReplayBuffer), use +`examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_t2va_lora_v1.sh`. + MiniMax H3 t2va requires an explicit named `aspect_ratio` (one of `21:9/16:9/4:3/1:1/3:4/9:16`); the launch script sets `16:9` and explicit `height`/`width` control the actual canvas (must be multiples of 32). @@ -259,6 +262,9 @@ DATA_DIR=/path/to/parquet \ bash examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_fl2va_lora.sh ``` +For CUDA V1 sync, use +`examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_fl2va_lora_v1.sh`. + The latest vLLM-Omni contract requires 4–15 seconds at 24 FPS. The launcher's `NUM_FRAMES=96` is aligned by vLLM-Omni to the next valid `17n+5` boundary. Sampling edges must be multiples of 32 (the H3 pipeline silently floors @@ -348,6 +354,9 @@ VAL_REF_IMAGE_SHORT_EDGE=1024 \ bash examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_ref2va_lora.sh ``` +For CUDA V1 sync, use +`examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_ref2va_lora_v1.sh`. + The H3 Agent Loop keeps the user prompt token-ID-native while vLLM-Omni adds the reference presentation. The default `MAX_PROMPT_EMBEDS=12288` is both the prompt-embedding limit and the fixed transport limit for each video/audio diff --git a/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_fl2va_lora_v1.sh b/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_fl2va_lora_v1.sh new file mode 100644 index 000000000..540e24622 --- /dev/null +++ b/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_fl2va_lora_v1.sh @@ -0,0 +1,176 @@ +#!/usr/bin/env bash +# MiniMax H3 FL2VA (image-conditioned) DiffusionNFT LoRA recipe +# (V1 trainer: TransferQueue + ReplayBuffer + sync mode). +# +# This is the v1 counterpart of run_minimax_h3_fl2va_lora.sh. It uses +# `verl_omni.trainer.main_diffusion_v1`, which selects +# `PolicyGradientDiffusionTrainerV1Sync` via `trainer.v1.trainer_mode=sync`. +# TransferQueue is force-enabled inside the runner, so it does not need to be +# set on the CLI. GPU layout stays colocated (same as v0). +# +# Reference (legacy v0 script): +# verl-omni/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_fl2va_lora.sh +set -euo pipefail + +export WANDB_MODE=${WANDB_MODE:-offline} +export WANDB_RUN_ID=${WANDB_RUN_ID:-minimax_h3_fl2va_lora_v1} +export WANDB_RESUME=${WANDB_RESUME:-allow} +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} + +: "${DATA_DIR:?Set DATA_DIR to the parquet directory produced by prepare_data.py}" +if [[ -z "${MODEL_PATH:-}" || ! -d "$MODEL_PATH/FL2VA" || ! -d "$MODEL_PATH/transformer" ]]; then + echo "MODEL_PATH must point to a MiniMax-H3 repo root containing FL2VA/ (fused rollout) and transformer/ (diffusers actor) (got: '${MODEL_PATH:-}')" >&2 + exit 1 +fi + +N_GPUS=${N_GPUS:-8} +ROLLOUT_TP=${ROLLOUT_TP:-4} +ROLLOUT_N=${ROLLOUT_N:-16} +HEIGHT=${HEIGHT:-288} +WIDTH=${WIDTH:-448} +VAL_HEIGHT=${VAL_HEIGHT:-576} +VAL_WIDTH=${VAL_WIDTH:-928} +NUM_FRAMES=${NUM_FRAMES:-96} +INFER_STEPS=${INFER_STEPS:-10} +MAX_PROMPT_EMBEDS=${MAX_PROMPT_EMBEDS:-1024} +FRAME_INDICES=${FRAME_INDICES:-'[0]'} +ACTOR_ATTN_BACKEND=${ACTOR_ATTN_BACKEND:-_flash_3_varlen_hub} +ROLLOUT_ATTN_BACKEND=${ROLLOUT_ATTN_BACKEND:-FLASH_ATTN_3_HUB} +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-1000} + +script_path=$(readlink -f "$0") +script_name=$(basename "$script_path" .sh) +repo_root=$(dirname "$script_path") +while [[ "$repo_root" != "/" && ! -f "$repo_root/LICENSE" ]]; do + repo_root=$(dirname "$repo_root") +done +output_dir=${OUTPUT_DIR:-$(dirname "$script_path")/outputs/$script_name} +checkpoint_dir=$output_dir/checkpoints +run_timestamp=$(date +"%Y%m%d_%H%M") +log_file=$output_dir/logs/$run_timestamp/${NODE_RANK:-0}.log +mkdir -p "$checkpoint_dir" "$(dirname "$log_file")" + +lora_warmstart_arg=() +if [[ -n "${LORA_WARMSTART_PATH:-}" ]]; then + lora_warmstart_arg=(actor_rollout_ref.model.lora_adapter_path=$LORA_WARMSTART_PATH) +fi + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + algorithm.trainer_type=direct_preference \ + algorithm.sample_source=online \ + algorithm.adv_mode=continuous \ + algorithm.timestep_fraction=1.0 \ + algorithm.old_policy_decay_schedule=delayed_linear_to_0_999 \ + algorithm.old_policy_update_interval=2 \ + data.train_files="$DATA_DIR/train.parquet" \ + data.val_files="$DATA_DIR/test.parquet" \ + data.train_batch_size=32 \ + data.val_max_samples=128 \ + data.max_prompt_length=4096 \ + data.truncation=error \ + data.seed=42 \ + actor_rollout_ref.model.path="$MODEL_PATH/FL2VA" \ + actor_rollout_ref.model.tokenizer_path="$MODEL_PATH/FL2VA/tokenizer" \ + actor_rollout_ref.model.config_path="$MODEL_PATH/transformer" \ + +actor_rollout_ref.model.architecture=MiniMaxH3Pipeline \ + actor_rollout_ref.model.external_lib=verl_omni.pipelines.minimax_h3_diffusion_nft \ + actor_rollout_ref.model.algorithm=diffusion_nft \ + actor_rollout_ref.model.model_type=diffusion_nft_model \ + actor_rollout_ref.model.attn_backend="$ACTOR_ATTN_BACKEND" \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + "${lora_warmstart_arg[@]}" \ + actor_rollout_ref.model.policy_state_adapters='["default","old"]' \ + actor_rollout_ref.model.target_modules='["to_q","to_k","to_v","to_out.0","ff.net.0.proj","ff.net.2"]' \ + actor_rollout_ref.model.fsdp_layer_prefixes="['transformer_blocks.','token_refiner.refiner_blocks.']" \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.optim.lr=3e-4 \ + actor_rollout_ref.actor.optim.weight_decay=1e-4 \ + actor_rollout_ref.actor.optim.betas="[0.9,0.999]" \ + actor_rollout_ref.actor.optim.override_optimizer_config="{eps: 1e-8}" \ + actor_rollout_ref.actor.optim.clip_grad=1.0 \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.actor.diffusion_loss.loss_mode=diffusion_nft \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.actor.diffusion_loss.mix_beta=0.1 \ + actor_rollout_ref.actor.diffusion_loss.ref_kl_coef=0.0001 \ + actor_rollout_ref.actor.diffusion_loss.adv_clip_max=5.0 \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm_omni \ + actor_rollout_ref.rollout.max_num_seqs=1 \ + actor_rollout_ref.rollout.rollout_attn_backend="$ROLLOUT_ATTN_BACKEND" \ + actor_rollout_ref.rollout.rollout_adapter=old \ + actor_rollout_ref.rollout.tensor_model_parallel_size="$ROLLOUT_TP" \ + actor_rollout_ref.rollout.n="$ROLLOUT_N" \ + actor_rollout_ref.rollout.seed=42 \ + actor_rollout_ref.rollout.agent.num_workers=$((N_GPUS / ROLLOUT_TP)) \ + actor_rollout_ref.rollout.agent.default_agent_loop=minimax_h3_diffusion_single_turn_agent \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.calculate_log_probs=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.enable_cpu_offload=True \ + actor_rollout_ref.rollout.max_prompt_embed_length="$MAX_PROMPT_EMBEDS" \ + actor_rollout_ref.rollout.pipeline.task=fl2va \ + actor_rollout_ref.rollout.pipeline.frame_indices="$FRAME_INDICES" \ + actor_rollout_ref.rollout.pipeline.height="$HEIGHT" \ + actor_rollout_ref.rollout.pipeline.width="$WIDTH" \ + actor_rollout_ref.rollout.pipeline.num_frames="$NUM_FRAMES" \ + actor_rollout_ref.rollout.pipeline.frame_rate=24.0 \ + actor_rollout_ref.rollout.pipeline.num_inference_steps="$INFER_STEPS" \ + actor_rollout_ref.rollout.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.pipeline.max_sequence_length="$MAX_PROMPT_EMBEDS" \ + actor_rollout_ref.rollout.pipeline.video_flow_shift=12.0 \ + +actor_rollout_ref.rollout.pipeline.output_type=pt \ + actor_rollout_ref.rollout.val_kwargs.pipeline.height="$VAL_HEIGHT" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.width="$VAL_WIDTH" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_frames="$NUM_FRAMES" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.frame_rate=24.0 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=40 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.true_cfg_scale=1.0 \ + +actor_rollout_ref.rollout.val_kwargs.pipeline.output_type=pt \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + reward.reward_model.enable=False \ + reward.num_workers=1 \ + reward.custom_reward_function.path=pkg://verl_omni.reward_loop.reward_manager.multi \ + reward.custom_reward_function.name=_multi_reward_placeholder \ + reward.reward_manager.name=MultiVisualRewardManager \ + reward.reward_manager.module.path=pkg://verl_omni.reward_loop.reward_manager \ + "+reward.reward_functions.clap.path=$repo_root/verl_omni/utils/reward_score/clap.py" \ + '+reward.reward_functions.clap.name=compute_score' \ + '+reward.reward_functions.clap.weight=1.0' \ + '+reward.reward_functions.clap.device=cuda:0' \ + '+reward.reward_functions.clap.model_name_or_path=laion/larger_clap_general' \ + "+reward.reward_functions.imagebind.path=$repo_root/verl_omni/utils/reward_score/imagebind.py" \ + '+reward.reward_functions.imagebind.name=compute_score' \ + '+reward.reward_functions.imagebind.weight=1.0' \ + '+reward.reward_functions.imagebind.device=cuda:1' \ + '+reward.reward_functions.imagebind.model_name_or_path=.checkpoints/imagebind_huge.pth' \ + '+reward.reward_functions.imagebind.mode=audio_video' \ + reward.aggregation=weighted_sum \ + trainer.logger='["console","tensorboard","wandb"]' \ + trainer.project_name=diffusion_nft \ + trainer.experiment_name=minimax_h3_fl2va_lora_v1 \ + trainer.default_local_dir="$checkpoint_dir" \ + trainer.validation_data_dir="$output_dir/validation_data" \ + trainer.rollout_data_dir="$output_dir/rollout_data" \ + trainer.rollout_data_save_freq=10 \ + trainer.log_val_generations=8 \ + trainer.video_fps=24 \ + trainer.val_before_train=True \ + trainer.n_gpus_per_node="$N_GPUS" \ + trainer.nnodes=1 \ + trainer.save_freq=5 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 \ + trainer.total_training_steps="$TOTAL_TRAINING_STEPS" \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync \ + "$@" diff --git a/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_ref2va_lora_v1.sh b/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_ref2va_lora_v1.sh new file mode 100644 index 000000000..42d07e499 --- /dev/null +++ b/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_ref2va_lora_v1.sh @@ -0,0 +1,183 @@ +#!/usr/bin/env bash +# MiniMax H3 Ref2VA (multi-reference) DiffusionNFT LoRA recipe +# (V1 trainer: TransferQueue + ReplayBuffer + sync mode). +# References may mix images, videos and standalone audio (up to twelve files). +# +# This is the v1 counterpart of run_minimax_h3_ref2va_lora.sh. It uses +# `verl_omni.trainer.main_diffusion_v1`, which selects +# `PolicyGradientDiffusionTrainerV1Sync` via `trainer.v1.trainer_mode=sync`. +# TransferQueue is force-enabled inside the runner, so it does not need to be +# set on the CLI. GPU layout stays colocated (same as v0). +# +# Reference (legacy v0 script): +# verl-omni/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_ref2va_lora.sh +set -euo pipefail + +export WANDB_MODE=${WANDB_MODE:-online} +export WANDB_RUN_ID=${WANDB_RUN_ID:-minimax_h3_ref2va_lora_v1} +export WANDB_RESUME=${WANDB_RESUME:-allow} +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} + +: "${DATA_DIR:?Set DATA_DIR to the parquet directory produced by prepare_ref2va_data.py}" +if [[ -z "${MODEL_PATH:-}" || ! -d "$MODEL_PATH/Ref2VA" || ! -d "$MODEL_PATH/transformer_ref" ]]; then + echo "MODEL_PATH must point to a MiniMax-H3 repo root containing Ref2VA/ (fused rollout) and transformer_ref/ (diffusers actor) (got: '${MODEL_PATH:-}')" >&2 + exit 1 +fi + +N_GPUS=${N_GPUS:-8} +ROLLOUT_TP=${ROLLOUT_TP:-4} +TEXT_ENCODER_TP=${TEXT_ENCODER_TP:-$ROLLOUT_TP} +ROLLOUT_N=${ROLLOUT_N:-16} +HEIGHT=${HEIGHT:-288} +WIDTH=${WIDTH:-448} +VAL_HEIGHT=${VAL_HEIGHT:-576} +VAL_WIDTH=${VAL_WIDTH:-928} +NUM_FRAMES=${NUM_FRAMES:-96} +INFER_STEPS=${INFER_STEPS:-10} +MAX_PROMPT_EMBEDS=${MAX_PROMPT_EMBEDS:-12288} +REF_IMAGE_SHORT_EDGE=${REF_IMAGE_SHORT_EDGE:-2048} +VAL_REF_IMAGE_SHORT_EDGE=${VAL_REF_IMAGE_SHORT_EDGE:-$REF_IMAGE_SHORT_EDGE} +export REF_IMAGE_SHORT_EDGE +ACTOR_ATTN_BACKEND=${ACTOR_ATTN_BACKEND:-_flash_3_varlen_hub} +ROLLOUT_ATTN_BACKEND=${ROLLOUT_ATTN_BACKEND:-FLASH_ATTN_3_HUB} +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-1000} + +script_path=$(readlink -f "$0") +script_name=$(basename "$script_path" .sh) +repo_root=$(dirname "$script_path") +while [[ "$repo_root" != "/" && ! -f "$repo_root/LICENSE" ]]; do + repo_root=$(dirname "$repo_root") +done +output_dir=${OUTPUT_DIR:-$(dirname "$script_path")/outputs/$script_name} +checkpoint_dir=$output_dir/checkpoints +run_timestamp=$(date +"%Y%m%d_%H%M") +log_file=$output_dir/logs/$run_timestamp/${NODE_RANK:-0}.log +mkdir -p "$checkpoint_dir" "$(dirname "$log_file")" + +lora_warmstart_arg=() +if [[ -n "${LORA_WARMSTART_PATH:-}" ]]; then + lora_warmstart_arg=(actor_rollout_ref.model.lora_adapter_path=$LORA_WARMSTART_PATH) +fi + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + algorithm.trainer_type=direct_preference \ + algorithm.sample_source=online \ + algorithm.adv_mode=continuous \ + algorithm.timestep_fraction=1.0 \ + algorithm.old_policy_decay_schedule=delayed_linear_to_0_999 \ + algorithm.old_policy_update_interval=2 \ + data.train_files="$DATA_DIR/train.parquet" \ + data.val_files="$DATA_DIR/test.parquet" \ + data.train_batch_size=32 \ + data.val_max_samples=128 \ + data.max_prompt_length=4096 \ + data.truncation=error \ + data.seed=42 \ + actor_rollout_ref.model.path="$MODEL_PATH/Ref2VA" \ + actor_rollout_ref.model.tokenizer_path="$MODEL_PATH/Ref2VA/tokenizer" \ + actor_rollout_ref.model.config_path="$MODEL_PATH/transformer_ref" \ + +actor_rollout_ref.model.architecture=MiniMaxH3Pipeline \ + actor_rollout_ref.model.external_lib=verl_omni.pipelines.minimax_h3_diffusion_nft \ + actor_rollout_ref.model.algorithm=diffusion_nft \ + actor_rollout_ref.model.model_type=diffusion_nft_model \ + actor_rollout_ref.model.attn_backend="$ACTOR_ATTN_BACKEND" \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + "${lora_warmstart_arg[@]}" \ + actor_rollout_ref.model.policy_state_adapters='["default","old"]' \ + actor_rollout_ref.model.target_modules='["to_q","to_k","to_v","to_out.0","ff.net.0.proj","ff.net.2"]' \ + actor_rollout_ref.model.fsdp_layer_prefixes="['transformer_blocks.','token_refiner.refiner_blocks.']" \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.optim.lr=3e-4 \ + actor_rollout_ref.actor.optim.weight_decay=1e-4 \ + actor_rollout_ref.actor.optim.betas="[0.9,0.999]" \ + actor_rollout_ref.actor.optim.override_optimizer_config="{eps: 1e-8}" \ + actor_rollout_ref.actor.optim.clip_grad=1.0 \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.actor.diffusion_loss.loss_mode=diffusion_nft \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.actor.diffusion_loss.mix_beta=0.1 \ + actor_rollout_ref.actor.diffusion_loss.ref_kl_coef=0.0001 \ + actor_rollout_ref.actor.diffusion_loss.adv_clip_max=5.0 \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm_omni \ + actor_rollout_ref.rollout.max_num_seqs=1 \ + actor_rollout_ref.rollout.rollout_attn_backend="$ROLLOUT_ATTN_BACKEND" \ + actor_rollout_ref.rollout.rollout_adapter=old \ + actor_rollout_ref.rollout.tensor_model_parallel_size="$ROLLOUT_TP" \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.text_encoder_tp_size="$TEXT_ENCODER_TP" \ + actor_rollout_ref.rollout.n="$ROLLOUT_N" \ + actor_rollout_ref.rollout.seed=42 \ + actor_rollout_ref.rollout.agent.num_workers=$((N_GPUS / ROLLOUT_TP)) \ + actor_rollout_ref.rollout.agent.default_agent_loop=minimax_h3_diffusion_single_turn_agent \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.calculate_log_probs=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.enable_layerwise_offload=True \ + actor_rollout_ref.rollout.max_prompt_embed_length="$MAX_PROMPT_EMBEDS" \ + actor_rollout_ref.rollout.pipeline.task=ref2va \ + actor_rollout_ref.rollout.pipeline.height="$HEIGHT" \ + actor_rollout_ref.rollout.pipeline.width="$WIDTH" \ + actor_rollout_ref.rollout.pipeline.num_frames="$NUM_FRAMES" \ + actor_rollout_ref.rollout.pipeline.frame_rate=24.0 \ + actor_rollout_ref.rollout.pipeline.num_inference_steps="$INFER_STEPS" \ + actor_rollout_ref.rollout.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.pipeline.max_sequence_length="$MAX_PROMPT_EMBEDS" \ + actor_rollout_ref.rollout.pipeline.reference_image_short_edge="$REF_IMAGE_SHORT_EDGE" \ + actor_rollout_ref.rollout.pipeline.video_flow_shift=12.0 \ + +actor_rollout_ref.rollout.pipeline.output_type=pt \ + actor_rollout_ref.rollout.val_kwargs.pipeline.height="$VAL_HEIGHT" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.width="$VAL_WIDTH" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_frames="$NUM_FRAMES" \ + actor_rollout_ref.rollout.val_kwargs.pipeline.frame_rate=24.0 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=40 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.reference_image_short_edge="$VAL_REF_IMAGE_SHORT_EDGE" \ + +actor_rollout_ref.rollout.val_kwargs.pipeline.output_type=pt \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + reward.reward_model.enable=False \ + reward.num_workers=1 \ + reward.custom_reward_function.path=pkg://verl_omni.reward_loop.reward_manager.multi \ + reward.custom_reward_function.name=_multi_reward_placeholder \ + reward.reward_manager.name=MultiVisualRewardManager \ + reward.reward_manager.module.path=pkg://verl_omni.reward_loop.reward_manager \ + "+reward.reward_functions.clap.path=$repo_root/verl_omni/utils/reward_score/clap.py" \ + '+reward.reward_functions.clap.name=compute_score' \ + '+reward.reward_functions.clap.weight=1.0' \ + '+reward.reward_functions.clap.device=cuda:0' \ + '+reward.reward_functions.clap.model_name_or_path=laion/larger_clap_general' \ + "+reward.reward_functions.imagebind.path=$repo_root/verl_omni/utils/reward_score/imagebind.py" \ + '+reward.reward_functions.imagebind.name=compute_score' \ + '+reward.reward_functions.imagebind.weight=1.0' \ + '+reward.reward_functions.imagebind.device=cuda:1' \ + '+reward.reward_functions.imagebind.model_name_or_path=.checkpoints/imagebind_huge.pth' \ + '+reward.reward_functions.imagebind.mode=audio_video' \ + reward.aggregation=weighted_sum \ + trainer.logger='["console","tensorboard","wandb"]' \ + trainer.project_name=diffusion_nft \ + trainer.experiment_name=minimax_h3_ref2va_lora_v1 \ + trainer.default_local_dir="$checkpoint_dir" \ + trainer.validation_data_dir="$output_dir/validation_data" \ + trainer.rollout_data_dir="$output_dir/rollout_data" \ + trainer.rollout_data_save_freq=10 \ + trainer.rollout_data_max_samples=8 \ + trainer.log_val_generations=8 \ + trainer.video_fps=24 \ + trainer.val_before_train=True \ + trainer.n_gpus_per_node="$N_GPUS" \ + trainer.nnodes=1 \ + trainer.save_freq=5 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 \ + trainer.total_training_steps="$TOTAL_TRAINING_STEPS" \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync \ + "$@" diff --git a/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_t2va_lora_v1.sh b/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_t2va_lora_v1.sh new file mode 100644 index 000000000..b41bcd071 --- /dev/null +++ b/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_t2va_lora_v1.sh @@ -0,0 +1,181 @@ +#!/usr/bin/env bash +# MiniMax-H3 text-to-audio-video (t2va) DiffusionNFT LoRA recipe +# (V1 trainer: TransferQueue + ReplayBuffer + sync mode). +# +# This is the v1 counterpart of run_minimax_h3_t2va_lora.sh. It uses +# `verl_omni.trainer.main_diffusion_v1`, which selects +# `PolicyGradientDiffusionTrainerV1Sync` via `trainer.v1.trainer_mode=sync`. +# TransferQueue is force-enabled inside the runner, so it does not need to be +# set on the CLI. GPU layout stays colocated (same as v0). +# +# Reference (legacy v0 script): +# verl-omni/examples/diffusionnft_trainer/minimax_h3/run_minimax_h3_t2va_lora.sh +set -x + +export WANDB_MODE=${WANDB_MODE:-offline} + +WORKSPACE=${WORKSPACE:-$HOME} +MODEL_PATH=${MODEL_PATH:-} +DATA_DIR=${DATA_DIR:-$WORKSPACE/data/vid_prompt/verl_omni} +NUM_GPUS=${NUM_GPUS:-8} +ROLLOUT_TP=${ROLLOUT_TP:-2} +ROLLOUT_N=${ROLLOUT_N:-16} +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-1000} +HEIGHT=${HEIGHT:-256} +WIDTH=${WIDTH:-384} +NUM_FRAMES=${NUM_FRAMES:-121} +INFER_STEPS=${INFER_STEPS:-10} +VAL_HEIGHT=${VAL_HEIGHT:-512} +VAL_WIDTH=${VAL_WIDTH:-768} +ACTOR_ATTN_BACKEND=${ACTOR_ATTN_BACKEND:-_flash_3_varlen_hub} +ROLLOUT_ATTN_BACKEND=${ROLLOUT_ATTN_BACKEND:-FLASH_ATTN_3_HUB} + +if [[ -z "$MODEL_PATH" || ! -d "$MODEL_PATH/FL2VA" || ! -d "$MODEL_PATH/transformer" ]]; then + echo "MODEL_PATH must point to a local MiniMax-H3 repo root containing FL2VA/ and transformer/ (got: '${MODEL_PATH:-}')" >&2 + exit 1 +fi + +train_path=$DATA_DIR/train.parquet +test_path=$DATA_DIR/test.parquet + +script_path=$(readlink -f "$0") +script_name=$(basename "$script_path" .sh) +repo_root=$(dirname "$script_path") +while [[ "$repo_root" != "/" && ! -f "$repo_root/LICENSE" ]]; do + repo_root=$(dirname "$repo_root") +done +if [[ ! -f "$repo_root/LICENSE" ]]; then + echo "Unable to locate repo root from $script_path: no LICENSE found" >&2 + exit 1 +fi + +output_dir=${OUTPUT_DIR:-$repo_root/outputs/$script_name} +checkpoint_dir=$output_dir/checkpoints +run_timestamp=$(date +"%Y%m%d_%H%M") +log_file=$output_dir/logs/$run_timestamp/${NODE_RANK:-0}.log +mkdir -p "$checkpoint_dir" "$(dirname "$log_file")" +exec > >(tee -a "$log_file") 2>&1 + +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 + +h3_lora_targets="['to_q','to_k','to_v','to_out.0','ff.net.0.proj','ff.net.2']" + +lora_warmstart_arg=() +if [[ -n "${LORA_WARMSTART_PATH:-}" ]]; then + lora_warmstart_arg=(actor_rollout_ref.model.lora_adapter_path=$LORA_WARMSTART_PATH) +fi + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + data.train_files=$train_path \ + data.val_files=$test_path \ + data.train_batch_size=32 \ + data.val_max_samples=128 \ + data.max_prompt_length=1024 \ + data.truncation=error \ + data.seed=42 \ + algorithm.trainer_type=direct_preference \ + algorithm.sample_source=online \ + algorithm.timestep_fraction=1.0 \ + algorithm.old_policy_decay_schedule=delayed_linear_to_0_999 \ + algorithm.old_policy_update_interval=2 \ + algorithm.adv_mode=continuous \ + actor_rollout_ref.model.path=$MODEL_PATH/FL2VA \ + actor_rollout_ref.model.config_path=$MODEL_PATH/transformer \ + +actor_rollout_ref.model.architecture=MiniMaxH3Pipeline \ + actor_rollout_ref.model.algorithm=diffusion_nft \ + actor_rollout_ref.model.model_type=diffusion_nft_model \ + actor_rollout_ref.model.attn_backend=$ACTOR_ATTN_BACKEND \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + "${lora_warmstart_arg[@]}" \ + actor_rollout_ref.model.policy_state_adapters='["default","old"]' \ + actor_rollout_ref.model.target_modules="$h3_lora_targets" \ + actor_rollout_ref.model.fsdp_layer_prefixes="['transformer_blocks.','token_refiner.refiner_blocks.']" \ + '+actor_rollout_ref.actor.fsdp_config.wrap_policy.transformer_layer_cls_to_wrap=[MiniMaxH3TransformerBlock,MiniMaxH3TokenRefinerBlock]' \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.optim.lr=3e-4 \ + actor_rollout_ref.actor.optim.weight_decay=1e-4 \ + actor_rollout_ref.actor.optim.betas="[0.9,0.999]" \ + actor_rollout_ref.actor.optim.override_optimizer_config="{eps: 1e-8}" \ + actor_rollout_ref.actor.optim.clip_grad=1.0 \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.actor.diffusion_loss.loss_mode=diffusion_nft \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.actor.diffusion_loss.mix_beta=0.1 \ + actor_rollout_ref.actor.diffusion_loss.ref_kl_coef=0.0001 \ + actor_rollout_ref.actor.diffusion_loss.adv_clip_max=5.0 \ + actor_rollout_ref.actor.use_kl_loss=False \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm_omni \ + actor_rollout_ref.rollout.max_num_seqs=1 \ + actor_rollout_ref.rollout.rollout_attn_backend=$ROLLOUT_ATTN_BACKEND \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP \ + actor_rollout_ref.rollout.n=$ROLLOUT_N \ + actor_rollout_ref.rollout.seed=42 \ + actor_rollout_ref.rollout.agent.num_workers=$((NUM_GPUS / ROLLOUT_TP)) \ + actor_rollout_ref.rollout.agent.default_agent_loop=minimax_h3_diffusion_single_turn_agent \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.calculate_log_probs=False \ + actor_rollout_ref.rollout.rollout_adapter=old \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.rollout.pipeline.aspect_ratio=16:9 \ + actor_rollout_ref.rollout.pipeline.height=$HEIGHT \ + actor_rollout_ref.rollout.pipeline.width=$WIDTH \ + actor_rollout_ref.rollout.pipeline.num_frames=$NUM_FRAMES \ + actor_rollout_ref.rollout.pipeline.frame_rate=24.0 \ + actor_rollout_ref.rollout.pipeline.num_inference_steps=$INFER_STEPS \ + actor_rollout_ref.rollout.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.pipeline.max_sequence_length=1024 \ + actor_rollout_ref.rollout.pipeline.video_flow_shift=12.0 \ + +actor_rollout_ref.rollout.pipeline.output_type=pt \ + actor_rollout_ref.rollout.val_kwargs.pipeline.height=$VAL_HEIGHT \ + actor_rollout_ref.rollout.val_kwargs.pipeline.width=$VAL_WIDTH \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_frames=$NUM_FRAMES \ + actor_rollout_ref.rollout.val_kwargs.pipeline.frame_rate=24.0 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=40 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.true_cfg_scale=1.0 \ + +actor_rollout_ref.rollout.val_kwargs.pipeline.output_type=pt \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + reward.num_workers=1 \ + reward.reward_model.enable=False \ + reward.custom_reward_function.path=pkg://verl_omni.reward_loop.reward_manager.multi \ + reward.custom_reward_function.name=_multi_reward_placeholder \ + reward.reward_manager.name=MultiVisualRewardManager \ + reward.reward_manager.module.path=pkg://verl_omni.reward_loop.reward_manager \ + "+reward.reward_functions.clap.path=$repo_root/verl_omni/utils/reward_score/clap.py" \ + '+reward.reward_functions.clap.name=compute_score' \ + '+reward.reward_functions.clap.weight=1.0' \ + '+reward.reward_functions.clap.device=cuda:0' \ + '+reward.reward_functions.clap.model_name_or_path=laion/larger_clap_general' \ + "+reward.reward_functions.imagebind.path=$repo_root/verl_omni/utils/reward_score/imagebind.py" \ + '+reward.reward_functions.imagebind.name=compute_score' \ + '+reward.reward_functions.imagebind.weight=1.0' \ + '+reward.reward_functions.imagebind.device=cuda:1' \ + '+reward.reward_functions.imagebind.model_name_or_path=.checkpoints/imagebind_huge.pth' \ + '+reward.reward_functions.imagebind.mode=audio_video' \ + reward.aggregation=weighted_sum \ + trainer.logger='["console","tensorboard","wandb"]' \ + trainer.project_name=diffusion_nft \ + trainer.experiment_name=minimax_h3_t2va_lora_v1 \ + trainer.default_local_dir=$checkpoint_dir \ + trainer.validation_data_dir=$output_dir/validation_data \ + trainer.rollout_data_dir=$output_dir/rollout_data \ + trainer.rollout_data_save_freq=10 \ + trainer.log_val_generations=8 \ + trainer.video_fps=24 \ + trainer.val_before_train=True \ + trainer.n_gpus_per_node=$NUM_GPUS \ + trainer.nnodes=1 \ + trainer.save_freq=5 \ + trainer.max_actor_ckpt_to_keep=5 \ + trainer.test_freq=10 \ + trainer.total_epochs=15 \ + trainer.total_training_steps=$TOTAL_TRAINING_STEPS \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync "$@" diff --git a/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_npu_v1.sh b/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_npu_v1.sh new file mode 100644 index 000000000..f358dd244 --- /dev/null +++ b/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_npu_v1.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash +# Qwen-Image DiffusionNFT LoRA RL on Ascend NPU (V1 trainer: TransferQueue + ReplayBuffer + sync mode). +# +# This is the v1 counterpart of run_qwen_image_ocr_lora_npu.sh. It uses +# `verl_omni.trainer.main_diffusion_v1`, which selects +# `PolicyGradientDiffusionTrainerV1Sync` via `trainer.v1.trainer_mode=sync`. +# TransferQueue is force-enabled inside the runner, so it does not need to be +# set on the CLI. Attention / TP layout stays on the v0 NPU script. +# +# Reference (legacy v0 script): +# verl-omni/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_npu.sh +set -x +export VERL_DATAPROTO_SERIALIZATION_METHOD=numpy + +# Set WORKSPACE to any writable directory; defaults to $HOME +WORKSPACE=${WORKSPACE:-$HOME} + +ocr_train_path=$WORKSPACE/data/ocr/train.parquet +ocr_test_path=$WORKSPACE/data/ocr/test.parquet + +model_name=Qwen/Qwen-Image +reward_model_name=Qwen/Qwen3-VL-8B-Instruct +reward_function_path=verl_omni/utils/reward_score/genrm_ocr.py + +NUM_GPUS_ACTOR_ROLLOUT_REWARD=16 +ROLLOUT_TP=2 +REWARD_TP=4 +IMAGE_RESOLUTION=512 + +ENGINE=vllm_omni +REWARD_ENGINE=vllm + + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + data.train_files="$ocr_train_path" \ + data.val_files="$ocr_test_path" \ + data.train_max_samples=7200 \ + data.train_batch_size=24 \ + data.max_prompt_length=256 \ + actor_rollout_ref.model.attn_backend='_native_npu' \ + actor_rollout_ref.model.algorithm=diffusion_nft \ + actor_rollout_ref.model.model_type=diffusion_nft_model \ + actor_rollout_ref.model.path=$model_name \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + actor_rollout_ref.model.policy_state_adapters='["default","old"]' \ + actor_rollout_ref.model.target_modules="['to_q','to_k','to_v','to_out.0','add_q_proj','add_k_proj','add_v_proj','to_add_out','img_mlp.net.0.proj','img_mlp.net.2','txt_mlp.net.0.proj','txt_mlp.net.2']" \ + actor_rollout_ref.actor.optim.lr=3e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.0001 \ + actor_rollout_ref.actor.ppo_mini_batch_size=12 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=12 \ + actor_rollout_ref.actor.diffusion_loss.loss_mode=diffusion_nft \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.actor.diffusion_loss.mix_beta=0.1 \ + actor_rollout_ref.actor.diffusion_loss.ref_kl_coef=0.0001 \ + actor_rollout_ref.actor.diffusion_loss.adv_clip_max=5.0 \ + 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.rollout.tensor_model_parallel_size=$ROLLOUT_TP \ + actor_rollout_ref.rollout.rollout_attn_backend=TORCH_SDPA \ + 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.calculate_log_probs=False \ + actor_rollout_ref.rollout.rollout_adapter=old \ + actor_rollout_ref.rollout.pipeline.num_inference_steps=10 \ + actor_rollout_ref.rollout.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.pipeline.height=$IMAGE_RESOLUTION \ + actor_rollout_ref.rollout.pipeline.width=$IMAGE_RESOLUTION \ + actor_rollout_ref.rollout.pipeline.max_sequence_length=256 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=40 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \ + algorithm.trainer_type=direct_preference \ + algorithm.sample_source=online \ + algorithm.timestep_fraction=1.0 \ + algorithm.old_policy_decay_schedule=delayed_linear_to_0_999 \ + algorithm.old_policy_update_interval=2 \ + algorithm.adv_mode=continuous \ + reward.num_workers=$((NUM_GPUS_ACTOR_ROLLOUT_REWARD / REWARD_TP)) \ + reward.reward_model.enable=True \ + reward.reward_model.model_path=$reward_model_name \ + reward.reward_model.rollout.enforce_eager=False \ + reward.reward_model.rollout.name=$REWARD_ENGINE \ + reward.reward_model.rollout.tensor_model_parallel_size=$REWARD_TP \ + reward.custom_reward_function.path=$reward_function_path \ + reward.custom_reward_function.name=compute_score_ocr \ + trainer.logger='["console", "tensorboard"]' \ + trainer.project_name=diffusion_nft \ + trainer.experiment_name=qwen_image_ocr_lora_npu_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=60 \ + trainer.test_freq=20 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=300 \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync "$@" + diff --git a/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_v1.sh b/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_v1.sh new file mode 100644 index 000000000..cd4567e29 --- /dev/null +++ b/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora_v1.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# Qwen-Image DiffusionNFT LoRA RL (V1 trainer: TransferQueue + ReplayBuffer + sync mode). +# +# This is the v1 counterpart of run_qwen_image_ocr_lora.sh. It uses +# `verl_omni.trainer.main_diffusion_v1`, which selects +# `PolicyGradientDiffusionTrainerV1Sync` via `trainer.v1.trainer_mode=sync`. +# TransferQueue is force-enabled inside the runner, so it does not need to be +# set on the CLI. +# +# Reference (legacy v0 script): +# verl-omni/examples/diffusionnft_trainer/qwen_image/run_qwen_image_ocr_lora.sh +set -x + +# Set WORKSPACE to any writable directory; defaults to $HOME +WORKSPACE=${WORKSPACE:-$HOME} + +ocr_train_path=$WORKSPACE/data/ocr/train.parquet +ocr_test_path=$WORKSPACE/data/ocr/test.parquet + +model_name=Qwen/Qwen-Image +reward_model_name=Qwen/Qwen3-VL-8B-Instruct +reward_function_path=verl_omni/utils/reward_score/genrm_ocr.py + +NUM_GPUS_ACTOR_ROLLOUT_REWARD=4 +ROLLOUT_TP=1 +REWARD_TP=4 +IMAGE_RESOLUTION=512 + +ENGINE=vllm_omni +REWARD_ENGINE=vllm + + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + data.train_files=$ocr_train_path \ + data.val_files=$ocr_test_path \ + data.train_max_samples=7200 \ + data.train_batch_size=24 \ + data.max_prompt_length=256 \ + actor_rollout_ref.model.algorithm=diffusion_nft \ + actor_rollout_ref.model.model_type=diffusion_nft_model \ + actor_rollout_ref.model.path=$model_name \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + actor_rollout_ref.model.policy_state_adapters='["default","old"]' \ + actor_rollout_ref.model.target_modules="['to_q','to_k','to_v','to_out.0','add_q_proj','add_k_proj','add_v_proj','to_add_out','img_mlp.net.0.proj','img_mlp.net.2','txt_mlp.net.0.proj','txt_mlp.net.2']" \ + actor_rollout_ref.actor.optim.lr=3e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.0001 \ + actor_rollout_ref.actor.ppo_mini_batch_size=12 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=12 \ + actor_rollout_ref.actor.diffusion_loss.loss_mode=diffusion_nft \ + actor_rollout_ref.actor.diffusion_loss.clip_ratio=1e-5 \ + actor_rollout_ref.actor.diffusion_loss.mix_beta=0.1 \ + actor_rollout_ref.actor.diffusion_loss.ref_kl_coef=0.0001 \ + actor_rollout_ref.actor.diffusion_loss.adv_clip_max=5.0 \ + 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.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.calculate_log_probs=False \ + actor_rollout_ref.rollout.rollout_adapter=old \ + actor_rollout_ref.rollout.pipeline.num_inference_steps=10 \ + actor_rollout_ref.rollout.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.pipeline.height=$IMAGE_RESOLUTION \ + actor_rollout_ref.rollout.pipeline.width=$IMAGE_RESOLUTION \ + actor_rollout_ref.rollout.pipeline.max_sequence_length=256 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=40 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \ + algorithm.trainer_type=direct_preference \ + algorithm.sample_source=online \ + algorithm.timestep_fraction=1.0 \ + algorithm.old_policy_decay_schedule=delayed_linear_to_0_999 \ + algorithm.old_policy_update_interval=2 \ + algorithm.adv_mode=continuous \ + reward.num_workers=$((NUM_GPUS_ACTOR_ROLLOUT_REWARD / REWARD_TP)) \ + reward.reward_model.enable=True \ + reward.reward_model.model_path=$reward_model_name \ + reward.reward_model.rollout.name=$REWARD_ENGINE \ + reward.reward_model.rollout.tensor_model_parallel_size=$REWARD_TP \ + reward.custom_reward_function.path=$reward_function_path \ + reward.custom_reward_function.name=compute_score_ocr \ + trainer.logger='["console", "wandb"]' \ + trainer.project_name=diffusion_nft \ + trainer.experiment_name=qwen_image_ocr_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=60 \ + trainer.test_freq=20 \ + trainer.total_epochs=1 \ + trainer.total_training_steps=300 \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync "$@"