[trainer] feat: support BAGEL non-lora training with pickscore reward

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

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)

examples/flowgrpo_trainer/bagel/bagel.md → 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 |

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 "$@"

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)

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.

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)

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:

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``),