[bugfix]: stabilize GenRL reward and PPO training

AGENTS.md

@@ -54,3 +54,30 @@ This repository is agent-friendly. Before doing any work, read:
 If you are exploring a new procedure that has no existing SOP, document your
 progress in `.agents/exploration/` and flag it for review at the end of your
 session.
+
+## Context Efficiency
+
+### Subagent Discipline
+
+**Context-aware delegation:**
+ - Under ~50k context: prefer inline work for tasks under ~5 tool calls.
+ - Over ~50k context: prefer subagents for self-contained tasks, even simple ones — the per-call token tax on large contexts adds up fast.
+
+When using subagents, include output rules: "Final response under 2000 characters. List outcomes, not process."
+Never call TaskOutput twice for the same subagent. If it times out, increase the timeout — don't re-read.
+
+### File Reading
+Read files with purpose. Before reading a file, know what you're looking for.
+Use Grep to locate relevant sections before reading entire large files.
+Never re-read a file you've already read in this session.
+For files over 500 lines, use offset/limit to read only the relevant section.
+
+### Responses
+Don't echo back file contents you just read — the user can see them.
+Don't narrate tool calls ("Let me read the file..." / "Now I'll edit..."). Just do it.
+Keep explanations proportional to complexity. Simple changes need one sentence, not three paragraphs.
+
+**Tables — STRICT RULES (apply everywhere, always):**
+- Markdown tables: use minimum separator (`|-|-|`). Never pad with repeated hyphens (`|---|---|`).
+- NEVER use box-drawing / ASCII-art tables with characters like `┌`, `┬`, `─`, `│`, `└`, `┘`, `├`, `┤`, `┼`. These are completely banned.
+- No exceptions. Not for "clarity", not for alignment, not for terminal output.

examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml

