diff --git a/docs/contributing/integrating_a_diffusion_model.md b/docs/contributing/integrating_a_diffusion_model.md index c63e2855..09b2a741 100644 --- a/docs/contributing/integrating_a_diffusion_model.md +++ b/docs/contributing/integrating_a_diffusion_model.md @@ -201,7 +201,15 @@ class MyModel(DiffusionModelBase): scheduler_inputs, step): ... ``` -### 3.1 `build_scheduler` and `set_timesteps` +### 3.1 (Optional) `configure_trainable_params` + +Override this hook to selectively set ``requires_grad`` for non-LoRA +full-weight training. The engine calls it after module build, before +FSDP wrapping, when ``lora_rank=0``. When LoRA is enabled this hook +is **not** called — ``requires_grad`` is managed by the LoRA adapter +instead. The default is a no-op (all params trainable). + +### 3.2 `build_scheduler` and `set_timesteps` Reuse [`FlowMatchSDEDiscreteScheduler`](../../verl_omni/pipelines/schedulers/flow_match_sde.py) @@ -212,7 +220,7 @@ Compute `image_seq_len` and `mu` exactly as the upstream diffusers pipeline does. If they drift, the training-time noise schedule will not match deployment. -### 3.2 `prepare_model_inputs` +### 3.3 `prepare_model_inputs` This method receives the **full** batched tensors for the entire denoising trajectory (`latents` of shape `(B, T, ...)`, `timesteps` of @@ -233,7 +241,7 @@ inputs. The typical steps are: The dict keys must match the kwargs of the diffusers transformer class verbatim — the FSDP engine calls `module(**model_inputs)`. -### 3.3 `forward_and_sample_previous_step` +### 3.4 `forward_and_sample_previous_step` Call the transformer once for the positive prompt; if CFG is active, call it again for the negative prompt and combine them. Always finish with diff --git a/docs/contributing/integrating_a_non_diffusers_model.md b/docs/contributing/integrating_a_non_diffusers_model.md index 3fcc0ebe..2f6970c8 100644 --- a/docs/contributing/integrating_a_non_diffusers_model.md +++ b/docs/contributing/integrating_a_non_diffusers_model.md @@ -271,7 +271,28 @@ When `build_module()` returns a non-`None` value, the FSDP engine uses it directly. When it returns `None` (the default), the engine falls back to `diffusers.AutoModel.from_pretrained`. -### 3.2 Implement `prepare_model_inputs` +### 3.2 (Optional) `configure_trainable_params` + +Override this hook to selectively set ``requires_grad`` for non-LoRA +full-weight training. The engine calls it after module build, before +FSDP wrapping, when ``lora_rank=0``. When LoRA is enabled this hook +is **not** called — ``requires_grad`` is managed by the LoRA adapter +instead. The default is a no-op (all params trainable). Note that +mixed ``requires_grad`` requires ``strategy=fsdp2``. + +Example from the BAGEL integration (trains only ``moe_gen``, casts to fp32): + +```python +@classmethod +def configure_trainable_params(cls, module, model_config): + for name, param in module.named_parameters(): + param.requires_grad = "moe_gen" in name + for name, param in module.named_parameters(): + if param.requires_grad: + param.data = param.data.to(torch.float32) +``` + +### 3.3 Implement `prepare_model_inputs` This classmethod receives the training module, model config, latents, timesteps, prompt embeddings (plus masks), and the micro-batch from the @@ -281,7 +302,7 @@ signature. Note: non-diffusers models often ignore the ``prompt_embeds*`` parameters and read token IDs directly from ``micro_batch`` instead (see the BAGEL integration's ``_prompt_token_ids_to_batch`` for an example). -### 3.3 Implement `forward_and_sample_previous_step` +### 3.4 Implement `forward_and_sample_previous_step` This follows the same pattern as the diffusers guide. If your model supports classifier-free guidance, implement multi-branch forwarding @@ -291,7 +312,7 @@ a 3-branch CFG with sigma-interval gating as a reference. The return signature is always ``(log_prob, prev_sample_mean, std_dev_t, sqrt_dt)``. -### 3.4 Implement the Scheduler +### 3.5 Implement the Scheduler Non-diffusers models use ``FlowMatchSDEDiscreteScheduler`` just like diffusers models. If your model needs custom sigma schedules (e.g. @@ -414,6 +435,9 @@ checklist to verify your implementation against it: ### Training adapter (`diffusers_training_adapter.py`) - [ ] `@DiffusionModelBase.register("OmniBagelForConditionalGeneration", algorithm="flow_grpo")` - [ ] `build_module()` returns `BagelForTraining.from_pretrained(...)` +- [ ] (Optional) `configure_trainable_params()` — selectively freezes / sets + ``requires_grad`` for non-LoRA full-weight training (e.g. train only the + generation pathway, keep understanding frozen). - [ ] `build_scheduler()` and `set_timesteps()` with shifted sigmas - [ ] `prepare_model_inputs()` reads ``prompts`` and ``attention_mask`` from micro-batch (standard tensors, no extra field needed) diff --git a/examples/flowgrpo_trainer/bagel/bagel.md b/examples/flowgrpo_trainer/bagel/README.md similarity index 76% rename from examples/flowgrpo_trainer/bagel/bagel.md rename to examples/flowgrpo_trainer/bagel/README.md index e19a839b..b6d2b6f3 100644 --- a/examples/flowgrpo_trainer/bagel/bagel.md +++ b/examples/flowgrpo_trainer/bagel/README.md @@ -84,7 +84,7 @@ python3 examples/flowgrpo_trainer/data_process/bagel_pickscore.py \ This produces ``$WORKSPACE/data/pickscore/bagel/train.parquet`` and ``test.parquet``. -### Run training +### Run LoRA training ```bash bash examples/flowgrpo_trainer/bagel/run_bagel_pickscore_lora.sh @@ -97,6 +97,34 @@ Key configuration differences from OCR: (``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 | ``run_bagel_pickscore_lora.sh`` | ``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 | diff --git a/examples/flowgrpo_trainer/bagel/run_bagel_pickscore.sh b/examples/flowgrpo_trainer/bagel/run_bagel_pickscore.sh new file mode 100644 index 00000000..24f703aa --- /dev/null +++ b/examples/flowgrpo_trainer/bagel/run_bagel_pickscore.sh @@ -0,0 +1,92 @@ +# Bagel full-weight RL, vllm_omni rollout (FlowGRPO) with PickScore reward +# +# Non-LoRA counterpart of run_bagel_pickscore_lora.sh, following the official +# flow_grpo `pickscore_bagel` config (use_lora=False). Only the generation +# (moe_gen) pathway is trained; the understanding pathway is frozen via the +# `configure_trainable_params` hook in the BAGEL training adapter. +# +# Uses FSDP2 (strategy=fsdp2) which natively supports mixed requires_grad +# (frozen understanding + trainable generation pathway) and reshards layer +# params after forward, reducing peak memory during gradient checkpointing. +# +# 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 \ + 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.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.strategy=fsdp2 \ + 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=3 \ + 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 \ + 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 "$@" diff --git a/verl_omni/pipelines/bagel_flow_grpo/bagel_model.py b/verl_omni/pipelines/bagel_flow_grpo/bagel_model.py index e3790e5a..ffb143f0 100644 --- a/verl_omni/pipelines/bagel_flow_grpo/bagel_model.py +++ b/verl_omni/pipelines/bagel_flow_grpo/bagel_model.py @@ -347,12 +347,12 @@ def __init__(self, config: BagelTrainingConfig): self.input_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = RotaryEmbedding(config.head_dim, theta=config.rope_theta) def forward( self, hidden_states: Tensor, - cos: Tensor, - sin: Tensor, + position_ids: Tensor, text_mask: Tensor, latent_mask: Tensor, L_ctx: int = 0, @@ -362,8 +362,7 @@ def forward( Args: hidden_states: ``(B, L, D)`` input sequence. - cos: RoPE cosine embedding. - sin: RoPE sine embedding. + position_ids: ``(B, L)`` long tensor of position indices. text_mask: Bool mask — True for text pathway. latent_mask: Bool mask — True for gen pathway. L_ctx: Text context length for causal split. @@ -372,6 +371,8 @@ def forward( Returns: Output of shape ``(B, L, D)``. """ + cos, sin = self.rotary_emb(position_ids) + text_idx = text_mask.nonzero(as_tuple=True) latent_idx = latent_mask.nonzero(as_tuple=True) @@ -474,7 +475,6 @@ def __init__(self, config: BagelTrainingConfig): self.layers = nn.ModuleList([BagelMoTLayer(config) for _ in range(config.num_hidden_layers)]) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.norm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = RotaryEmbedding(config.head_dim, theta=config.rope_theta) self.time_embedder = TimestepEmbedder(config.hidden_size) self.vae2llm = nn.Linear(config.patch_latent_dim, config.hidden_size) @@ -561,7 +561,6 @@ def forward( position_ids = torch.cat([ctx_pos, img_pos]).unsqueeze(0).expand(B, -1) else: position_ids = torch.zeros(1, L_total, dtype=torch.long, device=dev).expand(B, -1) - cos, sin = self.rotary_emb(position_ids) # Key padding mask: zero-padded text tokens in uneven micro-batches # must not attend to image queries. ``None`` keeps the flash backend. @@ -574,10 +573,12 @@ def forward( # 7. Transformer layers (split attention: text causal + image full) for layer in self.layers: - def _layer_fn(seq, cos_, sin_, text_mask_, latent_mask_, kpm, *, _layer=layer): - return _layer(seq, cos_, sin_, text_mask_, latent_mask_, L_ctx, key_padding_mask=kpm) + def _layer_fn(seq, pos_ids, text_mask_, latent_mask_, kpm, *, _layer=layer): + return _layer(seq, pos_ids, text_mask_, latent_mask_, L_ctx, key_padding_mask=kpm) - sequence = self._checkpointed_call(_layer_fn, sequence, cos, sin, text_mask, latent_mask, key_padding_mask) + sequence = self._checkpointed_call( + _layer_fn, sequence, position_ids, text_mask, latent_mask, key_padding_mask + ) # 8. Final norm with MoT routing normed = sequence.new_zeros(sequence.shape) 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 3a18e260..e88e39f6 100644 --- a/verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py +++ b/verl_omni/pipelines/bagel_flow_grpo/diffusers_training_adapter.py @@ -63,6 +63,17 @@ def configure_train_mode(cls, module): if hasattr(layer_inner, "self_attn"): layer_inner.self_attn.training = False + @classmethod + def configure_trainable_params(cls, module, model_config): + """Freeze all params except the generation (``moe_gen``) pathway.""" + for name, param in module.named_parameters(): + param.requires_grad = "moe_gen" in name + + # cast all trainable parameters to fp32 + for name, param in module.named_parameters(): + if param.requires_grad: + param.data = param.data.to(torch.float32) + @classmethod def build_scheduler(cls, model_config: DiffusionModelConfig): # Build on GPU so scheduler buffers are comparable with cuda timesteps in FSDP forward. diff --git a/verl_omni/pipelines/bagel_flow_grpo/vllm_omni_rollout_adapter.py b/verl_omni/pipelines/bagel_flow_grpo/vllm_omni_rollout_adapter.py index fe8ee02d..edcc525f 100644 --- a/verl_omni/pipelines/bagel_flow_grpo/vllm_omni_rollout_adapter.py +++ b/verl_omni/pipelines/bagel_flow_grpo/vllm_omni_rollout_adapter.py @@ -25,7 +25,7 @@ import os import random from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Iterable, Optional import torch from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -233,6 +233,30 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.scheduler = _BagelSchedulerAdapter(inner) logger.info("BagelPipelineWithLogProb: SDE scheduler enabled for RL rollouts") + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights, routing by name prefix. + + Weight-sync from the actor uses ``transformer.`` prefix with separate + q/k/v projections; the parent's ``AutoWeightsLoader`` cannot map these + to the rollout's fused ``qkv_proj``. Delegate such weights to + ``language_model.load_weights`` which handles stacked-param remapping. + Other weights (initial checkpoint load) defer to the parent. + """ + actor_weights: list[tuple[str, torch.Tensor]] = [] + checkpoint_weights: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + if name.startswith("transformer."): + actor_weights.append((f"model.{name[len('transformer.') :]}", tensor)) + else: + checkpoint_weights.append((name, tensor)) + + loaded: set[str] = set() + if checkpoint_weights: + loaded |= super().load_weights(checkpoint_weights) + if actor_weights: + loaded |= self.language_model.load_weights(actor_weights) + return loaded + def _decode_token_prompt(self, token_ids: Any) -> str | None: """Decode BAGEL token IDs to a cleaned prompt text string.""" token_list = _to_token_list(token_ids) diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index 017d3cda..a76801d5 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -102,6 +102,20 @@ def configure_train_mode(cls, module: torch.nn.Module) -> None: """Hook called after ``module.train()`` for architecture-specific overrides.""" return + @classmethod + def configure_trainable_params( + cls, + module: torch.nn.Module, + model_config: DiffusionModelConfig, + ) -> None: + """Hook called after module build to set ``requires_grad`` on trainable params. + + Args: + module: The loaded model module (pre-FSDP). + model_config: The ``DiffusionModelConfig``. + """ + return + @classmethod @abstractmethod def build_scheduler(cls, model_config: DiffusionModelConfig) -> SchedulerMixin: diff --git a/verl_omni/pipelines/non_diffusers_model_base.py b/verl_omni/pipelines/non_diffusers_model_base.py index cbbc552c..048f4ea1 100644 --- a/verl_omni/pipelines/non_diffusers_model_base.py +++ b/verl_omni/pipelines/non_diffusers_model_base.py @@ -26,7 +26,7 @@ import json import logging import os -from abc import ABC, abstractmethod +from abc import ABCMeta, abstractmethod from typing import Callable import torch @@ -37,7 +37,7 @@ __all__ = ["NonDiffusersModelBase"] -class NonDiffusersModelBase(nn.Module, ABC): +class NonDiffusersModelBase(nn.Module, metaclass=ABCMeta): """ABC for non-diffusers models used in verl-omni FSDP training. Provides LoRA/PEFT adapter lifecycle (required by ``LoRAAdapterMixin``), diff --git a/verl_omni/workers/engine/fsdp/diffusers_impl.py b/verl_omni/workers/engine/fsdp/diffusers_impl.py index 611ea6cc..3839f25f 100644 --- a/verl_omni/workers/engine/fsdp/diffusers_impl.py +++ b/verl_omni/workers/engine/fsdp/diffusers_impl.py @@ -58,6 +58,7 @@ from verl.workers.engine.fsdp.utils import create_device_mesh, get_sharding_strategy from verl.workers.engine.utils import enable_full_determinism, prepare_micro_batches +from verl_omni.pipelines.model_base import DiffusionModelBase from verl_omni.pipelines.utils import ( build_scheduler, forward, @@ -202,8 +203,6 @@ def _build_module_from_registry(self, torch_dtype: torch.dtype) -> Optional[torc Returns ``None`` if the registry has no custom loader, so the caller falls back to ``diffusers.AutoModel``. """ - from verl_omni.pipelines.model_base import DiffusionModelBase - model_cls = DiffusionModelBase.get_class(self.model_config) module = model_cls.build_module(self.model_config, torch_dtype) if module is None: @@ -448,6 +447,9 @@ def _build_model_optimizer(self): # Apply LoRA adapters if low-rank adaptation is enabled if self._is_lora: module = self._build_lora_module(module) + else: + # configure trainable parameters for non-lora training + DiffusionModelBase.get_class(self.model_config).configure_trainable_params(module, self.model_config) if self.use_ulysses_sp: sp_size = self.ulysses_sequence_parallel_size @@ -1266,8 +1268,6 @@ def __init__(self, engine: DiffusersFSDPEngine, **kwargs): super().__init__(engine=engine, mode="train", **kwargs) def __enter__(self): - from verl_omni.pipelines.model_base import DiffusionModelBase - assert isinstance(self.engine, DiffusersFSDPEngine) super().__enter__() self.engine.module.train()