update

Updated the number of GPUs used for training from 1 to 4.

.github/workflows/fastvideo-kernel-publish.yml

@@ -156,23 +156,8 @@ jobs:
           
           # Fix the wheel to be manylinux compliant
           pip install auditwheel
-          # Point auditwheel at torch libs, but do not vendor them into the wheel.
-          TORCH_LIB_DIR=$(python - <<'PY'
-          import os
-          import torch
-
-          print(os.path.join(os.path.dirname(torch.__file__), "lib"))
-          PY
-          )
-          export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${LD_LIBRARY_PATH}"
           # Target manylinux_2_35 (Ubuntu 22.04 native)
-          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist \
-            --exclude libtorch_cuda.so \
-            --exclude libtorch_cpu.so \
-            --exclude libtorch.so \
-            --exclude libc10.so \
-            --exclude libc10_cuda.so \
-            --exclude libtorch_python.so
+          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist
           # Move fixed wheels back to dist for upload consistency
           rm dist/*.whl
           mv fixed_dist/*.whl dist/

examples/training/finetune/hunyuan15_t2v/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/hunyuan15_t2v/preprocess_hunyuan_data_t2v.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+DATASET_PATH="data/crush-smol"
+OUTPUT_DIR="data/crush-smol_processed_t2v_hunyuan15/"
+
+torchrun --nproc_per_node=$GPU_NUM --master_port=29513 \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+    --model_path $MODEL_PATH \
+    --mode preprocess \
+    --workload_type t2v \
+    --preprocess.video_loader_type torchvision \
+    --preprocess.dataset_type merged \
+    --preprocess.dataset_path $DATASET_PATH \
+    --preprocess.dataset_output_dir $OUTPUT_DIR \
+    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.dataloader_num_workers 0 \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 61 \
+    --preprocess.train_fps 15 \
+    --preprocess.samples_per_file 8 \
+    --preprocess.flush_frequency 8 \
+    --preprocess.video_length_tolerance_range 5

examples/training/finetune/hunyuan_t2v/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/hunyuan_t2v/preprocess_hunyuan_data_t2v.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="hunyuanvideo-community/HunyuanVideo"
+DATASET_PATH="data/crush-smol"
+OUTPUT_DIR="data/crush-smol_processed_t2v_hunyuan/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+    --model_path $MODEL_PATH \
+    --mode preprocess \
+    --workload_type t2v \
+    --preprocess.video_loader_type torchvision \
+    --preprocess.dataset_type merged \
+    --preprocess.dataset_path $DATASET_PATH \
+    --preprocess.dataset_output_dir $OUTPUT_DIR \
+    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.dataloader_num_workers 0 \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 77 \
+    --preprocess.train_fps 16 \
+    --preprocess.samples_per_file 8 \
+    --preprocess.flush_frequency 8 \
+    --preprocess.video_length_tolerance_range 5

examples/training/rl/finetune_t2v_grpo.sh

@@ -1,129 +0,0 @@
-#!/bin/bash
-
-# Change to FastVideo root directory (3 levels up from this script)
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-FASTVIDEO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
-cd "$FASTVIDEO_ROOT"
-
-# Add FastVideo root to PYTHONPATH so Python can find the fastvideo package
-export PYTHONPATH="$FASTVIDEO_ROOT${PYTHONPATH:+:$PYTHONPATH}"
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-RL_DATASET_DIR="data/ocr/"  # Path to RL prompt dataset directory (should contain train.txt and test.txt)
-VALIDATION_DATASET_FILE="$SCRIPT_DIR/validation.json"
-NUM_GPUS=1
-
-# use GPU 3
-export CUDA_VISIBLE_DEVICES=3
-
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_t2v_grpo"
-  --output_dir "checkpoints/wan_t2v_grpo"
-  --max_train_steps 5000
-  --train_batch_size 4
-  # --train_sp_batch_size 4
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 5
-  --num_height 240
-  --num_width 416
-  --num_frames 33
-  --lora_rank 32
-  --lora_training True
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS 
-  --sp_size $NUM_GPUS 
-  --tp_size $NUM_GPUS
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
-  # --use-fsdp-inference False
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments (for RL prompt dataset)
-dataset_args=(
-  --data_path $RL_DATASET_DIR  # Used as fallback if rl_dataset_path not set
-  --rl_dataset_path $RL_DATASET_DIR  # RL prompt dataset directory
-  --rl_dataset_type "text"  # "text" or "geneval"
-  --rl_num_image_per_prompt 4  # k parameter (number of samples per prompt)
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation True
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 5
-  --validation_sampling_steps "50" 
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 5e-5
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 10
-  --training_state_checkpointing_steps 10
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# RL-specific arguments
-rl_args=(
-  --inference_mode False
-  --rl_mode True
-  --rl_algorithm "grpo"
-  --rl_kl_beta 0.004  # KL regularization coefficient
-  --rl_policy_clip_range 0.2  # Policy clipping range for GRPO
-  --rl_kl_reward 0.0  # KL reward coefficient (typically 0)
-  --rl_global_std False  # Use per-prompt std (recommended for GRPO)
-  --rl_per_prompt_stat_tracking True  # Enable per-prompt stat tracking
-  --rl_warmup_steps 0  # Number of warmup steps (SFT before RL)
-  --reward-models "{\"paddle_ocr\": 1.0}" # use video_ocr reward function
-)
-
-# CFG arguments
-cfg_args=(
-  --guidance_scale 1.0 # use guidance_scale > 1.0 to enable CFG
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0  # No CFG during training (CFG used in sampling)
-  --dit_precision "fp32"
-  # --dit_precision "bf16"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-  # --resume_from_checkpoint "checkpoints/wan_t2v_grpo/checkpoint-XXX"
-  --enable-gradient-checkpointing-type "full"
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-  --master_port 29501 \
-    "$FASTVIDEO_ROOT/fastvideo/training/wan_rl_training_pipeline.py" \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${rl_args[@]}" \
-    "${miscellaneous_args[@]}"

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu

@@ -4,7 +4,6 @@
 #include <torch/all.h>
 #include <torch/python.h>
 #include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
 #include "common/common.hpp"
 #include "norm/layernorm.hpp"
 
@@ -15,6 +14,10 @@ auto layer_norm(
   std::optional<at::Tensor const> const B,
   std::optional<at::Tensor> Output
 ) {
+  using ElementIn = float;
+  using ElementOut = float;
+  using ElementWeight = float;
+
   int64_t const m = Input.size(0);
   int64_t const n = Input.size(1);
   torch::Device const input_device = Input.device();
@@ -23,70 +26,31 @@ auto layer_norm(
     Output.emplace(
       torch::empty(
         {m, n},
-        torch::TensorOptions().device(input_device).dtype(Input.scalar_type())
+        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
       )
     );
   }
 
-  TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(),
-              "Output dtype must match Input dtype. Got Output=",
-              Output.value().scalar_type(), ", Input=", Input.scalar_type());
-  if (W.has_value()) {
-    TORCH_CHECK(W.value().scalar_type() == Input.scalar_type(),
-                "W dtype must match Input dtype. Got W=",
-                W.value().scalar_type(), ", Input=", Input.scalar_type());
-  }
-  if (B.has_value()) {
-    TORCH_CHECK(B.value().scalar_type() == Input.scalar_type(),
-                "B dtype must match Input dtype. Got B=",
-                B.value().scalar_type(), ", Input=", Input.scalar_type());
-  }
 
   void *Iptr = Input.data_ptr();
   void *Wptr = W.has_value() ? W.value().data_ptr() : nullptr;
   void *Bptr = B.has_value() ? B.value().data_ptr() : nullptr;
   void *Optr = Output.value().data_ptr();
 
-  auto stream = at::cuda::getCurrentCUDAStream().stream();
-  if (Input.scalar_type() == at::kHalf) {
-    using ElementIn = cutlass::half_t;
-    using ElementOut = cutlass::half_t;
-    using ElementWeight = cutlass::half_t;
-    BOOL_SWITCH(B.has_value(), BIAS, [&]{
-      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-        CONFIG_SWITCH(n, [&]{
-          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
-        });
+  BOOL_SWITCH(B.has_value(), BIAS, [&]{
+    BOOL_SWITCH(W.has_value(), AFFINE, [&]{
+    CONFIG_SWITCH(n, [&]{
+    layernorm<
+    ElementIn, ElementOut, ElementWeight,
+    AFFINE, BIAS,
+    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA> (
+      Iptr, Wptr, Bptr,
+      Optr, eps, m, n,
+      at::cuda::getCurrentCUDAStream().stream()
+        );
       });
     });
-  } else if (Input.scalar_type() == at::kBFloat16) {
-    using ElementIn = cutlass::bfloat16_t;
-    using ElementOut = cutlass::bfloat16_t;
-    using ElementWeight = cutlass::bfloat16_t;
-    BOOL_SWITCH(B.has_value(), BIAS, [&]{
-      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-        CONFIG_SWITCH(n, [&]{
-          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
-        });
-      });
-    });
-  } else if (Input.scalar_type() == at::kFloat) {
-    using ElementIn = float;
-    using ElementOut = float;
-    using ElementWeight = float;
-    BOOL_SWITCH(B.has_value(), BIAS, [&]{
-      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-        CONFIG_SWITCH(n, [&]{
-          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
-        });
-      });
-    });
-  } else {
-    TORCH_CHECK(false, "Unsupported dtype for layer_norm_cuda: ", Input.scalar_type());
-  }
+  });
   
     
 

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp

@@ -68,22 +68,9 @@ public:
     // mean reduction
     float u = _reduce_sum(x, shared_data) / params.n;
 
-    // IMPORTANT:
-    // Loader pads out-of-range lanes with 0. That is OK for the sum, but after
-    // subtracting mean, those padded lanes become -u and would incorrectly
-    // contribute to the variance. Mask them back to 0 before variance reduction.
-    // We launch exactly NumThrPerCta threads for a 1xMaxHiddenSize tile,
-    // so each thread is responsible for a contiguous chunk in N.
-    int thr_n_offset = tidx * NumElementPerThread;
     CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i) {
-      int idx = thr_n_offset + i;
-      if (idx < params.n) {
-        x[i] -= u;
-      } else {
-        x[i] = 0.f;
-      }
-    }
+    for (int i = 0; i < NumElementPerThread; ++i)
+      x[i] -= u;
 
     __syncthreads();
     // var reduction

fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu

@@ -4,7 +4,6 @@
 #include <torch/all.h>
 #include <torch/python.h>
 #include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
 #include <pybind11/pybind11.h>
 
 #include "common/common.hpp"
@@ -17,6 +16,10 @@ auto rms_norm(
   std::optional<at::Tensor>& Output
 ) {
   
+  using ElementIn = float;
+  using ElementOut = float;
+  using ElementWeight = float;
+  
   int64_t const m = Input.size(0);
   int64_t const n = Input.size(1);
   torch::Device const input_device = Input.device();
@@ -25,51 +28,27 @@ auto rms_norm(
     Output.emplace(
       torch::empty(
         {m, n},
-        torch::TensorOptions().device(input_device).dtype(Input.scalar_type())
+        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
       )
     );
   }
 
-  TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(),
-              "Output dtype must match Input dtype. Got Output=",
-              Output.value().scalar_type(), ", Input=", Input.scalar_type());
-  if (Weight.has_value()) {
-    TORCH_CHECK(Weight.value().scalar_type() == Input.scalar_type(),
-                "Weight dtype must match Input dtype. Got Weight=",
-                Weight.value().scalar_type(), ", Input=", Input.scalar_type());
-  }
-
   void *Iptr = Input.data_ptr();
   void *Wptr = Weight.has_value() ? Weight.value().data_ptr() : nullptr;
   void *Optr = Output.value().data_ptr();
 
-  if (Input.scalar_type() == at::kHalf) {
-    using ElementIn = cutlass::half_t;
-    using ElementOut = cutlass::half_t;
-    using ElementWeight = cutlass::half_t;
-    CONFIG_SWITCH(n, [&]{
-      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
-    });
-  } else if (Input.scalar_type() == at::kBFloat16) {
-    using ElementIn = cutlass::bfloat16_t;
-    using ElementOut = cutlass::bfloat16_t;
-    using ElementWeight = cutlass::bfloat16_t;
-    CONFIG_SWITCH(n, [&]{
-      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
-    });
-  } else if (Input.scalar_type() == at::kFloat) {
-    using ElementIn = float;
-    using ElementOut = float;
-    using ElementWeight = float;
-    CONFIG_SWITCH(n, [&]{
-      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
-    });
-  } else {
-    TORCH_CHECK(false, "Unsupported dtype for rms_norm_cuda: ", Input.scalar_type());
-  }
+
+  CONFIG_SWITCH(n, [&]{
+    rmsnorm<
+    ElementIn, ElementOut, ElementWeight,
+    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA
+    > (
+      Iptr, Wptr,
+      Optr, 
+      eps, m, n, 
+      at::cuda::getCurrentCUDAStream().stream()
+    );
+});
   
 
   return Output;

fastvideo-kernel/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "fastvideo-kernel"
-version = "0.2.4"
+version = "0.2.3"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 from typing import Tuple
 
 import torch
@@ -24,12 +23,6 @@ def _is_sm90() -> bool:
     return major == 9 and minor == 0
 
 
-def _force_triton() -> bool:
-    # Force Triton even on SM90 and even if the compiled extension is available.
-    # Useful for CI / debugging / parity testing.
-    return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1"
-
-
 def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Pure-torch (no triton) conversion:
@@ -288,7 +281,7 @@ def block_sparse_attn(
     - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today).
     """
     block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
-    if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
+    if _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
         return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes)
     # Triton path: generally assumes q/k/v share the same padded length
     if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]:

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -46,22 +46,14 @@ def sliding_tile_attention(
     flag = shape_map[seq_shape]
 
     for head_idx, (t, h, w) in enumerate(window_size):
-        # Per-head slices are not contiguous in the batch dimension when batch>1
-        # (they keep the original head-stride). The TK kernel assumes contiguous
-        # [B, H, S, D] layout, so we materialize a contiguous [B,1,S,D] view.
-        q_h = q[:, head_idx:head_idx + 1].contiguous()
-        k_h = k[:, head_idx:head_idx + 1].contiguous()
-        v_h = v[:, head_idx:head_idx + 1].contiguous()
-        o_h = torch.empty_like(q_h)
         sta_fwd(
-            q_h, k_h,
-            v_h, o_h,
+            q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
+            v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
             t, h, w, text_length, False, has_text, flag
         )
-        output[:, head_idx:head_idx + 1] = o_h
 
     if has_text:
-        sta_fwd(q.contiguous(), k.contiguous(), v.contiguous(), output, 3, 3, 3, text_length, True, True, flag)
+        sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
 
     return output[:, :, :seq_length]
 

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.4"
+__version__ = "0.2.3"

fastvideo/configs/models/encoders/clip.py

@@ -44,6 +44,15 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
     _fsdp_shard_conditions: list = field(
         default_factory=lambda: [_is_transformer_layer, _is_embeddings])
 
+    tokenizer_kwargs: dict = field(
+        default_factory=lambda: {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": 77,
+            "return_tensors": "pt",
+        }
+    )
+
 
 @dataclass
 class CLIPVisionArchConfig(ImageEncoderArchConfig):

fastvideo/configs/models/encoders/llama.py

@@ -56,6 +56,15 @@ class LlamaArchConfig(TextEncoderArchConfig):
         default_factory=lambda:
         [_is_transformer_layer, _is_embeddings, _is_final_norm])
 
+    tokenizer_kwargs: dict = field(
+        default_factory=lambda: {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": 256,
+            "return_tensors": "pt",
+        }
+    )
+
 
 @dataclass
 class LlamaConfig(TextEncoderConfig):

fastvideo/dataset/__init__.py

@@ -8,7 +8,6 @@ from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset,
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                          TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
-from fastvideo.dataset.rl_prompt_dataset import build_rl_prompt_dataloader
 
 
 def getdataset(args) -> VideoCaptionMergedDataset:
@@ -48,6 +47,5 @@ def gettextdataset(args) -> TextDataset:
 
 __all__ = [
     "build_parquet_map_style_dataloader", "ValidationDataset",
-    "VideoCaptionMergedDataset", "TextDataset",
-    "build_rl_prompt_dataloader"
+    "VideoCaptionMergedDataset", "TextDataset"
 ]

fastvideo/dataset/dataloader/record_schema.py

@@ -72,9 +72,16 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
     """Create a record for the Parquet dataset with CLIP features."""
     records = basic_t2v_record_creator(batch)
 
-    assert len(
-        batch.image_embeds) == 1, "image embedding should be a single tensor"
-    image_embeds = batch.image_embeds[0]
+    # Adapt for model doesn't have image encoder, e.g., Hunyuan
+    if len(batch.image_embeds) == 0:
+        image_embeds = None
+    elif len(batch.image_embeds) == 1:
+        image_embeds = batch.image_embeds[0]
+    else:
+        raise ValueError(
+            "Unexpected number of image_embeds in batch: expected 0 or 1, got {}".format(
+                len(batch.image_embeds)))
+        
     image_latent = batch.image_latent
     pil_image = batch.pil_image
 

fastvideo/dataset/rl_prompt_dataset.py

@@ -1,174 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import torch
-from torch.utils.data import Dataset, DataLoader, Sampler
-import json
-import os
-
-
-class TextPromptDataset(Dataset):
-    """Dataset for loading text prompts from a simple text file (one prompt per line)."""
-    
-    def __init__(self, dataset, split='train'):
-        self.file_path = os.path.join(dataset, f'{split}.txt')
-        with open(self.file_path, 'r') as f:
-            self.prompts = [line.strip() for line in f.readlines()]
-        
-    def __len__(self):
-        return len(self.prompts)
-    
-    def __getitem__(self, idx):
-        return {"prompt": self.prompts[idx], "metadata": {}}
-
-    @staticmethod
-    def collate_fn(examples):
-        prompts = [example["prompt"] for example in examples]
-        metadatas = [example["metadata"] for example in examples]
-        return prompts, metadatas
-
-
-class GenevalPromptDataset(Dataset):
-    """Dataset for loading prompts with metadata from JSONL files (e.g., GenEval format)."""
-    
-    def __init__(self, dataset, split='train'):
-        self.file_path = os.path.join(dataset, f'{split}_metadata.jsonl')
-        with open(self.file_path, 'r', encoding='utf-8') as f:
-            self.metadatas = [json.loads(line) for line in f]
-            self.prompts = [item['prompt'] for item in self.metadatas]
-        
-    def __len__(self):
-        return len(self.prompts)
-    
-    def __getitem__(self, idx):
-        return {"prompt": self.prompts[idx], "metadata": self.metadatas[idx]}
-
-    @staticmethod
-    def collate_fn(examples):
-        prompts = [example["prompt"] for example in examples]
-        metadatas = [example["metadata"] for example in examples]
-        return prompts, metadatas
-
-
-class KRepeatSampler(Sampler):
-    """Sampler that repeats each sample k times, ensuring synchronized random selection. For single-node training, set num_replicas=1 and rank=0."""
-    
-    def __init__(self, dataset, batch_size, k, num_replicas, rank, seed=0):
-        self.dataset = dataset
-        self.batch_size = batch_size        # Batch size per GPU/card
-        self.k = k                          # Number of repetitions per sample
-        self.num_replicas = num_replicas    # Total number of GPUs/cards
-        self.rank = rank                    # Current GPU/card rank
-        self.seed = seed                    # Random seed for synchronization
-        
-        # Calculate the number of unique samples needed for each iteration
-        self.total_samples = self.num_replicas * self.batch_size
-        assert self.total_samples % self.k == 0, f"k can not div n*b, k{k}-num_replicas{num_replicas}-batch_size{batch_size}"
-        self.m = self.total_samples // self.k  # different number of samples
-        self.step = 0
-
-    def __iter__(self):
-        while True:
-            # Generate a deterministic random sequence to ensure all cards are synchronized
-            g = torch.Generator()
-            g.manual_seed(self.seed + self.step)
-            
-            # Randomly select m unique samples
-            indices = torch.randperm(len(self.dataset), generator=g)[:self.m].tolist()
-            
-            # Repeat each sample k times to generate a total of n*b samples
-            repeated_indices = [idx for idx in indices for _ in range(self.k)]
-            
-            # Shuffle the order to ensure even distribution
-            shuffled_indices = torch.randperm(len(repeated_indices), generator=g).tolist()
-            shuffled_samples = [repeated_indices[i] for i in shuffled_indices]
-            
-            # Split samples among all cards
-            per_card_samples = []
-            for i in range(self.num_replicas):
-                start = i * self.batch_size
-                end = start + self.batch_size
-                per_card_samples.append(shuffled_samples[start:end])
-            
-            # Return the sample indices for the current card
-            yield per_card_samples[self.rank]
-        
-    def __len__(self):
-        return len(self.dataset) // self.batch_size
-    
-    def set_step(self, step):
-        """Used to synchronize the random state for different epochs."""
-        self.step = step
-
-
-def build_rl_prompt_dataloader(
-    dataset_path: str,
-    dataset_type: str = "text",
-    split: str = "train",
-    train_batch_size: int = 8,
-    test_batch_size: int = 8,
-    k: int = 1,
-    seed: int = 42,
-    train_num_workers: int = 1,
-    test_num_workers: int = 8,
-    num_replicas: int = 1,
-    rank: int = 0,
-) -> tuple[DataLoader, DataLoader]:
-    """
-    Factory function to create train and test dataloaders for RL prompt datasets.
-    
-    Args:
-        dataset_path: Path to dataset directory
-        dataset_type: "text" for TextPromptDataset or "geneval" for GenevalPromptDataset
-        split: Dataset split ("train" or "test")
-        train_batch_size: Batch size per GPU for training
-        test_batch_size: Batch size for testing
-        k: Number of times to repeat each sample (num_image_per_prompt)
-        seed: Random seed for sampler synchronization
-        train_num_workers: Number of workers for training dataloader
-        test_num_workers: Number of workers for test dataloader
-        num_replicas: Number of replicas (default 1 for single-node)
-        rank: Rank of current process (default 0 for single-node)
-    
-    Returns:
-        Tuple of (train_dataloader, test_dataloader)
-    """
-    # Create datasets based on type
-    if dataset_type == "text":
-        train_dataset = TextPromptDataset(dataset_path, 'train')
-        test_dataset = TextPromptDataset(dataset_path, 'test')
-        collate_fn = TextPromptDataset.collate_fn
-    elif dataset_type == "geneval":
-        train_dataset = GenevalPromptDataset(dataset_path, 'train')
-        test_dataset = GenevalPromptDataset(dataset_path, 'test')
-        collate_fn = GenevalPromptDataset.collate_fn
-    else:
-        raise ValueError(f"Unknown dataset_type: {dataset_type}. Must be 'text' or 'geneval'")
-    
-    # Create infinite-loop training sampler
-    train_sampler = KRepeatSampler(
-        dataset=train_dataset,
-        batch_size=train_batch_size,
-        k=k,
-        num_replicas=num_replicas,
-        rank=rank,
-        seed=seed
-    )
-    
-    # Create training dataloader with batch_sampler (infinite loop)
-    train_dataloader = DataLoader(
-        train_dataset,
-        batch_sampler=train_sampler,
-        num_workers=train_num_workers,
-        collate_fn=collate_fn,
-    )
-    
-    # Create standard test dataloader
-    test_dataloader = DataLoader(
-        test_dataset,
-        batch_size=test_batch_size,
-        collate_fn=collate_fn,
-        shuffle=False,
-        num_workers=test_num_workers,
-    )
-    
-    return train_dataloader, test_dataloader, train_dataset, test_dataset
-

fastvideo/dataset/utils.py

@@ -30,12 +30,13 @@ def get_torch_tensors_from_row_dict(row_dict,
     """
     return_dict = {}
     for key in keys:
-        shape, bytes = None, None
+        shape, bytes, dtype_str = None, None, None
         if isinstance(key, tuple):
             for k in key:
                 try:
                     shape = row_dict[f"{k}_shape"]
                     bytes = row_dict[f"{k}_bytes"]
+                    dtype_str = row_dict.get(f"{k}_dtype", 'float32')
                 except KeyError:
                     continue
             key = key[0]
@@ -44,13 +45,25 @@ def get_torch_tensors_from_row_dict(row_dict,
         else:
             shape = row_dict[f"{key}_shape"]
             bytes = row_dict[f"{key}_bytes"]
+            dtype_str = row_dict.get(f"{key}_dtype", 'float32')
+
+        # Convert dtype string to numpy dtype
+        if 'float16' in dtype_str:
+            np_dtype = np.float16
+        elif 'float32' in dtype_str:
+            np_dtype = np.float32
+        elif 'float64' in dtype_str:
+            np_dtype = np.float64
+        else:
+            # Default to float32 if dtype is unrecognized
+            np_dtype = np.float32
 
         # TODO (peiyuan): read precision
         if key == 'text_embedding' and (rng.random()
                                         if rng else random.random()) < cfg_rate:
             data = np.zeros((512, 4096), dtype=np.float32)
         else:
-            data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
+            data = np.frombuffer(bytes, dtype=np_dtype).reshape(shape).copy()
         data = torch.from_numpy(data)
         if len(data.shape) == 3:
             B, L, D = data.shape
@@ -141,21 +154,35 @@ def collate_rows_from_parquet_schema(rows,
             # Get tensor data from row using the existing helper function pattern
             shape_key = f"{tensor_name}_shape"
             bytes_key = f"{tensor_name}_bytes"
+            dtype_key = f"{tensor_name}_dtype"
 
             if shape_key in row and bytes_key in row:
                 shape = row[shape_key]
                 bytes_data = row[bytes_key]
+                
+                # Get dtype from row, default to float32 for backward compatibility
+                dtype_str = row.get(dtype_key, 'float32')
+                # Convert dtype string to numpy dtype
+                if 'float16' in dtype_str:
+                    np_dtype = np.float16
+                elif 'float32' in dtype_str:
+                    np_dtype = np.float32
+                elif 'float64' in dtype_str:
+                    np_dtype = np.float64
+                else:
+                    # Default to float32 if dtype is unrecognized
+                    np_dtype = np.float32
 
                 if len(bytes_data) == 0:
                     tensor = torch.zeros(0, dtype=torch.bfloat16)
                 else:
-                    # Convert bytes to tensor using float32 as default
+                    # Convert bytes to tensor using the correct dtype
                     if tensor_name == 'text_embedding' and (rng.random(
                     ) if rng else random.random()) < cfg_rate:
                         data = np.zeros((512, 4096), dtype=np.float32)
                     else:
                         data = np.frombuffer(
-                            bytes_data, dtype=np.float32).reshape(shape).copy()
+                            bytes_data, dtype=np_dtype).reshape(shape).copy()
                     tensor = torch.from_numpy(data)
                     # if len(data.shape) == 3:
                     #     B, L, D = tensor.shape

fastvideo/distributed/communication_op.py

@@ -53,6 +53,9 @@ def sequence_model_parallel_all_gather_with_unpad(
         Tensor: Gathered and unpadded tensor
     """
 
+    # NCCL all_gather expects contiguous inputs.
+    if not input_.is_contiguous():
+        input_ = input_.contiguous()
     # First gather across all ranks
     gathered = get_sp_group().all_gather(input_, dim)
 

fastvideo/fastvideo_args.py

@@ -133,7 +133,7 @@ class FastVideoArgs:
     # CPU offload parameters
     dit_cpu_offload: bool = True
     use_fsdp_inference: bool = False
-    dit_layerwise_offload: bool = True
+    dit_layerwise_offload: bool = False
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
     vae_cpu_offload: bool = True
@@ -740,271 +740,6 @@ def get_current_fastvideo_args() -> FastVideoArgs:
     return _current_fastvideo_args
 
 
-@dataclasses.dataclass
-class RLArgs:
-    """
-    Reinforcement Learning (RL) specific arguments
-    """
-    # ============================================================================
-    # SHARED RL CONFIGURATION
-    rl_mode: bool = False # Enable RL training mode
-    rl_algorithm: str = "grpo"  # RL algorithm to use: "grpo", "ppo", "dpo"
-
-    # Trajectory collection
-    num_rollouts: int = 4  # Number of rollouts to collect per training step
-    rollout_steps: str = "20,30"  # Random intermediate steps for sampling (comma-separated)
-    noise_injection_min: int = 10  # Minimum timestep for noise injection
-    noise_injection_max: int = 40  # Maximum timestep for noise injection
-    use_sde_sampling: bool = True  # Use SDE sampling (Flow-GRPO-Fast)
-    num_denoising_steps: int = 2  # Number of denoising steps per trajectory (1-2 for fast)
-
-    # Advantage estimation
-    gamma: float = 0.99  # Discount factor for returns
-    lambda_param: float = 0.95  # GAE lambda parameter
-    use_gae: bool = True  # Use Generalized Advantage Estimation
-    normalize_advantages: bool = True  # Normalize advantages before policy update
-
-    # Reward models
-    reward_models: dict[str, float] = field(default_factory=lambda: {"dummy": 1.0})  # reward models (names, weight)
-    value_model_path: str = ""  # Path to value model (can be empty to train from scratch)
-    value_model_share_backbone: bool = False  # Share transformer backbone between policy and value
-
-    # Training schedule
-    warmup_steps: int = 1000  # Collect SFT-style data before starting RL
-    collect_on_policy: bool = True  # Collect fresh rollouts each step (on-policy)
-    timestep_fraction: float = 0.99  # Fraction of timesteps to train on
-    num_inner_epochs: int = 1  # Number of inner epochs per outer epoch
-
-    # KL regularization
-    kl_beta: float = 0.004  # KL loss coefficient (GRPO uses KL loss, DPO uses larger beta)
-    kl_reward: float = 0.0  # KL reward coefficient (alternative to KL loss, typically 0)
-
-    # SFT integration
-    sft_weight: float = 0.0  # SFT loss weight for supervised learning in RL training
-    sft_batch_size: int = 3  # Batch size for SFT data
-
-    # CFG
-    guidance_scale = 1.0 # use guidance_scale > 1.0 to enable CFG
-
-    # Statistics tracking
-    global_std: bool = False  # Use global std across all samples vs per-group std
-    per_prompt_stat_tracking: bool = True  # Track statistics per prompt
-
-    # Training options
-    use_diffusion_loss: bool = True  # Use diffusion loss in training
-
-    # ============================================================================
-    # GRPO-SPECIFIC CONFIGURATION
-
-    # Policy optimization
-    grpo_policy_clip_range: float = 0.001  # PPO-style clipping range for policy ratio
-    grpo_value_clip_range: float = 0.2  # Value function clipping range
-    grpo_num_policy_epochs: int = 1  # Number of policy update epochs (GRPO typically uses 1)
-    grpo_num_value_epochs: int = 1  # Number of value function update epochs
-    grpo_target_kl: float = 0.01  # Target KL divergence for early stopping
-    grpo_entropy_coef: float = 0.0  # Entropy coefficient for exploration
-    grpo_value_loss_coef: float = 0.5  # Value loss coefficient
-
-    # GRPO-Guard safety mechanisms
-    grpo_use_grpo_guard: bool = True  # Enable GRPO-Guard safety mechanisms
-    grpo_ratio_norm_correction: bool = True  # RatioNorm: correct importance ratio bias
-    grpo_gradient_reweighting: bool = True  # Reweight gradients across denoising steps
-    grpo_max_importance_ratio: float = 10.0  # Clip importance ratios above this value
-
-    # ============================================================================
-    # DPO-SPECIFIC CONFIGURATION
-
-    dpo_beta: float = 100.0  # DPO regularization parameter (typically much larger than GRPO beta)
-    dpo_ref_update_step: int = 10000000  # Reference model update frequency for OnlineDPO
-    dpo_label_smoothing: float = 0.0  # Label smoothing for DPO loss
-
-    @staticmethod
-    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
-        """Add RL-specific CLI arguments to the parser."""
-        # RL (Reinforcement Learning) arguments
-        parser.add_argument("--rl-mode",
-                            action=StoreBoolean,
-                            help="Enable RL training mode")
-        parser.add_argument("--rl-algorithm",
-                            type=str,
-                            default=RLArgs.rl_algorithm,
-                            choices=["grpo", "ppo", "dpo"],
-                            help="RL algorithm to use (grpo, ppo, dpo)")
-
-        # Trajectory collection (Flow-GRPO-Fast)
-        parser.add_argument("--rl-num-rollouts",
-                            type=int,
-                            default=RLArgs.num_rollouts,
-                            help="Number of rollouts to collect per training step")
-        parser.add_argument("--rl-rollout-steps",
-                            type=str,
-                            default=RLArgs.rollout_steps,
-                            help="Random intermediate steps for sampling (comma-separated)")
-        parser.add_argument("--rl-noise-injection-min",
-                            type=int,
-                            default=RLArgs.noise_injection_min,
-                            help="Minimum timestep for noise injection")
-        parser.add_argument("--rl-noise-injection-max",
-                            type=int,
-                            default=RLArgs.noise_injection_max,
-                            help="Maximum timestep for noise injection")
-        parser.add_argument("--rl-use-sde-sampling",
-                            action=StoreBoolean,
-                            help="Use SDE sampling (Flow-GRPO-Fast)")
-        parser.add_argument("--rl-num-denoising-steps",
-                            type=int,
-                            default=RLArgs.num_denoising_steps,
-                            help="Number of denoising steps per trajectory (1-2 for fast)")
-
-        # Advantage estimation
-        parser.add_argument("--rl-gamma",
-                            type=float,
-                            default=RLArgs.gamma,
-                            help="Discount factor for returns")
-        parser.add_argument("--rl-lambda",
-                            type=float,
-                            default=RLArgs.lambda_param,
-                            help="GAE lambda parameter")
-        parser.add_argument("--rl-use-gae",
-                            action=StoreBoolean,
-                            help="Use Generalized Advantage Estimation")
-        parser.add_argument("--rl-normalize-advantages",
-                            action=StoreBoolean,
-                            help="Normalize advantages before policy update")
-
-        # Policy optimization (GRPO/PPO)
-        parser.add_argument("--rl-policy-clip-range",
-                            type=float,
-                            default=RLArgs.grpo_policy_clip_range,
-                            dest="grpo_policy_clip_range",  # Map to RLArgs field name
-                            help="PPO-style clipping range for policy ratio")
-        parser.add_argument("--rl-value-clip-range",
-                            type=float,
-                            default=RLArgs.grpo_value_clip_range,
-                            help="Value function clipping range")
-        parser.add_argument("--rl-num-policy-epochs",
-                            type=int,
-                            default=RLArgs.grpo_num_policy_epochs,
-                            help="Number of policy update epochs (GRPO typically uses 1)")
-        parser.add_argument("--rl-num-value-epochs",
-                            type=int,
-                            default=RLArgs.grpo_num_value_epochs,
-                            help="Number of value function update epochs")
-        parser.add_argument("--rl-target-kl",
-                            type=float,
-                            default=RLArgs.grpo_target_kl,
-                            help="Target KL divergence for early stopping")
-        parser.add_argument("--rl-entropy-coef",
-                            type=float,
-                            default=RLArgs.grpo_entropy_coef,
-                            help="Entropy coefficient for exploration")
-        parser.add_argument("--rl-value-loss-coef",
-                            type=float,
-                            default=RLArgs.grpo_value_loss_coef,
-                            help="Value loss coefficient")
-
-        # GRPO-Guard (safety mechanisms)
-        parser.add_argument("--rl-use-grpo-guard",
-                            action=StoreBoolean,
-                            help="Enable GRPO-Guard safety mechanisms")
-        parser.add_argument("--rl-ratio-norm-correction",
-                            action=StoreBoolean,
-                            help="RatioNorm: correct importance ratio bias")
-        parser.add_argument("--rl-gradient-reweighting",
-                            action=StoreBoolean,
-                            help="Reweight gradients across denoising steps")
-        parser.add_argument("--rl-max-importance-ratio",
-                            type=float,
-                            default=RLArgs.grpo_max_importance_ratio,
-                            help="Clip importance ratios above this value")
-
-        # Reward models
-        parser.add_argument("--reward-models",
-                            type=str,
-                            default='{"dummy": 1.0}',
-                            help="Reward models as JSON dict (e.g., '{\"video_ocr\": 1.0, \"pickscore\": 0.5}')")
-        parser.add_argument("--value-model-path",
-                            type=str,
-                            default=RLArgs.value_model_path,
-                            help="Path to value model (can be empty to train from scratch)")
-        parser.add_argument("--value-model-share-backbone",
-                            action=StoreBoolean,
-                            help="Share transformer backbone between policy and value")
-
-        # Training schedule
-        parser.add_argument("--rl-warmup-steps",
-                            type=int,
-                            default=RLArgs.warmup_steps,
-                            help="Collect SFT-style data before starting RL")
-        parser.add_argument("--rl-collect-on-policy",
-                            action=StoreBoolean,
-                            help="Collect fresh rollouts each step (on-policy)")
-        parser.add_argument("--rl-timestep-fraction",
-                            type=float,
-                            default=RLArgs.timestep_fraction,
-                            help="Fraction of timesteps to train on")
-        parser.add_argument("--rl-num-inner-epochs",
-                            type=int,
-                            default=RLArgs.num_inner_epochs,
-                            help="Number of inner epochs per outer epoch")
-
-        # KL regularization
-        parser.add_argument("--rl-kl-beta",
-                            type=float,
-                            default=RLArgs.kl_beta,
-                            dest="kl_beta",  # Map CLI arg to RLArgs field name
-                            help="KL loss coefficient (GRPO uses KL loss, DPO uses larger beta)")
-        parser.add_argument("--rl-kl-reward",
-                            type=float,
-                            default=RLArgs.kl_reward,
-                            help="KL reward coefficient (alternative to KL loss, typically 0)")
-
-        # SFT integration
-        parser.add_argument("--rl-sft-weight",
-                            type=float,
-                            default=RLArgs.sft_weight,
-                            help="SFT loss weight for supervised learning in RL training")
-        parser.add_argument("--rl-sft-batch-size",
-                            type=int,
-                            default=RLArgs.sft_batch_size,
-                            help="Batch size for SFT data")
-
-        # CFG settings
-        parser.add_argument("--guidance-scale",
-                            type=float,
-                            default=1.0,
-                            help="Guidance scale for CFG")
-
-        # Statistics tracking
-        parser.add_argument("--rl-global-std",
-                            action=StoreBoolean,
-                            help="Use global std across all samples vs per-group std")
-        parser.add_argument("--rl-per-prompt-stat-tracking",
-                            action=StoreBoolean,
-                            help="Track statistics per prompt")
-
-        # Training options
-        parser.add_argument("--rl-use-diffusion-loss",
-                            action=StoreBoolean,
-                            help="Use diffusion loss in training")
-
-        # DPO-specific
-        parser.add_argument("--dpo-beta",
-                            type=float,
-                            default=RLArgs.dpo_beta,
-                            help="DPO regularization parameter (typically much larger than GRPO beta)")
-        parser.add_argument("--dpo-ref-update-step",
-                            type=int,
-                            default=RLArgs.dpo_ref_update_step,
-                            help="Reference model update frequency for OnlineDPO")
-        parser.add_argument("--dpo-label-smoothing",
-                            type=float,
-                            default=RLArgs.dpo_label_smoothing,
-                            help="Label smoothing for DPO loss")
-
-        return parser
-
-
 @dataclasses.dataclass
 class TrainingArgs(FastVideoArgs):
     """
@@ -1017,11 +752,6 @@ class TrainingArgs(FastVideoArgs):
     num_height: int = 0
     num_width: int = 0
     num_frames: int = 0
-    
-    # RL dataset configuration (for RL prompt datasets)
-    rl_dataset_path: str = ""  # Path to RL prompt dataset directory (defaults to data_path if not set)
-    rl_dataset_type: str = "text"  # "text" or "geneval"
-    rl_num_image_per_prompt: int = 4  # k parameter for KRepeatSampler (num_image_per_prompt)
 
     train_batch_size: int = 0
     num_latent_t: int = 0
@@ -1132,9 +862,6 @@ class TrainingArgs(FastVideoArgs):
     last_step_only: bool = False  # Only use the last timestep for training
     context_noise: int = 0  # Context noise level for cache updates
 
-    # Nested RL configuration
-    rl_args: RLArgs = dataclasses.field(default_factory=RLArgs)
-
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
         provided_args = clean_cli_args(args)
@@ -1159,25 +886,6 @@ class TrainingArgs(FastVideoArgs):
                 kwargs[attr] = WorkloadType.from_string(
                     workload_type_value) if isinstance(
                         workload_type_value, str) else workload_type_value
-            elif attr == 'rl_args':
-                # Construct nested RLArgs from CLI arguments
-                rl_kwargs = {}
-                for rl_field in dataclasses.fields(RLArgs):
-                    rl_attr = rl_field.name
-                    if hasattr(args, rl_attr):
-                        value = getattr(args, rl_attr)
-                        # Special handling for reward_models: parse JSON string to dict
-                        if rl_attr == 'reward_models' and isinstance(value, str):
-                            rl_kwargs[rl_attr] = json.loads(value) if value else {}
-                        else:
-                            rl_kwargs[rl_attr] = value
-                    else:
-                        # Use default value from RLArgs
-                        if rl_field.default_factory is not dataclasses.MISSING:
-                            rl_kwargs[rl_attr] = rl_field.default_factory()
-                        elif rl_field.default is not dataclasses.MISSING:
-                            rl_kwargs[rl_attr] = rl_field.default
-                kwargs[attr] = RLArgs(**rl_kwargs)
             # Use getattr with default value from the dataclass for potentially missing attributes
             else:
                 # Get the field to check its default value
@@ -1207,26 +915,11 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--data-path",
                             type=str,
                             required=True,
-                            help="Path to parquet files (or RL prompt dataset directory for RL training)")
+                            help="Path to parquet files")
         parser.add_argument("--dataloader-num-workers",
                             type=int,
                             required=True,
                             help="Number of workers for dataloader")
-        
-        # RL dataset arguments (optional, defaults to data_path)
-        parser.add_argument("--rl-dataset-path",
-                            type=str,
-                            default="",
-                            help="Path to RL prompt dataset directory (defaults to --data-path if not set)")
-        parser.add_argument("--rl-dataset-type",
-                            type=str,
-                            default="text",
-                            choices=["text", "geneval"],
-                            help="RL dataset type: 'text' for TextPromptDataset or 'geneval' for GenevalPromptDataset")
-        parser.add_argument("--rl-num-image-per-prompt",
-                            type=int,
-                            default=4,
-                            help="Number of times to repeat each prompt (k parameter for KRepeatSampler)")
         parser.add_argument("--num-height",
                             type=int,
                             required=True,
@@ -1591,9 +1284,6 @@ class TrainingArgs(FastVideoArgs):
                             default=TrainingArgs.context_noise,
                             help="Context noise level for cache updates")
 
-        # RL (Reinforcement Learning) arguments
-        RLArgs.add_cli_args(parser)
-
         return parser
 
 

fastvideo/hooks/hooks.py

@@ -1,122 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import functools
-from typing import Any
-from torch import nn
-
-
-class ForwardHook:
-    """
-    Base class for forward hooks.
-    Hooks are used in the way:
-        modified_args, modified_kwargs = hook.pre_forward(module, *args, **kwargs)
-        output = module.forward(*modified_args, **modified_kwargs)
-        modified_output = hook.post_forward(module, output)
-    """
-
-    @classmethod
-    def name(cls) -> str:
-        raise NotImplementedError
-
-    def on_attach(self, module: nn.Module):  # noqa: B027
-        """Called once when the hook is attached to the module."""
-        pass
-
-    def on_detach(self, module: nn.Module):  # noqa: B027
-        """
-        Called once when the hook is detached from the module.
-        Note: this function is not guaranteed to be called if the module is
-              deleted before the hook is detached.
-        """
-        pass
-
-    def pre_forward(self, module: nn.Module, *args,
-                    **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]:
-        """Called before the module's forward method is executed."""
-        return args, kwargs
-
-    def post_forward(self, module: nn.Module, output: Any) -> Any:
-        """Called after the module's forward method is executed."""
-        return output
-
-
-class ModuleHookManager:
-    module_hook_attribute = "_hook_manager"
-
-    def __init__(self, module: nn.Module):
-        self.module = module
-        self.forward_hooks: dict[str, ForwardHook] = {}
-        self.original_forward = module.forward
-
-    @classmethod
-    def get_from(cls, module: nn.Module) -> "ModuleHookManager | None":
-        if hasattr(module, cls.module_hook_attribute):
-            return getattr(module, cls.module_hook_attribute)
-        return None
-
-    @classmethod
-    def get_from_or_default(cls, module: nn.Module) -> "ModuleHookManager":
-        if not hasattr(module, cls.module_hook_attribute):
-            setattr(module, cls.module_hook_attribute, cls(module))
-
-            def forward_hook_wrapper(mod: nn.Module, *args, **kwargs):
-                manager: ModuleHookManager = getattr(mod,
-                                                     cls.module_hook_attribute)
-                for hook in manager.forward_hooks.values():
-                    args, kwargs = hook.pre_forward(mod, *args, **kwargs)
-                output = manager.original_forward(*args, **kwargs)
-                for hook in reversed(manager.forward_hooks.values()):
-                    output = hook.post_forward(mod, output)
-                return output
-
-            module.forward = functools.partial(forward_hook_wrapper, module)
-
-        return getattr(module, cls.module_hook_attribute)
-
-    @staticmethod
-    def remove_from_manager(module: nn.Module) -> None:
-        if hasattr(module, ModuleHookManager.module_hook_attribute):
-            manager: ModuleHookManager = getattr(
-                module, ModuleHookManager.module_hook_attribute)
-            module.forward = manager.original_forward
-            delattr(module, ModuleHookManager.module_hook_attribute)
-
-    def _check_manager_attached(self) -> None:
-        if not hasattr(self.module, self.module_hook_attribute):
-            raise ValueError("ModuleHookManager is not attached to the module.")
-        if getattr(self.module, self.module_hook_attribute) is not self:
-            raise ValueError(
-                "ModuleHookManager attached to the module is different.")
-
-    def append_forward_hook(self, hook: ForwardHook):
-        self._check_manager_attached()
-        if hook.name() in self.forward_hooks:
-            raise ValueError(
-                f"Hook with name {hook.name()} is already registered.")
-        # after python 3.7, dicts maintain insertion order
-        self.forward_hooks[hook.name()] = hook
-        hook.on_attach(self.module)
-
-    def replace_forward_hook(self,
-                             hook_name: str,
-                             new_hook: ForwardHook,
-                             run_on_attach: bool = True):
-        self._check_manager_attached()
-        if hook_name not in self.forward_hooks:
-            raise ValueError(f"No hook with name {hook_name} found.")
-        old_hook = self.forward_hooks[hook_name]
-        if run_on_attach:
-            old_hook.on_detach(self.module)
-        self.forward_hooks[hook_name] = new_hook
-        new_hook.on_attach(self.module)
-
-    def remove_forward_hook(self, hook_name: str, run_detach: bool = True):
-        self._check_manager_attached()
-        if hook_name not in self.forward_hooks:
-            raise ValueError(f"No hook with name {hook_name} found.")
-        if run_detach:
-            self.forward_hooks[hook_name].on_detach(self.module)
-        del self.forward_hooks[hook_name]
-
-    def get_forward_hook(self, hook_name: str) -> ForwardHook | None:
-        return self.forward_hooks.get(hook_name, None)

fastvideo/hooks/layerwise_offload.py

@@ -1,164 +0,0 @@
-from contextlib import contextmanager
-from typing import Any
-import torch
-from torch import nn
-from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def _tensor_placeholder(tensor: torch.Tensor,
-                        device: torch.device) -> torch.Tensor:
-    """Create a rank-preserving empty placeholder on the specified device."""
-    shape = (0, ) if tensor.ndim <= 0 else (0, ) * tensor.ndim
-    return torch.empty(shape, device=device, dtype=tensor.dtype)
-
-
-class LayerwiseOffloadState:
-
-    def __init__(
-        self,
-        async_copy_stream: torch.cuda.Stream,
-        device: torch.device,
-        next_state: "LayerwiseOffloadState | None" = None,
-    ) -> None:
-        self.async_copy_stream = async_copy_stream
-        self.next_state = next_state
-        self.gpu_named_parameters: dict[str, torch.Tensor] = {}
-        self.cpu_named_parameters: dict[str, torch.Tensor] = {}
-        self.module_ref: nn.Module = None  # type: ignore
-        self.device: torch.device = device
-
-    def _will_offload(self, name: str) -> bool:
-        return True
-
-    @torch.compiler.disable
-    def on_init(self, module: nn.Module):
-        self.module_ref = module
-        for name, param in self.module_ref.named_parameters():
-            if self._will_offload(name):
-                self.cpu_named_parameters[name] = (
-                    param.data.detach().to("cpu").pin_memory())
-                param.data = _tensor_placeholder(param.data, self.device)
-
-    @torch.compiler.disable
-    def wait_and_replace_params(self):
-        torch.cuda.current_stream().wait_stream(self.async_copy_stream)
-        # now gpu_named_parameters are ready
-        for name, param in self.module_ref.named_parameters():
-            if not self._will_offload(name):
-                continue
-            if name not in self.gpu_named_parameters:
-                # first load with blocking load
-                self.gpu_named_parameters[name] = self.cpu_named_parameters[
-                    name].to(self.device)
-            param.data = self.gpu_named_parameters[name]
-
-    @torch.compiler.disable
-    def prefetch_params(self):
-        compute_stream = torch.cuda.current_stream()
-        with torch.cuda.stream(self.async_copy_stream):
-            for name, param in self.module_ref.named_parameters():
-                if not self._will_offload(name):
-                    continue
-                assert name not in self.gpu_named_parameters
-                gpu_param = self.cpu_named_parameters[name].to(
-                    self.device, non_blocking=True)
-                gpu_param.record_stream(
-                    compute_stream
-                )  # ensure tensor will not be freed until forward is completed
-                self.gpu_named_parameters[name] = gpu_param
-
-    @torch.compiler.disable
-    def release_gpu_params(self):
-        for name, param in self.module_ref.named_parameters():
-            if self._will_offload(name):
-                param.data = _tensor_placeholder(param.data, self.device)
-                del self.gpu_named_parameters[name]
-        assert len(self.gpu_named_parameters) == 0
-
-
-class LayerwiseOffloadHook(ForwardHook):
-    """A hook that enables layerwise CPU offloading during forward pass."""
-
-    def __init__(self, state: LayerwiseOffloadState) -> None:
-        self.state = state
-
-    def on_attach(self, module: nn.Module):
-        self.state.on_init(module)  # pyright: ignore
-
-    def on_detach(self, module: nn.Module):
-        named_parameters = dict(module.named_parameters())
-        for name, cpu_tensor in self.state.cpu_named_parameters.items():
-            if name not in self.state.gpu_named_parameters:
-                if name in named_parameters:
-                    named_parameters[name].data = cpu_tensor.to(
-                        device=self.state.device)
-                else:
-                    logger.warning(
-                        "Parameter {} not found in module during detachment.",
-                        name,
-                    )
-
-    @classmethod
-    def name(cls) -> str:
-        return "LayerwiseOffloadHook"
-
-    def pre_forward(self, module: nn.Module, *args, **kwargs):
-        self.state.wait_and_replace_params()  # pyright: ignore
-        if self.state.next_state is not None:
-            self.state.next_state.prefetch_params()  # pyright: ignore
-        return args, kwargs
-
-    def post_forward(self, module: torch.nn.Module, output: Any):
-        self.state.release_gpu_params()  # pyright: ignore
-        return output
-
-    @contextmanager
-    def mutate_params_scope(self):
-        try:
-            # load params to GPU and keep them there
-            self.state.wait_and_replace_params()  # pyright: ignore
-            yield
-        finally:
-            # instead of releasing, we should overwrite the original params since they have been modified
-            self.state.cpu_named_parameters.clear()
-            self.state.gpu_named_parameters.clear()
-            self.state.on_init(self.state.module_ref)  # pyright: ignore
-
-
-def enable_layerwise_offload(model: nn.Module, is_replace: bool = False):
-    if torch.cuda.is_available():
-        device = torch.device("cuda", torch.cuda.current_device())
-    else:
-        logger.warning(
-            "CUDA is not available. Layerwise offloading is disabled.")
-        return
-    state_list = []
-    async_stream = torch.cuda.Stream()
-    for name, submodule in model.named_children():
-        if isinstance(submodule, nn.ModuleList):
-            for idx, module_entry in enumerate(submodule):
-                state = LayerwiseOffloadState(async_copy_stream=async_stream,
-                                              device=device)
-                state_list.append(state)
-                hook_mgr = ModuleHookManager.get_from_or_default(module_entry)
-                hook = LayerwiseOffloadHook(state)
-                if is_replace:
-                    existing_hook = hook_mgr.forward_hooks.get(hook.name())
-                    if existing_hook is not None:
-                        hook_mgr.replace_forward_hook(hook.name(), hook)
-                    else:
-                        raise AssertionError(
-                            f"Expect hook exists in {name} for replacement.")
-                else:
-                    hook_mgr.append_forward_hook(hook)
-            break
-    if len(state_list) == 0:
-        raise ValueError(
-            "No nn.ModuleList found in the model for layerwise offloading.")
-
-    # circular linking of states
-    for i in range(len(state_list)):
-        state_list[i].next_state = state_list[(i + 1) % len(state_list)]

fastvideo/layers/rotary_embedding.py

@@ -194,14 +194,14 @@ class RotaryEmbedding(CustomOp):
         cos, sin = cos_sin.chunk(2, dim=-1)
 
         query_shape = query.shape
-        query = query.view(num_tokens, -1, self.head_size)
+        query = query.reshape(num_tokens, -1, self.head_size)
         query_rot = query[..., :self.rotary_dim]
         query_pass = query[..., self.rotary_dim:]
         query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
         query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
 
         key_shape = key.shape
-        key = key.view(num_tokens, -1, self.head_size)
+        key = key.reshape(num_tokens, -1, self.head_size)
         key_rot = key[..., :self.rotary_dim]
         key_pass = key[..., self.rotary_dim:]
         key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)

fastvideo/models/dits/hunyuanvideo15.py

@@ -690,12 +690,15 @@ class SingleTokenRefiner(nn.Module):
         timestep_aware_representations = self.t_embedder(t)
 
         # Get context-aware representations
-        original_dtype = x.dtype
+        target_dtype = self.c_embedder.fc_in.weight.dtype
+        if x.dtype != target_dtype:
+            x = x.to(dtype=target_dtype)
         if mask is None:
             context_aware_representations = x.mean(dim=1)
         else:
-            mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
-            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+            mask_float = mask.to(dtype=target_dtype).unsqueeze(-1)  # [B, L, 1]
+            context_aware_representations = (
+                x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
 
         context_aware_representations = self.c_embedder(
             context_aware_representations)
@@ -850,4 +853,4 @@ class FinalLayer(nn.Module):
         scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
         x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         x, _ = self.linear(x)
-        return x
+        return x

fastvideo/models/dits/wanvideo.py

@@ -734,8 +734,18 @@ class WanTransformer3DModel(CachableDiT):
                         block, hidden_states, encoder_hidden_states,
                         timestep_proj, freqs_cis, attention_mask)
             else:
-                for block in self.blocks:
-                    hidden_states = block(hidden_states, encoder_hidden_states,
+                offload_mgr = getattr(self, "_layerwise_offload_manager", None)
+                use_offload = offload_mgr is not None and getattr(offload_mgr, "enabled", False)
+
+                for i, block in enumerate(self.blocks):
+                    scope = offload_mgr.layer_scope(
+                        prefetch_layer_idx=i + 1 if i + 1 < len(self.blocks) else None,
+                        release_layer_idx=i,
+                        non_blocking=True,
+                    ) if use_offload else nullcontext()
+                    
+                    with scope:
+                        hidden_states = block(hidden_states, encoder_hidden_states,
                                               timestep_proj, freqs_cis, attention_mask)
             # if teacache is enabled, we need to cache the original hidden states
 

fastvideo/models/loader/component_loader.py

@@ -35,7 +35,7 @@ from fastvideo.models.loader.weight_utils import (
 )
 from fastvideo.models.registry import ModelRegistry
 from fastvideo.utils import PRECISION_TO_TYPE, is_pin_memory_available
-from fastvideo.hooks.layerwise_offload import enable_layerwise_offload
+from fastvideo.models.layerwise_offload import LayerwiseOffloadManager
 
 logger = init_logger(__name__)
 
@@ -280,7 +280,7 @@ class TextEncoderLoader(ComponentLoader):
         target_device: torch.device,
         fastvideo_args: FastVideoArgs,
         dtype: str = "fp16",
-        use_text_encoder_override: bool = False,  # prevent subclasses from misusing
+        use_text_encoder_override: bool = False, # prevent subclasses from misusing
     ):
         use_cpu_offload = (
             fastvideo_args.text_encoder_cpu_offload
@@ -297,10 +297,7 @@ class TextEncoderLoader(ComponentLoader):
             )
 
         # Set quantization config if specified
-        if (
-            use_text_encoder_override
-            and fastvideo_args.override_text_encoder_quant is not None
-        ):
+        if use_text_encoder_override and fastvideo_args.override_text_encoder_quant is not None:
             if fastvideo_args.override_text_encoder_safetensors is None:
                 raise ValueError(
                     "override_text_encoder_quant is set but override_text_encoder_safetensors is None"
@@ -317,10 +314,7 @@ class TextEncoderLoader(ComponentLoader):
                 model: TextEncoder = model_cls(model_config)  # type: ignore
 
             weights_to_load = {name for name, _ in model.named_parameters()}
-            if (
-                use_text_encoder_override
-                and fastvideo_args.override_text_encoder_safetensors is not None
-            ):
+            if use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None:
                 loaded_weights: set[str] = model.load_weights(
                     safetensors_weights_iterator(
                         [fastvideo_args.override_text_encoder_safetensors],
@@ -348,6 +342,7 @@ class TextEncoderLoader(ComponentLoader):
 
             if use_cpu_offload:
                 pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available()
+                pin_cpu_memory = False
                 # Disable FSDP for MPS as it's not compatible
                 if current_platform.is_mps():
                     logger.info(
@@ -678,8 +673,35 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
-        if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload:
-            enable_layerwise_offload(model)
+        if fastvideo_args.dit_layerwise_offload and hasattr(model, "blocks"):
+            # Check if this is a Wan model (only Wan models support layerwise offload)
+            is_wan_model = "Wan" in cls_name
+            if not is_wan_model:
+                logger.warning(
+                    "Layerwise offload is currently only supported for Wan models. "
+                    "Model class '%s' does not support layerwise offload. "
+                    "Disabling layerwise offload for this model.",
+                    cls_name
+                )
+            else:
+                try:
+                    num_layers = len(getattr(model, "blocks"))
+                except TypeError:
+                    num_layers = None
+                if isinstance(num_layers, int) and num_layers > 0:
+                    # Ensure model is on the correct device (CUDA) before initializing manager
+                    # This ensures non-managed parameters (embeddings, final norms) are on GPU
+                    model = model.to(get_local_torch_device())
+                    mgr = LayerwiseOffloadManager(
+                        model,
+                        module_list_attr="blocks",
+                        num_layers=num_layers,
+                        enabled=True,
+                        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
+                        auto_initialize=True,
+                    )
+                    setattr(model, "_layerwise_offload_manager", mgr)
+
         return model
 
 

fastvideo/models/loader/fsdp_load.py

@@ -108,6 +108,7 @@ def maybe_load_fsdp_model(
 
     if use_fsdp:
         pin_cpu_memory = pin_cpu_memory and is_pin_memory_available()
+        pin_cpu_memory = False
         world_size = hsdp_replicate_dim * hsdp_shard_dim
         if not training_mode and not fsdp_inference:
             hsdp_replicate_dim = world_size
@@ -153,7 +154,6 @@ def maybe_load_fsdp_model(
                 f"Unexpected param or buffer {n} on meta device.")
         # Avoid unintended computation graph accumulation during inference
         if isinstance(p, torch.nn.Parameter):
-            
             p.requires_grad = False
 
     compile_in_loader = enable_torch_compile and training_mode

fastvideo/models/vaes/hunyuan15vae.py

@@ -114,11 +114,14 @@ class HunyuanVideo15AttnBlock(nn.Module):
         """
         seq_len = n_frame * n_hw
         mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
+        
         for i in range(seq_len):
             i_frame = i // n_hw
             mask[i, : (i_frame + 1) * n_hw] = 0
         if batch_size is not None:
-            mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+            # mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+            mask = mask.unsqueeze(0).unsqueeze(1).expand(batch_size, 1, -1, -1)
+
         return mask
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -131,6 +134,7 @@ class HunyuanVideo15AttnBlock(nn.Module):
         value = self.to_v(x)
 
         batch_size, channels, frames, height, width = query.shape
+        
 
         query = query.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
         key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()

fastvideo/models/vaes/hunyuanvae.py

@@ -92,6 +92,11 @@ class HunyuanVAEAttention(nn.Module):
         value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
 
         # Perform scaled dot-product attention
+        # If attention_mask is provided, it needs to be expanded for multi-head attention
+        if attention_mask is not None:
+            # Expand mask from [batch_size, seq_len, seq_len] to [batch_size, num_heads, seq_len, seq_len]
+            attention_mask = attention_mask.unsqueeze(1).expand(-1, self.heads, -1, -1)
+        
         hidden_states = F.scaled_dot_product_attention(query,
                                                        key,
                                                        value,

fastvideo/pipelines/composed_pipeline_base.py

@@ -201,6 +201,8 @@ class ComposedPipelineBase(ABC):
             # fwd, bwd, and other operations' precision.
             assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
 
+        logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
+
         pipe = cls(model_path,
                    fastvideo_args,
                    required_config_modules=required_config_modules,

fastvideo/pipelines/lora_pipeline.py

@@ -1,9 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from collections import defaultdict
 from collections.abc import Hashable
-from contextlib import nullcontext
 from typing import Any
-from collections.abc import Generator
 
 import torch
 import torch.distributed as dist
@@ -14,13 +12,8 @@ from torch.distributed.tensor import DTensor
 
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.hooks.hooks import ModuleHookManager
-from fastvideo.hooks.layerwise_offload import LayerwiseOffloadHook
-from fastvideo.layers.lora.linear import (
-    BaseLayerWithLoRA,
-    get_lora_layer,
-    replace_submodule,
-)
+from fastvideo.layers.lora.linear import (BaseLayerWithLoRA, get_lora_layer,
+                                          replace_submodule)
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.utils import get_param_names_mapping
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -29,89 +22,17 @@ from fastvideo.utils import maybe_download_lora
 logger = init_logger(__name__)
 
 
-def _get_hook_ctx(module: nn.Module | None):
-    if module is None:
-        return nullcontext()
-    hook_mgr = ModuleHookManager.get_from(module)
-    if hook_mgr is not None:
-        offload_hook = hook_mgr.forward_hooks.get(LayerwiseOffloadHook.name())
-        if offload_hook is not None:
-            return offload_hook.mutate_params_scope()  # type: ignore
-    return nullcontext()
-
-
-def _named_module_by_prefix(
-    module: nn.Module, prefixes: list[str]
-) -> list[tuple[str | None, list[tuple[str, nn.Module]]]]:
-    none_list: list[tuple[str, nn.Module]] = []
-    prefix_list: list[tuple[str, list[tuple[str, nn.Module]]]] = [
-        (prefix, []) for prefix in prefixes
-    ]
-    for name, submodule in module.named_modules():
-        for cur_prefix, cur_list in prefix_list:
-            # we should exclude e.g. block.1 and block.12.attn
-            if name.startswith(cur_prefix + "."):
-                cur_list.append((name, submodule))
-                break
-        else:
-            none_list.append((name, submodule))
-    return prefix_list + [(None, none_list)]  # type: ignore
-
-
-class LoRAModelLayers:
-
-    def __init__(self, block_list: list[tuple[str, nn.Module]]) -> None:
-        # block_name -> {layer_name -> layer}
-        self.block_to_lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {}
-        # layer_name -> block_name
-        self.lora_layers_to_block: dict[str, str | None] = {}
-        self.other_lora_layers: dict[str, BaseLayerWithLoRA] = {}
-        self.block_mapping = dict(block_list)
-
-    def add_lora_layer(self, block_name: str | None, layer_name: str,
-                       layer: BaseLayerWithLoRA):
-        if block_name is None:
-            self.other_lora_layers[layer_name] = layer
-            self.lora_layers_to_block[layer_name] = None
-        else:
-            if block_name not in self.block_to_lora_layers:
-                self.block_to_lora_layers[block_name] = {}
-            self.block_to_lora_layers[block_name][layer_name] = layer
-            self.lora_layers_to_block[layer_name] = block_name
-
-    def all_lora_layers(
-        self, ) -> Generator[tuple[str, BaseLayerWithLoRA], Any, None]:
-        for block_layers in self.block_to_lora_layers.values():
-            for name, layer in block_layers.items():
-                yield name, layer
-        for name, layer in self.other_lora_layers.items():
-            yield name, layer
-
-    def lora_layers_by_block(
-        self,
-    ) -> Generator[
-            tuple[nn.Module | None, dict[str, BaseLayerWithLoRA]],
-            Any,
-            None,
-    ]:
-        for block_name, layers in self.block_to_lora_layers.items():
-            yield self.block_mapping[block_name], layers
-        yield None, self.other_lora_layers
-
-
 class LoRAPipeline(ComposedPipelineBase):
     """
     Pipeline that supports injecting LoRA adapters into the diffusion transformer.
     TODO: support training.
     """
-
     lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(
         dict
     )  # state dicts of loaded lora adapters (includes lora_A, lora_B, and lora_alpha)
     cur_adapter_name: str = ""
     cur_adapter_path: str = ""
-    # model_name -> layers
-    lora_layers: dict[str, LoRAModelLayers] = {}
+    lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {}
     fastvideo_args: FastVideoArgs | TrainingArgs
     exclude_lora_layers: dict[str, list[str]] = {}
     device: torch.device = get_local_torch_device()
@@ -127,10 +48,10 @@ class LoRAPipeline(ComposedPipelineBase):
         self.device = get_local_torch_device()
         # build list of trainable transformers
         for transformer_name in self.trainable_transformer_names:
-            if (transformer_name in self.modules
-                    and self.modules[transformer_name] is not None):
-                self.trainable_transformer_modules[transformer_name] = (
-                    self.modules[transformer_name])
+            if transformer_name in self.modules and self.modules[
+                    transformer_name] is not None:
+                self.trainable_transformer_modules[
+                    transformer_name] = self.modules[transformer_name]
             # check for transformer_2 in case of Wan2.2 MoE or fake_score_transformer_2
             if transformer_name.endswith("_2"):
                 raise ValueError(
@@ -138,23 +59,19 @@ class LoRAPipeline(ComposedPipelineBase):
                 )
 
             secondary_transformer_name = transformer_name + "_2"
-            if (secondary_transformer_name in self.modules
-                    and self.modules[secondary_transformer_name] is not None):
+            if secondary_transformer_name in self.modules and self.modules[
+                    secondary_transformer_name] is not None:
                 self.trainable_transformer_modules[
                     secondary_transformer_name] = self.modules[
                         secondary_transformer_name]
 
-        logger.info(
-            "trainable_transformer_modules: %s",
-            self.trainable_transformer_modules.keys(),
-        )
+        logger.info("trainable_transformer_modules: %s",
+                    self.trainable_transformer_modules.keys())
 
-        for (
-                transformer_name,
-                transformer_module,
-        ) in self.trainable_transformer_modules.items():
-            self.exclude_lora_layers[transformer_name] = (
-                transformer_module.config.arch_config.exclude_lora_layers)
+        for transformer_name, transformer_module in self.trainable_transformer_modules.items(
+        ):
+            self.exclude_lora_layers[
+                transformer_name] = transformer_module.config.arch_config.exclude_lora_layers
         self.lora_target_modules = self.fastvideo_args.lora_target_modules
         self.lora_path = self.fastvideo_args.lora_path
         self.lora_nickname = self.fastvideo_args.lora_nickname
@@ -166,33 +83,20 @@ class LoRAPipeline(ComposedPipelineBase):
                 self.fastvideo_args.lora_alpha = self.fastvideo_args.lora_rank
             self.lora_rank = self.fastvideo_args.lora_rank  # type: ignore
             self.lora_alpha = self.fastvideo_args.lora_alpha  # type: ignore
-            logger.info(
-                "Using LoRA training with rank %d and alpha %d",
-                self.lora_rank,
-                self.lora_alpha,
-            )
+            logger.info("Using LoRA training with rank %d and alpha %d",
+                        self.lora_rank, self.lora_alpha)
             if self.lora_target_modules is None:
                 self.lora_target_modules = [
-                    "q_proj",
-                    "k_proj",
-                    "v_proj",
-                    "o_proj",
-                    "to_q",
-                    "to_k",
-                    "to_v",
-                    "to_out",
-                    "to_qkv",
-                    "to_gate_compress",
+                    "q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k",
+                    "to_v", "to_out", "to_qkv", "to_gate_compress"
                 ]
                 logger.info(
                     "Using default lora_target_modules for all transformers: %s",
-                    self.lora_target_modules,
-                )
+                    self.lora_target_modules)
             else:
                 logger.warning(
                     "Using custom lora_target_modules for all transformers, which may not be intended: %s",
-                    self.lora_target_modules,
-                )
+                    self.lora_target_modules)
 
             self.convert_to_lora_layers()
         # Inference
@@ -200,8 +104,7 @@ class LoRAPipeline(ComposedPipelineBase):
             self.convert_to_lora_layers()
             self.set_lora_adapter(
                 self.lora_nickname,  # type: ignore
-                self.lora_path,
-            )  # type: ignore
+                self.lora_path)  # type: ignore
 
     def is_target_layer(self, module_name: str) -> bool:
         if self.lora_target_modules is None:
@@ -211,9 +114,9 @@ class LoRAPipeline(ComposedPipelineBase):
 
     def set_trainable(self) -> None:
 
-        def set_lora_grads(lora_layers: LoRAModelLayers,
+        def set_lora_grads(lora_layers: dict[str, BaseLayerWithLoRA],
                            device_mesh: DeviceMesh):
-            for name, layer in lora_layers.all_lora_layers():
+            for name, layer in lora_layers.items():
                 layer.lora_A.requires_grad_(True)
                 layer.lora_B.requires_grad_(True)
                 layer.base_layer.requires_grad_(False)
@@ -228,15 +131,10 @@ class LoRAPipeline(ComposedPipelineBase):
             super().set_trainable()
             return
 
-        device_mesh = init_device_mesh(
-            "cuda",
-            (dist.get_world_size(), 1),
-            mesh_dim_names=["fake", "replicate"],
-        )
-        for (
-                transformer_name,
-                transformer_module,
-        ) in self.trainable_transformer_modules.items():
+        device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1),
+                                       mesh_dim_names=["fake", "replicate"])
+        for transformer_name, transformer_module in self.trainable_transformer_modules.items(
+        ):
             transformer_module.train()
             transformer_module.requires_grad_(False)
             if transformer_name in self.lora_layers:
@@ -253,71 +151,32 @@ class LoRAPipeline(ComposedPipelineBase):
         if self.lora_initialized:
             return
         self.lora_initialized = True
-        for (
-                transformer_name,
-                transformer_module,
-        ) in self.trainable_transformer_modules.items():
+        for transformer_name, transformer_module in self.trainable_transformer_modules.items(
+        ):
             converted_count = 0
-            # init bookkeeping structures
             if transformer_name not in self.lora_layers:
-                # get block list
-                block_list = []
-                for name, submodule in transformer_module.named_children():
-                    if isinstance(submodule, nn.ModuleList):
-                        block_list = [(f"{name}.{i}", m)
-                                      for i, m in enumerate(submodule)]
-                        break
-                self.lora_layers[transformer_name] = LoRAModelLayers(block_list)
+                self.lora_layers[transformer_name] = {}
             logger.info("Converting %s to LoRA Transformer", transformer_name)
-            # scan every module and convert to LoRA layer if applicable
+            for name, layer in transformer_module.named_modules():
+                if not self.is_target_layer(name):
+                    continue
 
-            for block_name, block_modules in _named_module_by_prefix(
-                    transformer_module,
-                    list(self.lora_layers[transformer_name].block_mapping),
-            ):
-                if block_name is not None and (
-                        not self.fastvideo_args.training_mode
-                        and self.fastvideo_args.dit_layerwise_offload):
-                    scope_ctx = _get_hook_ctx(
-                        self.lora_layers[transformer_name].
-                        block_mapping[block_name])
-                else:
-                    scope_ctx = nullcontext()
-                with scope_ctx:
-                    for name, layer in block_modules:
-                        if not self.is_target_layer(name):
-                            continue
+                excluded = False
+                for exclude_layer in self.exclude_lora_layers[transformer_name]:
+                    if exclude_layer in name:
+                        excluded = True
+                        break
+                if excluded:
+                    continue
 
-                        excluded = False
-                        for exclude_layer in self.exclude_lora_layers[
-                                transformer_name]:
-                            if exclude_layer in name:
-                                excluded = True
-                                break
-                        if excluded:
-                            continue
-
-                        layer = get_lora_layer(
-                            layer,
-                            lora_rank=self.lora_rank,
-                            lora_alpha=self.lora_alpha,
-                            training_mode=self.training_mode,
-                        )
-                        if layer is not None:
-                            block_name_split = name.split(".", 2)
-                            if len(block_name_split) > 2:
-                                block_name = (block_name_split[0] + "." +
-                                              block_name_split[1])
-                            else:
-                                block_name = None
-                            if (block_name
-                                    not in self.lora_layers[transformer_name].
-                                    block_mapping):
-                                block_name = None
-                            self.lora_layers[transformer_name].add_lora_layer(
-                                block_name, name, layer)
-                            replace_submodule(transformer_module, name, layer)
-                            converted_count += 1
+                layer = get_lora_layer(layer,
+                                       lora_rank=self.lora_rank,
+                                       lora_alpha=self.lora_alpha,
+                                       training_mode=self.training_mode)
+                if layer is not None:
+                    self.lora_layers[transformer_name][name] = layer
+                    replace_submodule(transformer_module, name, layer)
+                    converted_count += 1
             logger.info("Converted %d layers to LoRA layers", converted_count)
 
     def set_lora_adapter(self,
@@ -350,7 +209,7 @@ class LoRAPipeline(ComposedPipelineBase):
 
             # Extract alpha values and weights in a single pass
             to_merge_params: defaultdict[Hashable,
-                                         dict[Any, Any]] = (defaultdict(dict))
+                                         dict[Any, Any]] = defaultdict(dict)
             for name, weight in lora_state_dict.items():
                 # Extract weights (lora_A, lora_B, and lora_alpha)
                 name = name.replace("diffusion_model.", "")
@@ -364,14 +223,13 @@ class LoRAPipeline(ComposedPipelineBase):
                     target_name, _, _ = param_names_mapping_fn(layer_name)
                     # Store alpha alongside weights with same target_name base
                     alpha_key = target_name + ".lora_alpha"
-                    self.lora_adapters[lora_nickname][alpha_key] = (
-                        weight.item()
-                        if weight.numel() == 1 else float(weight.mean()))
+                    self.lora_adapters[lora_nickname][alpha_key] = weight.item(
+                    ) if weight.numel() == 1 else float(weight.mean())
                     continue
 
                 name, _, _ = lora_param_names_mapping_fn(name)
-                target_name, merge_index, num_params_to_merge = (
-                    param_names_mapping_fn(name))
+                target_name, merge_index, num_params_to_merge = param_names_mapping_fn(
+                    name)
                 # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
                 # see param mapping in HunyuanVideoArchConfig
                 if merge_index is not None and "lora_B" in name:
@@ -403,84 +261,45 @@ class LoRAPipeline(ComposedPipelineBase):
 
         # Merge the new adapter
         adapted_count = 0
-        for (
-                transformer_name,
-                transformer_lora_layers,
-        ) in self.lora_layers.items():
-            for (
-                    module,
-                    layers,
-            ) in transformer_lora_layers.lora_layers_by_block():
-                with _get_hook_ctx(module):
-                    for name, layer in layers.items():
-                        lora_A_name = name + ".lora_A"
-                        lora_B_name = name + ".lora_B"
-                        lora_alpha_name = name + ".lora_alpha"
-                        if (lora_A_name in self.lora_adapters[lora_nickname]
-                                and lora_B_name
-                                in self.lora_adapters[lora_nickname]):
-                            # Get alpha value for this layer (defaults to None if not present)
-                            lora_A = self.lora_adapters[lora_nickname][
-                                lora_A_name]
-                            lora_B = self.lora_adapters[lora_nickname][
-                                lora_B_name]
-                            # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
-                            alpha = (self.lora_adapters[lora_nickname].get(
-                                lora_alpha_name) if adapter_updated else None)
-                            try:
-                                layer.set_lora_weights(
-                                    lora_A,
-                                    lora_B,
-                                    lora_alpha=alpha,
-                                    training_mode=self.fastvideo_args.
-                                    training_mode,
-                                    lora_path=lora_path,
-                                )
-                            except Exception as e:
-                                logger.error(
-                                    "Error setting LoRA weights for layer %s: %s",
-                                    name,
-                                    str(e),
-                                )
-                                raise e
-                            adapted_count += 1
-                        else:
-                            if rank == 0:
-                                logger.warning(
-                                    "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
-                                    lora_path,
-                                    name,
-                                )
-                            layer.disable_lora = True
-        logger.info(
-            "Rank %d: LoRA adapter %s applied to %d layers",
-            rank,
-            lora_path,
-            adapted_count,
-        )
+        for transformer_name, transformer_lora_layers in self.lora_layers.items(
+        ):
+            for name, layer in transformer_lora_layers.items():
+                lora_A_name = name + ".lora_A"
+                lora_B_name = name + ".lora_B"
+                lora_alpha_name = name + ".lora_alpha"
+                if lora_A_name in self.lora_adapters[lora_nickname]\
+                    and lora_B_name in self.lora_adapters[lora_nickname]:
+                    # Get alpha value for this layer (defaults to None if not present)
+                    lora_A = self.lora_adapters[lora_nickname][lora_A_name]
+                    lora_B = self.lora_adapters[lora_nickname][lora_B_name]
+                    # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
+                    alpha = self.lora_adapters[lora_nickname].get(
+                        lora_alpha_name) if adapter_updated else None
+
+                    layer.set_lora_weights(
+                        lora_A,
+                        lora_B,
+                        lora_alpha=alpha,
+                        training_mode=self.fastvideo_args.training_mode,
+                        lora_path=lora_path)
+                    adapted_count += 1
+                else:
+                    if rank == 0:
+                        logger.warning(
+                            "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
+                            lora_path, name)
+                    layer.disable_lora = True
+        logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank,
+                    lora_path, adapted_count)
 
     def merge_lora_weights(self) -> None:
-        for (
-                transformer_name,
-                transformer_lora_layers,
-        ) in self.lora_layers.items():
-            for (
-                    module,
-                    layers,
-            ) in transformer_lora_layers.lora_layers_by_block():
-                with _get_hook_ctx(module):
-                    for name, layer in layers.items():
-                        layer.merge_lora_weights()
+        for transformer_name, transformer_lora_layers in self.lora_layers.items(
+        ):
+            for name, layer in transformer_lora_layers.items():
+                layer.merge_lora_weights()
 
     def unmerge_lora_weights(self) -> None:
-        for (
-                transformer_name,
-                transformer_lora_layers,
-        ) in self.lora_layers.items():
-            for (
-                    module,
-                    layers,
-            ) in transformer_lora_layers.lora_layers_by_block():
-                with _get_hook_ctx(module):
-                    for name, layer in layers.items():
-                        layer.unmerge_lora_weights()
+        for transformer_name, transformer_lora_layers in self.lora_layers.items(
+        ):
+            for name, layer in transformer_lora_layers.items():
+                layer.unmerge_lora_weights()

fastvideo/pipelines/pipeline_batch_info.py

@@ -67,21 +67,6 @@ class ForwardBatch:
     execution, allowing methods to update specific components without needing
     to manage numerous individual parameters.
     """
-
-    @dataclass
-    class RLData:
-        """RL-specific data collection options and outputs."""
-        enabled: bool = False
-        collect_log_probs: bool = True
-        collect_kl: bool = False
-        kl_reward: float = 0.0
-        store_trajectory: bool = True
-        keep_trajectory_on_cpu: bool = False
-        log_probs: torch.Tensor | None = None
-        kl: torch.Tensor | None = None
-        trajectory_latents: torch.Tensor | None = None
-        trajectory_timesteps: torch.Tensor | None = None
-
     # TODO(will): double check that args are separate from fastvideo_args
     # properly. Also maybe think about providing an abstraction for pipeline
     # specific arguments.
@@ -212,9 +197,6 @@ class ForwardBatch:
     logging_info: PipelineLoggingInfo = field(
         default_factory=PipelineLoggingInfo)
 
-    # RL data collection
-    rl_data: "ForwardBatch.RLData" = field(default_factory=RLData)
-
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 
@@ -285,36 +267,6 @@ class TrainingBatch:
     latent_vis_dict: dict[str, Any] = field(default_factory=dict)
     fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
 
-    # RL/GRPO-specific attributes
-    reward_scores: torch.Tensor | None = None  # Computed rewards from reward models
-    log_probs: torch.Tensor | None = None  # Current policy log probabilities [B, num_steps] or [B]
-    old_log_probs: torch.Tensor | None = None  # Old policy log probs (for importance ratio) [B, num_steps] or [B]
-    advantages: torch.Tensor | None = None  # GAE advantages [B, num_steps] or [B]
-    returns: torch.Tensor | None = None  # TD returns (advantages + values) [B, num_steps] or [B]
-    values: torch.Tensor | None = None  # Value function predictions [B]
-    old_values: torch.Tensor | None = None  # Old value predictions (for clipping) [B]
-
-    # GRPO sampling-specific attributes
-    kl: torch.Tensor | None = None  # KL divergences from sampling [B, num_steps] (if kl_reward > 0)
-    prompt_ids: torch.Tensor | None = None  # Prompt token IDs for stat tracking [B, seq_len]
-    prompt_embeds: torch.Tensor | None = None  # Prompt embeddings used in sampling [B, seq_len, hidden_dim]
-    negative_prompt_embeds: torch.Tensor | None = None  # Negative prompt embeddings for CFG [B, seq_len, hidden_dim]
-
-    # RL loss components
-    policy_loss: float = 0.0  # GRPO/PPO policy loss
-    value_loss: float = 0.0  # Value function loss
-    kl_divergence: float = 0.0  # KL(new_policy || old_policy)
-    importance_ratio: float = 1.0  # exp(log_prob - old_log_prob)
-    clip_fraction: float = 0.0  # Fraction of ratios that were clipped
-
-    # RL metrics
-    advantage_mean: float = 0.0  # Mean advantage (should be ~0 after normalization)
-    advantage_std: float = 1.0  # Std of advantages
-    reward_mean: float = 0.0  # Mean reward across batch
-    reward_std: float = 0.0  # Std of rewards
-    value_mean: float = 0.0  # Mean value prediction
-    entropy: float = 0.0  # Policy entropy (for exploration)
-
 
 @dataclass
 class PreprocessBatch(ForwardBatch):

fastvideo/pipelines/preprocess/hunyuan/hunyuan_preprocess_pipelins.py

@@ -0,0 +1,59 @@
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.preprocess.preprocess_stages import (
+    TextTransformStage, VideoTransformStage)
+from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
+                                        TextEncodingStage)
+from fastvideo.pipelines.stages.image_encoding import ImageVAEEncodingStage
+        
+
+class PreprocessPipelineT2V(ComposedPipelineBase):
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "text_encoder_2", "tokenizer_2", "vae"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        assert fastvideo_args.preprocess_config is not None
+        
+        self.add_stage(
+            stage_name="text_transform_stage",
+            stage=TextTransformStage(
+                cfg_uncondition_drop_rate=fastvideo_args.preprocess_config.training_cfg_rate,
+                seed=fastvideo_args.preprocess_config.seed,
+            )
+        )
+        
+        text_encoders = [
+            self.get_module("text_encoder"),
+            self.get_module("text_encoder_2")
+        ]
+        tokenizers = [
+            self.get_module("tokenizer"),
+            self.get_module("tokenizer_2")
+        ]
+        
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=text_encoders,
+                tokenizers=tokenizers,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_transform_stage",
+            stage=VideoTransformStage(
+                train_fps=fastvideo_args.preprocess_config.train_fps,
+                num_frames=fastvideo_args.preprocess_config.num_frames,
+                max_height=fastvideo_args.preprocess_config.max_height,
+                max_width=fastvideo_args.preprocess_config.max_width,
+                do_temporal_sample=fastvideo_args.preprocess_config.do_temporal_sample,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_encoding_stage",
+            stage=EncodingStage(vae=self.get_module("vae"))
+        )
+
+EntryClass = [PreprocessPipelineT2V]

fastvideo/pipelines/preprocess/hunyuan15/hunyuan_preprocess_pipelins.py

@@ -0,0 +1,59 @@
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.preprocess.preprocess_stages import (
+    TextTransformStage, VideoTransformStage)
+from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
+                                        TextEncodingStage)
+from fastvideo.pipelines.stages.image_encoding import ImageVAEEncodingStage
+        
+
+class PreprocessPipelineT2V(ComposedPipelineBase):
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "text_encoder_2", "tokenizer_2", "vae"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        assert fastvideo_args.preprocess_config is not None
+        
+        self.add_stage(
+            stage_name="text_transform_stage",
+            stage=TextTransformStage(
+                cfg_uncondition_drop_rate=fastvideo_args.preprocess_config.training_cfg_rate,
+                seed=fastvideo_args.preprocess_config.seed,
+            )
+        )
+        
+        text_encoders = [
+            self.get_module("text_encoder"),
+            self.get_module("text_encoder_2")
+        ]
+        tokenizers = [
+            self.get_module("tokenizer"),
+            self.get_module("tokenizer_2")
+        ]
+        
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=text_encoders,
+                tokenizers=tokenizers,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_transform_stage",
+            stage=VideoTransformStage(
+                train_fps=fastvideo_args.preprocess_config.train_fps,
+                num_frames=fastvideo_args.preprocess_config.num_frames,
+                max_height=fastvideo_args.preprocess_config.max_height,
+                max_width=fastvideo_args.preprocess_config.max_width,
+                do_temporal_sample=fastvideo_args.preprocess_config.do_temporal_sample,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_encoding_stage",
+            stage=EncodingStage(vae=self.get_module("vae"))
+        )
+
+EntryClass = [PreprocessPipelineT2V]

fastvideo/pipelines/preprocess/v1_preprocessing_new.py

@@ -5,10 +5,6 @@ from fastvideo.logger import init_logger
 from fastvideo.utils import FlexibleArgumentParser
 from fastvideo.workflow.workflow_base import WorkflowBase
 
-import os
-
-os.environ["MASTER_PORT"] = "29513"
-
 logger = init_logger(__name__)
 
 

fastvideo/pipelines/stages/denoising.py

@@ -4,14 +4,11 @@ Denoising stage for diffusion pipelines.
 """
 
 import inspect
-import math
 import weakref
 from collections.abc import Iterable
-from contextlib import nullcontext
 from typing import Any
 
 import torch
-from diffusers.utils.torch_utils import randn_tensor
 from tqdm.auto import tqdm
 
 from fastvideo.attention import get_attn_backend
@@ -55,84 +52,6 @@ except ImportError:
 logger = init_logger(__name__)
 
 
-def sde_step_with_logprob(
-    scheduler,
-    model_output: torch.FloatTensor,
-    timestep: float | torch.FloatTensor,
-    sample: torch.FloatTensor,
-    prev_sample: torch.FloatTensor | None = None,
-    generator: torch.Generator | None = None,
-    deterministic: bool = False,
-    return_pixel_log_prob: bool = False,
-    return_dt_and_std_dev_t: bool = False
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, ...]:
-    """
-    Predict the sample from the previous timestep by reversing the SDE and
-    compute log probabilities for the transition.
-    """
-    if isinstance(timestep, torch.Tensor):
-        if timestep.ndim == 0:
-            timestep = timestep.unsqueeze(0)
-        step_indices = [
-            scheduler.index_for_timestep(t.item()) for t in timestep
-        ]
-    else:
-        step_indices = [scheduler.index_for_timestep(timestep)]
-
-    prev_step_indices = [step + 1 for step in step_indices]
-
-    sigmas = scheduler.sigmas.to(sample.device, sample.dtype)
-    sigma = sigmas[step_indices].view(-1, 1, 1, 1, 1)
-    sigma_prev = sigmas[prev_step_indices].view(-1, 1, 1, 1, 1)
-    sigma_max = sigmas[0].item()
-    sigma_min = sigmas[-1].item()
-
-    dt = sigma_prev - sigma
-
-    std_dev_t = sigma_min + (sigma_max - sigma_min) * sigma
-    prev_sample_mean = (sample * (1 + std_dev_t**2 / (2 * sigma) * dt) +
-                        model_output * (1 + std_dev_t**2 * (1 - sigma) /
-                                        (2 * sigma)) * dt)
-
-    if prev_sample is not None and generator is not None:
-        raise ValueError(
-            "Cannot pass both generator and prev_sample. Please make sure that either `generator` or"
-            " `prev_sample` stays `None`.")
-
-    if prev_sample is None:
-        variance_noise = randn_tensor(
-            model_output.shape,
-            generator=generator,
-            device=model_output.device,
-            dtype=model_output.dtype,
-        )
-        sqrt_dt = torch.sqrt(-1 * dt)
-        prev_sample = prev_sample_mean + std_dev_t * sqrt_dt * variance_noise
-    else:
-        sqrt_dt = torch.sqrt(-1 * dt)
-
-    if deterministic:
-        prev_sample = sample + dt * model_output
-        sqrt_dt = torch.sqrt(-1 * dt)
-
-    if return_pixel_log_prob:
-        raise NotImplementedError(
-            "Pixel-level log prob is not supported in this helper.")
-
-    std_dev_sqrt_dt = std_dev_t * sqrt_dt
-    log_prob = (
-        -((prev_sample.detach() - prev_sample_mean)**2) /
-        (2 *
-         (std_dev_sqrt_dt**2)) - torch.log(std_dev_sqrt_dt + 1e-8) - torch.log(
-             torch.sqrt(2 * torch.as_tensor(math.pi, device=sample.device))))
-
-    log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
-
-    if return_dt_and_std_dev_t:
-        return prev_sample, log_prob, prev_sample_mean, std_dev_t, sqrt_dt
-    return prev_sample, log_prob, prev_sample_mean, std_dev_t * sqrt_dt
-
-
 class DenoisingStage(PipelineStage):
     """
     Stage for running the denoising loop in diffusion pipelines.
@@ -284,11 +203,9 @@ class DenoisingStage(PipelineStage):
         else:
             boundary_timestep = None
         latent_model_input = latents.to(target_dtype)
-        rl_data = batch.rl_data if batch.rl_data and batch.rl_data.enabled else None
+        assert latent_model_input.shape[0] == 1, "only support batch size 1"
 
         if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
-            assert latent_model_input.shape[
-                0] == 1, "TI2V task only supports batch size 1"
             # TI2V directly replaces the first frame of the latent with
             # the image latent instead of appending along the channel dim
             assert batch.image_latent is None, "TI2V task should not have image latents"
@@ -326,12 +243,6 @@ class DenoisingStage(PipelineStage):
         # Initialize lists for ODE trajectory
         trajectory_timesteps: list[torch.Tensor] = []
         trajectory_latents: list[torch.Tensor] = []
-        rl_timesteps: list[torch.Tensor] = []
-        rl_latents: list[torch.Tensor] = []
-        rl_log_probs: list[torch.Tensor] = []
-        rl_kl: list[torch.Tensor] = []
-        if rl_data is not None and rl_data.store_trajectory:
-            rl_latents.append(latents)
 
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -418,24 +329,6 @@ class DenoisingStage(PipelineStage):
                     1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
                     is not None else None)
 
-                def run_transformer(model, encoder_hidden_states, cond_kwargs,
-                                    is_cfg_negative: bool):
-                    batch.is_cfg_negative = is_cfg_negative
-                    with set_forward_context(
-                            current_timestep=i,
-                            attn_metadata=attn_metadata,
-                            forward_batch=batch,
-                    ):
-                        return model(
-                            latent_model_input,
-                            encoder_hidden_states,
-                            t_expand,
-                            guidance=guidance_expand,
-                            **image_kwargs,
-                            **cond_kwargs,
-                            **action_kwargs,
-                        )
-
                 # Predict noise residual
                 with torch.autocast(device_type="cuda",
                                     dtype=target_dtype,
@@ -497,13 +390,40 @@ class DenoisingStage(PipelineStage):
                     # support torch dynamo compilation. They pass in
                     # attn_metadata, vllm_config, and num_tokens. We can pass in
                     # fastvideo_args or training_args, and attn_metadata.
-                    noise_pred = run_transformer(current_model, prompt_embeds,
-                                                 pos_cond_kwargs, False)
+                    batch.is_cfg_negative = False
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=attn_metadata,
+                            forward_batch=batch,
+                            # fastvideo_args=fastvideo_args
+                    ):
+                        # Run transformer
+                        noise_pred = current_model(
+                            latent_model_input,
+                            prompt_embeds,
+                            t_expand,
+                            guidance=guidance_expand,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                            **action_kwargs,
+                        )
 
                     if batch.do_classifier_free_guidance:
-                        noise_pred_uncond = run_transformer(
-                            current_model, neg_prompt_embeds, neg_cond_kwargs,
-                            True)
+                        batch.is_cfg_negative = True
+                        with set_forward_context(
+                                current_timestep=i,
+                                attn_metadata=attn_metadata,
+                                forward_batch=batch,
+                        ):
+                            noise_pred_uncond = current_model(
+                                latent_model_input,
+                                neg_prompt_embeds,
+                                t_expand,
+                                guidance=guidance_expand,
+                                **image_kwargs,
+                                **neg_cond_kwargs,
+                                **action_kwargs,
+                            )
 
                         noise_pred_text = noise_pred
                         noise_pred = noise_pred_uncond + current_guidance_scale * (
@@ -518,58 +438,11 @@ class DenoisingStage(PipelineStage):
                                 guidance_rescale=batch.guidance_rescale,
                             )
                     # Compute the previous noisy sample
-                    prev_latents = latents
                     latents = self.scheduler.step(noise_pred,
                                                   t,
                                                   latents,
                                                   **extra_step_kwargs,
                                                   return_dict=False)[0]
-                    if rl_data is not None:
-                        if rl_data.collect_log_probs:
-                            _, log_prob, prev_latents_mean, std_dev_t, _ = sde_step_with_logprob(
-                                self.scheduler,
-                                noise_pred.float(),
-                                t,
-                                prev_latents.float(),
-                                prev_sample=latents.float(),
-                                deterministic=False,
-                                return_dt_and_std_dev_t=True,
-                            )
-                            rl_log_probs.append(log_prob)
-
-                        if rl_data.collect_kl and rl_data.kl_reward > 0:
-                            adapter_ctx = nullcontext()
-                            if hasattr(current_model, "disable_adapter"):
-                                adapter_ctx = current_model.disable_adapter()
-                            with adapter_ctx:
-                                noise_pred_ref = run_transformer(
-                                    current_model, prompt_embeds,
-                                    pos_cond_kwargs, False)
-                                if batch.do_classifier_free_guidance:
-                                    noise_pred_uncond_ref = run_transformer(
-                                        current_model, neg_prompt_embeds,
-                                        neg_cond_kwargs, True)
-                                    noise_pred_text_ref = noise_pred_ref
-                                    noise_pred_ref = noise_pred_uncond_ref + current_guidance_scale * (
-                                        noise_pred_text_ref -
-                                        noise_pred_uncond_ref)
-                            _, _, prev_latents_mean_ref, std_dev_t_ref, _ = sde_step_with_logprob(
-                                self.scheduler,
-                                noise_pred_ref.float(),
-                                t,
-                                prev_latents.float(),
-                                prev_sample=latents.float(),
-                                deterministic=False,
-                                return_dt_and_std_dev_t=True,
-                            )
-                            if not torch.allclose(std_dev_t, std_dev_t_ref):
-                                logger.warning(
-                                    "std_dev_t mismatch in RL KL computation at step %s",
-                                    i)
-                            kl = (prev_latents_mean -
-                                  prev_latents_mean_ref)**2 / (2 * std_dev_t**2)
-                            kl = kl.mean(dim=tuple(range(1, kl.ndim)))
-                            rl_kl.append(kl)
                     if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
                         latents = latents.squeeze(0)
                         latents = (1. - mask2[0]) * z + mask2[0] * latents
@@ -579,15 +452,6 @@ class DenoisingStage(PipelineStage):
                 if batch.return_trajectory_latents:
                     trajectory_timesteps.append(t)
                     trajectory_latents.append(latents)
-                if rl_data is not None:
-                    rl_timesteps.append(t)
-                    if rl_data.store_trajectory:
-                        rl_latents.append(latents)
-                    if rl_data.collect_kl and rl_data.kl_reward <= 0:
-                        rl_kl.append(
-                            torch.zeros(latents.shape[0],
-                                        device=latents.device,
-                                        dtype=latents.dtype))
 
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
@@ -608,25 +472,6 @@ class DenoisingStage(PipelineStage):
         if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
             batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
             batch.trajectory_latents = trajectory_tensor.cpu()
-        if rl_data is not None:
-            if rl_timesteps:
-                rl_data.trajectory_timesteps = torch.stack(rl_timesteps, dim=0)
-                if rl_data.keep_trajectory_on_cpu:
-                    rl_data.trajectory_timesteps = rl_data.trajectory_timesteps.cpu(
-                    )
-            if rl_data.store_trajectory and rl_latents:
-                rl_data.trajectory_latents = torch.stack(rl_latents, dim=1)
-                if rl_data.keep_trajectory_on_cpu:
-                    rl_data.trajectory_latents = rl_data.trajectory_latents.cpu(
-                    )
-            if rl_log_probs:
-                rl_data.log_probs = torch.stack(rl_log_probs, dim=1)
-                if rl_data.keep_trajectory_on_cpu:
-                    rl_data.log_probs = rl_data.log_probs.cpu()
-            if rl_kl:
-                rl_data.kl = torch.stack(rl_kl, dim=1)
-                if rl_data.keep_trajectory_on_cpu:
-                    rl_data.kl = rl_data.kl.cpu()
 
         # Update batch with final latents
         batch.latents = latents

fastvideo/tests/hooks/test_hooks_api.py

@@ -1,60 +0,0 @@
-from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
-from torch import nn
-from typing import Any
-import torch
-
-
-class EventHook(ForwardHook):
-    def __init__(self, content: str, event_list: list[str]):
-        self.content = content
-        self.event_list = event_list
-
-    def name(self) -> str:
-        return f"EventHook_{self.content}"
-
-    def pre_forward(self, module: nn.Module, *args, **kwargs):
-        print(
-            f"[{self.content}] Pre-forward called with args[0].shape: {args[0].shape}"
-        )
-        self.event_list.append(f"[pre]{self.content}")
-        return args, kwargs
-
-    def post_forward(self, module: nn.Module, output: Any):
-        print(
-            f"[{self.content}] Post-forward called with outputs.shape: {output.shape}"
-        )
-        self.event_list.append(f"[post]{self.content}")
-        return output
-
-
-def test_hook_execution_order():
-    """Test that hooks are executed in the correct order: LIFO for pre-hooks, FIFO for post-hooks."""
-    # Create a simple model
-    model = nn.Linear(10, 20)
-
-    # Create event list to track hook execution order
-    events = []
-
-    # Create and push hooks in order: A then B
-
-    manager = ModuleHookManager.get_from_or_default(model)
-
-    hook_a = EventHook("A", events)
-    hook_b = EventHook("B", events)
-
-    manager.append_forward_hook(hook_a)
-    manager.append_forward_hook(hook_b)
-
-    # Perform a forward pass
-    input_tensor = torch.randn(2, 10)
-    model(input_tensor)
-
-    # Verify the execution order is [pre_a, pre_b, post_b, post_a]
-    # Pre-hooks should be FILO (First In Last Out): A then B
-    # Post-hooks should be LIFO (Last In First Out): B then A
-    expected_events = ["[pre]A", "[pre]B", "[post]B", "[post]A"]
-
-    assert events == expected_events, (
-        f"Expected {expected_events}, but got {events}"
-    )
-    print(f"✓ Hook execution order test passed: {events}")

fastvideo/tests/hooks/test_layerwise_offload.py

@@ -1,395 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-import torch
-import torch.nn as nn
-
-from fastvideo.hooks.layerwise_offload import (
-    LayerwiseOffloadHook,
-    enable_layerwise_offload,
-)
-from fastvideo.hooks.hooks import ModuleHookManager
-
-
-class SimpleBlock(nn.Module):
-    """A simple block with linear layers for testing."""
-
-    def __init__(self, hidden_size: int, dtype: torch.dtype = torch.float32):
-        super().__init__()
-        self.linear1 = nn.Linear(hidden_size, hidden_size, dtype=dtype)
-        self.linear2 = nn.Linear(hidden_size, hidden_size, dtype=dtype)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.linear1(x)
-        x = torch.relu(x)
-        x = self.linear2(x)
-        return x
-
-
-class SimpleModelWithModuleList(nn.Module):
-    """A simple model with ModuleList for testing layerwise offloading."""
-
-    def __init__(
-        self,
-        num_blocks: int = 4,
-        hidden_size: int = 128,
-        dtype: torch.dtype = torch.float32,
-    ):
-        super().__init__()
-        self.blocks = nn.ModuleList(
-            [SimpleBlock(hidden_size, dtype=dtype) for _ in range(num_blocks)]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        for block in self.blocks:
-            x = block(x)
-        return x
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_basic():
-    """Test basic functionality of layerwise offloading."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    batch_size = 2
-    seq_len = 16
-    num_blocks = 4
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Get reference output without offloading
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-    with torch.no_grad():
-        reference_output = model(input_tensor.clone())
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Verify parameters are offloaded to CPU
-    for block in model.blocks:
-        for param in block.parameters():
-            # Parameters should be placeholder tensors (empty)
-            assert param.numel() == 0, (
-                "Parameters should be offloaded (empty tensors)"
-            )
-
-    # Run forward pass with offloading
-    with torch.no_grad():
-        offloaded_output = model(input_tensor.clone())
-
-    # Check output correctness
-    assert torch.allclose(
-        reference_output, offloaded_output, rtol=1e-4, atol=1e-5
-    ), "Output with offloading should match reference output"
-
-    print(
-        f" Layerwise offload basic test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}"
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_bf16():
-    """Test layerwise offloading with bfloat16 precision."""
-    device = torch.device("cuda")
-    hidden_size = 256
-    batch_size = 1
-    seq_len = 32
-    num_blocks = 3
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.bfloat16
-    ).to(device)
-
-    # Get reference output without offloading
-    input_tensor = torch.randn(
-        batch_size, seq_len, hidden_size, device=device, dtype=torch.bfloat16
-    )
-    with torch.no_grad():
-        reference_output = model(input_tensor.clone())
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Run forward pass with offloading
-    with torch.no_grad():
-        offloaded_output = model(input_tensor.clone())
-
-    # Check output correctness (looser tolerance for bf16)
-    assert torch.allclose(
-        reference_output, offloaded_output, rtol=1e-2, atol=1e-3
-    ), "Output with offloading should match reference output for bf16"
-
-    print(
-        f" Layerwise offload bf16 test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}"
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_multiple_forward_passes():
-    """Test that layerwise offloading works correctly across multiple forward passes."""
-    device = torch.device("cuda")
-    hidden_size = 64
-    batch_size = 2
-    seq_len = 8
-    num_blocks = 3
-    num_iterations = 5
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Get reference outputs without offloading
-    torch.manual_seed(42)
-    reference_outputs = []
-    for i in range(num_iterations):
-        input_tensor = torch.randn(
-            batch_size, seq_len, hidden_size, device=device
-        )
-        with torch.no_grad():
-            reference_outputs.append(model(input_tensor.clone()))
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Run multiple forward passes with offloading
-    torch.manual_seed(42)
-    for i in range(num_iterations):
-        input_tensor = torch.randn(
-            batch_size, seq_len, hidden_size, device=device
-        )
-        with torch.no_grad():
-            offloaded_output = model(input_tensor.clone())
-
-        # Check output correctness for each iteration
-        assert torch.allclose(
-            reference_outputs[i], offloaded_output, rtol=1e-4, atol=1e-5
-        ), f"Output mismatch in iteration {i}"
-
-    print(
-        f" Multiple forward passes test passed for {num_iterations} iterations"
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_parameter_integrity():
-    """Test that parameters are correctly restored during forward pass."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    num_blocks = 2
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Store original parameter values per block
-    original_params_per_block = []
-    for block in model.blocks:
-        block_params = {}
-        for name, param in block.named_parameters():
-            block_params[name] = param.data.clone()
-        original_params_per_block.append(block_params)
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Create hook managers and verify they exist
-    for block_idx, block in enumerate(model.blocks):
-        manager = ModuleHookManager.get_from(block)
-        assert manager is not None, "Hook manager should be attached to blocks"
-        hook: LayerwiseOffloadHook | None = manager.get_forward_hook(
-            "LayerwiseOffloadHook"
-        )
-        assert hook is not None, "LayerwiseOffloadHook should be registered"
-
-        # Verify parameters are stored in CPU
-        state = hook.state
-        assert len(state.cpu_named_parameters) > 0, (
-            "CPU parameters should be stored"
-        )
-
-        # Verify CPU parameters match original values for this block
-        original_params = original_params_per_block[block_idx]
-        for name, cpu_param in state.cpu_named_parameters.items():
-            assert name in original_params, (
-                f"CPU parameter {name} not found in original params for block {block_idx}"
-            )
-            assert torch.allclose(
-                cpu_param.cpu(), original_params[name].cpu(), rtol=1e-5, atol=1e-6
-            ), f"CPU parameter {name} should match original for block {block_idx}"
-
-    print(" Parameter integrity test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_no_modulelist_error():
-    """Test that enabling offload on a model without ModuleList raises an error."""
-
-    class ModelWithoutModuleList(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.linear = nn.Linear(128, 128)
-
-        def forward(self, x):
-            return self.linear(x)
-
-    model = ModelWithoutModuleList().to("cuda")
-
-    with pytest.raises(
-        ValueError,
-        match="No nn.ModuleList found in the model for layerwise offloading",
-    ):
-        enable_layerwise_offload(model)
-
-    print(" No ModuleList error test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_memory_reduction():
-    """Test that layerwise offloading reduces GPU memory usage."""
-    device = torch.device("cuda")
-    hidden_size = 512
-    num_blocks = 8
-    batch_size = 1
-    seq_len = 64
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Measure initial GPU memory
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-    with torch.no_grad():
-        _ = model(input_tensor)
-
-    memory_without_offload = torch.cuda.max_memory_allocated()
-
-    # Reset model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Enable offloading
-    enable_layerwise_offload(model)
-
-    # Measure GPU memory with offloading
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-    with torch.no_grad():
-        _ = model(input_tensor)
-
-    memory_with_offload = torch.cuda.max_memory_allocated()
-
-    # Memory with offload should be less (parameters are offloaded)
-    # Note: This is a weak check as memory usage depends on many factors
-    print(
-        f"Memory without offload: {memory_without_offload / 1024**2:.2f} MB, "
-        f"with offload: {memory_with_offload / 1024**2:.2f} MB"
-    )
-
-    # We expect some reduction but this test is more informational
-    assert memory_with_offload < memory_without_offload * 1.5, (
-        "Memory usage should not increase significantly with offloading"
-    )
-
-    print(" Memory reduction test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_gradient_disabled():
-    """Test that layerwise offloading works correctly with gradients disabled."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    batch_size = 2
-    seq_len = 16
-    num_blocks = 3
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-    model.eval()
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-
-    # Forward pass should work without gradients
-    with torch.no_grad():
-        output = model(input_tensor)
-
-    assert output.requires_grad is False, "Output should not require gradients"
-    assert output.shape == (
-        batch_size,
-        seq_len,
-        hidden_size,
-    ), "Output shape should match input"
-
-    print(" Gradient disabled test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_different_batch_sizes():
-    """Test layerwise offloading with different batch sizes."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    seq_len = 16
-    num_blocks = 3
-
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Get reference model without offloading
-    reference_model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-    reference_model.load_state_dict(model.state_dict())
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Test different batch sizes
-    for batch_size in [1, 2, 4, 8]:
-        input_tensor = torch.randn(
-            batch_size, seq_len, hidden_size, device=device
-        )
-
-        with torch.no_grad():
-            reference_output = reference_model(input_tensor.clone())
-            offloaded_output = model(input_tensor.clone())
-
-        assert torch.allclose(
-            reference_output, offloaded_output, rtol=1e-4, atol=1e-5
-        ), f"Output mismatch for batch_size={batch_size}"
-
-    print(" Different batch sizes test passed")
-
-
-if __name__ == "__main__":
-    # Run tests manually for debugging
-    if torch.cuda.is_available():
-        print("Running layerwise offloading tests...")
-        test_layerwise_offload_basic()
-        test_layerwise_offload_bf16()
-        test_layerwise_offload_multiple_forward_passes()
-        test_layerwise_offload_parameter_integrity()
-        test_layerwise_offload_no_modulelist_error()
-        test_layerwise_offload_memory_reduction()
-        test_layerwise_offload_gradient_disabled()
-        test_layerwise_offload_different_batch_sizes()
-        print("\n All tests passed!")
-    else:
-        print("CUDA not available, skipping tests")

fastvideo/tests/inference/lora/test_lora_inference_similarity.py

@@ -93,11 +93,8 @@ def test_merge_lora_weights(model_id):
 
     lora_nickname = lora_config["lora_nickname"]
     lora_path = lora_config["lora_path"]
-    # When layerwise offload is enabled, placeholder tensors cannot be compared directly.
     args = FastVideoArgs.from_kwargs(
         model_path=model_id,
-        dit_layerwise_offload=False,
-        use_fsdp_inference=True,
         dit_cpu_offload=True,
         dit_precision="bf16",
     )

fastvideo/tests/nightly/test_e2e_hunyuan15_overfit_single_sample.py

@@ -0,0 +1,186 @@
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+
+from huggingface_hub import snapshot_download
+from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+
+NUM_NODES = "1"
+MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+
+DATA_DIR = "data"
+LOCAL_RAW_DATA_DIR = Path(DATA_DIR) / "cats15"
+LOCAL_PREPROCESSED_DATA_DIR = Path(DATA_DIR) / "cats_processed_t2v_hunyuan15"
+LOCAL_OUTPUT_DIR = Path(DATA_DIR) / "outputs_hunyuan15"
+
+NUM_GPUS_PER_NODE_PREPROCESSING = "1"
+NUM_GPUS_PER_NODE_TRAINING = "8"
+
+# entrypoints (adjust to what hunyuan15 scripts use)
+PREPROCESS_ENTRY = ["-m", "fastvideo.pipelines.preprocess.v1_preprocessing_new"]
+TRAIN_ENTRY_FILE_PATH = "fastvideo/training/hunyuan15_training_pipeline.py"  # change if hunyuan15 uses a different file
+
+LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "training_dataset", "worker_0", "worker_0")
+LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt.json")
+
+
+def download_data():
+    os.makedirs(DATA_DIR, exist_ok=True)
+    snapshot_download(
+        repo_id="wlsaidhi/cats-overfit-merged",
+        local_dir=str(LOCAL_RAW_DATA_DIR),
+        repo_type="dataset",
+        resume_download=True,
+        token=os.environ.get("HF_TOKEN"),
+    )
+
+    # normalize dataset layout like your hunyuan test
+    video_dir = LOCAL_RAW_DATA_DIR / "video"
+    videos_dir = LOCAL_RAW_DATA_DIR / "videos"
+    if video_dir.exists() and not videos_dir.exists():
+        video_dir.rename(videos_dir)
+
+    src_val = LOCAL_RAW_DATA_DIR / "validation_prompt_1_sample.json"
+    shutil.copy2(src_val, LOCAL_VALIDATION_DATASET_FILE)
+
+    src_v2c = LOCAL_RAW_DATA_DIR / "videos2caption_1_sample.json"
+    shutil.copy2(src_v2c, LOCAL_RAW_DATA_DIR / "videos2caption.json")
+
+
+def run_preprocessing():
+    if LOCAL_PREPROCESSED_DATA_DIR.exists():
+        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.getcwd()
+
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
+        *PREPROCESS_ENTRY,
+        "--model-path", MODEL_PATH,
+        "--mode", "preprocess",
+        "--workload-type", "t2v",
+        "--preprocess.video_loader_type", "torchvision",
+        "--preprocess.dataset_type", "merged",
+        "--preprocess.dataset_path", str(LOCAL_RAW_DATA_DIR),
+        "--preprocess.dataset_output_dir", str(LOCAL_PREPROCESSED_DATA_DIR),
+        "--preprocess.preprocess_video_batch_size", "1",
+        "--preprocess.dataloader_num_workers", "0",
+        "--preprocess.max_height", "480",
+        "--preprocess.max_width", "832",
+        "--preprocess.num_frames", "61",
+        "--preprocess.train_fps", "24",
+        "--preprocess.samples_per_file", "1",
+        "--preprocess.flush_frequency", "1",
+        "--preprocess.video_length_tolerance_range", "5",
+
+        # IMPORTANT: add a dtype/precision flag here if available in your codebase,
+        # so saved arrays are fp16/fp32 instead of bf16.
+        # Example (ONLY if supported by args):
+        # "--preprocess.save_dtype", "fp16",
+        # "--preprocess.embedding_dtype", "fp16",
+    ]
+
+    subprocess.run(cmd, check=True, env=env)
+
+
+def run_training():
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.getcwd()
+
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
+        TRAIN_ENTRY_FILE_PATH,
+        "--model_path", MODEL_PATH,
+        "--inference_mode", "False",  # Required: must be False for training mode
+        "--pretrained_model_name_or_path", MODEL_PATH,
+        "--data_path", LOCAL_TRAINING_DATA_DIR,
+        "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
+
+        "--train_batch_size", "1",
+        "--num_latent_t", "16",  # (61-1)//4 + 1 = 16 for 61 frames with temporal_compression=4
+        "--max_train_steps", "901",
+        "--learning_rate", "1e-5",
+        "--gradient_accumulation_steps", "1",  # Required: default 0 causes empty training loop
+
+        # If your pipeline produces bf16 artifacts that later go to numpy/parquet,
+        # switch mixed_precision to fp16 to avoid bf16 -> numpy issues.
+        "--mixed_precision", "fp16",
+
+        # Distributed training configuration - must match nproc_per_node
+        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
+        "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
+        "--tp_size", "1",
+        "--hsdp_replicate_dim", "1",
+        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
+        "--train_sp_batch_size", "1",
+
+        # Dataloader
+        "--dataloader_num_workers", "4",
+
+        # Validation & logging
+        "--log_validation",
+        "--validation_steps", "100",
+        "--validation_sampling_steps", "50",
+        "--validation_guidance_scale", "6.0",  # Match Hunyuan15_480P_SamplingParam default
+
+        # Checkpointing
+        "--weight_only_checkpointing_steps", "6000",
+        "--training_state_checkpointing_steps", "6000",
+        "--checkpoints_total_limit", "3",
+
+        # Optimizer settings
+        "--weight_decay", "0.01",
+        "--max_grad_norm", "1.0",
+
+        # Training config
+        "--ema_start_step", "0",
+        "--training_cfg_rate", "0.0",
+        "--dit_precision", "fp32",
+        "--enable_gradient_checkpointing_type", "full",
+
+        # Output
+        "--output_dir", str(LOCAL_OUTPUT_DIR),
+        "--tracker_project_name", "hunyuan15_finetune_overfit_ci",
+        "--num_height", "480",
+        "--num_width", "832",
+        "--num_frames", "61",
+        "--embedded_cfg_scale", "6.0",
+        "--num_euler_timesteps", "50",
+    ]
+
+    subprocess.run(cmd, check=True, env=env)
+
+
+def test_e2e_hunyuan15_overfit_single_sample():
+    os.environ["WANDB_MODE"] = "online"
+    os.environ["WANDB_API_KEY"] = "8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+
+    # download_data()
+    # run_preprocessing()
+    run_training()
+
+    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_hy15.mp4")
+    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
+
+    assert os.path.exists(reference_video_file)
+    assert os.path.exists(final_validation_video_file)
+
+    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
+        reference_video_file,
+        final_validation_video_file,
+        use_ms_ssim=True,
+    )
+    assert max_ssim > 0.5, f"Max SSIM is below threshold: {max_ssim}"
+
+
+if __name__ == "__main__":
+    test_e2e_hunyuan15_overfit_single_sample()

fastvideo/tests/nightly/test_e2e_hunyuan_overfit_single_sample.py

@@ -0,0 +1,218 @@
+import os
+from pathlib import Path
+from huggingface_hub import snapshot_download
+import shutil
+import subprocess
+import sys
+from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+
+# Example of using Hunyuan preprocessing and training pipeline
+
+# Import the training pipeline
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+
+NUM_NODES = "1"
+MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
+
+# preprocessing
+DATA_DIR = "data"
+LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats"))
+NUM_GPUS_PER_NODE_PREPROCESSING = "1"
+PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocessing_new.py"
+
+LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_processed_t2v_hunyuan"))
+
+
+# training
+NUM_GPUS_PER_NODE_TRAINING = "4"
+TRAINING_ENTRY_FILE_PATH = "fastvideo/training/hunyuan_training_pipeline.py"
+# New preprocessing pipeline creates files in training_dataset/worker_0/worker_0/
+LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "training_dataset", "worker_0", "worker_0")
+LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt.json")
+LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs_hunyuan"))
+
+def download_data():
+    # create the data dir if it doesn't exist
+    data_dir = Path(DATA_DIR)
+
+    print(f"Creating data directory at {data_dir}")
+    os.makedirs(data_dir, exist_ok=True)
+
+    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
+    try:
+        result = snapshot_download(
+            repo_id="wlsaidhi/cats-overfit-merged",
+            local_dir=str(LOCAL_RAW_DATA_DIR),
+            repo_type="dataset",
+            resume_download=True,
+            token=os.environ.get("HF_TOKEN"),  # In case authentication is needed
+        )
+        print(f"Download completed successfully. Files downloaded to: {result}")
+        
+        # Verify the download
+        if not LOCAL_RAW_DATA_DIR.exists():
+            raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
+            
+        # List downloaded files
+        print("Downloaded files:")
+        for file in LOCAL_RAW_DATA_DIR.rglob("*"):
+            if file.is_file():
+                print(f"  - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
+        
+        # Rename video directory if needed (dataset has 'video' but preprocessing expects 'videos')
+        video_dir = os.path.join(LOCAL_RAW_DATA_DIR, "video")
+        videos_dir = os.path.join(LOCAL_RAW_DATA_DIR, "videos")
+        if os.path.exists(video_dir) and not os.path.exists(videos_dir):
+            print(f"Renaming video directory to videos...")
+            os.rename(video_dir, videos_dir)
+            
+        # Copy validation file to expected name
+        source_validation_file = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.json")
+        target_validation_file = LOCAL_VALIDATION_DATASET_FILE
+        
+        if os.path.exists(source_validation_file):
+            print(f"Copying {source_validation_file} to {target_validation_file}...")
+            shutil.copy2(source_validation_file, target_validation_file)
+        else:
+            raise FileNotFoundError(f"Source validation file not found: {source_validation_file}")
+        
+        # Override videos2caption.json with the 1-sample version for preprocessing
+        # The new preprocessing pipeline automatically reads videos2caption.json from the dataset path
+        source_videos2caption = os.path.join(LOCAL_RAW_DATA_DIR, "videos2caption_1_sample.json")
+        target_videos2caption = os.path.join(LOCAL_RAW_DATA_DIR, "videos2caption.json")
+        
+        if os.path.exists(source_videos2caption):
+            print(f"Overriding videos2caption.json with 1-sample version...")
+            shutil.copy2(source_videos2caption, target_videos2caption)
+        else:
+            raise FileNotFoundError(f"Source videos2caption file not found: {source_videos2caption}")
+        
+    except Exception as e:
+        print(f"Error during download: {str(e)}")
+        raise
+
+
+def run_preprocessing():
+    # remove the local_preprocessed_data_dir if it exists
+    if LOCAL_PREPROCESSED_DATA_DIR.exists():
+        print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
+        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
+    env = os.environ.copy()
+    env['PYTHONPATH'] = os.getcwd()
+    # Run torchrun command using the new preprocessing pipeline
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
+        "-m", "fastvideo.pipelines.preprocess.v1_preprocessing_new",
+        "--model-path", MODEL_PATH,
+        "--mode", "preprocess",
+        "--workload-type", "t2v",
+        "--preprocess.video_loader_type", "torchvision",
+        "--preprocess.dataset_type", "merged",
+        "--preprocess.dataset_path", str(LOCAL_RAW_DATA_DIR),
+        "--preprocess.dataset_output_dir", str(LOCAL_PREPROCESSED_DATA_DIR),
+        "--preprocess.preprocess_video_batch_size", "1",
+        "--preprocess.dataloader_num_workers", "0",
+        "--preprocess.max_height", "480",
+        "--preprocess.max_width", "832",
+        "--preprocess.num_frames", "77",
+        "--preprocess.train_fps", "16",
+        "--preprocess.samples_per_file", "1",
+        "--preprocess.flush_frequency", "1",
+        "--preprocess.video_length_tolerance_range", "5",
+    ]
+
+    process = subprocess.run(cmd, check=True, env=env)
+
+
+def run_training():
+    env = os.environ.copy()
+    env['PYTHONPATH'] = os.getcwd()
+
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
+        TRAINING_ENTRY_FILE_PATH,
+        "--model_path", MODEL_PATH,
+        "--inference_mode", "False",
+        "--pretrained_model_name_or_path", MODEL_PATH,
+        "--data_path", LOCAL_TRAINING_DATA_DIR,
+        "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
+        "--train_batch_size", "1",
+        "--num_latent_t", "8",
+        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
+        "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
+        "--tp_size", "1",
+        "--hsdp_replicate_dim", "1",
+        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
+        "--train_sp_batch_size", "1",
+        "--dataloader_num_workers", "10",
+        "--gradient_accumulation_steps", "1",
+        "--max_train_steps", "901",
+        "--learning_rate", "1e-5",
+        "--mixed_precision", "bf16",
+        "--weight_only_checkpointing_steps", "6000",
+        "--training_state_checkpointing_steps", "6000",
+        "--validation_steps", "100",
+        "--validation_sampling_steps", "50",
+        "--log_validation",
+        "--checkpoints_total_limit", "3",
+        "--ema_start_step", "0",
+        "--training_cfg_rate", "0.0",
+        "--output_dir", str(LOCAL_OUTPUT_DIR),
+        "--tracker_project_name", "hunyuan_finetune_overfit_ci",
+        "--num_height", "480",
+        "--num_width", "832",
+        "--num_frames", "81",
+        "--validation_guidance_scale", "1.0",
+        "--embedded_cfg_scale", "6.0",
+        "--num_euler_timesteps", "50",
+        "--multi_phased_distill_schedule", "4000-1",
+        "--weight_decay", "0.01",
+        "--not_apply_cfg_solver",
+        "--dit_precision", "fp32",
+        "--max_grad_norm", "1.0",
+        "--flow_shift", "7",
+    ]
+
+    print(f"Running training with command: {cmd}")
+    process = subprocess.run(cmd, check=True, env=env)
+
+
+def test_e2e_hunyuan_overfit_single_sample():
+    os.environ["WANDB_MODE"] = "online"
+
+    #download_data()
+    run_preprocessing()
+    run_training()
+
+    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
+    print(f"reference_video_file: {reference_video_file}")
+    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
+    print(f"final_validation_video_file: {final_validation_video_file}")
+
+
+    # Ensure both files exist
+    assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
+    assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
+
+    #Compute SSIM
+    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
+        reference_video_file,
+        final_validation_video_file,
+        use_ms_ssim=True  # Using MS-SSIM for better quality assessment
+    )
+
+    print("\n===== SSIM Results for Step 900 Validation (Hunyuan) =====")
+    print(f"Mean MS-SSIM: {mean_ssim:.4f}")
+    print(f"Min MS-SSIM: {min_ssim:.4f}")
+    print(f"Max MS-SSIM: {max_ssim:.4f}")
+
+    assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
+
+
+
+if __name__ == "__main__":
+    test_e2e_hunyuan_overfit_single_sample()

fastvideo/tests/ssim/test_inference_similarity.py

@@ -19,12 +19,16 @@ if "A40" in device_name:
     device_reference_folder = "A40" + device_reference_folder_suffix
 elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
+elif "H100" in device_name or "NVIDIA H100" in device_name:
+    device_reference_folder = "H100" + device_reference_folder_suffix
+
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")
     # raise ValueError(f"Unsupported device for ssim tests: {device_name}")
 
 # Base parameters from the shell script
+
 HUNYUAN_PARAMS = {
     "num_gpus": 4,
     "model_path": "FastVideo/FastHunyuan-diffusers",
@@ -234,7 +238,6 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
         "tp_size": BASE_PARAMS["tp_size"],
         "use_fsdp_inference": True,
         "dit_cpu_offload": False,
-        "dit_layerwise_offload": False,
     }
     if BASE_PARAMS.get("vae_sp"):
         init_kwargs["vae_sp"] = True

fastvideo/tests/ssim/test_matrixgame_similarity.py

@@ -88,7 +88,6 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id):
     init_kwargs = {
         "num_gpus": BASE_PARAMS["num_gpus"],
         "use_fsdp_inference": True,
-        "dit_layerwise_offload": False,
         "dit_cpu_offload": False,
         "vae_cpu_offload": False,
         "text_encoder_cpu_offload": True,

fastvideo/tests/ssim/test_turbodiffusion_similarity.py

@@ -218,7 +218,6 @@ def test_turbodiffusion_i2v_inference_similarity(prompt, model_id):
         # Keep both transformers in VRAM - avoids CPU RAM bottleneck
         "dit_cpu_offload": False,
         "use_fsdp_inference": True,
-        "dit_layerwise_offload": False,
     }
 
     generation_kwargs = {

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -6,11 +6,6 @@ import json
 from huggingface_hub import snapshot_download
 import torch
 
-# Ensure backend selection happens during import-time initialization
-os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
-# Force VSA to use Triton implementation even on H100 / when CUDA extension is available
-# os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
-
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
 from fastvideo.training.wan_training_pipeline import main
@@ -24,6 +19,8 @@ h200_reference_wandb_summary_file = "fastvideo/tests/training/VSA/h200_reference
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "2"
 
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
@@ -121,10 +118,10 @@ def test_distributed_training():
     wandb_summary = json.load(open(summary_file))
 
     fields_and_thresholds = {
-        'avg_step_time': 3,
-        'grad_norm': 0.2,
-        'step_time': 2.5,
-        'train_loss': 0.04
+        'avg_step_time': 1.0,
+        'grad_norm': 0.1,
+        'step_time': 1.0,
+        'train_loss': 0.02
     }
 
     failures = []

fastvideo/tests/transformers/test_wanvideo.py

@@ -17,11 +17,10 @@ from fastvideo.configs.models.dits import WanVideoConfig
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 
 
-
 logger = init_logger(__name__)
 
 os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29701"
+os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
@@ -122,24 +121,4 @@ def test_wan_transformer():
     assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
 
     # Check if outputs are similar (allowing for small numerical differences)
-    max_diff = torch.max(torch.abs(output1 - output2))
-    mean_diff = torch.mean(torch.abs(output1 - output2))
-    logger.info("Max Diff: %s", max_diff.item())
-    logger.info("Mean Diff: %s", mean_diff.item())
-    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
-    # mean diff
-    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
-
-
-if __name__ == "__main__":
-    from fastvideo.distributed import (
-        cleanup_dist_env_and_memory,
-        maybe_init_distributed_environment_and_model_parallel,
-    )
-    # Allow running this test file directly without pytest.
-    maybe_init_distributed_environment_and_model_parallel(1, 1)
-    try:
-        test_wan_transformer()
-        logger.info("test_wan_transformer finished successfully.")
-    finally:
-        cleanup_dist_env_and_memory()
+    assert_close(output1, output2, atol=1e-1, rtol=1e-2)

fastvideo/training/__init__.py

@@ -1,12 +1,5 @@
 from .distillation_pipeline import DistillationPipeline
 from .training_pipeline import TrainingPipeline
 from .wan_training_pipeline import WanTrainingPipeline
-from fastvideo.training.rl import RLPipeline, create_rl_pipeline
 
-__all__ = [
-    "TrainingPipeline",
-    "WanTrainingPipeline",
-    "DistillationPipeline",
-    "RLPipeline",
-    "create_rl_pipeline",
-]
+__all__ = ["TrainingPipeline", "WanTrainingPipeline", "DistillationPipeline"]

fastvideo/training/hunyuan15_training_pipeline.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.training.training_pipeline import TrainingPipeline
+
+from fastvideo.pipelines.basic.hunyuan15.hunyuan15_pipeline import HunyuanVideo15Pipeline
+
+logger = init_logger(__name__)
+
+
+class Hunyuan15TrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for HunyuanVideo-1.5.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        if training_args.enable_gradient_checkpointing_type is None:
+            training_args.enable_gradient_checkpointing_type = "full"
+        super().initialize_training_pipeline(training_args)
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        pass
+        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+        #     shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        # reserved for future refactors
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline (Hunyuan15)...")
+        args_copy = deepcopy(training_args)
+        args_copy.inference_mode = True
+
+        validation_pipeline = HunyuanVideo15Pipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={
+                # reuse the training transformer weights for validation sampling
+                "transformer": self.get_module("transformer"),
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True,
+            # dit_layerwise_offload=True,
+            use_fsdp_inference=True,
+        )
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting Hunyuan15 training pipeline...")
+
+    pipeline = Hunyuan15TrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path,
+        args=args,
+    )
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.utils import FlexibleArgumentParser
+
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/hunyuan_training_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines.basic.hunyuan.hunyuan_pipeline import HunyuanVideoPipeline
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.utils import is_vsa_available
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+class HunyuanTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Hunyuan.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        validation_pipeline = HunyuanVideoPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={
+                "transformer": self.get_module("transformer"),
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = HunyuanTrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/rl/__init__.py

@@ -1,6 +0,0 @@
-from .rl_pipeline import RLPipeline, create_rl_pipeline
-
-__all__ = [
-    "RLPipeline",
-    "create_rl_pipeline",
-]

fastvideo/training/rl/rewards/__init__.py

@@ -1,11 +0,0 @@
-from .rewards import (
-    create_reward_models,
-    MultiRewardAggregator,
-    ValueModel
-)
-
-__all__ = [
-    "create_reward_models",
-    "MultiRewardAggregator",
-    "ValueModel",
-]

fastvideo/training/rl/rewards/base.py

@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-    Abstract base class for VIDEO reward models.
-
-    All VIDEO reward models should inherit from this class and implement
-    the compute_reward() method.
-
-    IMPORTANT: Reward models must process FULL VIDEO SEQUENCES, not individual frames.
-    Input shape is [B, T, C, H, W] where T is the temporal (frame) dimension.
-
-    For video-specific rewards, consider:
-    - Temporal coherence across frames
-    - Motion quality and smoothness
-    - Video-text alignment (not just frame-text)
-    - Multi-frame aesthetic quality
-    """
-
-from typing import Any
-
-from abc import ABC, abstractmethod
-import torch
-import torch.nn as nn
-
-class BaseRewardModel(ABC, nn.Module):
-    def __init__(self, model_path: str | None = None, device: str = "cuda"):
-        super().__init__()
-        self.model_path = model_path
-        self.device = device
-
-    @abstractmethod
-    def compute_reward(
-        self,
-        videos: torch.Tensor,  # [B, T, C, H, W] decoded video sequences
-        prompts: list[str] | None,  # Text prompts
-        **kwargs: Any
-    ) -> torch.Tensor:
-        """
-        Compute rewards for generated VIDEO sequences.
-
-        IMPORTANT: This method must process the FULL temporal sequence [B, T, C, H, W].
-        Do NOT evaluate individual frames independently and average.
-
-        Args:
-            videos: Decoded video tensors [B, T, C, H, W] in range [0, 1]
-                   B = batch size
-                   T = number of frames (temporal dimension)
-                   C = channels (typically 3 for RGB)
-                   H, W = height, width
-            prompts: List of text prompts (length B) describing each video
-            **kwargs: Additional model-specific arguments
-
-        Returns:
-            rewards: Tensor of shape [B] with reward scores for each video sequence
-
-        Example:
-            >>> videos = torch.rand(4, 17, 3, 256, 256)  # 4 videos, 17 frames each
-            >>> prompts = ["A cat jumping", "A dog running", ...]
-            >>> rewards = model.compute_reward(videos, prompts)
-        """
-        raise NotImplementedError("Subclasses must implement compute_reward()")
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(model_path={self.model_path})"

fastvideo/training/rl/rewards/ocr.py

@@ -1,206 +0,0 @@
-from paddleocr import PaddleOCR
-import torch
-import numpy as np
-from Levenshtein import distance
-from typing import Any
-from PIL import Image
-
-from fastvideo.training.rl.rewards.base import BaseRewardModel
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class OcrScorerVideo(BaseRewardModel):
-    """
-    OCR reward model for multi-frame video OCR evaluation.
-
-    This model evaluates multiple frames across the video sequence,
-    sampling frames at a specified interval and averaging the OCR scores.
-    """
-
-    def __init__(self,
-                 model_path: str | None = None,
-                 device: str = "cpu",
-                 frame_interval: int = 4):
-        """
-        OCR reward calculator for videos
-
-        Args:
-            model_path: Not used for PaddleOCR (kept for BaseRewardModel compatibility)
-            device: Device string (used to determine use_gpu if not explicitly set)
-            frame_interval: Sample every Nth frame (default: 4)
-        """
-        super().__init__(model_path=model_path, device=device)
-
-        self.frame_interval = frame_interval
-        self.ocr = PaddleOCR(
-            use_angle_cls=False,
-            lang="en",
-            use_gpu=False,
-            show_log=False  # Disable unnecessary log output
-        )
-
-        logger.info("Initialized OcrScorerVideo (device=%s, frame_interval=%d)",
-                    device, frame_interval)
-
-    def _process_single_video(self, video_tensor: torch.Tensor,
-                              prompt: str) -> float:
-        """
-        Process a single video tensor and return its OCR reward.
-
-        Args:
-            video_tensor: Video tensor of shape [C, T, H, W]
-            prompt: Text prompt containing target OCR text in quotes
-
-        Returns:
-            Average reward across positive-scoring frames
-        """
-        # Extract target text from prompt
-        try:
-            target_text = prompt.split('"')[1].replace(' ', '').lower()
-        except IndexError:
-            logger.warning("Failed to extract quoted text from prompt: %s",
-                           prompt)
-            target_text = prompt.replace(' ', '').lower()
-
-        if not target_text:
-            return 0.0
-
-        # video_tensor is [C, T, H, W]
-        C, T, H, W = video_tensor.shape
-
-        # Convert to numpy and move to CPU if needed
-        video_np = video_tensor.detach().cpu().numpy()
-
-        # Convert from [C, T, H, W] to [T, H, W, C] for easier frame extraction
-        video_np = np.transpose(video_np, (1, 2, 3, 0))  # [T, H, W, C]
-        logger.info(f"in ocr 1.5, video_np[0][0]: {video_np[0][0]}")
-
-        # Normalize to [0, 255] uint8 if needed
-        if video_np.max() <= 1.0:
-            video_np = (video_np * 255).astype(np.uint8)
-        else:
-            video_np = video_np.astype(np.uint8)
-
-        frame_rewards = []
-
-        # Sample frames at specified interval
-        for frame_idx in range(0, T, self.frame_interval):
-            frame = video_np[frame_idx]  # [H, W, C]
-            logger.info(f"in ocr 2, frame.shape: {frame.shape}")
-            # Run OCR
-            try:
-                result = self.ocr.ocr(frame, cls=False)
-                logger.info(f"in ocr 3, result: {result}")
-                if result and result[0]:
-                    recognized_text = "".join(
-                        [line[1][0] for line in result[0] if line[1][1] > 0])
-                else:
-                    recognized_text = ""
-            except Exception as e:
-                logger.info("OCR failed on frame %d: %s", frame_idx, str(e))
-                recognized_text = ''
-
-            logger.info(f"in ocr 4, recognized_text: {recognized_text}")
-
-            recognized_text = recognized_text.replace(' ', '').lower()
-            if target_text in recognized_text:
-                dist = 0
-            else:
-                dist = distance(recognized_text, target_text)
-            dist = min(dist, len(target_text))
-            reward = 1.0 - dist / len(target_text)
-
-            logger.info(f"in ocr 5, reward: {reward}")
-            if reward > 0:
-                frame_rewards.append(reward)
-
-        logger.info(f"in ocr 6, frame_rewards: {frame_rewards}")
-
-        return sum([reward / len(frame_rewards)
-                    for reward in frame_rewards]) if frame_rewards else 0.0
-
-    @torch.no_grad()
-    def compute_reward(self, videos: torch.Tensor, prompts: list[str],
-                       **kwargs: Any) -> torch.Tensor:
-        """
-        Calculate OCR reward by evaluating sampled frames across the video.
-
-        Args:
-            videos: Video tensor of shape [B, C, T, H, W]
-                   B = batch size
-                   C = channels (typically 3 for RGB)
-                   T = number of frames (temporal dimension)
-                   H, W = height, width
-            prompts: List of text prompts containing target OCR text in quotes (length B)
-            **kwargs: Additional arguments
-
-        Returns:
-            Reward tensor [B] with averaged OCR similarity scores across frames
-        """
-        # Ensure videos is a torch tensor with correct shape
-        assert isinstance(
-            videos,
-            torch.Tensor), f"videos must be torch.Tensor, got {type(videos)}"
-        assert videos.ndim == 5, f"videos must have 5 dimensions [B, C, T, H, W], got shape {videos.shape}"
-
-        logger.info(f"in ocr 1, videos.shape: {videos.shape}")
-
-        B, C, T, H, W = videos.shape
-        assert len(
-            prompts
-        ) == B, f"Number of prompts ({len(prompts)}) must match batch size ({B})"
-
-        rewards = []
-        for b in range(B):
-            # Extract single video: [C, T, H, W]
-            video = videos[b]
-            reward = self._process_single_video(video, prompts[b])
-            rewards.append(reward)
-
-        logger.info(f"in ocr 7, rewards: {rewards}")
-
-        rewards = torch.tensor(rewards, dtype=torch.float32, device=self.device)
-
-        logger.info(f"in ocr 8, rewards: {rewards}")
-
-        # Check for NaN or Inf values
-        if torch.isnan(rewards).any() or torch.isinf(rewards).any():
-            logger.warning(
-                "NaN or Inf detected in OCR rewards, returning zero tensor")
-            return torch.zeros_like(rewards)
-
-        return rewards
-
-
-if __name__ == "__main__":
-    example_image_path = "flowgrpo_cmd.png"
-    example_image = Image.open(example_image_path)
-    example_prompt = '/f1ow_grpo$'
-
-    # Convert image to RGB if needed
-    if example_image.mode != 'RGB':
-        example_image = example_image.convert('RGB')
-
-    # Convert PIL Image to numpy array [H, W, C]
-    image_np = np.array(example_image)
-
-    # Normalize to [0, 1] range and convert to float32
-    image_np = image_np.astype(np.float32) / 255.0
-
-    # Convert to torch tensor and reshape: [H, W, C] -> [C, H, W]
-    image_tensor = torch.from_numpy(image_np).permute(2, 0, 1)
-
-    # Add temporal dimension: [C, H, W] -> [C, T, H, W] where T=1
-    video_tensor = image_tensor.unsqueeze(1)  # [C, 1, H, W]
-
-    # Add batch dimension: [C, T, H, W] -> [B, C, T, H, W] where B=1
-    video_tensor = video_tensor.unsqueeze(0)  # [1, C, 1, H, W]
-
-    # Instantiate scorer
-    scorer = OcrScorerVideo(device="cpu")
-
-    # Call compute_reward method with video tensor
-    reward = scorer.compute_reward(video_tensor, [example_prompt])
-    print(f"OCR Reward: {reward.item()}")

fastvideo/training/rl/rewards/rewards.py

@@ -1,338 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Base infrastructure for VIDEO reward models in RL/GRPO training.
-
-IMPORTANT: This module is designed exclusively for VIDEO generation models.
-All reward models must operate on video sequences [B, T, C, H, W], not single frames.
-
-This module provides:
-1. Multi-reward aggregation for video
-2. Value model wrapper
-3. Integration with FastVideo video generation infrastructure
-"""
-
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from fastvideo.logger import init_logger
-from fastvideo.training.rl.rewards.ocr import OcrScorerVideo
-from fastvideo.training.rl.rewards.base import BaseRewardModel
-logger = init_logger(__name__)
-
-
-class MultiRewardAggregator(nn.Module):
-    """
-    Aggregates multiple reward models with configurable weights.
-
-    This implements the multi-reward aggregation strategy from flow_grpo,
-    allowing combination of different reward signals (aesthetic quality,
-    text-video alignment, compositional understanding, etc.)
-    """
-
-    def __init__(
-        self,
-        reward_models: list[BaseRewardModel],
-        reward_weights: list[float] | None = None,
-        normalize_rewards: bool = True
-    ):
-        """
-        Initialize multi-reward aggregator.
-
-        Args:
-            reward_models: List of reward model instances
-            reward_weights: Weights for each reward model (default: uniform)
-            normalize_rewards: Whether to normalize rewards before aggregation
-        """
-        super().__init__()
-        self.reward_models = nn.ModuleList(reward_models)
-
-        if reward_weights is None:
-            reward_weights = [1.0 / len(reward_models)] * len(reward_models)
-
-        assert len(reward_weights) == len(reward_models), \
-            f"Number of weights ({len(reward_weights)}) must match number of models ({len(reward_models)})"
-
-        assert abs(sum(reward_weights) - 1.0) < 1e-6, \
-            f"Reward weights must sum to 1.0, got {sum(reward_weights)}"
-
-        self.reward_weights = reward_weights
-        self.normalize_rewards = normalize_rewards
-
-        logger.info(
-            "Initialized MultiRewardAggregator with %d models: %s",
-            len(reward_models),
-            [(type(m).__name__, w) for m, w in zip(reward_models, reward_weights, strict=False)]
-        )
-
-    def compute_reward(
-        self,
-        videos: torch.Tensor,
-        prompts: list[str],
-        return_individual: bool = False,
-        **kwargs: Any
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
-        """
-        Compute aggregated reward from multiple models.
-
-        Args:
-            videos: Decoded video tensors [B, C, T, H, W]
-            prompts: List of text prompts
-            return_individual: If True, return dict with individual rewards
-            **kwargs: Additional arguments passed to reward models
-
-        Returns:
-            If return_individual=False: aggregated_rewards [B]
-            If return_individual=True: dict with "aggregated" and individual model rewards
-        """
-        batch_size = videos.shape[0]
-        individual_rewards: dict[str, torch.Tensor] = {}
-
-        # Collect rewards from all models
-        all_rewards = []
-        for i, (model, weight) in enumerate(zip(self.reward_models, self.reward_weights, strict=False)):
-            reward = model.compute_reward(videos, prompts, **kwargs)
-            assert reward.shape == (batch_size,), \
-                f"Reward model {i} returned shape {reward.shape}, expected ({batch_size},)"
-
-            # Optionally normalize individual rewards
-            if self.normalize_rewards:
-                reward = (reward - reward.mean()) / (reward.std() + 1e-8)
-
-            individual_rewards[f"reward_{type(model).__name__}"] = reward
-            all_rewards.append(weight * reward)
-
-        # Aggregate with weights
-        aggregated = sum(all_rewards)
-
-        if return_individual:
-            individual_rewards["aggregated"] = aggregated
-            return individual_rewards
-
-        return aggregated
-
-    def __repr__(self) -> str:
-        models_str = ", ".join([
-            f"{type(m).__name__}(w={w:.3f})"
-            for m, w in zip(self.reward_models, self.reward_weights, strict=False)
-        ])
-        return f"MultiRewardAggregator({models_str})"
-
-
-class ValueModel(nn.Module):
-    """
-    Value function model wrapper for RL training.
-
-    The value model can either:
-    1. Share the transformer backbone with the policy (memory efficient)
-    2. Use a separate transformer (more flexible)
-
-    For now, this is a placeholder that will be expanded based on
-    the chosen architecture strategy.
-    """
-
-    def __init__(
-        self,
-        transformer: nn.Module,
-        share_backbone: bool = False,
-        hidden_size: int | None = None
-    ):
-        """
-        Initialize value model.
-
-        Args:
-            transformer: Transformer model (policy or separate)
-            share_backbone: Whether to share backbone with policy
-            hidden_size: Hidden size for value head (inferred if None)
-        """
-        super().__init__()
-        self.transformer = transformer
-        self.share_backbone = share_backbone
-
-        # Value head will be added later based on transformer architecture
-        # For now, just store the transformer reference
-        logger.info(
-            "Initialized ValueModel (share_backbone=%s)",
-            share_backbone
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.Tensor,
-        **kwargs: Any
-    ) -> torch.Tensor:
-        """
-        Forward pass to compute value predictions.
-
-        Args:
-            hidden_states: Latent states [B, C, T, H, W]
-            encoder_hidden_states: Text embeddings [B, L, D]
-            timestep: Timesteps [B]
-            **kwargs: Additional transformer arguments
-
-        Returns:
-            values: Value predictions [B]
-        """
-        # TODO: Implement value prediction
-        # For now, return dummy values
-        batch_size = hidden_states.shape[0]
-        return torch.zeros(batch_size, device=hidden_states.device)
-
-
-class DummyRewardModel(BaseRewardModel):
-    """
-    Dummy VIDEO reward model for testing and development.
-
-    Returns random rewards in the range [0, 1] for VIDEO inputs.
-    This is a placeholder for testing the RL pipeline before real video reward models
-    are implemented.
-
-    NOTE: This does NOT actually evaluate video quality - it's just for testing!
-    """
-
-    def __init__(self, mean: float = 0.5, std: float = 0.1):
-        super().__init__(model_path=None)
-        self.mean = mean
-        self.std = std
-        logger.info("Initialized DummyRewardModel (VIDEO) - mean=%.2f, std=%.2f", mean, std)
-        logger.warning(
-            "DummyRewardModel is for TESTING ONLY - does not evaluate actual video quality!"
-        )
-
-    def compute_reward(
-        self,
-        videos: torch.Tensor,  # [B, T, C, H, W]
-        prompts: list[str],
-        **kwargs: Any
-    ) -> torch.Tensor:
-        """
-        Return random rewards for testing.
-
-        Args:
-            videos: Video sequences [B, T, C, H, W]
-            prompts: Text prompts
-
-        Returns:
-            Random rewards [B] in range [0, 1]
-        """
-        batch_size = videos.shape[0]
-        num_frames = videos.shape[1]
-
-        logger.debug(
-            "DummyRewardModel processing %d videos with %d frames each",
-            batch_size,
-            num_frames
-        )
-
-        # Generate random rewards (not based on actual video content!)
-        rewards = torch.randn(batch_size, device=videos.device) * self.std + self.mean
-        return rewards.clamp(0.0, 1.0)
-
-    def load_model(self) -> None:
-        """No model to load for dummy."""
-        pass
-
-
-def create_reward_models(
-    reward_models: dict,
-    device: str = "cuda"
-) -> MultiRewardAggregator:
-    """
-    Factory function to create VIDEO reward models from configuration strings.
-
-    IMPORTANT: Only creates VIDEO reward models. Image-only reward models
-    (PickScore, ImageReward, GenEval, etc.) are NOT supported.
-
-    Args:
-        reward_models: dictionary of reward model names to weights
-                      Example: {"paddle_ocr": 0.5, "video_score": 0.5}
-        device: Device to load models on
-
-    Returns:
-        MultiRewardAggregator with loaded VIDEO reward models
-
-    Supported VIDEO Reward Types:
-        - "paddle_ocr": PaddleOCR multi-frame video text recognition
-        - "video_score": Video aesthetic quality (multi-frame) - TODO
-        - "video_text_alignment": CLIP-based video-text similarity - TODO
-        - "temporal_coherence": Frame-to-frame consistency - TODO
-        - "motion_quality": Motion smoothness and realism - TODO
-        - "dummy": Random rewards for testing (VIDEO-aware)
-
-    NOT Supported (Image-Only):
-        - "pickscore": Image aesthetic (use "video_score" instead)
-        - "imagereward": Image quality (use "video_score" instead)
-        - "geneval": Image compositional (no video equivalent yet)
-        - Any single-frame reward models
-
-    Example:
-        >>> models = create_reward_models(
-        ...     reward_models={
-        ...         "paddle_ocr": 0.5,
-        ...         "video_text_alignment": 0.5
-        ...     },
-        ...     device="cuda"
-        ... )
-    """
-
-
-    assert reward_models, "No reward models specified. Please select at least 1 reward model"
-
-    types = [t.strip() for t in reward_models.keys()]
-    weights = list(reward_models.values())
-
-    assert len(types) == len(weights), \
-        f"Number of models ({len(types)}) must match number of weights ({len(weights)})"
-
-    # Create reward models based on types
-    models_list: list[BaseRewardModel] = []
-    for reward_type in types:
-        if reward_type == "dummy":
-            model = DummyRewardModel()
-
-        elif reward_type == "paddle_ocr":
-            logger.info("Creating PaddleOCR reward model")
-            model = OcrScorerVideo(device=device)
-
-        elif reward_type == "video_score":
-            # TODO: Implement VideoScore reward model (Phase 2)
-            logger.warning(
-                "VideoScore reward not implemented yet, using DummyRewardModel"
-            )
-            model = DummyRewardModel()
-        elif reward_type == "video_text_alignment":
-            # TODO: Implement VideoTextAlignment reward model (Phase 2)
-            logger.warning(
-                "VideoTextAlignment reward not implemented yet, using DummyRewardModel"
-            )
-            model = DummyRewardModel()
-        elif reward_type == "temporal_coherence":
-            # TODO: Implement TemporalCoherence reward model (Phase 2)
-            logger.warning(
-                "TemporalCoherence reward not implemented yet, using DummyRewardModel"
-            )
-            model = DummyRewardModel()
-        elif reward_type == "motion_quality":
-            # TODO: Implement MotionQuality reward model (Phase 2)
-            logger.warning(
-                "MotionQuality reward not implemented yet, using DummyRewardModel"
-            )
-            model = DummyRewardModel()
-        else:
-            logger.warning(
-                "Unknown VIDEO reward type '%s', using DummyRewardModel",
-                reward_type
-            )
-            model = DummyRewardModel()
-
-        models_list.append(model)
-
-    logger.info(
-        "Created MultiRewardAggregator with %d VIDEO reward models",
-        len(models_list)
-    )
-
-    return MultiRewardAggregator(models_list, weights, normalize_rewards=True)

fastvideo/training/rl/rl_utils.py

@@ -1,385 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Utility functions for RL/GRPO training.
-"""
-
-from typing import Any
-
-import torch
-import torch.nn.functional as F
-
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def compute_gae(
-    rewards: torch.Tensor,
-    values: torch.Tensor,
-    next_values: torch.Tensor,
-    dones: torch.Tensor | None = None,
-    gamma: float = 0.99,
-    lambda_: float = 0.95
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Compute Generalized Advantage Estimation (GAE-lambda).
-
-    GAE reduces variance in advantage estimation while allowing some bias.
-    This is a key component of modern policy gradient methods like PPO and GRPO.
-
-    Args:
-        rewards: Rewards at each step [B, T] or [B]
-        values: Value predictions at each step [B, T] or [B]
-        next_values: Value predictions at next step [B, T] or [B]
-        dones: Episode termination flags [B, T] or [B] (1 if done, 0 otherwise)
-        gamma: Discount factor
-        lambda_: GAE lambda parameter (0=TD(0), 1=Monte Carlo)
-
-    Returns:
-        advantages: GAE advantages [B, T] or [B]
-        returns: TD(lambda) returns [B, T] or [B]
-
-    Reference:
-        Schulman et al. "High-Dimensional Continuous Control Using Generalized Advantage Estimation"
-        https://arxiv.org/abs/1506.02438
-    """
-    if dones is None:
-        dones = torch.zeros_like(rewards)
-
-    # Compute TD residuals: delta_t = r_t + gamma * V(s_{t+1}) - V(s_t)
-    deltas = rewards + gamma * next_values * (1.0 - dones) - values
-
-    # If single step (no time dimension), return directly
-    if deltas.dim() == 1:
-        advantages = deltas
-        returns = advantages + values
-        return advantages, returns
-
-    # Multi-step: compute GAE recursively
-    batch_size, num_steps = deltas.shape
-    advantages = torch.zeros_like(deltas)
-    gae = torch.zeros(batch_size, device=deltas.device)
-
-    # Backward pass to compute GAE
-    for t in reversed(range(num_steps)):
-        gae = deltas[:, t] + gamma * lambda_ * (1.0 - dones[:, t]) * gae
-        advantages[:, t] = gae
-
-    # Returns are advantages + values
-    returns = advantages + values
-
-    return advantages, returns
-
-
-def normalize_advantages(
-    advantages: torch.Tensor,
-    epsilon: float = 1e-8
-) -> torch.Tensor:
-    """
-    Normalize advantages to have zero mean and unit variance.
-
-    This is a common practice in PPO and GRPO to stabilize training.
-
-    Args:
-        advantages: Raw advantages [B, ...]
-        epsilon: Small constant for numerical stability
-
-    Returns:
-        normalized_advantages: Normalized advantages [B, ...]
-    """
-    mean = advantages.mean()
-    std = advantages.std()
-    return (advantages - mean) / (std + epsilon)
-
-#TODO(jiali): refactor into algorithm
-def compute_grpo_policy_loss(
-    log_probs: torch.Tensor,
-    old_log_probs: torch.Tensor,
-    advantages: torch.Tensor,
-    clip_range: float = 0.2,
-    use_ratio_norm: bool = True,
-    max_importance_ratio: float = 10.0
-) -> tuple[torch.Tensor, dict[str, Any]]:
-    """
-    Compute GRPO policy loss with importance sampling and clipping.
-
-    This implements the core GRPO objective with safety mechanisms from GRPO-Guard:
-    - Importance ratio clipping (PPO-style)
-    - RatioNorm correction (GRPO-Guard)
-    - Ratio clamping for extreme values
-
-    Args:
-        log_probs: Log probabilities from current policy [B]
-        old_log_probs: Log probabilities from old policy [B]
-        advantages: Advantages [B]
-        clip_range: Clipping range for importance ratios
-        use_ratio_norm: Apply RatioNorm correction (GRPO-Guard)
-        max_importance_ratio: Maximum importance ratio before clamping
-
-    Returns:
-        loss: Policy loss (scalar)
-        info: Dictionary with diagnostic information
-
-    Reference:
-        - PPO: Schulman et al. "Proximal Policy Optimization Algorithms"
-        - GRPO-Guard: RatioNorm and gradient reweighting
-    """
-    # Compute importance ratio: r_t = pi_new(a|s) / pi_old(a|s)
-    log_ratio = log_probs - old_log_probs
-    ratio = torch.exp(log_ratio)
-
-    # Clamp extreme ratios for numerical stability
-    ratio = torch.clamp(ratio, 1.0 / max_importance_ratio, max_importance_ratio)
-
-    # RatioNorm correction (GRPO-Guard)
-    # Corrects bias in importance sampling when ratio >> 1
-    if use_ratio_norm:
-        ratio_mean = ratio.mean()
-        ratio = ratio / (ratio_mean + 1e-8)
-
-    # Clipped surrogate objective
-    ratio_clipped = torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
-    surrogate1 = ratio * advantages
-    surrogate2 = ratio_clipped * advantages
-    policy_loss = -torch.min(surrogate1, surrogate2).mean()
-
-    # Compute diagnostics
-    with torch.no_grad():
-        # Clip fraction: how often ratios were clipped
-        clip_fraction = ((ratio < 1.0 - clip_range) | (ratio > 1.0 + clip_range)).float().mean()
-
-        # KL divergence (approximate)
-        kl_div = log_ratio.mean()
-
-        # Importance ratio stats
-        importance_ratio_mean = ratio.mean()
-        importance_ratio_std = ratio.std()
-
-    info = {
-        "policy_loss": policy_loss.item(),
-        "clip_fraction": clip_fraction.item(),
-        "kl_divergence": kl_div.item(),
-        "importance_ratio_mean": importance_ratio_mean.item(),
-        "importance_ratio_std": importance_ratio_std.item(),
-    }
-
-    return policy_loss, info
-
-
-def compute_value_loss(
-    values: torch.Tensor,
-    returns: torch.Tensor,
-    old_values: torch.Tensor | None = None,
-    clip_range: float = 0.2,
-    use_clipping: bool = True
-) -> tuple[torch.Tensor, dict[str, Any]]:
-    """
-    Compute value function loss with optional clipping.
-
-    Args:
-        values: Value predictions from current model [B]
-        returns: Target returns (from GAE) [B]
-        old_values: Value predictions from old model [B] (for clipping)
-        clip_range: Clipping range for value updates
-        use_clipping: Whether to use clipped value loss (PPO-style)
-
-    Returns:
-        loss: Value loss (scalar)
-        info: Dictionary with diagnostic information
-    """
-    # Standard MSE loss
-    value_loss_unclipped = F.mse_loss(values, returns, reduction="none")
-
-    # Clipped value loss (PPO-style)
-    if use_clipping and old_values is not None:
-        values_clipped = old_values + torch.clamp(
-            values - old_values,
-            -clip_range,
-            clip_range
-        )
-        value_loss_clipped = F.mse_loss(values_clipped, returns, reduction="none")
-        value_loss = torch.max(value_loss_unclipped, value_loss_clipped).mean()
-    else:
-        value_loss = value_loss_unclipped.mean()
-
-    # Compute diagnostics
-    with torch.no_grad():
-        explained_variance = 1.0 - (returns - values).var() / (returns.var() + 1e-8)
-
-    info = {
-        "value_loss": value_loss.item(),
-        "explained_variance": explained_variance.item(),
-        "value_mean": values.mean().item(),
-        "value_std": values.std().item(),
-    }
-
-    return value_loss, info
-
-
-def compute_policy_entropy(log_probs: torch.Tensor) -> torch.Tensor:
-    """
-    Compute policy entropy for exploration bonus.
-
-    Args:
-        log_probs: Log probabilities [B]
-
-    Returns:
-        entropy: Mean entropy across batch (scalar)
-    """
-    # For continuous actions: H = -log_prob (assuming Gaussian)
-    # For discrete: H = -sum(p * log(p))
-    # Here we use a simple approximation
-    entropy = -log_probs.mean()
-    return entropy
-
-
-def apply_gradient_reweighting(
-    gradients: torch.Tensor,
-    timesteps: torch.Tensor,
-    num_train_timesteps: int = 1000
-) -> torch.Tensor:
-    """
-    Apply GRPO-Guard gradient reweighting across denoising steps.
-
-    This reweights gradients based on the timestep to balance learning
-    across different noise levels.
-
-    Args:
-        gradients: Gradients to reweight [B, ...]
-        timesteps: Timesteps at which gradients were computed [B]
-        num_train_timesteps: Total number of training timesteps
-
-    Returns:
-        reweighted_gradients: Reweighted gradients [B, ...]
-    """
-    # Compute timestep weights (higher weight for later timesteps)
-    # This is a simple linear weighting, can be made more sophisticated
-    timestep_weights = 1.0 + (timesteps.float() / num_train_timesteps)
-    timestep_weights = timestep_weights.view(-1, *([1] * (gradients.dim() - 1)))
-
-    return gradients * timestep_weights
-
-
-def sample_random_timesteps(
-    batch_size: int,
-    min_timestep: int,
-    max_timestep: int,
-    device: torch.device,
-    generator: torch.Generator | None = None
-) -> torch.Tensor:
-    """
-    Sample random timesteps for noise injection (Flow-GRPO-Fast).
-
-    Args:
-        batch_size: Number of samples
-        min_timestep: Minimum timestep
-        max_timestep: Maximum timestep
-        device: Device for tensor
-        generator: Random generator for reproducibility
-
-    Returns:
-        timesteps: Random timesteps [B]
-    """
-    if generator is not None:
-        timesteps = torch.randint(
-            min_timestep,
-            max_timestep + 1,
-            (batch_size,),
-            device=device,
-            generator=generator
-        )
-    else:
-        timesteps = torch.randint(
-            min_timestep,
-            max_timestep + 1,
-            (batch_size,),
-            device=device
-        )
-
-    return timesteps
-
-
-def compute_reward_statistics(
-    rewards: torch.Tensor
-) -> dict[str, float]:
-    """
-    Compute statistics for reward distribution.
-
-    Args:
-        rewards: Reward values [B]
-
-    Returns:
-        stats: Dictionary with mean, std, min, max
-    """
-    return {
-        "reward_mean": rewards.mean().item(),
-        "reward_std": rewards.std().item(),
-        "reward_min": rewards.min().item(),
-        "reward_max": rewards.max().item(),
-    }
-
-
-def check_early_stopping(
-    kl_divergence: float,
-    target_kl: float
-) -> bool:
-    """
-    Check if training should stop early based on KL divergence.
-
-    Args:
-        kl_divergence: Current KL divergence
-        target_kl: Target KL threshold
-
-    Returns:
-        should_stop: True if KL exceeds target
-    """
-    if kl_divergence > target_kl:
-        logger.warning(
-            "Early stopping triggered: KL divergence %.4f > target %.4f",
-            kl_divergence,
-            target_kl
-        )
-        return True
-    return False
-
-
-def compute_log_probs_from_model_output(
-    model_output: torch.Tensor,
-    target: torch.Tensor,
-    noise_level: float = 0.1
-) -> torch.Tensor:
-    """
-    Compute log probabilities from model predictions.
-
-    For diffusion models, we approximate log probabilities using the
-    negative squared error (assuming Gaussian likelihood).
-
-    Args:
-        model_output: Model predictions [B, C, T, H, W]
-        target: Target values [B, C, T, H, W]
-        noise_level: Assumed noise level (std) for Gaussian likelihood
-
-    Returns:
-        log_probs: Log probabilities [B]
-    """
-    # Compute mean squared error per sample
-    mse = ((model_output - target) ** 2).flatten(1).mean(dim=1)
-
-    # Log probability under Gaussian: log p(x) = -0.5 * (x - mu)^2 / sigma^2 + const
-    log_probs = -0.5 * mse / (noise_level ** 2)
-
-    return log_probs
-
-
-def check_for_nan_inf(tensor: torch.Tensor, name: str) -> None:
-    """
-    Check tensor for NaN or Inf values and raise error if found.
-
-    Args:
-        tensor: Tensor to check
-        name: Name for error message
-    """
-    if torch.isnan(tensor).any():
-        raise ValueError(f"{name} contains NaN values")
-    if torch.isinf(tensor).any():
-        raise ValueError(f"{name} contains Inf values")

fastvideo/training/rl/stat_tracking.py

@@ -1,189 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Per-prompt statistics tracking for GRPO training.
-
-This module ports the PerPromptStatTracker from FlowGRPO to FastVideo.
-It tracks reward statistics per unique prompt and computes normalized advantages.
-
-Ported from:
-- flow_grpo/flow_grpo/stat_tracking.py
-
-Key adaptations:
-1. Uses FastVideo's logging instead of print statements
-2. Works with single GPU (no distributed logic)
-3. Supports numpy arrays and torch tensors
-"""
-
-import numpy as np
-from typing import Union
-
-import torch
-
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class PerPromptStatTracker:
-    """
-    Tracks reward statistics per unique prompt for advantage normalization.
-    
-    This class maintains running statistics (mean, std) for each unique prompt
-    and computes normalized advantages using either per-prompt or global statistics.
-    
-    Used in GRPO training to normalize advantages within groups of samples
-    generated from the same prompt, which helps stabilize training when different
-    prompts have different reward scales.
-    """
-
-    def __init__(self, global_std: bool = False):
-        """
-        Initialize the per-prompt stat tracker.
-        
-        Args:
-            global_std: If True, use global std across all rewards for normalization.
-                       If False, use per-prompt std (default, recommended for GRPO).
-        """
-        self.global_std = global_std
-        self.stats: dict[str, list] = {}  # Maps prompt -> list of rewards
-        self.history_prompts: set[int] = set()  # Set of hashed prompts seen
-
-    def update(
-        self,
-        prompts: Union[list[str], np.ndarray],
-        rewards: Union[list[float], np.ndarray, torch.Tensor],
-        type: str = 'grpo'
-    ) -> np.ndarray:
-        """
-        Update statistics and compute normalized advantages.
-        
-        Args:
-            prompts: List or array of prompt strings (one per sample)
-            rewards: Array or tensor of reward values (one per sample)
-            type: Advantage computation type:
-                - 'grpo': Normalize by (reward - mean) / std (default)
-                - 'rwr': Return rewards as-is (reward-weighted regression)
-                - 'sft': Binary advantages (1 for max, 0 otherwise)
-                - 'dpo': DPO-style advantages (1 for max, -1 for min)
-        
-        Returns:
-            advantages: Normalized advantages array [num_samples] or [num_samples, ...]
-                       Shape matches rewards shape
-        """
-        # Convert to numpy arrays
-        prompts = np.array(prompts)
-        if isinstance(rewards, torch.Tensor):
-            rewards = rewards.detach().cpu().numpy()
-        rewards = np.array(rewards, dtype=np.float64)
-        
-        # Ensure rewards are 1D (one reward per sample)
-        # FlowGRPO expects rewards to be aggregated per sample
-        if rewards.ndim > 1:
-            # If multi-dimensional, flatten or take mean
-            # For [B, num_steps] shape, we typically want one reward per sample
-            # So we take the mean across timesteps
-            if rewards.ndim == 2:
-                # Assume shape is [B, num_steps] - take mean across timesteps
-                rewards = rewards.mean(axis=1)
-            else:
-                # Flatten and take mean for higher dimensions
-                rewards = rewards.reshape(len(prompts), -1).mean(axis=1)
-        
-        # Ensure prompts and rewards have matching lengths
-        assert len(prompts) == len(rewards), \
-            f"Prompts ({len(prompts)}) and rewards ({len(rewards)}) must have same length"
-        
-        unique_prompts = np.unique(prompts)
-        advantages = np.zeros_like(rewards, dtype=np.float64)
-        
-        # First pass: collect rewards for each prompt
-        for prompt in unique_prompts:
-            prompt_mask = prompts == prompt
-            prompt_rewards = rewards[prompt_mask]
-            
-            # Store rewards in stats
-            if prompt not in self.stats:
-                self.stats[prompt] = []
-            self.stats[prompt].extend(prompt_rewards.tolist())
-            self.history_prompts.add(hash(prompt))
-        
-        # Second pass: compute statistics and advantages
-        for prompt in unique_prompts:
-            prompt_mask = prompts == prompt
-            prompt_rewards = rewards[prompt_mask]
-            
-            # Stack all historical rewards for this prompt
-            if len(self.stats[prompt]) > 0:
-                all_prompt_rewards = np.array(self.stats[prompt])
-            else:
-                all_prompt_rewards = prompt_rewards
-            
-            # Compute mean and std
-            mean = np.mean(all_prompt_rewards, axis=0, keepdims=True)
-            
-            if self.global_std:
-                # Use global std across all rewards
-                std = np.std(rewards, axis=0, keepdims=True) + 1e-4
-            else:
-                # Use per-prompt std
-                std = np.std(all_prompt_rewards, axis=0, keepdims=True) + 1e-4
-            
-            # Compute advantages based on type
-            if type == 'grpo':
-                # GRPO: normalize by (reward - mean) / std
-                advantages[prompt_mask] = (prompt_rewards - mean) / std
-            elif type == 'rwr':
-                # Reward-weighted regression: use rewards as-is
-                advantages[prompt_mask] = prompt_rewards
-            elif type == 'sft':
-                # Supervised fine-tuning: binary (1 for max, 0 otherwise)
-                max_reward = np.max(prompt_rewards)
-                advantages[prompt_mask] = (prompt_rewards == max_reward).astype(np.float64)
-            elif type == 'dpo':
-                # DPO-style: 1 for max, -1 for min
-                prompt_rewards_tensor = torch.tensor(prompt_rewards)
-                max_idx = torch.argmax(prompt_rewards_tensor)
-                min_idx = torch.argmin(prompt_rewards_tensor)
-                
-                # If all rewards are the same, use first two indices
-                if max_idx == min_idx:
-                    min_idx = torch.tensor(0)
-                    max_idx = torch.tensor(1) if len(prompt_rewards_tensor) > 1 else torch.tensor(0)
-                
-                result = torch.zeros_like(prompt_rewards_tensor, dtype=torch.float64)
-                result[max_idx] = 1.0
-                result[min_idx] = -1.0
-                advantages[prompt_mask] = result.numpy()
-            else:
-                raise ValueError(f"Unknown advantage type: {type}. Must be one of: 'grpo', 'rwr', 'sft', 'dpo'")
-        
-        return advantages
-
-    def get_stats(self) -> tuple[float, int]:
-        """
-        Get statistics about tracked prompts.
-        
-        Returns:
-            avg_group_size: Average number of samples per unique prompt
-            history_prompts: Number of unique prompts seen (across all updates)
-        """
-        if not self.stats:
-            avg_group_size = 0.0
-        else:
-            total_samples = sum(len(v) for v in self.stats.values())
-            avg_group_size = total_samples / len(self.stats)
-        
-        history_prompts = len(self.history_prompts)
-        
-        return avg_group_size, history_prompts
-
-    def clear(self) -> None:
-        """
-        Clear all statistics (but keep history_prompts for tracking).
-        
-        This is typically called after each epoch to reset per-epoch statistics
-        while maintaining a record of all prompts seen during training.
-        """
-        self.stats = {}
-        logger.debug("Cleared per-prompt statistics (kept %d unique prompts in history)", 
-                    len(self.history_prompts))

fastvideo/training/rl/wan_grpo_utils.py

@@ -1,877 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-GRPO utilities for Wan model in FastVideo.
-
-This module ports the SDE step and pipeline functions from FlowGRPO to work with
-FastVideo's scheduler and pipeline interfaces.
-
-Ported from:
-- flow_grpo/flow_grpo/diffusers_patch/wan_pipeline_with_logprob.py
-
-Key adaptations:
-1. Uses FastVideo's FlowUniPCMultistepScheduler instead of diffusers' UniPCMultistepScheduler
-2. Works with FastVideo's WanPipeline (ComposedPipelineBase) instead of diffusers' WanPipeline
-3. Direct module access via pipeline.get_module() instead of pipeline attributes
-4. Simplified prompt encoding (direct text encoder usage instead of pipeline stages)
-"""
-
-import math
-import time
-from typing import Any
-
-import torch
-from tqdm import tqdm
-from diffusers.utils.torch_utils import randn_tensor
-
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
-    FlowUniPCMultistepScheduler)
-from fastvideo.utils import get_compute_dtype
-
-# for test_wan_transformer2
-import os
-from diffusers import WanTransformer3DModel
-from fastvideo.utils import maybe_download_model
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.configs.models.dits import WanVideoConfig
-from fastvideo.models.loader.component_loader import TransformerLoader
-
-logger = init_logger(__name__)
-
-
-def test_wan_transformer():
-    BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-    MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                      local_dir=os.path.join(
-                                          'data', BASE_MODEL_PATH))
-    TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(
-                             dit_config=WanVideoConfig(),
-                             dit_precision=precision_str))
-    args.device = device
-
-    loader = TransformerLoader()
-    model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
-
-    model1 = WanTransformer3DModel.from_pretrained(
-        TRANSFORMER_PATH, device=device,
-        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
-
-    total_params = sum(p.numel() for p in model1.parameters())
-    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
-    weight_sum_model1 = sum(
-        p.to(torch.float64).sum().item() for p in model1.parameters())
-    # Also calculate mean for more stable comparison
-    weight_mean_model1 = weight_sum_model1 / total_params
-    logger.info("Model 1 weight sum: %s", weight_sum_model1)
-    logger.info("Model 1 weight mean: %s", weight_mean_model1)
-
-    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
-    total_params_model2 = sum(p.numel() for p in model2.parameters())
-    weight_sum_model2 = sum(
-        p.to(torch.float64).sum().item() for p in model2.parameters())
-    # Also calculate mean for more stable comparison
-    weight_mean_model2 = weight_sum_model2 / total_params_model2
-    logger.info("Model 2 weight sum: %s", weight_sum_model2)
-    logger.info("Model 2 weight mean: %s", weight_mean_model2)
-
-    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
-    logger.info("Weight sum difference: %s", weight_sum_diff)
-    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
-    logger.info("Weight mean difference: %s", weight_mean_diff)
-
-    # Set both models to eval mode
-    model1 = model1.eval()
-    model2 = model2.eval()
-
-    # Create identical inputs for both models
-    batch_size = 1
-    seq_len = 30
-
-    # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size,
-                                16,
-                                21,
-                                160,
-                                90,
-                                device=device,
-                                dtype=precision)
-
-    # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size,
-                                        seq_len + 1,
-                                        4096,
-                                        device=device,
-                                        dtype=precision)
-
-    # Timestep
-    timestep = torch.tensor([500], device=device, dtype=precision)
-
-    forward_batch = ForwardBatch(data_type="dummy", )
-
-    with torch.amp.autocast('cuda', dtype=precision):
-        output1 = model1(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            timestep=timestep,
-            return_dict=False,
-        )[0]
-        with set_forward_context(
-                current_timestep=0,
-                attn_metadata=None,
-                forward_batch=forward_batch,
-        ):
-            output2 = model2(hidden_states=hidden_states,
-                             encoder_hidden_states=encoder_hidden_states,
-                             timestep=timestep)
-
-    # Check if outputs have the same shape
-    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
-    assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
-
-    # Check if outputs are similar (allowing for small numerical differences)
-    max_diff = torch.max(torch.abs(output1 - output2))
-    mean_diff = torch.mean(torch.abs(output1 - output2))
-    logger.info("Max Diff: %s", max_diff.item())
-    logger.info("Mean Diff: %s", mean_diff.item())
-    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
-    # mean diff
-    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
-    '''
-    INFO 01-19 22:53:46 [wan_grpo_utils.py:74] Model 1 weight sum: 395834.3506456231████                                                    | 1/2 [00:00<00:00,  7.84it/s]
-    INFO 01-19 22:53:46 [wan_grpo_utils.py:75] Model 1 weight mean: 0.0002789536598289884
-    INFO 01-19 22:53:47 [wan_grpo_utils.py:83] Model 2 weight sum: 395834.3506456231
-    INFO 01-19 22:53:47 [wan_grpo_utils.py:84] Model 2 weight mean: 0.0002789536598289884
-    INFO 01-19 22:53:47 [wan_grpo_utils.py:87] Weight sum difference: 0.0
-    INFO 01-19 22:53:47 [wan_grpo_utils.py:89] Weight mean difference: 0.0
-    INFO 01-19 22:53:54 [wan_grpo_utils.py:145] Max Diff: 0.08203125
-    INFO 01-19 22:53:54 [wan_grpo_utils.py:146] Mean Diff: 0.01129150390625
-    '''
-
-
-def test_wan_transformer2(model2):
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-
-    logger.info("loading model1 transformer weight")
-    BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-    MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                      local_dir=os.path.join(
-                                          'data', BASE_MODEL_PATH))
-    TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-    model1 = WanTransformer3DModel.from_pretrained(
-        TRANSFORMER_PATH,
-        device=device,
-        torch_dtype=precision,
-    ).to(device, dtype=precision).requires_grad_(False)
-
-    total_params = sum(p.numel() for p in model1.parameters())
-    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
-    weight_sum_model1 = sum(
-        p.to(torch.float64).sum().item() for p in model1.parameters())
-    # Also calculate mean for more stable comparison
-    weight_mean_model1 = weight_sum_model1 / total_params
-    logger.info("Model 1 weight sum: %s", weight_sum_model1)
-    logger.info("Model 1 weight mean: %s", weight_mean_model1)
-
-    # Calculate weight sum for model2 (converting to float64 to avoid overflow)
-    total_params_model2 = sum(p.numel() for p in model2.parameters())
-    weight_sum_model2 = sum(
-        p.to(torch.float64).sum().item() for p in model2.parameters())
-    # Also calculate mean for more stable comparison
-    weight_mean_model2 = weight_sum_model2 / total_params_model2
-    logger.info("Model 2 weight sum: %s", weight_sum_model2)
-    logger.info("Model 2 weight mean: %s", weight_mean_model2)
-
-    weight_sum_diff = abs(weight_sum_model1 - weight_sum_model2)
-    logger.info("Weight sum difference: %s", weight_sum_diff)
-    weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2)
-    logger.info("Weight mean difference: %s", weight_mean_diff)
-
-    # Set both models to eval mode
-    model1 = model1.eval()
-    model2 = model2.eval()
-
-    # Create identical inputs for both models
-    batch_size = 1
-    seq_len = 30
-
-    # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(
-        batch_size,
-        16,
-        21,
-        160,
-        90,
-        device=device,
-        dtype=precision,
-    )
-
-    # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(
-        batch_size,
-        seq_len + 1,
-        4096,
-        device=device,
-        dtype=precision,
-    )
-
-    # Timestep
-    timestep = torch.tensor([500], device=device, dtype=precision)
-
-    forward_batch = ForwardBatch(data_type="dummy", )
-
-    with torch.amp.autocast("cuda", dtype=precision):
-        output1 = model1(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            timestep=timestep,
-            return_dict=False,
-        )[0]
-        with set_forward_context(
-                current_timestep=0,
-                attn_metadata=None,
-                forward_batch=forward_batch,
-        ):
-            output2 = model2(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                timestep=timestep,
-            )
-
-    # Print basic stats for debugging (cast to float32 for stability)
-    out1 = output1.detach().float()
-    out2 = output2.detach().float()
-    logger.info(
-        "output1 stats: min=%s max=%s mean=%s std=%s",
-        out1.min().item(),
-        out1.max().item(),
-        out1.mean().item(),
-        out1.std(unbiased=False).item(),
-    )
-    logger.info(
-        "output2 stats: min=%s max=%s mean=%s std=%s",
-        out2.min().item(),
-        out2.max().item(),
-        out2.mean().item(),
-        out2.std(unbiased=False).item(),
-    )
-
-    # Check if outputs have the same shape
-    assert (output1.shape == output2.shape
-            ), f"Output shapes don't match: {output1.shape} vs {output2.shape}"
-    assert (output1.dtype == output2.dtype
-            ), f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
-
-    # Check if outputs are similar (allowing for small numerical differences)
-    max_diff = torch.max(torch.abs(output1 - output2))
-    mean_diff = torch.mean(torch.abs(output1 - output2))
-    logger.info("Max Diff: %s", max_diff.item())
-    logger.info("Mean Diff: %s", mean_diff.item())
-    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
-    # mean diff
-    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
-    '''
-    when --dit_precision "bf16", use_fsdp hardcoded to False:
-    INFO 01-19 22:01:24 [wan_grpo_utils.py:65] Model 1 weight sum: 395834.3506456231████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  3.25it/s]
-    INFO 01-19 22:01:24 [wan_grpo_utils.py:66] Model 1 weight mean: 0.0002789536598289884
-    INFO 01-19 22:01:24 [wan_grpo_utils.py:75] Model 2 weight sum: 395125.463677882
-    INFO 01-19 22:01:24 [wan_grpo_utils.py:76] Model 2 weight mean: 0.0002739000890162162
-    INFO 01-19 22:01:24 [wan_grpo_utils.py:79] Weight sum difference: 708.8869677411276
-    INFO 01-19 22:01:24 [wan_grpo_utils.py:81] Weight mean difference: 5.053570812772192e-06
-    INFO 01-19 22:01:32 [wan_grpo_utils.py:139] output1 stats: min=-2.28125 max=1.921875 mean=-0.16638492047786713 std=0.458170622587204
-    INFO 01-19 22:01:32 [wan_grpo_utils.py:146] output2 stats: min=-2.296875 max=1.90625 mean=-0.166452556848526 std=0.4579130709171295
-    INFO 01-19 22:01:32 [wan_grpo_utils.py:165] Max Diff: 0.08984375
-    INFO 01-19 22:01:32 [wan_grpo_utils.py:166] Mean Diff: 0.0120849609375
-
-    when --dit_precision "fp32", use_fsdp not changed:
-
-    '''
-
-
-def sde_step_with_logprob(
-    scheduler: FlowUniPCMultistepScheduler,
-    model_output: torch.FloatTensor,
-    timestep: float | torch.FloatTensor,
-    sample: torch.FloatTensor,
-    prev_sample: torch.FloatTensor | None = None,
-    generator: torch.Generator | None = None,
-    deterministic: bool = False,
-    return_pixel_log_prob: bool = False,
-    return_dt_and_std_dev_t: bool = False
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, ...]:
-    """
-    Predict the sample from the previous timestep by reversing the SDE. 
-    This function propagates the flow process from the learned model outputs 
-    (most often the predicted velocity) and computes log probabilities.
-
-    Ported from FlowGRPO's sde_step_with_logprob to work with FastVideo's
-    FlowUniPCMultistepScheduler.
-
-    Args:
-        scheduler: FastVideo FlowUniPCMultistepScheduler instance
-        model_output: The direct output from learned flow model
-        timestep: The current discrete timestep in the diffusion chain
-        sample: A current instance of a sample created by the diffusion process
-        prev_sample: Optional previous sample (if provided, used instead of sampling)
-        generator: Optional random number generator
-        deterministic: If True, no noise is added (deterministic sampling)
-        return_pixel_log_prob: If True, return pixel-level log probabilities (not used)
-        return_dt_and_std_dev_t: If True, return dt and std_dev_t separately
-
-    Returns:
-        If return_dt_and_std_dev_t=True:
-            (prev_sample, log_prob, prev_sample_mean, std_dev_t, sqrt_dt)
-        Otherwise:
-            (prev_sample, log_prob, prev_sample_mean, std_dev_t * sqrt_dt)
-    """
-
-    # # Convert all variables to fp32 for numerical stability
-    # model_output = model_output.float()
-    # sample = sample.float()
-    # if prev_sample is not None:
-    #     prev_sample = prev_sample.float()
-
-    # Get step indices for current and previous timesteps
-    # Handle both single timestep and batch of timesteps
-    if isinstance(timestep, torch.Tensor):
-        if timestep.ndim == 0:
-            timestep = timestep.unsqueeze(0)
-        step_indices = [
-            scheduler.index_for_timestep(t.item()) for t in timestep
-        ]
-    else:
-        step_indices = [scheduler.index_for_timestep(timestep)]
-
-    prev_step_indices = [step + 1 for step in step_indices]
-
-    # Move sigmas to sample device
-    sigmas = scheduler.sigmas.to(sample.device)
-    # myregion debug: hardcode sigmas to flow_grpo's
-    sigmas = torch.Tensor([
-        0.9997, 0.9824, 0.9639, 0.9441, 0.9227, 0.8996, 0.8746, 0.8475, 0.8178,
-        0.7853, 0.7496, 0.7102, 0.6663, 0.6173, 0.5621, 0.4997, 0.4283, 0.3459,
-        0.2498, 0.1362, 0.0000
-    ]).to(sample.device, sample.dtype)
-    # end region
-
-    # Get sigma values for current and previous steps
-    sigma = sigmas[step_indices].view(-1, 1, 1, 1, 1)
-    sigma_prev = sigmas[prev_step_indices].view(-1, 1, 1, 1, 1)
-    sigma_max = sigmas[0].item()  # First sigma (highest)
-    sigma_min = sigmas[-1].item()  # Last sigma (lowest)
-
-    dt = sigma_prev - sigma
-
-    # myregion debug
-    print(f"[DEBUG]: sigma_max: {sigma_max}, sigma_min: {sigma_min}, dt: {dt}")
-    print(f"[DEBUG]: in sde_step_with_logprob(), timestep: {timestep}")
-    print(f"[DEBUG]: in sde_step_with_logprob(), sigmas: {sigmas}")
-    print(f"[DEBUG]: in sde_step_with_logprob(), step_indices: {step_indices}")
-    print(
-        f"[DEBUG]: in sde_step_with_logprob(), prev_step_indices: {prev_step_indices}"
-    )
-    '''
-    [DEBUG]: in sde_step_with_logprob(), timestep: tensor([428, 428, 428, 428], device='cuda:0')
-    [DEBUG]: in sde_step_with_logprob(), sigmas: tensor([0.9999, 0.9826, 0.9642, 0.9443, 0.9230, 0.8999, 0.8749, 0.8477, 0.8181,
-            0.7856, 0.7499, 0.7104, 0.6665, 0.6175, 0.5624, 0.4999, 0.4285, 0.3461,
-            0.2499, 0.1363, 0.0000], device='cuda:0')
-    [DEBUG]: in sde_step_with_logprob(), step_indices: [16, 16, 16, 16]
-    [DEBUG]: in sde_step_with_logprob(), prev_step_indices: [17, 17, 17, 17]
-
-    DEBUG]: in sde_step_with_logprob(), timestep: tensor([249], device='cuda:0')███████████████▎       | 18/20 [00:04<00:00,  3.87step/s, step_time=0.26s, timestep=346.0]
-    [DEBUG]: in sde_step_with_logprob(), sigmas: tensor([0.9999, 0.9826, 0.9642, 0.9443, 0.9230, 0.8999, 0.8749, 0.8477, 0.8181,
-            0.7856, 0.7499, 0.7104, 0.6665, 0.6175, 0.5624, 0.4999, 0.4285, 0.3461,
-            0.2499, 0.1363, 0.0000], device='cuda:0')
-    [DEBUG]: in sde_step_with_logprob(), step_indices: [18]
-    [DEBUG]: in sde_step_with_logprob(), prev_step_indices: [19]
-
-    [DEBUG]: in sde_step_with_logprob(), timestep: tensor([617, 617, 617, 617], device='cuda:0')
-    [DEBUG]: in sde_step_with_logprob(), sigmas: tensor([0.9999, 0.9826, 0.9642, 0.9443, 0.9230, 0.8999, 0.8749, 0.8477, 0.8181,
-            0.7856, 0.7499, 0.7104, 0.6665, 0.6175, 0.5624, 0.4999, 0.4285, 0.3461,
-            0.2499, 0.1363, 0.0000], device='cuda:0')
-    [DEBUG]: in sde_step_with_logprob(), step_indices: [13, 13, 13, 13]
-    [DEBUG]: in sde_step_with_logprob(), prev_step_indices: [14, 14, 14, 14]
-    '''
-    # endregion
-
-    # Compute std_dev_t and prev_sample_mean using SDE formulation
-    std_dev_t = sigma_min + (sigma_max - sigma_min) * sigma
-    prev_sample_mean = (sample * (1 + std_dev_t**2 / (2 * sigma) * dt) +
-                        model_output * (1 + std_dev_t**2 * (1 - sigma) /
-                                        (2 * sigma)) * dt)
-
-    if prev_sample is not None and generator is not None:
-        raise ValueError(
-            "Cannot pass both generator and prev_sample. Please make sure that either `generator` or"
-            " `prev_sample` stays `None`.")
-
-    # Sample prev_sample if not provided
-    if prev_sample is None:
-        variance_noise = randn_tensor(
-            model_output.shape,
-            generator=generator,
-            device=model_output.device,
-            dtype=model_output.dtype,
-        )
-        sqrt_dt = torch.sqrt(-1 * dt)  # dt is negative (going backwards)
-        prev_sample = prev_sample_mean + std_dev_t * sqrt_dt * variance_noise
-    else:
-        sqrt_dt = torch.sqrt(-1 * dt)
-
-    # No noise is added during evaluation (deterministic)
-    if deterministic:
-        prev_sample = sample + dt * model_output
-        sqrt_dt = torch.sqrt(-1 * dt)
-
-    # Compute log probability: log p(prev_sample | sample, model_output)
-    # Assuming Gaussian distribution: N(prev_sample_mean, (std_dev_t * sqrt_dt)^2)
-    std_dev_sqrt_dt = std_dev_t * sqrt_dt
-    log_prob = (
-        -((prev_sample.detach() - prev_sample_mean)**2) /
-        (2 * (std_dev_sqrt_dt**2)) - torch.log(
-            std_dev_sqrt_dt + 1e-8)  # Add small epsilon for numerical stability
-        - torch.log(
-            torch.sqrt(2 * torch.as_tensor(math.pi, device=sample.device))))
-
-    # Mean along all but batch dimension
-    log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
-
-    if return_dt_and_std_dev_t:
-        return prev_sample, log_prob, prev_sample_mean, std_dev_t, sqrt_dt
-    return prev_sample, log_prob, prev_sample_mean, std_dev_t * sqrt_dt
-
-
-def wan_pipeline_with_logprob(
-    pipeline,
-    prompt: str | list[str] = None,
-    negative_prompt: str | list[str] = None,
-    height: int = 480,
-    width: int = 832,
-    num_frames: int = 81,
-    num_inference_steps: int = 50,
-    guidance_scale: float = 5.0,
-    num_videos_per_prompt: int | None = 1,
-    generator: torch.Generator | list[torch.Generator] | None = None,
-    latents: torch.Tensor | None = None,
-    prompt_embeds: torch.Tensor | None = None,
-    negative_prompt_embeds: torch.Tensor | None = None,
-    output_type: str | None = "pt",
-    return_dict: bool = False,
-    attention_kwargs: dict[str, Any] | None = None,
-    max_sequence_length: int = 512,
-    deterministic: bool = False,
-    kl_reward: float = 0.0,
-    return_pixel_log_prob: bool = False,
-) -> tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor],
-           list[torch.Tensor], torch.Tensor | None]:
-    """
-    Wan pipeline with log probability computation for GRPO training.
-
-    Ported from FlowGRPO's wan_pipeline_with_logprob to work with FastVideo's WanPipeline.
-    This function generates videos and computes log probabilities at each denoising step.
-
-    Args:
-        pipeline: FastVideo WanPipeline instance
-        prompt: Text prompt(s) for generation
-        negative_prompt: Negative prompt(s) for classifier-free guidance
-        height: Height of generated video
-        width: Width of generated video
-        num_frames: Number of frames in generated video
-        num_inference_steps: Number of denoising steps
-        guidance_scale: Classifier-free guidance scale
-        num_videos_per_prompt: Number of videos to generate per prompt
-        generator: Random generator for reproducibility
-        latents: Optional initial latents
-        prompt_embeds: Optional pre-computed prompt embeddings
-        negative_prompt_embeds: Optional pre-computed negative prompt embeddings
-        output_type: Output type ("pt" for PyTorch tensor, "np" for numpy, "latent" for latents only)
-        return_dict: Whether to return dict (not used, always returns tuple)
-        attention_kwargs: Optional attention kwargs
-        max_sequence_length: Maximum sequence length for text encoding
-        deterministic: If True, use deterministic sampling (no noise)
-        kl_reward: KL reward coefficient (if > 0, computes KL divergence)
-        return_pixel_log_prob: If True, return pixel-level log probabilities (not used)
-
-    Returns:
-        Tuple of:
-            - video: Generated video tensor [B, C, T, H, W] or latents if output_type="latent"
-            - all_latents: List of latents at each step [num_steps+1] of shape [B, C, T, H, W]
-            - all_log_probs: List of log probabilities at each step [num_steps] of shape [B]
-            - all_kl: List of KL divergences at each step [num_steps] of shape [B] (if kl_reward > 0)
-            - prompt_ids: Tokenized prompt IDs [B, seq_len] (None if prompt_embeds were provided)
-    """
-    # Get device from transformer
-    transformer = pipeline.get_module("transformer")
-
-    # myregion debug: test transformer output
-    logger.info("testing transformer, running test_wan_transformer2")
-    test_wan_transformer()
-    # test_wan_transformer2(transformer)
-    # endregion
-
-    # hardcode dtype for debug
-    # transformer_dtype = torch.float32
-    # use get_compute_dtype() to get dtype based on mixed precision
-    transformer_dtype = get_compute_dtype()
-    logger.info(f"[DEBUG]: transformer_dtype: {transformer_dtype}")
-
-    # Get scheduler and other modules
-    scheduler = pipeline.get_module("scheduler")
-    vae = pipeline.get_module("vae")
-
-    # Determine batch size
-    if prompt is not None and isinstance(prompt, str):
-        batch_size = 1
-    elif prompt is not None and isinstance(prompt, list):
-        batch_size = len(prompt)
-    elif prompt_embeds is not None:
-        batch_size = prompt_embeds.shape[0]
-    else:
-        raise ValueError("Either prompt or prompt_embeds must be provided")
-
-    # Encode prompts if not provided
-    prompt_ids = None
-    if prompt_embeds is None:
-        # Encode prompts directly using text encoder and tokenizer
-        # This is a simplified encoding - for full pipeline encoding, use TextEncodingStage
-        text_encoder = pipeline.get_module("text_encoder")
-        tokenizer = pipeline.get_module("tokenizer")
-
-        # Normalize to list
-        if isinstance(prompt, str):
-            prompts_list = [prompt]
-        else:
-            prompts_list = prompt
-
-        # Tokenize prompts
-        text_inputs = tokenizer(prompts_list,
-                                padding="max_length",
-                                max_length=max_sequence_length,
-                                truncation=True,
-                                return_tensors="pt").to(pipeline.device)
-
-        # Store prompt_ids for return
-        prompt_ids = text_inputs["input_ids"]
-
-        # Encode with text encoder
-        with torch.no_grad():
-            outputs = text_encoder(
-                text_inputs["input_ids"],
-                attention_mask=text_inputs["attention_mask"],
-                output_hidden_states=True,
-            )
-            # Get last hidden state (Wan typically uses last hidden state)
-            prompt_embeds = outputs.last_hidden_state
-
-        # Encode negative prompts if CFG is enabled
-        if guidance_scale > 1.0:
-            if negative_prompt is None:
-                negative_prompt = [""] * len(prompts_list)
-            elif isinstance(negative_prompt, str):
-                negative_prompt = [negative_prompt]
-
-            neg_text_inputs = tokenizer(negative_prompt,
-                                        padding="max_length",
-                                        max_length=max_sequence_length,
-                                        truncation=True,
-                                        return_tensors="pt").to(pipeline.device)
-
-            with torch.no_grad():
-                neg_outputs = text_encoder(
-                    neg_text_inputs["input_ids"],
-                    attention_mask=neg_text_inputs["attention_mask"],
-                    output_hidden_states=True,
-                )
-                negative_prompt_embeds = neg_outputs.last_hidden_state
-        else:
-            negative_prompt_embeds = None
-
-# myregion Debug: Print shapes of prompt embeddings
-        logger.info(
-            f"After encoding - prompt_embeds shape: {prompt_embeds.shape if prompt_embeds is not None else None}"
-        )
-        logger.info(
-            f"After encoding - negative_prompt_embeds shape: {negative_prompt_embeds.shape if negative_prompt_embeds is not None else None}"
-        )
-        logger.info(
-            f"After encoding - prompt_embeds dtype: {prompt_embeds.dtype if prompt_embeds is not None else None}"
-        )
-        logger.info(
-            f"After encoding - negative_prompt_embeds dtype: {negative_prompt_embeds.dtype if negative_prompt_embeds is not None else None}"
-        )
-        '''
-        INFO 01-17 05:31:13 [wan_grpo_utils.py:290] After encoding - prompt_embeds shape: torch.Size([4, 512, 4096])
-        INFO 01-17 05:31:13 [wan_grpo_utils.py:291] After encoding - negative_prompt_embeds shape: None
-        INFO 01-17 05:31:13 [wan_grpo_utils.py:292] After encoding - prompt_embeds dtype: torch.float32
-        INFO 01-17 05:31:13 [wan_grpo_utils.py:293] After encoding - negative_prompt_embeds dtype: None
-        '''
-# endregion
-# logger.info("wan_pipeline_with_logprob's transformer class type: %s", type(transformer))
-# logger.info("Variables in transformer: %s", str(dir(transformer)))
-    prompt_embeds = prompt_embeds.to(transformer_dtype)
-    if negative_prompt_embeds is not None:
-        negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
-
-    # Prepare timesteps
-    scheduler.set_timesteps(num_inference_steps, device=pipeline.device)
-    timesteps = scheduler.timesteps
-
-    # Prepare latent variables
-    num_channels_latents = transformer.config.in_channels
-    vae = pipeline.get_module("vae")
-    # Get VAE scale factors
-    vae_scale_factor_spatial = vae.spatial_compression_ratio
-    vae_scale_factor_temporal = vae.temporal_compression_ratio
-
-    if latents is None:
-        # Generate random latents
-        # Note: num_frames in latents accounts for temporal compression
-        num_latent_frames = (num_frames - 1) // vae_scale_factor_temporal + 1
-        latents_shape = (
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            num_latent_frames,
-            height // vae_scale_factor_spatial,
-            width // vae_scale_factor_spatial,
-        )
-        if generator is not None:
-            if isinstance(generator, list):
-                latents = [
-                    torch.randn(
-                        latents_shape[1:],
-                        generator=gen,
-                        device=pipeline.device,
-                        dtype=transformer_dtype,
-                    ) for gen in generator
-                ]
-                latents = torch.stack(latents, dim=0)
-            else:
-                latents = torch.randn(
-                    latents_shape,
-                    generator=generator,
-                    device=pipeline.device,
-                    dtype=transformer_dtype,
-                )
-        else:
-            latents = torch.randn(latents_shape,
-                                  device=pipeline.device,
-                                  dtype=transformer_dtype)
-    else:
-        latents = latents.to(device=pipeline.device, dtype=transformer_dtype)
-
-
-# myregion Debug: Print latents shape, dtype, and value range
-    logger.info("=" * 80)
-    logger.info("Latents Debug Information:")
-    logger.info(f"  Shape: {latents.shape}")
-    logger.info(f"  Dtype: {latents.dtype}")
-    logger.info(f"  Min value: {latents.min().item():.6f}")
-    logger.info(f"  Max value: {latents.max().item():.6f}")
-    logger.info(f"  Mean value: {latents.mean().item():.6f}")
-    logger.info(f"  Std value: {latents.std().item():.6f}")
-    logger.info(f"  Device: {latents.device}")
-    logger.info("=" * 80)
-    '''
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:355] ================================================================================
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:356] Latents Debug Information:
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:357]   Shape: torch.Size([4, 16, 9, 30, 52])
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:358]   Dtype: torch.bfloat16
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:359]   Min value: -4.500000
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:360]   Max value: 4.656250
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:361]   Mean value: 0.000111
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:362]   Std value: 1.000000
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:363]   Device: cuda:0
-    INFO 01-17 07:41:33 [wan_grpo_utils.py:364] ================================================================================
-    '''
-    # endregion
-
-    all_latents = [latents]
-    all_log_probs = []
-    all_kl = []
-
-    # myregion Debug
-    logger.info("Tensor type issue debugging:")
-    logger.info(f"latents: {type(latents)}")
-    logger.info(f"prompt_embeds: {type(prompt_embeds)}")
-    logger.info(
-        f"[DEBUG]: before denoising loop: type(timesteps): {type(timesteps)}")
-    logger.info(
-        f"[DEBUG]: before denoising loop: timesteps.shape: {timesteps.shape}")
-    # endregion
-
-    # Progress bar for denoising loop
-    progress_bar = tqdm(enumerate(timesteps),
-                        total=len(timesteps),
-                        desc="Denoising steps",
-                        unit="step")
-
-    for i, t in progress_bar:
-        step_start_time = time.time()
-        latents_ori = latents.clone()
-        timestep = t.expand(latents.shape[0]) if isinstance(
-            t, torch.Tensor) else torch.tensor([t] * latents.shape[0],
-                                               device=pipeline.device)
-
-        logger.info(
-            f"[DEBUG]: before set_forward_context: current_timestep=i:{i}")
-        # Predict noise with transformer
-        with set_forward_context(
-                current_timestep=t.item(),
-                attn_metadata=None,
-                forward_batch=None,
-        ):
-            noise_pred = transformer(
-                hidden_states=latents,
-                timestep=timestep,
-                encoder_hidden_states=prompt_embeds,
-                attention_kwargs=attention_kwargs,
-                return_dict=False,
-            )[0]
-        noise_pred = noise_pred.to(prompt_embeds.dtype)
-
-        # Classifier-free guidance
-        if guidance_scale > 1.0:
-            with set_forward_context(
-                    current_timestep=i,
-                    attn_metadata=None,
-                    forward_batch=None,
-            ):
-                noise_uncond = transformer(
-                    hidden_states=latents,
-                    timestep=timestep,
-                    encoder_hidden_states=negative_prompt_embeds,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-            noise_pred = noise_uncond + guidance_scale * (noise_pred -
-                                                          noise_uncond)
-
-        # SDE step with log probability
-        latents, log_prob, prev_latents_mean, std_dev_t = sde_step_with_logprob(
-            scheduler,
-            noise_pred,  #.float(),
-            t.unsqueeze(0) if isinstance(t, torch.Tensor) else t,
-            latents,  #.float(),
-            deterministic=deterministic,
-            return_pixel_log_prob=return_pixel_log_prob)
-        # sde_step_with_logprob returns fp32
-        # latents = latents.to(transformer_dtype)
-        prev_latents = latents.clone()
-
-        all_latents.append(latents)
-        all_log_probs.append(log_prob)
-
-        # Compute KL divergence if kl_reward > 0 (for KL reward in sampling)
-        if kl_reward > 0 and not deterministic:
-            # Use reference model (disable adapter if using LoRA)
-            latent_model_input_ref = torch.cat(
-                [latents_ori] * 2) if guidance_scale > 1.0 else latents_ori
-            with set_forward_context(
-                    current_timestep=i,
-                    attn_metadata=None,
-                    forward_batch=None,
-            ):
-                with transformer.disable_adapter() if hasattr(
-                        transformer, 'disable_adapter') else torch.no_grad():
-                    noise_pred_ref = transformer(
-                        hidden_states=latent_model_input_ref,
-                        timestep=timestep,
-                        encoder_hidden_states=prompt_embeds,
-                        attention_kwargs=attention_kwargs,
-                        return_dict=False,
-                    )[0]
-            noise_pred_ref = noise_pred_ref.to(prompt_embeds.dtype)
-
-            # Perform guidance for reference model
-            if guidance_scale > 1.0:
-                noise_pred_uncond_ref, noise_pred_text_ref = noise_pred_ref.chunk(
-                    2)
-                noise_pred_ref = noise_pred_uncond_ref + guidance_scale * (
-                    noise_pred_text_ref - noise_pred_uncond_ref)
-
-            # Compute reference log prob
-            _, ref_log_prob, ref_prev_latents_mean, ref_std_dev_t = sde_step_with_logprob(
-                scheduler,
-                noise_pred_ref.float(),
-                t.unsqueeze(0) if isinstance(t, torch.Tensor) else t,
-                latents_ori.float(),
-                prev_sample=prev_latents.float(),
-                deterministic=deterministic,
-            )
-
-            # Compute KL divergence: KL = (mean_diff)^2 / (2 * std^2)
-            assert torch.allclose(
-                std_dev_t, ref_std_dev_t
-            ), "std_dev_t should match between current and reference"
-            kl = (prev_latents_mean - ref_prev_latents_mean)**2 / (2 *
-                                                                   std_dev_t**2)
-            kl = kl.mean(dim=tuple(range(1, kl.ndim)))
-            all_kl.append(kl)
-        else:
-            # No KL reward, set to zero
-            all_kl.append(torch.zeros(len(latents), device=latents.device))
-
-        # Update progress bar with timing information
-        step_time = time.time() - step_start_time
-        progress_bar.set_postfix({
-            "step_time":
-            f"{step_time:.2f}s",
-            "timestep":
-            f"{t.item() if isinstance(t, torch.Tensor) else t:.1f}"
-        })
-
-    # Decode latents to video if needed
-    if output_type != "latent":
-        latents = latents.to(vae.dtype)
-
-        # Apply VAE normalization (Wan VAE specific)
-        # Wan VAE requires denormalization before decoding
-        if hasattr(vae, 'config') and hasattr(vae.config,
-                                              'latents_mean') and hasattr(
-                                                  vae.config, 'latents_std'):
-            # Get z_dim from config or VAE
-            z_dim = getattr(vae.config, 'z_dim', latents.shape[1])
-            latents_mean = (torch.tensor(vae.config.latents_mean,
-                                         device=latents.device,
-                                         dtype=latents.dtype).view(
-                                             1, z_dim, 1, 1, 1))
-            latents_std = (
-                1.0 / torch.tensor(vae.config.latents_std,
-                                   device=latents.device,
-                                   dtype=latents.dtype).view(1, z_dim, 1, 1, 1))
-            latents = latents / latents_std + latents_mean
-        elif hasattr(vae, 'latents_mean') and hasattr(vae, 'latents_std'):
-            # Alternative: check if latents_mean/std are direct attributes
-            z_dim = latents.shape[1]
-            latents_mean = (torch.tensor(vae.latents_mean,
-                                         device=latents.device,
-                                         dtype=latents.dtype).view(
-                                             1, z_dim, 1, 1, 1))
-            latents_std = (1.0 / torch.tensor(
-                vae.latents_std, device=latents.device,
-                dtype=latents.dtype).view(1, z_dim, 1, 1, 1))
-            latents = latents / latents_std + latents_mean
-
-        # Decode using VAE
-        with torch.no_grad():
-            video = vae.decode(latents.float(), return_dict=False)[0]
-            # VAE.decode returns tensor directly (not tuple)
-
-            # Postprocess video: convert from [-1, 1] to [0, 1]
-            # FastVideo VAE typically outputs in [-1, 1] range
-            video = (video / 2 + 0.5).clamp(0, 1)
-    else:
-        video = latents
-
-    return video, all_latents, all_log_probs, all_kl, prompt_ids

fastvideo/training/training_pipeline.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import asdict
+import inspect
 import math
 import os
 import time
@@ -174,18 +175,25 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 last_epoch=self.init_steps - 1,
             )
 
-        if not self.training_args.rl_args.rl_mode:
-            self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
-                training_args.data_path,
-                training_args.train_batch_size,
-                parquet_schema=self.train_dataset_schema,
-                num_data_workers=training_args.dataloader_num_workers,
-                cfg_rate=training_args.training_cfg_rate,
-                drop_last=True,
-                text_padding_length=training_args.pipeline_config.
-                text_encoder_configs[0].arch_config.
-                text_len,  # type: ignore[attr-defined]
-                seed=self.seed)
+        text_padding_length = training_args.pipeline_config.text_encoder_configs[
+            0].arch_config.text_len  # type: ignore[attr-defined]
+        if not text_padding_length:
+            text_max_lengths = getattr(training_args.pipeline_config,
+                                       "text_encoder_max_lengths", None)
+            if text_max_lengths:
+                text_padding_length = text_max_lengths[0]
+            else:
+                text_padding_length = 512
+
+        self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
+            training_args.data_path,
+            training_args.train_batch_size,
+            parquet_schema=self.train_dataset_schema,
+            num_data_workers=training_args.dataloader_num_workers,
+            cfg_rate=training_args.training_cfg_rate,
+            drop_last=True,
+            text_padding_length=text_padding_length,
+            seed=self.seed)
 
         self.noise_scheduler = noise_scheduler
         if self.training_args.boundary_ratio is not None:
@@ -193,21 +201,19 @@ class TrainingPipeline(LoRAPipeline, ABC):
         else:
             self.boundary_timestep = None
 
-        if not self.training_args.rl_args.rl_mode:
-            logger.info("train_dataloader length: %s", len(self.train_dataloader))
+        logger.info("train_dataloader length: %s", len(self.train_dataloader))
         logger.info("train_sp_batch_size: %s",
                     training_args.train_sp_batch_size)
         logger.info("gradient_accumulation_steps: %s",
                     training_args.gradient_accumulation_steps)
         logger.info("sp_size: %s", training_args.sp_size)
 
-        if not self.training_args.rl_args.rl_mode:
-            self.num_update_steps_per_epoch = math.ceil(
-                len(self.train_dataloader) /
-                training_args.gradient_accumulation_steps * training_args.sp_size /
-                training_args.train_sp_batch_size)
-            self.num_train_epochs = math.ceil(training_args.max_train_steps /
-                                            self.num_update_steps_per_epoch)
+        self.num_update_steps_per_epoch = math.ceil(
+            len(self.train_dataloader) /
+            training_args.gradient_accumulation_steps * training_args.sp_size /
+            training_args.train_sp_batch_size)
+        self.num_train_epochs = math.ceil(training_args.max_train_steps /
+                                          self.num_update_steps_per_epoch)
 
         # TODO(will): is there a cleaner way to track epochs?
         self.current_epoch = 0
@@ -275,11 +281,27 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _normalize_dit_input(self,
                              training_batch: TrainingBatch) -> TrainingBatch:
         # TODO(will): support other models
+        # Automatically detect model type based on VAE class name
+        vae = self.get_module("vae")
+        vae_class_name = vae.__class__.__name__
+
+        # Map VAE class names to model types
+        if 'Hunyuan' in vae_class_name:
+            model_type = 'hunyuan'
+        elif 'Wan' in vae_class_name or 'AutoencoderKLCausal3D' in vae_class_name:
+            model_type = 'wan'
+        else:
+            # Default to checking for latents_mean attribute
+            if hasattr(vae, 'latents_mean'):
+                model_type = 'wan'
+            else:
+                model_type = 'hunyuan'
+
         with self.tracker.timed("timing/normalize_input"):
             training_batch.latents = normalize_dit_input(
-                'wan',
+                model_type,
                 training_batch.latents,
-                self.get_module("vae"),
+                vae,
             )
         return training_batch
 
@@ -431,6 +453,87 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #         device=training_batch.noisy_model_input.device,
         #         dtype=torch.bfloat16)
         current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+        if "return_dict" in input_kwargs:
+            try:
+                sig = inspect.signature(current_model.forward)
+            except (TypeError, ValueError):
+                sig = None
+            if sig is not None:
+                params = sig.parameters.values()
+                accepts_kwargs = any(param.kind == inspect.Parameter.VAR_KEYWORD
+                                     for param in params)
+                if ("return_dict" not in sig.parameters and not accepts_kwargs):
+                    input_kwargs = dict(input_kwargs)
+                    input_kwargs.pop("return_dict", None)
+        base_model = current_model
+        for attr in ("module", "_orig_mod"):
+            if hasattr(base_model, attr):
+                base_model = getattr(base_model, attr)
+        if hasattr(base_model, "_fsdp_wrapped_module"):
+            base_model = base_model._fsdp_wrapped_module
+            if hasattr(base_model, "module"):
+                base_model = base_model.module
+        model_config = getattr(base_model, "config", None)
+        if model_config is None:
+            model_config = getattr(current_model, "config", None)
+        if model_config is not None and hasattr(model_config,
+                                                "text_embed_2_dim"):
+            input_kwargs = dict(input_kwargs)
+            text_states = training_batch.encoder_hidden_states
+            text_mask = training_batch.encoder_attention_mask
+            if text_states is None or text_mask is None:
+                raise ValueError(
+                    "HunyuanVideo15 training requires text embeddings and masks"
+                )
+            if not isinstance(input_kwargs.get("encoder_hidden_states"),
+                              (list, tuple)):
+                batch_size = text_states.shape[0]
+                text_2_len = 256
+                if hasattr(self.training_args.pipeline_config,
+                           "text_encoder_max_lengths"):
+                    text_2_len = self.training_args.pipeline_config.text_encoder_max_lengths[
+                        1]
+                text_2_dim = model_config.text_embed_2_dim
+                text_2 = torch.zeros((batch_size, text_2_len, text_2_dim),
+                                     device=text_states.device,
+                                     dtype=text_states.dtype)
+                input_kwargs["encoder_hidden_states"] = [text_states, text_2]
+            if not isinstance(input_kwargs.get("encoder_attention_mask"),
+                              (list, tuple)):
+                batch_size = text_mask.shape[0]
+                text_2_len = 256
+                if hasattr(self.training_args.pipeline_config,
+                           "text_encoder_max_lengths"):
+                    text_2_len = self.training_args.pipeline_config.text_encoder_max_lengths[
+                        1]
+                text_mask_2 = torch.zeros((batch_size, text_2_len),
+                                          device=text_mask.device,
+                                          dtype=text_mask.dtype)
+                input_kwargs["encoder_attention_mask"] = [
+                    text_mask, text_mask_2
+                ]
+            if "encoder_hidden_states_image" not in input_kwargs:
+                image_embeds = training_batch.image_embeds
+                if image_embeds is None or image_embeds.numel() == 0:
+                    batch_size = text_states.shape[0]
+                    image_len = 729
+                    image_dim = model_config.image_embed_dim
+                    image_embeds = torch.zeros(
+                        (batch_size, image_len, image_dim),
+                        device=text_states.device,
+                        dtype=text_states.dtype,
+                    )
+                input_kwargs["encoder_hidden_states_image"] = [image_embeds]
+            hidden_states = input_kwargs.get("hidden_states")
+            if (isinstance(hidden_states, torch.Tensor)
+                    and hidden_states.shape[1] != model_config.in_channels):
+                batch_size, _, t, h, w = hidden_states.shape
+                video_latent = torch.zeros((batch_size, 1, t, h, w),
+                                           device=hidden_states.device,
+                                           dtype=hidden_states.dtype)
+                zeros_latent = torch.zeros_like(hidden_states)
+                input_kwargs["hidden_states"] = torch.cat(
+                    [hidden_states, video_latent, zeros_latent], dim=1)
 
         with self.tracker.timed("timing/forward_backward"), set_forward_context(
                 current_timestep=training_batch.current_timestep,
@@ -578,12 +681,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 round(num_trainable_params / 1e9, 3))
 
         # Set random seeds for deterministic training
-        if not self.training_args.rl_args.rl_mode:
-            self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
-                self.seed)
-        else:
-            self.noise_random_generator = torch.Generator(device=self.device).manual_seed(
-                self.seed)
+        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
+            self.seed)
         self.noise_gen_cuda = torch.Generator(
             device=current_platform.device_name).manual_seed(self.seed)
         self.validation_random_generator = torch.Generator(
@@ -657,7 +756,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 }
                 metrics["batch_size"] = int(training_batch.raw_latent_shape[0])
 
-                patch_t, patch_h, patch_w = self.training_args.pipeline_config.dit_config.patch_size
+                patch_size = self.training_args.pipeline_config.dit_config.patch_size
+                if isinstance(patch_size, tuple):
+                    patch_t, patch_h, patch_w = patch_size
+                else:
+                    patch_t = patch_size
+                    patch_h = patch_size
+                    patch_w = patch_size
                 seq_len = (training_batch.raw_latent_shape[2] // patch_t) * (
                     training_batch.raw_latent_shape[3] //
                     patch_h) * (training_batch.raw_latent_shape[4] // patch_w)
@@ -670,7 +775,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
                 metrics["hidden_dim"] = arch_config.hidden_size
                 metrics["num_layers"] = arch_config.num_layers
-                metrics["ffn_dim"] = arch_config.ffn_dim
+                ffn_dim = getattr(arch_config, "ffn_dim", None)
+                if ffn_dim is None:
+                    mlp_ratio = getattr(arch_config, "mlp_ratio", None)
+                    if mlp_ratio and arch_config.hidden_size:
+                        ffn_dim = int(arch_config.hidden_size * mlp_ratio)
+                if ffn_dim is not None:
+                    metrics["ffn_dim"] = ffn_dim
 
                 self.tracker.log(metrics, step)
             if step % self.training_args.training_state_checkpointing_steps == 0:
@@ -845,8 +956,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 step_captions.append(batch.prompt)
 
                 # Run validation inference
+                # Use the validation pipeline args to match its loaded modules.
+                validation_args = self.validation_pipeline.fastvideo_args
+                if validation_args is None:
+                    validation_args = training_args
                 output_batch = self.validation_pipeline.forward(
-                    batch, training_args)
+                    batch, validation_args)
                 samples = output_batch.output
 
                 if self.rank_in_sp_group != 0:

fastvideo/training/training_utils.py

@@ -849,7 +849,14 @@ def load_distillation_checkpoint(
 
 def normalize_dit_input(model_type, latents, vae) -> torch.Tensor:
     if model_type == "hunyuan_hf" or model_type == "hunyuan":
-        return latents * 0.476986
+        scaling_factor = vae.scaling_factor
+        if scaling_factor is None:
+            raise AttributeError(
+                "Hunyuan VAE must define scaling_factor for training.")
+        if isinstance(scaling_factor, torch.Tensor):
+            scaling_factor = scaling_factor.to(device=latents.device,
+                                               dtype=latents.dtype)
+        return latents * scaling_factor
     elif model_type == "wan":
         latents_mean = torch.tensor(vae.latents_mean)
         latents_std = 1.0 / torch.tensor(vae.latents_std)

fastvideo/training/wan_rl_training_pipeline.py

@@ -1,100 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import sys
-from copy import deepcopy
-
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
-    FlowUniPCMultistepScheduler)
-from fastvideo.pipelines.basic.wan.wan_pipeline import WanPipeline
-from fastvideo.training.rl.rl_pipeline import RLPipeline
-from fastvideo.utils import is_vsa_available
-
-vsa_available = is_vsa_available()
-
-logger = init_logger(__name__)
-
-
-class WanRLTrainingPipeline(RLPipeline):
-    """
-    A training pipeline for Wan with RL/GRPO support.
-    
-    This pipeline extends RLPipeline with Wan-specific initialization.
-    """
-    _required_config_modules = [
-        "scheduler", "transformer", "vae", "text_encoder", "tokenizer"
-    ]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def create_training_stages(self, training_args: TrainingArgs):
-        """
-        May be used in future refactors.
-        """
-        pass
-
-    def _create_inference_pipeline(self, training_args: TrainingArgs,
-                                   dit_cpu_offload: bool):
-        args_copy = deepcopy(training_args)
-        args_copy.inference_mode = True
-        loaded_modules = {
-            "transformer": self.get_module("transformer"),
-        }
-        transformer_2 = self.get_module("transformer_2", None)
-        if transformer_2 is not None:
-            loaded_modules["transformer_2"] = transformer_2
-        text_encoder = self.get_module("text_encoder", None)
-        if text_encoder is not None:
-            loaded_modules["text_encoder"] = text_encoder
-        tokenizer = self.get_module("tokenizer", None)
-        if tokenizer is not None:
-            loaded_modules["tokenizer"] = tokenizer
-        vae = self.get_module("vae", None)
-        if vae is not None:
-            loaded_modules["vae"] = vae
-
-        return WanPipeline.from_pretrained(
-            training_args.model_path,
-            args=args_copy,  # type: ignore
-            inference_mode=True,
-            loaded_modules=loaded_modules,
-            tp_size=training_args.tp_size,
-            sp_size=training_args.sp_size,
-            num_gpus=training_args.num_gpus,
-            pin_cpu_memory=training_args.pin_cpu_memory,
-            dit_cpu_offload=dit_cpu_offload)
-
-    def initialize_validation_pipeline(self, training_args: TrainingArgs):
-        logger.info("Initializing validation pipeline...")
-        self.validation_pipeline = self._create_inference_pipeline(
-            training_args, dit_cpu_offload=True)
-
-    def _build_sampling_pipeline(self, training_args: TrainingArgs):
-        return self._create_inference_pipeline(training_args,
-                                               dit_cpu_offload=False)
-
-
-def main(args) -> None:
-    logger.info("Starting RL training pipeline...")
-
-    pipeline = WanRLTrainingPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("RL training pipeline done")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    args.dit_cpu_offload = False
-    # Enable RL mode
-    args.rl_mode = True
-    main(args)

fastvideo/workflow/preprocess/components.py

@@ -200,8 +200,14 @@ class ParquetDatasetSaver:
             value = getattr(batch, key.name)
             if isinstance(value, list):
                 for idx in range(len(value)):
-                    if isinstance(value[idx], torch.Tensor):
-                        value[idx] = value[idx].cpu().numpy()
+                    # if isinstance(value[idx], torch.Tensor):
+                    #     value[idx] = value[idx].cpu().numpy()
+                    t = value[idx]
+                    if isinstance(t, torch.Tensor):
+                        if t.dtype == torch.bfloat16:
+                            t = t.float()   # or t.to(torch.float16) if you want smaller
+                        value[idx] = t.cpu().numpy()
+
             elif isinstance(value, torch.Tensor):
                 value = value.cpu().numpy()
                 setattr(batch, key.name, value)
@@ -288,10 +294,17 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
                                split=split)
         column_names = dataset.column_names
         # rename columns to match the schema
-        if "cap" in column_names:
+        if "cap" in column_names and "caption" not in column_names:
             dataset = dataset.rename_column("cap", "caption")
         if "path" in column_names:
             dataset = dataset.rename_column("path", "name")
+        
+        # Add num_frames column if it doesn't exist but fps and duration do
+        if 'num_frames' not in column_names and 'fps' in column_names and 'duration' in column_names:
+            def add_num_frames(item: dict[str, Any]) -> dict[str, Any]:
+                item['num_frames'] = int(item['fps'] * item['duration'])
+                return item
+            dataset = dataset.map(add_num_frames)
 
         dataset = dataset.filter(validator)
         dataset = dataset.shard(num_shards=get_world_size(),

kill.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+# Kill all GPU processes shown by nvidia-smi
+
+pids=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | tr -d ' ')
+
+if [ -z "$pids" ]; then
+    echo "No GPU processes found."
+    exit 0
+fi
+
+echo "Killing GPU processes: $pids"
+for pid in $pids; do
+    kill -9 "$pid" 2>/dev/null && echo "Killed PID $pid" || echo "Failed to kill PID $pid"
+done

pyproject.toml

@@ -63,7 +63,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.4",
+    "fastvideo-kernel==0.2.3",
     "wheel",
 
     # Training Dependencies

pyproject_other.toml

@@ -63,7 +63,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.4",
+    "fastvideo-kernel==0.2.3",
     "wheel",
 
     # Training Dependencies