@@ -0,0 +1,164 @@
+# GenRL / Video GRPO: Wan 2.1 T2V 1.3B with LongCat reweighting.
+#
+# Ported from GenRL/config/longcat.yaml.
+#
+# - Student: trainable full-parameter model by default.
+# - LoRA is still available via models.student.use_lora=true.
+# - Full fine-tuning with beta > 0 requires models.reference and much
+#   more memory; keep beta at 0.0 for the 4xH100 probe run.
+#
+# Usage:
+#   torchrun --nnodes=1 --nproc_per_node=4 \
+#     -m fastvideo.train.entrypoint.train \
+#     --config examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    # Set true for LoRA. LoRA can use method.beta with disable_adapter().
+    use_lora: false
+    lora_r: 128
+    lora_alpha: 64
+    lora_init_weights: gaussian
+    lora_path: null
+    lora_target_modules:
+      - to_k
+      - to_out
+      - to_q
+      - to_v
+      - ffn.fc_in
+      - ffn.fc_out
+    enable_gradient_checkpointing_type: full
+
+method:
+  _target_: fastvideo.train.methods.rl.genrl.GenRLMethod
+
+  # ---- Reward functions ----
+  reward_fn:
+    hpsv3_general: 1.0
+    hpsv3_percentile: 1.0
+    videoalign_mq: 1.0
+    videoalign_ta: 1.0
+  reward_module: null
+  reward_on_gpu: true
+
+  # ---- Data ----
+  prompt_dataset_path: GenRL/datasets/filtered_prompts
+  prompt_fn: filtered_prompts
+
+  # ---- Sampling ----
+  sample_batch_size: 4
+  eval_batch_size: 2
+  # Sample multiple rollout microbatches, average their PPO losses, then
+  # apply one optimizer update. This reduces reward/advantage variance.
+  num_batches_per_epoch: 4
+  accumulate_ppo_microbatches: true
+  eval_every_steps: 20
+  eval_num_batches: 1
+  eval_num_steps: 16
+  eval_guidance_scale: 4.5
+  num_inference_steps: 16
+  guidance_scale: 4.5
+  num_video_per_prompt: 4
+  noise_level: 1.0
+  sde_type: flow_sde
+  sde_window_size: 1
+  sde_window_range: [0, 6]
+  diffusion_clip: true
+  diffusion_clip_value: 0.45
+  kl_reward: 0
+  same_latent: true
+
+  # ---- Video dimensions ----
+  height: 480
+  width: 832
+  num_frames: 81
+
+  # ---- PPO training ----
+  train_batch_size: 4
+  num_inner_epochs: 1
+  clip_range: 1.0e-4
+  adv_clip_max: 5.0
+  # Official LoRA LongCat uses beta: 3.0e-4 with disable_adapter().
+  # For full fine-tuning on 4 H100s, avoid a second frozen Wan copy.
+  beta: 0.0
+  use_cfg: true
+  # Flash-GRPO-style temporal gradient rectification: avoid the large
+  # LongCat sigma/dt multiplier while debugging full-FT stability.
+  loss_reweighting: flash_tgr
+  loss_reweighting_clip: null
+  weight_advantages: true
+  # Match official GenRL PPO cadence. With sde_window_size: 1 this is
+  # equivalent to one optimizer step per sampled trajectory timestep.
+  optimizer_step_per_timestep: true
+  log_post_update_kl: true
+  max_grad_norm: 1.0
+  seed: 42
+
+  # ---- Advantage computation ----
+  per_prompt_stat_tracking: true
+  global_std: false
+  max_group_std: true
+
+training:
+  distributed:
+
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    # Full fine-tuning needs FSDP/HSDP sharding across all 4 GPUs.
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    # Not used by GenRL (prompt dataloaders are in method config)
+    # but required by the config parser.
+    data_path: ""
+    train_batch_size: 1
+    seed: 42
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    # Full-parameter visual GRPO is much more sensitive than LoRA.
+    # DanceGRPO reports 5e-6 to 2e-5 as the practical range.
+    learning_rate: 1.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 1.0e-4
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    # max_train_steps == num_epochs in GenRL terms.
+    max_train_steps: 100000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/genrl_longcat
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: VideoRL
+    # Leave blank so W&B auto-generates a unique display name per run.
+    run_name: ""
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  # Gradnorm call back Disabled; GenRLMethod clips internally.
+  ema:
+    decay: 0.9
+    start_iter: 0
+    update_interval: 8
+  log_rl_samples:
+    every_steps: 1
+    max_videos: 4
+    fps: 16
+
+pipeline:
+  flow_shift: 3.0

examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml

@@ -0,0 +1,113 @@
+# GenRL / Video GRPO: Wan 2.1 T2V 1.3B — OCR reward, full finetune, 4 GPUs.
+#
+# Usage:
+#   torchrun --nnodes=1 --nproc_per_node=4 \
+#     -m fastvideo.train.entrypoint.train \
+#     --config examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+
+method:
+  _target_: fastvideo.train.methods.rl.genrl.GenRLMethod
+
+  # ---- Reward functions ----
+  reward_fn:
+    video_ocr: 1.0
+  reward_module: null
+
+  # ---- Data ----
+  prompt_dataset_path: GenRL/datasets/ocr
+  prompt_fn: general_ocr
+
+  # ---- Sampling ----
+  sample_batch_size: 4
+  eval_batch_size: 2
+  num_batches_per_epoch: 1
+  num_inference_steps: 16
+  guidance_scale: 4.5
+  num_video_per_prompt: 4
+  noise_level: 1.0
+  sde_type: flow_sde
+  sde_window_size: 1
+  sde_window_range: [0, 6]
+  diffusion_clip: true
+  diffusion_clip_value: 0.45
+  kl_reward: 0
+  same_latent: true
+
+  # ---- Video dimensions ----
+  height: 480
+  width: 832
+  num_frames: 81
+
+  # ---- PPO training ----
+  train_batch_size: 4
+  num_inner_epochs: 1
+  clip_range: 1.0e-3
+  adv_clip_max: 5.0
+  # No frozen reference model is configured in this launch.
+  beta: 0.0
+  use_cfg: true
+  loss_reweighting: longcat
+  weight_advantages: false
+  max_grad_norm: 1.0
+  seed: 42
+
+  # ---- Advantage computation ----
+  per_prompt_stat_tracking: true
+  global_std: false
+  max_group_std: true
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: ""
+    train_batch_size: 1
+    seed: 42
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.999]
+    weight_decay: 1.0e-4
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 100000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/genrl_ocr
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: VideoRL
+    run_name: wan_2_1_t2v_1_3b_ocr
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 0.0  # Disabled; GenRLMethod clips internally.
+  ema:
+    decay: 0.9
+    start_iter: 0
+
+pipeline:
+  flow_shift: 3.0

examples/train/run.sh

@@ -55,7 +55,7 @@ python -m torch.distributed.run \
     --nproc_per_node "${NUM_GPUS}" \
     --master_addr "${MASTER_ADDR}" \
     --master_port "${MASTER_PORT}" \
-    fastvideo/train/entrypoint/train.py \
+    -m fastvideo.train.entrypoint.train \
     --config "${CONFIG}" \
     "$@" \
     2>&1 | tee "${LOG_FILE}"

examples/train/run_slurm.sh

@@ -107,7 +107,7 @@ srun torchrun \\
     --node_rank \$SLURM_PROCID \\
     --rdzv_backend=c10d \\
     --rdzv_endpoint="\$MASTER_ADDR:\$MASTER_PORT" \\
-    fastvideo/train/entrypoint/train.py \\
+    -m fastvideo.train.entrypoint.train \\
     --config ${CONFIG} \\
     --training.distributed.num_gpus ${TOTAL_GPUS} \\
     ${EXTRA_ARGS[*]:-}

fastvideo/layers/lora/linear.py

@@ -336,8 +336,24 @@ def forward(self, input_: torch.Tensor):
         output_parallel = self.base_layer.quant_method.apply(
             self.base_layer, input_parallel)
 
-        if self.set_lora:
-            output_parallel = self.apply_lora(output_parallel, input_parallel)
+        if not self.merged and not self.disable_lora:
+            lora_A = self.lora_A
+            lora_B = self.lora_B
+            assert lora_A is not None and lora_B is not None
+            if isinstance(lora_B, DTensor):
+                lora_B = lora_B.to_local()
+                lora_A = lora_A.to_local()
+
+            lora_A_sliced = self.slice_lora_a_weights(
+                lora_A.to(input_parallel, non_blocking=True))
+            lora_B_sliced = self.slice_lora_b_weights(
+                lora_B.to(output_parallel, non_blocking=True))
+            delta = input_parallel @ lora_A_sliced.T @ lora_B_sliced.T
+            if self.lora_alpha != self.lora_rank:
+                delta = delta * (
+                    self.lora_alpha / self.lora_rank  # type: ignore
+                )
+            output_parallel = output_parallel + delta
 
         if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
             output_ = tensor_model_parallel_all_reduce(output_parallel)

fastvideo/train/callbacks/__init__.py

@@ -8,6 +8,8 @@
     EMACallback, )
 from fastvideo.train.callbacks.grad_clip import (
     GradNormClipCallback, )
+from fastvideo.train.callbacks.log_rl_samples import (
+    LogRLSamplesCallback, )
 from fastvideo.train.callbacks.validation import (
     ValidationCallback, )
 
@@ -16,5 +18,6 @@
     "CallbackDict",
     "EMACallback",
     "GradNormClipCallback",
+    "LogRLSamplesCallback",
     "ValidationCallback",
 ]

fastvideo/train/callbacks/callback.py

@@ -24,6 +24,7 @@
     "grad_clip": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback",
     "validation": "fastvideo.train.callbacks.validation.ValidationCallback",
     "ema": "fastvideo.train.callbacks.ema.EMACallback",
+    "log_rl_samples": "fastvideo.train.callbacks.log_rl_samples.LogRLSamplesCallback",
 }
 
 
@@ -58,6 +59,7 @@ def on_before_optimizer_step(
         self,
         method: TrainingMethod,
         iteration: int = 0,
+        outputs: dict[str, Any] | None = None,
     ) -> None:
         pass