diff --git a/.gitignore b/.gitignore index 9f69c89a21..8ce8dd4f56 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,7 @@ env **/build/ **.pyc **.txt +!examples/train/requirements-genrl.txt *.log weights/ logs/ @@ -127,4 +128,5 @@ apps/dreamverse/web/.env.production.local .codex/ .sisyphus/ openspec/ +modal_train_genrl.py fastvideo/tests/ssim/reference_videos/** diff --git a/examples/train/README.md b/examples/train/README.md index b41b1b7847..ab166f8f1b 100644 --- a/examples/train/README.md +++ b/examples/train/README.md @@ -56,6 +56,55 @@ bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml \ --training.checkpoint.resume_from_checkpoint outputs/my_experiment/checkpoint-500 ``` +## GenRL HPSv3 + VideoAlign + +The GenRL reward runtime is vendored as normal FastVideo package files under +`fastvideo/train/methods/rl/reward/`; it is not a git submodule. The prompt +JSONL files and VideoReward checkpoint are runtime assets and should be +prepared before launch: + +```bash +uv pip install -e . +uv pip install -r examples/train/requirements-genrl.txt +python examples/train/prepare_genrl_assets.py +``` + +Before launching a long run, verify that both reward models load and can score +a dummy video: + +```bash +VIDEOALIGN_CHECKPOINT_PATH=.cache/VideoReward \ +FORCE_QWENVL_VIDEO_READER=opencv \ +python examples/train/prepare_genrl_assets.py --check-rewards +``` + +The helper writes: + +```bash +PROMPT_DATASET_PATH=.cache/genrl_filtered_prompts +VIDEOALIGN_CHECKPOINT_PATH=.cache/VideoReward +``` + +Launch the 4 GPU reproduction run: + +```bash +WANDB_MODE=online \ +WANDB_ENTITY= \ +NUM_GPUS=4 \ +VIDEOALIGN_CHECKPOINT_PATH=.cache/VideoReward \ +FORCE_QWENVL_VIDEO_READER=opencv \ +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +bash examples/train/run.sh \ + examples/train/configs/rl/wan/genrl_hpsv3_videoalign.yaml \ + --training.checkpoint.output_dir outputs/genrl_longcat +``` + +System packages used for the reproduced Modal run were `ffmpeg`, `libgl1`, +`libglib2.0-0`, `build-essential`, `ninja-build`, `cmake`, and `git-lfs`. +On CUDA 12.8 / Python 3.12, the Modal environment also used the PyTorch cu128 +wheels and the `flash_attn-2.8.3+cu128torch2.10` prebuilt wheel. If +FlashAttention-2 is unavailable, the VideoAlign wrapper falls back to SDPA. + ## W&B Logging Training metrics and validation videos are logged to diff --git a/examples/train/configs/README.md b/examples/train/configs/README.md index 450ad6833e..df621fd83b 100644 --- a/examples/train/configs/README.md +++ b/examples/train/configs/README.md @@ -7,6 +7,7 @@ configs/ ├── fine_tuning/ # Standard finetuning and DFSFT ├── distribution_matching/ # DMD2 and Self-Forcing ├── knowledge_distillation/ # KD from teacher to student +├── rl/ # RL methods such as DiffusionNFT and GenRL └── example.yaml # Annotated reference config with all fields ``` diff --git a/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml b/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml deleted file mode 100644 index 1b1c67db91..0000000000 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml +++ /dev/null @@ -1,113 +0,0 @@ -# GenRL / Video GRPO: Wan 2.1 T2V 1.3B — OCR reward, full finetune, 4 GPUs. -# -# Usage: -# torchrun --nnodes=1 --nproc_per_node=4 \ -# -m fastvideo.train.entrypoint.train \ -# --config examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml - -models: - student: - _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel - init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers - trainable: true - enable_gradient_checkpointing_type: full - -method: - _target_: fastvideo.train.methods.rl.genrl.GenRLMethod - - # ---- Reward functions ---- - reward_fn: - video_ocr: 1.0 - reward_module: null - - # ---- Data ---- - prompt_dataset_path: GenRL/datasets/ocr - prompt_fn: general_ocr - - # ---- Sampling ---- - sample_batch_size: 4 - eval_batch_size: 2 - num_batches_per_epoch: 1 - num_inference_steps: 16 - guidance_scale: 4.5 - num_video_per_prompt: 4 - noise_level: 1.0 - sde_type: flow_sde - sde_window_size: 1 - sde_window_range: [0, 6] - diffusion_clip: true - diffusion_clip_value: 0.45 - kl_reward: 0 - same_latent: true - - # ---- Video dimensions ---- - height: 480 - width: 832 - num_frames: 81 - - # ---- PPO training ---- - train_batch_size: 4 - num_inner_epochs: 1 - clip_range: 1.0e-3 - adv_clip_max: 5.0 - # No frozen reference model is configured in this launch. - beta: 0.0 - use_cfg: true - loss_reweighting: longcat - weight_advantages: false - max_grad_norm: 1.0 - seed: 42 - - # ---- Advantage computation ---- - per_prompt_stat_tracking: true - global_std: false - max_group_std: true - -training: - distributed: - num_gpus: 4 - sp_size: 1 - tp_size: 1 - hsdp_replicate_dim: 1 - hsdp_shard_dim: 4 - - data: - data_path: "" - train_batch_size: 1 - seed: 42 - num_height: 480 - num_width: 832 - num_frames: 81 - - optimizer: - learning_rate: 1.0e-4 - betas: [0.9, 0.999] - weight_decay: 1.0e-4 - lr_scheduler: constant - lr_warmup_steps: 0 - - loop: - max_train_steps: 100000 - gradient_accumulation_steps: 1 - - checkpoint: - output_dir: outputs/genrl_ocr - training_state_checkpointing_steps: 100 - checkpoints_total_limit: 3 - - tracker: - project_name: VideoRL - run_name: wan_2_1_t2v_1_3b_ocr - - model: - enable_gradient_checkpointing_type: full - -callbacks: - grad_clip: - max_grad_norm: 0.0 # Disabled; GenRLMethod clips internally. - ema: - decay: 0.9 - start_iter: 0 - -pipeline: - flow_shift: 3.0 diff --git a/examples/train/configs/genrl_wan2.1_t2v_1.3B_hpsv3_videoalign.yaml b/examples/train/configs/rl/wan/genrl_hpsv3_videoalign.yaml similarity index 84% rename from examples/train/configs/genrl_wan2.1_t2v_1.3B_hpsv3_videoalign.yaml rename to examples/train/configs/rl/wan/genrl_hpsv3_videoalign.yaml index 73e3caacf2..9dccd793ec 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_hpsv3_videoalign.yaml +++ b/examples/train/configs/rl/wan/genrl_hpsv3_videoalign.yaml @@ -7,10 +7,20 @@ # - Full fine-tuning with beta > 0 requires models.reference and much # more memory; keep beta at 0.0 for the 4xH100 probe run. # +# Reproduction-critical setup for the reported 4xH100 probe run: +# - Use NUM_GPUS=4 so the distributed rollout/grouping shape matches. +# - Prepare the exact GenRL prompt JSONL files and VideoReward checkpoint: +# python examples/train/prepare_genrl_assets.py +# - Set VIDEOALIGN_CHECKPOINT_PATH to the prepared KwaiVGI/VideoReward snapshot. +# - Set FORCE_QWENVL_VIDEO_READER=opencv if torchvision video IO is missing. +# # Usage: -# torchrun --nnodes=1 --nproc_per_node=4 \ -# -m fastvideo.train.entrypoint.train \ -# --config examples/train/configs/genrl_wan2.1_t2v_1.3B_hpsv3_videoalign.yaml +# NUM_GPUS=4 \ +# VIDEOALIGN_CHECKPOINT_PATH=.cache/VideoReward \ +# FORCE_QWENVL_VIDEO_READER=opencv \ +# bash examples/train/run.sh \ +# examples/train/configs/rl/wan/genrl_hpsv3_videoalign.yaml \ +# --training.checkpoint.output_dir /path/to/outputs/genrl_hpsv3_videoalign models: student: @@ -45,7 +55,7 @@ method: reward_on_gpu: true # ---- Data ---- - prompt_dataset_path: GenRL/datasets/filtered_prompts + prompt_dataset_path: .cache/genrl_filtered_prompts prompt_fn: filtered_prompts # ---- Sampling ---- diff --git a/examples/train/prepare_genrl_assets.py b/examples/train/prepare_genrl_assets.py new file mode 100644 index 0000000000..18026681cd --- /dev/null +++ b/examples/train/prepare_genrl_assets.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepare GenRL prompt and reward assets for example training runs.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import shutil +import subprocess +from pathlib import Path + +MIN_TRAIN_PROMPTS = 4 +GENRL_REPO = "https://github.com/ModelTC/GenRL.git" +GENRL_PROMPT_FILES = ("train.json", "test.json") +VIDEOREWARD_REPO = "KwaiVGI/VideoReward" + + +def _run(cmd: list[str], cwd: Path | None = None) -> None: + subprocess.run(cmd, cwd=cwd, check=True) + + +def _is_nonempty_dir(path: Path) -> bool: + return path.is_dir() and any(path.iterdir()) + + +def _ensure_genrl_sparse_checkout(genrl_cache_dir: Path) -> Path: + """Fetch only the GenRL prompt JSONL files into an ignored cache checkout.""" + if not genrl_cache_dir.exists(): + _run( + [ + "git", + "clone", + "--depth", + "1", + "--filter=blob:none", + "--sparse", + GENRL_REPO, + str(genrl_cache_dir), + ] + ) + elif not (genrl_cache_dir / ".git").exists(): + if _is_nonempty_dir(genrl_cache_dir): + raise RuntimeError( + f"{genrl_cache_dir} exists but is not a git checkout. " + "Pass --genrl-cache-dir to use a different cache directory." + ) + genrl_cache_dir.rmdir() + return _ensure_genrl_sparse_checkout(genrl_cache_dir) + + _run( + ["git", "sparse-checkout", "set", "datasets/filtered_prompts"], + cwd=genrl_cache_dir, + ) + _run(["git", "lfs", "install"], cwd=genrl_cache_dir) + _run( + [ + "git", + "lfs", + "pull", + "-I", + "datasets/filtered_prompts/*", + ], + cwd=genrl_cache_dir, + ) + return genrl_cache_dir / "datasets" / "filtered_prompts" + + +def prepare_genrl_prompts(prompt_dir: Path, genrl_cache_dir: Path) -> Path: + source_prompt_dir = _ensure_genrl_sparse_checkout(genrl_cache_dir) + prompt_dir.mkdir(parents=True, exist_ok=True) + for filename in GENRL_PROMPT_FILES: + shutil.copy2(source_prompt_dir / filename, prompt_dir / filename) + + validate_prompt_file( + prompt_dir / "train.json", + min_prompts=MIN_TRAIN_PROMPTS, + ) + validate_prompt_file(prompt_dir / "test.json", min_prompts=1) + return prompt_dir + + +def validate_prompt_file(path: Path, min_prompts: int) -> None: + if not path.exists(): + raise FileNotFoundError( + f"Missing {path}. Expected GenRL filtered_prompts JSONL files. " + "Run `python examples/train/prepare_genrl_assets.py`." + ) + + prompt_count = 0 + saw_content = False + with path.open(encoding="utf-8") as f: + for line_no, raw_line in enumerate(f, start=1): + line = raw_line.strip() + if not line: + continue + if not saw_content and line.startswith("version https://git-lfs.github.com"): + raise RuntimeError( + f"{path} is a Git LFS pointer, not real prompt JSON. " + "Install git-lfs and rerun this script." + ) + saw_content = True + try: + item = json.loads(line) + except json.JSONDecodeError as exc: + raise RuntimeError(f"Invalid JSON in {path} at line {line_no}: {exc}") from exc + if not isinstance(item, dict): + raise RuntimeError( + f"Expected JSON object in {path} at line {line_no}; " + f"got {type(item).__name__}." + ) + if item.get("prompt"): + prompt_count += 1 + + if prompt_count < min_prompts: + raise RuntimeError( + f"{path} has {prompt_count} usable prompts; expected at least " + f"{min_prompts}." + ) + + +def prepare_video_reward(videoalign_dir: Path) -> Path: + if has_video_reward_checkpoint(videoalign_dir): + return videoalign_dir + + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise ImportError( + f"huggingface_hub is required to download {VIDEOREWARD_REPO}. " + "Install FastVideo dependencies, then rerun this script." + ) from exc + + snapshot_download( + repo_id=VIDEOREWARD_REPO, + repo_type="model", + local_dir=str(videoalign_dir), + local_dir_use_symlinks=False, + ) + if not has_video_reward_checkpoint(videoalign_dir): + raise RuntimeError( + f"Downloaded {VIDEOREWARD_REPO}, but no VideoReward checkpoint " + f"was found under {videoalign_dir}." + ) + return videoalign_dir + + +def has_video_reward_checkpoint(root: Path) -> bool: + model_config = root / "model_config.json" + if not model_config.exists(): + return False + if (root / "model.pth").exists(): + return True + if ((root / "adapter_model.safetensors").exists() + and (root / "non_lora_state_dict.pth").exists()): + return True + for checkpoint in root.glob("checkpoint-*"): + if (checkpoint / "model.pth").exists(): + return True + if ( + (checkpoint / "adapter_model.safetensors").exists() + and (checkpoint / "non_lora_state_dict.pth").exists() + ): + return True + return False + + +def check_reward_runtime(device: str = "auto") -> None: + """Run the same lightweight reward-model preflight used by Modal.""" + import numpy as np + import torch + + selected_device = device + if selected_device == "auto": + selected_device = "cuda:0" if torch.cuda.is_available() else "cpu" + + print("=== GenRL reward preflight ===", flush=True) + import peft + import transformers + import torchvision + + print( + "Versions: " + f"torch={torch.__version__} " + f"torchvision={torchvision.__version__} " + f"transformers={transformers.__version__} " + f"peft={peft.__version__}", + flush=True, + ) + + from fastvideo.train.methods.rl.reward.hpsv3 import ( + _HPSV3_INFERENCERS, + hpsv3_general_score, + hpsv3_percentile_score, + set_hpsv3_device, + ) + + torch_device = torch.device(selected_device) + dummy_video = np.zeros((1, 1, 224, 224, 3), dtype=np.uint8) + for name, factory in ( + ("HPSv3-general", hpsv3_general_score), + ("HPSv3-percentile", hpsv3_percentile_score), + ): + reward = factory(torch_device) + scores, _ = reward(dummy_video, ["preflight prompt"], {}) + value = float(scores["avg"].detach().cpu()[0]) + print(f"{name} preflight score: {value:.4f}", flush=True) + + set_hpsv3_device("cpu") + _HPSV3_INFERENCERS.clear() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + from fastvideo.train.methods.rl.reward.videoalign import ( + _VIDEOALIGN_INFERENCERS, + set_videoalign_device, + videoalign_mq_score, + videoalign_ta_score, + ) + + dummy_video = np.zeros((1, 8, 224, 224, 3), dtype=np.uint8) + for name, factory in ( + ("VideoAlign-MQ", videoalign_mq_score), + ("VideoAlign-TA", videoalign_ta_score), + ): + reward = factory(torch_device) + scores, _ = reward(dummy_video, ["preflight prompt"], {}) + value = float(scores["avg"].detach().cpu()[0]) + print(f"{name} preflight score: {value:.4f}", flush=True) + + set_videoalign_device("cpu") + _VIDEOALIGN_INFERENCERS.clear() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + print("=== GenRL reward preflight OK ===", flush=True) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Prepare assets for GenRL HPSv3 + VideoAlign training." + ) + parser.add_argument( + "--prompt-dir", + type=Path, + default=Path(".cache/genrl_filtered_prompts"), + help="Directory that will contain train.json and test.json.", + ) + parser.add_argument( + "--genrl-cache-dir", + type=Path, + default=Path(".cache/GenRL"), + help="Ignored sparse checkout cache used to fetch only GenRL filtered prompts.", + ) + parser.add_argument( + "--videoalign-dir", + type=Path, + default=Path(".cache/VideoReward"), + help=f"Directory containing or receiving {VIDEOREWARD_REPO}.", + ) + parser.add_argument( + "--check-rewards", + action="store_true", + help="After preparing assets, load HPSv3 and VideoAlign on a dummy video.", + ) + parser.add_argument( + "--reward-device", + default="auto", + help="Device for --check-rewards: auto, cpu, cuda, or cuda:.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + prompt_dir = prepare_genrl_prompts(args.prompt_dir, args.genrl_cache_dir) + videoalign_dir = prepare_video_reward(args.videoalign_dir) + os.environ.setdefault("VIDEOALIGN_CHECKPOINT_PATH", str(videoalign_dir)) + print("GenRL assets ready.") + print(f"PROMPT_DATASET_PATH={prompt_dir}") + print(f"VIDEOALIGN_CHECKPOINT_PATH={videoalign_dir}") + if args.check_rewards: + check_reward_runtime(args.reward_device) + + +if __name__ == "__main__": + main() diff --git a/examples/train/requirements-genrl.txt b/examples/train/requirements-genrl.txt new file mode 100644 index 0000000000..a3754d000e --- /dev/null +++ b/examples/train/requirements-genrl.txt @@ -0,0 +1,18 @@ +# Dependency pins used for the reproduced GenRL HPSv3 + VideoAlign run. +# +# Install after the editable FastVideo install so these reward-stack versions +# match the Modal environment used for the reference run. +accelerate==1.0.1 +datasets==3.6.0 +diffusers==0.33.1 +fire>=0.7.0 +matplotlib==3.10.3 +numpy==1.26.4 +peft==0.15.0 +prettytable==3.8.0 +qwen-vl-utils==0.0.11 +safetensors==0.5.3 +scipy==1.15.2 +timm==1.0.15 +transformers==4.57.3 +trl==0.8.6 diff --git a/fastvideo/train/methods/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py index 89a467c223..98b9d6ac8d 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -689,7 +689,20 @@ def _ppo_train( samples[key] = samples[key][row_idx, perms] # Batch into micro-batches. + if total_batch_size % self._num_batches_per_epoch != 0: + raise ValueError( + "GenRL PPO batch size must be divisible by " + "num_batches_per_epoch after advantage filtering. Got " + f"total_batch_size={total_batch_size}, " + f"num_batches_per_epoch={self._num_batches_per_epoch}." + ) micro = (total_batch_size // self._num_batches_per_epoch) + if micro <= 0: + raise ValueError( + "GenRL PPO microbatch size must be positive. Got " + f"total_batch_size={total_batch_size}, " + f"num_batches_per_epoch={self._num_batches_per_epoch}." + ) batched = {k: v.reshape(-1, micro, *v.shape[1:]) for k, v in samples.items()} batched_list = [dict(zip(batched, x, strict=False)) for x in zip(*batched.values(), strict=False)] diff --git a/fastvideo/train/methods/rl/reward/HPSv3/__init__.py b/fastvideo/train/methods/rl/reward/HPSv3/__init__.py index a77d9a91f9..df97d809bd 100644 --- a/fastvideo/train/methods/rl/reward/HPSv3/__init__.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/__init__.py @@ -1,7 +1,7 @@ """Vendored runtime subset of HPSv3. Source: https://github.com/MizzenAI/HPSv3 -Commit: bd0c5fcb5f587617b0169c07222ab78d01e2f3c2 +Commit: a2eb2ef2c7b5d91a566347a5825cf6d872122149 Purpose: Runtime reward inference integration for FastVideo GenRL. This is temporary minimal vendoring for PR integration. It is expected to be diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/__init__.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/__init__.py index e0f779e31b..1ec1f41712 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/__init__.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/__init__.py @@ -1,12 +1,9 @@ """Vendored HPSv3 runtime package. Source: https://github.com/MizzenAI/HPSv3 -Commit: bd0c5fcb5f587617b0169c07222ab78d01e2f3c2 +Commit: a2eb2ef2c7b5d91a566347a5825cf6d872122149 Purpose: Runtime reward inference integration for FastVideo GenRL. -This is temporary minimal vendoring for PR integration. It is expected to be -cleaned up and normalized later. - Porting rules: - Include only files required by the runtime import closure used by FastVideo. - When an upstream file is required, copy the entire file faithfully. diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/data_collator_qwen.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/data_collator_qwen.py index 9d4de8fc8b..a982ccd881 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/data_collator_qwen.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/data_collator_qwen.py @@ -41,7 +41,6 @@ class QWen2VLDataCollator: - def __init__( self, processor, @@ -69,25 +68,28 @@ def _clean_message( remove unnecessary keys from message(very very necessary) """ message_list = [] - for text, image in zip(texts, images, strict=False): - out_message = [{ - "role": - "user", - "content": [ - { - "type": "image", - "image": image, - "min_pixels": min_pixels, - "max_pixels": max_pixels, - }, - { - "type": - "text", - "text": (INSTRUCTION.format(text_prompt=text) + - prompt_with_special_token if use_special_tokens else prompt_without_special_token), - }, - ], - }] + for text, image in zip(texts, images): + out_message = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": image, + "min_pixels": min_pixels, + "max_pixels": max_pixels, + }, + { + "type": "text", + "text": ( + INSTRUCTION.format(text_prompt=text) + prompt_with_special_token + if use_special_tokens + else prompt_without_special_token + ), + }, + ], + } + ] message_list.append(out_message) @@ -104,8 +106,9 @@ def _pad_sequence(self, sequences, attention_mask, max_len, padding_side="right" pad_len = max_len - sequences.shape[1] padding = (0, pad_len) if padding_side == "right" else (pad_len, 0) - sequences_padded = torch.nn.functional.pad(sequences, padding, "constant", - self.processor.tokenizer.pad_token_id) + sequences_padded = torch.nn.functional.pad( + sequences, padding, "constant", self.processor.tokenizer.pad_token_id + ) attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, "constant", 0) return sequences_padded, attention_mask_padded @@ -164,12 +167,12 @@ def __call__(self, inputs, with_instruction=True): # pdb.set_trace() max_len = max(batch_1["input_ids"].shape[1], batch_2["input_ids"].shape[1]) - batch_1["input_ids"], batch_1["attention_mask"] = self._pad_sequence(batch_1["input_ids"], - batch_1["attention_mask"], max_len, - "right") - batch_2["input_ids"], batch_2["attention_mask"] = self._pad_sequence(batch_2["input_ids"], - batch_2["attention_mask"], max_len, - "right") + batch_1["input_ids"], batch_1["attention_mask"] = self._pad_sequence( + batch_1["input_ids"], batch_1["attention_mask"], max_len, "right" + ) + batch_2["input_ids"], batch_2["attention_mask"] = self._pad_sequence( + batch_2["input_ids"], batch_2["attention_mask"], max_len, "right" + ) batch = { "batch_1": batch_1, diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/pairwise_dataset.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/pairwise_dataset.py index 577c6037ee..93de031092 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/pairwise_dataset.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/pairwise_dataset.py @@ -1,13 +1,13 @@ -import torch -from torch.utils.data import Dataset -import random import json import os +import random + +import torch +from torch.utils.data import Dataset from tqdm import tqdm class PairwiseOriginalDataset(Dataset): - def __init__( self, json_list, @@ -41,6 +41,7 @@ def __getitem__(self, idx): except Exception as e: print(f"Error processing sample at index {idx}: {e}") import traceback + traceback.print_exc() index = random.randint(0, len(self.samples) - 1) if index == idx: @@ -52,14 +53,14 @@ def get_single_item(self, idx): # Load image paths image_1 = sample["path1"] image_2 = sample["path2"] - assert os.path.exists(image_1) and os.path.exists(image_2), f'{image_1} or {image_2}' + assert os.path.exists(image_1) and os.path.exists(image_2), f"{image_1} or {image_2}" text_1 = sample["prompt"] text_2 = sample["prompt"] # Process Label if self.soft_label: choice_dist = sorted(sample["choice_dist"], reverse=True) - assert (torch.sum(torch.tensor(choice_dist)) > 0), "Choice distribution cannot be zero." + assert torch.sum(torch.tensor(choice_dist)) > 0, "Choice distribution cannot be zero." label = torch.tensor(choice_dist[0]) / torch.sum(torch.tensor(choice_dist)) else: label = torch.tensor(1).float() diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/utils.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/utils.py index 132655ecb3..1d4210a644 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/utils.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/utils.py @@ -6,6 +6,7 @@ import math import os import sys +import time import warnings from functools import lru_cache from io import BytesIO @@ -49,11 +50,13 @@ def floor_by_factor(number: int, factor: int) -> int: return math.floor(number / factor) * factor -def smart_resize(height: int, - width: int, - factor: int = IMAGE_FACTOR, - min_pixels: int = MIN_PIXELS, - max_pixels: int = MAX_PIXELS) -> tuple[int, int]: +def smart_resize( + height: int, + width: int, + factor: int = IMAGE_FACTOR, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS, +) -> tuple[int, int]: """ Rescales the image so that the following conditions are met: @@ -65,7 +68,8 @@ def smart_resize(height: int, """ if max(height, width) / min(height, width) > MAX_RATIO: raise ValueError( - f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}") + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) h_bar = max(factor, round_by_factor(height, factor)) w_bar = max(factor, round_by_factor(width, factor)) if h_bar * w_bar > max_pixels: @@ -80,9 +84,12 @@ def smart_resize(height: int, def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image: - image = ele["image"] if "image" in ele else ele["image_url"] + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] image_obj = None - if isinstance(image, Image.Image | torch.Tensor): + if isinstance(image, Image.Image) or isinstance(image, torch.Tensor): image_obj = image elif image.startswith("http://") or image.startswith("https://"): image_obj = Image.open(requests.get(image, stream=True).raw) @@ -112,18 +119,18 @@ def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACT if len(shape) == 4: if shape[1] in [1, 3]: # Likely [B, C, H, W] height, width = shape[2], shape[3] - image_mode = 'NCHW' + image_mode = "NCHW" elif shape[3] in [1, 3]: # Likely [B, H, W, C] height, width = shape[1], shape[2] - image_mode = 'NHWC' + image_mode = "NHWC" elif len(shape) == 3: if shape[0] in [1, 3]: # Likely [C, H, W] height, width = shape[1], shape[2] - image_mode = 'CHW' + image_mode = "CHW" elif shape[2] in [1, 3]: # Likely [H, W, C] height, width = shape[0], shape[1] - image_mode = 'HWC' + image_mode = "HWC" else: raise ValueError(f"Cannot determine tensor image format from shape {shape}") else: @@ -141,24 +148,36 @@ def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACT ) if isinstance(image, torch.Tensor): - if image_mode == 'NCHW': - image = transforms.functional.resize(image, [resized_height, resized_width], - interpolation=InterpolationMode.BICUBIC, - antialias=True) - elif image_mode == 'NHWC': - image = transforms.functional.resize(image.permute(0, 3, 1, 2), [resized_height, resized_width], - interpolation=InterpolationMode.BICUBIC, - antialias=True) - elif image_mode == 'CHW': + if image_mode == "NCHW": + image = transforms.functional.resize( + image, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ) + elif image_mode == "NHWC": + image = transforms.functional.resize( + image.permute(0, 3, 1, 2), + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ) + elif image_mode == "CHW": image = image.unsqueeze(0) # Add batch dimension - image = transforms.functional.resize(image, [resized_height, resized_width], - interpolation=InterpolationMode.BICUBIC, - antialias=True) - elif image_mode == 'HWC': + image = transforms.functional.resize( + image, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ) + elif image_mode == "HWC": image = image.permute(2, 0, 1).unsqueeze(0) # Add batch dimension and change to CHW - image = transforms.functional.resize(image, [resized_height, resized_width], - interpolation=InterpolationMode.BICUBIC, - antialias=True) + image = transforms.functional.resize( + image, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ) else: # If the image is a PIL Image, we resize it using PIL. @@ -172,7 +191,7 @@ def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACT def smart_nframes( ele: dict, total_frames: int, - video_fps: int | float, + video_fps: float, ) -> int: """calculate the number of frames for video used for model inputs. @@ -202,14 +221,15 @@ def smart_nframes( nframes = total_frames / video_fps * fps nframes = min(max(nframes, min_frames), max_frames) nframes = round_by_factor(nframes, FRAME_FACTOR) - if nframes > total_frames: - nframes = total_frames + nframes = min(nframes, total_frames) if not (nframes >= FRAME_FACTOR and nframes <= total_frames): raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") return nframes -def _read_video_torchvision(ele: dict, ) -> torch.Tensor: +def _read_video_torchvision( + ele: dict, +) -> torch.Tensor: """read video using torchvision.io.read_video Args: @@ -224,12 +244,10 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor: video_path = ele["video"] if version.parse(torchvision.__version__) < version.parse("0.19.0"): if "http://" in video_path or "https://" in video_path: - warnings.warn( - "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.", - stacklevel=2, - ) + warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.") if "file://" in video_path: video_path = video_path[7:] + st = time.time() video, audio, info = io.read_video( video_path, start_pts=ele.get("video_start", 0.0), @@ -240,10 +258,10 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor: total_frames, video_fps = video.size(0), info["video_fps"] # logger.info(f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") - if ele['sample_type'] == 'uniform': + if ele["sample_type"] == "uniform": nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() - elif ele['sample_type'] == 'multi_pts': + elif ele["sample_type"] == "multi_pts": frames_each_pts = 6 num_pts = 4 fps = 8 @@ -255,7 +273,7 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor: pts = torch.linspace(start_pt, end_pt, num_pts).round().long().tolist() idx = [] for pt in pts: - idx.extend(frames_idx[pt - frames_each_pts // 2:pt + frames_each_pts // 2]) + idx.extend(frames_idx[pt - frames_each_pts // 2 : pt + frames_each_pts // 2]) video = video[idx] return video @@ -267,7 +285,9 @@ def is_decord_available() -> bool: return importlib.util.find_spec("decord") is not None -def _read_video_decord(ele: dict, ) -> torch.Tensor: +def _read_video_decord( + ele: dict, +) -> torch.Tensor: """read video using decord.VideoReader Args: @@ -280,19 +300,21 @@ def _read_video_decord(ele: dict, ) -> torch.Tensor: torch.Tensor: the video tensor with shape (T, C, H, W). """ import decord + video_path = ele["video"] + st = time.time() vr = decord.VideoReader(video_path) # TODO: support start_pts and end_pts - if 'video_start' in ele or 'video_end' in ele: + if "video_start" in ele or "video_end" in ele: raise NotImplementedError("not support start_pts and end_pts in decord for now.") total_frames, video_fps = len(vr), vr.get_avg_fps() # logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") - if ele['sample_type'] == 'uniform': + if ele["sample_type"] == "uniform": nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) # nframes = max(nframes, 8) # import pdb; pdb.set_trace() idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() - elif ele['sample_type'] == 'multi_pts': + elif ele["sample_type"] == "multi_pts": frames_each_pts = 6 num_pts = 4 fps = 8 @@ -304,7 +326,7 @@ def _read_video_decord(ele: dict, ) -> torch.Tensor: pts = torch.linspace(start_pt, end_pt, num_pts).round().long().tolist() idx = [] for pt in pts: - idx.extend(frames_idx[pt - frames_each_pts // 2:pt + frames_each_pts // 2]) + idx.extend(frames_idx[pt - frames_each_pts // 2 : pt + frames_each_pts // 2]) video = vr.get_batch(idx).asnumpy() video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format return video @@ -339,7 +361,10 @@ def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | l min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS) total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS) - max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05)) + max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05), + ) max_pixels = ele.get("max_pixels", max_pixels) if "resized_height" in ele and "resized_width" in ele: resized_height, resized_width = smart_resize( @@ -362,21 +387,18 @@ def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | l antialias=True, ).float() return video - else: - assert isinstance(ele["video"], list | tuple) - process_info = ele.copy() - process_info.pop("type", None) - process_info.pop("video", None) - images = [ - fetch_image({ - "image": video_element, - **process_info - }, size_factor=image_factor) for video_element in ele["video"] - ] - nframes = ceil_by_factor(len(images), FRAME_FACTOR) - if len(images) < nframes: - images.extend([images[-1]] * (nframes - len(images))) - return images + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + images = [ + fetch_image({"image": video_element, **process_info}, size_factor=image_factor) + for video_element in ele["video"] + ] + nframes = ceil_by_factor(len(images), FRAME_FACTOR) + if len(images) < nframes: + images.extend([images[-1]] * (nframes - len(images))) + return images def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]: @@ -387,8 +409,12 @@ def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[di for message in conversation: if isinstance(message["content"], list): for ele in message["content"]: - if ("image" in ele or "image_url" in ele or "video" in ele - or ele["type"] in ("image", "image_url", "video")): + if ( + "image" in ele + or "image_url" in ele + or "video" in ele + or ele["type"] in ("image", "image_url", "video") + ): vision_infos.append(ele) return vision_infos diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/inference.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/inference.py index 62fd0ad4a6..ec8c2a2687 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/inference.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/inference.py @@ -1,27 +1,55 @@ import os from collections.abc import Mapping -import torch +from pathlib import Path + import huggingface_hub +import torch + +from .dataset.data_collator_qwen import ( + INSTRUCTION, + prompt_with_special_token, + prompt_without_special_token, +) from .dataset.utils import process_vision_info -from .dataset.data_collator_qwen import prompt_with_special_token, prompt_without_special_token, INSTRUCTION -from .utils.parser import ModelConfig, PEFTLoraConfig, TrainingConfig, DataConfig, parse_args_with_yaml from .train import create_model_and_processor -from pathlib import Path +from .utils.parser import ( + DataConfig, + ModelConfig, + PEFTLoraConfig, + TrainingConfig, + parse_args_with_yaml, +) _MODEL_CONFIG_PATH = Path(__file__).parent / "config/" class HPSv3RewardInferencer: - - def __init__(self, config_path=None, checkpoint_path=None, device='cuda', differentiable=False): + def __init__( + self, + config_path=None, + checkpoint_path=None, + device="cuda", + differentiable=False, + ): if config_path is None: - config_path = os.path.join(_MODEL_CONFIG_PATH, 'HPSv3_7B.yaml') + config_path = os.path.join(_MODEL_CONFIG_PATH, "HPSv3_7B.yaml") if checkpoint_path is None: - checkpoint_path = huggingface_hub.hf_hub_download("MizzenAI/HPSv3", 'HPSv3.safetensors', repo_type='model') - - (data_config, training_args, model_config, peft_lora_config), config_path = (parse_args_with_yaml( - (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig), config_path, is_train=False)) + checkpoint_path = huggingface_hub.hf_hub_download("MizzenAI/HPSv3", "HPSv3.safetensors", repo_type="model") + + ( + ( + data_config, + training_args, + model_config, + peft_lora_config, + ), + config_path, + ) = parse_args_with_yaml( + (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig), + config_path, + is_train=False, + ) training_args.output_dir = os.path.join(training_args.output_dir, config_path.split("/")[-1].split(".")[0]) model, processor, peft_config = create_model_and_processor( model_config=model_config, @@ -33,15 +61,36 @@ def __init__(self, config_path=None, checkpoint_path=None, device='cuda', differ self.device = device self.use_special_tokens = model_config.use_special_tokens - if checkpoint_path.endswith('.safetensors'): + if checkpoint_path.endswith(".safetensors"): import safetensors.torch + state_dict = safetensors.torch.load_file(checkpoint_path, device="cpu") else: state_dict = torch.load(checkpoint_path, map_location="cpu") if "model" in state_dict: state_dict = state_dict["model"] + before_rm_head = { + key: value.detach().cpu().clone() + for key, value in model.state_dict().items() + if "rm_head" in key + } model.load_state_dict(state_dict, strict=True) + after_rm_head = { + key: value.detach().cpu() + for key, value in model.state_dict().items() + if "rm_head" in key + } + unchanged = ( + before_rm_head + and set(before_rm_head) == set(after_rm_head) + and all(torch.equal(before_rm_head[key], after_rm_head[key]) for key in before_rm_head) + ) + if unchanged: + raise RuntimeError( + f"HPSv3 checkpoint {checkpoint_path} did not overwrite rm_head " + "weights. Refusing to score rewards with a randomly initialized head." + ) model.eval() self.model = model @@ -50,20 +99,21 @@ def __init__(self, config_path=None, checkpoint_path=None, device='cuda', differ self.model.to(self.device) self.data_config = data_config - def _pad_sequence(self, sequences, attention_mask, max_len, padding_side='right'): + def _pad_sequence(self, sequences, attention_mask, max_len, padding_side="right"): """ Pad the sequences to the maximum length. """ - assert padding_side in ['right', 'left'] + assert padding_side in ["right", "left"] if sequences.shape[1] >= max_len: return sequences, attention_mask pad_len = max_len - sequences.shape[1] - padding = (0, pad_len) if padding_side == 'right' else (pad_len, 0) + padding = (0, pad_len) if padding_side == "right" else (pad_len, 0) - sequences_padded = torch.nn.functional.pad(sequences, padding, 'constant', - self.processor.tokenizer.pad_token_id) - attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, 'constant', 0) + sequences_padded = torch.nn.functional.pad( + sequences, padding, "constant", self.processor.tokenizer.pad_token_id + ) + attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, "constant", 0) return sequences_padded, attention_mask_padded @@ -74,9 +124,9 @@ def _prepare_input(self, data): """ if isinstance(data, Mapping): return type(data)({k: self._prepare_input(v) for k, v in data.items()}) - elif isinstance(data, tuple | list): + if isinstance(data, (tuple, list)): return type(data)(self._prepare_input(v) for v in data) - elif isinstance(data, torch.Tensor): + if isinstance(data, torch.Tensor): kwargs = {"device": self.device} return data.to(**kwargs) return data @@ -95,26 +145,28 @@ def prepare_batch(self, image_paths, prompts): max_pixels = 256 * 28 * 28 min_pixels = 256 * 28 * 28 message_list = [] - for text, image in zip(prompts, image_paths, strict=False): - out_message = [{ - "role": - "user", - "content": [ - { - "type": "image", - "image": image, - "min_pixels": min_pixels, - "max_pixels": max_pixels, - }, - { - "type": - "text", - "text": - (INSTRUCTION.format(text_prompt=text) + - prompt_with_special_token if self.use_special_tokens else prompt_without_special_token), - }, - ], - }] + for text, image in zip(prompts, image_paths): + out_message = [ + { + "role": "user", + "content": [ + { + "type": "image", + "image": image, + "min_pixels": max_pixels, + "max_pixels": max_pixels, + }, + { + "type": "text", + "text": ( + INSTRUCTION.format(text_prompt=text) + prompt_with_special_token + if self.use_special_tokens + else prompt_without_special_token + ), + }, + ], + } + ] message_list.append(out_message) @@ -139,16 +191,16 @@ def reward(self, prompts, image_paths): if __name__ == "__main__": - config_path = 'config/inference/HPSv3_7B.yaml' - checkpoint_path = 'checkpoints/HPSv3_7B.pth' - device = 'cuda' + config_path = "config/inference/HPSv3_7B.yaml" + checkpoint_path = "checkpoints/HPSv3_7B.pth" + device = "cuda" dtype = torch.bfloat16 inferencer = HPSv3RewardInferencer(config_path, checkpoint_path, device=device) image_paths = ["assets/example1.png", "assets/example2.png"] prompts = [ "cute chibi anime cartoon fox, smiling wagging tail with a small cartoon heart above sticker", - "cute chibi anime cartoon fox, smiling wagging tail with a small cartoon heart above sticker" + "cute chibi anime cartoon fox, smiling wagging tail with a small cartoon heart above sticker", ] rewards = inferencer.reward(image_paths, prompts) print(rewards[0][0].item()) # miu and sigma. we select miu as the final output diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/differentiable_image_processor.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/differentiable_image_processor.py index 76d60bd9ce..2275e6a728 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/differentiable_image_processor.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/differentiable_image_processor.py @@ -40,7 +40,6 @@ import numpy as np import torch import torch.nn.functional as F - from transformers.image_processing_utils import BaseImageProcessor, BatchFeature from transformers.image_transforms import ( convert_to_rgb, @@ -67,6 +66,7 @@ logger = logging.get_logger(__name__) + if is_vision_available(): from PIL import Image @@ -82,13 +82,13 @@ def make_batched_images(images) -> list[list[ImageInput]]: Returns: list: A list of images. """ - if isinstance(images, list | tuple) and isinstance(images[0], list | tuple) and is_valid_image(images[0][0]): + if isinstance(images, (list, tuple)) and isinstance(images[0], (list, tuple)) and is_valid_image(images[0][0]): return [img for img_list in images for img in img_list] - elif isinstance(images, list | tuple) and is_valid_image(images[0]): + if isinstance(images, (list, tuple)) and is_valid_image(images[0]): return images - elif is_valid_image(images): + if is_valid_image(images): return [images] raise ValueError(f"Could not make batched images from {images}") @@ -96,13 +96,13 @@ def make_batched_images(images) -> list[list[ImageInput]]: # Copied from transformers.models.llava_next_video.image_processing_llava_next_video.make_batched_videos def make_batched_videos(videos) -> list[VideoInput]: - if isinstance(videos, list | tuple) and isinstance(videos[0], list | tuple) and is_valid_image(videos[0][0]): + if isinstance(videos, (list, tuple)) and isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0]): return videos - elif isinstance(videos, list | tuple) and is_valid_image(videos[0]): + if isinstance(videos, (list, tuple)) and is_valid_image(videos[0]): if isinstance(videos[0], Image.Image): return [videos] - elif len(videos[0].shape) == 4: + if len(videos[0].shape) == 4: return [list(video) for video in videos] elif is_valid_image(videos) and len(videos.shape) == 4: @@ -111,11 +111,13 @@ def make_batched_videos(videos) -> list[VideoInput]: raise ValueError(f"Could not make batched video from {videos}") -def smart_resize(height: int, - width: int, - factor: int = 28, - min_pixels: int = 56 * 56, - max_pixels: int = 14 * 14 * 4 * 1280): +def smart_resize( + height: int, + width: int, + factor: int = 28, + min_pixels: int = 56 * 56, + max_pixels: int = 14 * 14 * 4 * 1280, +): """Rescales the image so that the following conditions are met: 1. Both dimensions (height and width) are divisible by 'factor'. @@ -127,9 +129,10 @@ def smart_resize(height: int, """ if height < factor or width < factor: raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}") - elif max(height, width) / min(height, width) > 200: + if max(height, width) / min(height, width) > 200: raise ValueError( - f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}") + f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}" + ) h_bar = round(height / factor) * factor w_bar = round(width / factor) * factor if h_bar * w_bar > max_pixels: @@ -169,21 +172,26 @@ class Qwen2VLImageProcessor(BaseImageProcessor): max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`): The max pixels of the image to resize the image. patch_size (`int`, *optional*, defaults to 14): - The spatial patch size of the vision encoder. + The spacial patch size of the vision encoder. temporal_patch_size (`int`, *optional*, defaults to 2): The temporal patch size of the vision encoder. merge_size (`int`, *optional*, defaults to 2): The merge size of the vision encoder to llm encoder. """ - model_input_names = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"] + model_input_names = [ + "pixel_values", + "image_grid_thw", + "pixel_values_videos", + "video_grid_thw", + ] def __init__( self, do_resize: bool = True, resample: PILImageResampling = PILImageResampling.BICUBIC, do_rescale: bool = True, - rescale_factor: int | float = 1 / 255, + rescale_factor: float = 1 / 255, do_normalize: bool = True, image_mean: float | list[float] | None = None, image_std: float | list[float] | None = None, @@ -223,7 +231,7 @@ def _preprocess_differentiable( ): """ Differentiable version of image preprocessing using torch operations. - + Args: images: torch.Tensor of shape (B, C, H, W) or (C, H, W) Returns: @@ -250,16 +258,18 @@ def _preprocess_differentiable( max_pixels=self.max_pixels, ) # Use differentiable interpolation - image = F.interpolate(image.unsqueeze(0), - size=(resized_height, resized_width), - mode='bilinear', - align_corners=False).squeeze(0) + image = F.interpolate( + image.unsqueeze(0), + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + ).squeeze(0) if do_rescale: image = image * rescale_factor if do_normalize: - if isinstance(image_mean, list | tuple): + if isinstance(image_mean, (list, tuple)): mean = torch.tensor(image_mean, device=image.device, dtype=image.dtype).view(-1, 1, 1) std = torch.tensor(image_std, device=image.device, dtype=image.dtype).view(-1, 1, 1) else: @@ -279,7 +289,10 @@ def _preprocess_differentiable( # Reshape for patch extraction batch_size, channel, resized_height, resized_width = patches.shape grid_t = batch_size // self.temporal_patch_size - grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size + grid_h, grid_w = ( + resized_height // self.patch_size, + resized_width // self.patch_size, + ) # Differentiable patch extraction and reshaping patches = patches.view( @@ -294,8 +307,10 @@ def _preprocess_differentiable( self.patch_size, ) patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8) - flatten_patches = patches.reshape(grid_t * grid_h * grid_w, - channel * self.temporal_patch_size * self.patch_size * self.patch_size) + flatten_patches = patches.reshape( + grid_t * grid_h * grid_w, + channel * self.temporal_patch_size * self.patch_size * self.patch_size, + ) return flatten_patches, (grid_t, grid_h, grid_w) @@ -372,7 +387,8 @@ def _preprocess( if is_scaled_image(images[0]) and do_rescale: logger.warning_once( "It looks like you are trying to rescale already rescaled images. If the input" - " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.") + " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again." + ) if input_data_format is None: # We assume that all images have the same channel dimension format. input_data_format = infer_channel_dimension_format(images[0]) @@ -389,16 +405,23 @@ def _preprocess( min_pixels=self.min_pixels, max_pixels=self.max_pixels, ) - image = resize(image, - size=(resized_height, resized_width), - resample=resample, - input_data_format=input_data_format) + image = resize( + image, + size=(resized_height, resized_width), + resample=resample, + input_data_format=input_data_format, + ) if do_rescale: image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format) if do_normalize: - image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format) + image = self.normalize( + image=image, + mean=image_mean, + std=image_std, + input_data_format=input_data_format, + ) image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) processed_images.append(image) @@ -412,7 +435,10 @@ def _preprocess( patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1)) channel = patches.shape[1] grid_t = patches.shape[0] // self.temporal_patch_size - grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size + grid_h, grid_w = ( + resized_height // self.patch_size, + resized_width // self.patch_size, + ) patches = patches.reshape( grid_t, self.temporal_patch_size, @@ -425,8 +451,10 @@ def _preprocess( self.patch_size, ) patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8) - flatten_patches = patches.reshape(grid_t * grid_h * grid_w, - channel * self.temporal_patch_size * self.patch_size * self.patch_size) + flatten_patches = patches.reshape( + grid_t * grid_h * grid_w, + channel * self.temporal_patch_size * self.patch_size * self.patch_size, + ) return flatten_patches, (grid_t, grid_h, grid_w) @@ -442,10 +470,10 @@ def preprocess_tensor( ): """ Differentiable preprocessing method for torch tensors. - + Args: images: torch.Tensor of shape (B, C, H, W) or (C, H, W) - + Returns: dict containing: - pixel_values: torch.Tensor - processed patches @@ -468,7 +496,10 @@ def preprocess_tensor( image_std=image_std, ) - return {"pixel_values": patches, "image_grid_thw": torch.tensor(image_grid_thw, device=patches.device)} + return { + "pixel_values": patches, + "image_grid_thw": torch.tensor(image_grid_thw, device=patches.device), + } def preprocess( self, @@ -552,8 +583,10 @@ def preprocess( videos = make_batched_videos(videos) if images is not None and not valid_images(images): - raise ValueError("Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " - "torch.Tensor, tf.Tensor or jax.ndarray.") + raise ValueError( + "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " + "torch.Tensor, tf.Tensor or jax.ndarray." + ) validate_preprocess_arguments( rescale_factor=rescale_factor, @@ -610,6 +643,9 @@ def preprocess( vision_grid_thws.append(video_grid_thw) pixel_values = np.array(pixel_values) vision_grid_thws = np.array(vision_grid_thws) - data = {"pixel_values_videos": pixel_values, "video_grid_thw": vision_grid_thws} + data = { + "pixel_values_videos": pixel_values, + "video_grid_thw": vision_grid_thws, + } return BatchFeature(data=data, tensor_type=return_tensors) diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/qwen2vl_trainer.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/qwen2vl_trainer.py index c829ca0aad..1785180cf6 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/qwen2vl_trainer.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/qwen2vl_trainer.py @@ -1,35 +1,34 @@ -import os import math -import matplotlib.pyplot as plt +import os -import safetensors +import matplotlib.pyplot as plt import numpy as np +import safetensors import torch -import torch.nn as nn -import datasets -from torch.utils.data import Dataset, DataLoader +from ..utils.training_utils import get_peft_state_non_lora_maybe_zero_3 from peft import PeftModel +from torch import nn +from torch.utils.data import DataLoader, Dataset from transformers import Qwen2VLForConditionalGeneration from transformers.modeling_utils import PreTrainedModel -from transformers.trainer import TrainerCallback from transformers.trainer import ( - is_sagemaker_mp_enabled, - is_peft_available, - is_datasets_available, - WEIGHTS_NAME, - TRAINING_ARGS_NAME, - SAFE_WEIGHTS_NAME, PREFIX_CHECKPOINT_DIR, + SAFE_WEIGHTS_NAME, + TRAINING_ARGS_NAME, + WEIGHTS_NAME, + TrainerCallback, + is_datasets_available, + is_peft_available, + is_sagemaker_mp_enabled, logger, ) - from transformers.trainer_pt_utils import nested_detach from trl import RewardTrainer -from ..utils.training_utils import get_peft_state_non_lora_maybe_zero_3 +import datasets -class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration): +class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration): def __init__( self, config, @@ -57,7 +56,10 @@ def __init__( self.rm_head.add_module( f"layer_{layer}", nn.Sequential( - nn.Linear(rm_head_kwargs["hidden_size"], rm_head_kwargs["hidden_size"]), + nn.Linear( + rm_head_kwargs["hidden_size"], + rm_head_kwargs["hidden_size"], + ), nn.ReLU(), nn.Dropout(rm_head_kwargs.get("dropout", 0.1)), ), @@ -65,8 +67,11 @@ def __init__( else: self.rm_head.add_module( "output_layer", - nn.Linear(rm_head_kwargs["hidden_size"], output_dim, bias=rm_head_kwargs.get("bias", - False)), + nn.Linear( + rm_head_kwargs["hidden_size"], + output_dim, + bias=rm_head_kwargs.get("bias", False), + ), ) else: @@ -105,24 +110,25 @@ def forward( rope_deltas: torch.LongTensor | None = None, ): ## modified from the origin class Qwen2VLForConditionalGeneration - output_attentions = (output_attentions if output_attentions is not None else self.config.output_attentions) - output_hidden_states = (output_hidden_states - if output_hidden_states is not None else self.config.output_hidden_states) - return_dict = (return_dict if return_dict is not None else self.config.use_return_dict) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict # pdb.set_trace() if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) if pixel_values is not None: pixel_values = pixel_values.type(self.visual.get_dtype()) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - image_mask = ((input_ids == self.config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)) + image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds) image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - video_mask = ((input_ids == self.config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds)) + video_mask = (input_ids == self.config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds) video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) @@ -142,24 +148,26 @@ def forward( ) hidden_states = outputs[0] # [B, L, D] - with torch.autocast(device_type='cuda', dtype=torch.float32): + with torch.autocast(device_type="cuda", dtype=torch.float32): logits = self.rm_head(hidden_states) # [B, L, N] - batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] ## get sequence length if self.config.pad_token_id is None and batch_size != 1: raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") if self.config.pad_token_id is None: sequence_lengths = -1 + elif input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) else: - if input_ids is not None: - # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility - sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1) - sequence_lengths = sequence_lengths % input_ids.shape[-1] - sequence_lengths = sequence_lengths.to(logits.device) - else: - sequence_lengths = -1 + sequence_lengths = -1 ## get the last token's logits if self.reward_token == "last": @@ -167,7 +175,7 @@ def forward( elif self.reward_token == "mean": ## get the mean of all valid tokens' logits valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1) - pooled_logits = torch.stack([logits[i, :valid_lengths[i]].mean(dim=0) for i in range(batch_size)]) + pooled_logits = torch.stack([logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)]) elif self.reward_token == "special": # special_token_ids = self.tokenizer.convert_tokens_to_ids(self.special_tokens) # create a mask for special tokens @@ -201,6 +209,8 @@ def _convert_A_B_to_chosen_rejected( nontied_mask: [B, 1] (preference labels that is not tied) """ chosen_label = torch.ones_like(rewards_A, dtype=torch.int64).to(rewards_A.device) # [B, 1] + chosen_mask = chosen_label == 1 + rejected_mask = chosen_label != 1 rewards_chosen = rewards_A rewards_rejected = rewards_B @@ -208,8 +218,9 @@ def _convert_A_B_to_chosen_rejected( if tied_threshold is None: nontied_mask = torch.ones_like(chosen_label, dtype=torch.float32).to(rewards_A.device) else: - nontied_mask = (torch.abs((choice_dist[:, 0] - choice_dist[:, 1]) / torch.sum(choice_dist, dim=-1)) - > tied_threshold) + nontied_mask = ( + torch.abs((choice_dist[:, 0] - choice_dist[:, 1]) / torch.sum(choice_dist, dim=-1)) > tied_threshold + ) print(nontied_mask) return ( rewards_chosen, @@ -238,28 +249,29 @@ def on_step_end(self, args, state, control, **kwargs): model = kwargs.get("model") tokenizer = kwargs.get("tokenizer") - index_no_updates = torch.ones((len(tokenizer), ), dtype=torch.bool) + index_no_updates = torch.ones((len(tokenizer),), dtype=torch.bool) index_no_updates[self.special_token_ids] = False with torch.no_grad(): - model.get_input_embeddings().weight[index_no_updates] = (self.orig_embeds_params[index_no_updates]) + model.get_input_embeddings().weight[index_no_updates] = self.orig_embeds_params[index_no_updates] class VLMRewardTrainer(RewardTrainer): - - def __init__(self, - loss_type="regular", - loss_hyperparameters=None, - tied_threshold=None, - visualization_steps=500, - max_viz_samples=4, - *args, - **kwargs): + def __init__( + self, + loss_type="regular", + loss_hyperparameters={}, + tied_threshold=None, + visualization_steps=500, + max_viz_samples=4, + *args, + **kwargs, + ): super().__init__(*args, **kwargs) self.loss_type = loss_type self.tied_threshold = tied_threshold self.rewards_chosen_accumulated = [] self.rewards_rejected_accumulated = [] - self.loss_hyperparameters = loss_hyperparameters if loss_hyperparameters is not None else {} + self.loss_hyperparameters = loss_hyperparameters self.visualization_steps = visualization_steps self.max_viz_samples = max_viz_samples @@ -279,12 +291,20 @@ def get_eval_dataloader(self, eval_dataset: str | Dataset | None = None) -> Data # If we have persistent workers, don't do a fork bomb especially as eval datasets # don't change during training dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" - if (hasattr(self, "_eval_dataloaders") and dataloader_key in self._eval_dataloaders - and self.args.dataloader_persistent_workers): + if ( + hasattr(self, "_eval_dataloaders") + and dataloader_key in self._eval_dataloaders + and self.args.dataloader_persistent_workers + ): return self.accelerator.prepare(self._eval_dataloaders[dataloader_key]) - eval_dataset = (self.eval_dataset[eval_dataset] if isinstance(eval_dataset, str) else - eval_dataset if eval_dataset is not None else self.eval_dataset) + eval_dataset = ( + self.eval_dataset[eval_dataset] + if isinstance(eval_dataset, str) + else eval_dataset + if eval_dataset is not None + else self.eval_dataset + ) data_collator = self.data_collator @@ -355,107 +375,109 @@ def create_optimizer(self): optimizer_grouped_parameters = [ { "params": [ - p for n, p in opt_model.named_parameters() + p + for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in special_lr_parameters and p.requires_grad) ], - "weight_decay": - self.args.weight_decay, + "weight_decay": self.args.weight_decay, }, { "params": [ - p for n, p in opt_model.named_parameters() + p + for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in special_lr_parameters and p.requires_grad) ], - "weight_decay": - 0.0, + "weight_decay": 0.0, }, ] if visual_parameters: - optimizer_grouped_parameters.extend([ - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n in decay_parameters and n in visual_parameters and p.requires_grad) - ], - "weight_decay": - self.args.weight_decay, - "lr": - self.args.vision_lr, - }, - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n not in decay_parameters and n in visual_parameters and p.requires_grad) - ], - "weight_decay": - 0.0, - "lr": - self.args.vision_lr, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and n in visual_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + "lr": self.args.vision_lr, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and n in visual_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + "lr": self.args.vision_lr, + }, + ] + ) if merger_parameters: - optimizer_grouped_parameters.extend([ - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n in decay_parameters and n in merger_parameters and p.requires_grad) - ], - "weight_decay": - self.args.weight_decay, - "lr": - self.args.merger_lr, - }, - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n not in decay_parameters and n in merger_parameters and p.requires_grad) - ], - "weight_decay": - 0.0, - "lr": - self.args.merger_lr, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and n in merger_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + "lr": self.args.merger_lr, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and n in merger_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + "lr": self.args.merger_lr, + }, + ] + ) if rm_head_parameters: - optimizer_grouped_parameters.extend([ - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n in decay_parameters and n in rm_head_parameters and p.requires_grad) - ], - "weight_decay": - self.args.weight_decay, - "lr": - self.args.rm_head_lr, - }, - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n not in decay_parameters and n in rm_head_parameters and p.requires_grad) - ], - "weight_decay": - 0.0, - "lr": - self.args.rm_head_lr, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and n in rm_head_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + "lr": self.args.rm_head_lr, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and n in rm_head_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + "lr": self.args.rm_head_lr, + }, + ] + ) else: optimizer_grouped_parameters = [ { - "params": - [p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)], - "weight_decay": - self.args.weight_decay, + "params": [ + p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, }, { - "params": - [p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)], - "weight_decay": - 0.0, + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, }, ] @@ -464,14 +486,16 @@ def create_optimizer(self): special_token_embeddings.requires_grad = True - optimizer_grouped_parameters.extend([ - { - # "params": [p for n, p in opt_model.get_input_embeddings().named_parameters() if (p.requires_grad)], - "params": [special_token_embeddings], - "lr": self.args.special_token_lr, - "weight_decay": 0.0, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + # "params": [p for n, p in opt_model.get_input_embeddings().named_parameters() if (p.requires_grad)], + "params": [special_token_embeddings], + "lr": self.args.special_token_lr, + "weight_decay": 0.0, + }, + ] + ) optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model) @@ -484,8 +508,11 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): rewards_B = model(return_dict=True, **inputs["batch_2"])["logits"] # Log to TensorBoard for visualization - if (hasattr(self.state, 'global_step') and self.state.global_step % self.visualization_steps == 0 - and self.state.global_step > 0): + if ( + hasattr(self.state, "global_step") + and self.state.global_step % self.visualization_steps == 0 + and self.state.global_step > 0 + ): # Pass the original inputs which should contain the text prompts self._log_training_visualization(inputs, rewards_A, rewards_B) @@ -512,7 +539,7 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): loss = loss.mean() elif self.loss_type == "likelihood_displacement": # Bradley-Terry model - loss = -nn.functional.logsigmoid(rewards_chosen - self.loss_hyperparameters['tau'] * rewards_rejected) + loss = -nn.functional.logsigmoid(rewards_chosen - self.loss_hyperparameters["tau"] * rewards_rejected) out_mask = nontied_mask loss = loss * out_mask loss = loss.mean() @@ -529,8 +556,11 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): log_k = math.log(k) log_k2_sub_1 = math.log(k**2 - 1) bt_loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) - same_loss = (-nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) - - nn.functional.logsigmoid(rewards_rejected - rewards_chosen - log_k) - log_k2_sub_1) + same_loss = ( + -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) + - nn.functional.logsigmoid(rewards_rejected - rewards_chosen - log_k) + - log_k2_sub_1 + ) loss = bt_loss * nontied_mask.float() + same_loss * (1 - nontied_mask.float()) out_mask = torch.ones_like(nontied_mask, dtype=torch.float32).to(rewards_A.device) # [B, 1] loss = loss * out_mask @@ -566,8 +596,9 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): mean_z = mean_chosen - mean_rejected sigma_z = torch.sqrt(sigma_chosen**2 + sigma_rejected**2) - z_samples = torch.randn(batch_size, 1000).to(sigma_z.device).to( - torch.float16) * sigma_z.unsqueeze(1).repeat(1, 1000) + mean_z.unsqueeze(1).repeat(1, 1000) + z_samples = torch.randn(batch_size, 1000).to(sigma_z.device).to(torch.float16) * sigma_z.unsqueeze( + 1 + ).repeat(1, 1000) + mean_z.unsqueeze(1).repeat(1, 1000) loss = -torch.nn.functional.logsigmoid(z_samples).mean() else: raise NotImplementedError(f"Loss type {self.loss_type} not implemented.") @@ -620,13 +651,14 @@ def _log_training_visualization(self, inputs, rewards_A, rewards_B): try: # Get tensorboard writer from trainer writer = None - if hasattr(self, 'log_metrics') and hasattr(self.args, - 'report_to') and 'tensorboard' in self.args.report_to: + if hasattr(self, "log_metrics"): # Try to get the writer from the logger - from torch.utils.tensorboard import SummaryWriter - if not hasattr(self, '_tb_writer'): - self._tb_writer = SummaryWriter(log_dir=self.args.logging_dir) - writer = self._tb_writer + if hasattr(self.args, "report_to") and "tensorboard" in self.args.report_to: + from torch.utils.tensorboard import SummaryWriter + + if not hasattr(self, "_tb_writer"): + self._tb_writer = SummaryWriter(log_dir=self.args.logging_dir) + writer = self._tb_writer if writer is None: return @@ -644,34 +676,36 @@ def _log_training_visualization(self, inputs, rewards_A, rewards_B): score_B_val = float(score_B.mean()) if score_B.ndim > 0 else float(score_B) score_diff = score_A_val - score_B_val - writer.add_scalar(f'train_viz/sample_{i}/score_A', score_A_val, step) - writer.add_scalar(f'train_viz/sample_{i}/score_B', score_B_val, step) - writer.add_scalar(f'train_viz/sample_{i}/score_diff', score_diff, step) + writer.add_scalar(f"train_viz/sample_{i}/score_A", score_A_val, step) + writer.add_scalar(f"train_viz/sample_{i}/score_B", score_B_val, step) + writer.add_scalar(f"train_viz/sample_{i}/score_diff", score_diff, step) try: # Get image data from inputs - image_A = inputs['image_1'][i] if 'image_1' in inputs else None - image_B = inputs['image_2'][i] if 'image_2' in inputs else None + image_A = inputs["image_1"][i] if "image_1" in inputs else None + image_B = inputs["image_2"][i] if "image_2" in inputs else None # Get prompt text from the original batch (now properly stored) - prompt_A = inputs.get('text_1', ['Unknown prompt'])[i] if 'text_1' in inputs else 'Unknown prompt' + prompt_A = inputs.get("text_1", ["Unknown prompt"])[i] if "text_1" in inputs else "Unknown prompt" fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 8)) - fig.text(0.05, - 0.05, - f'Prompt:\n{prompt_A[:200]}{"..." if len(prompt_A) > 200 else ""}', - ha='left', - va='bottom', - fontsize=8, - wrap=True, - bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue", alpha=0.7)) + fig.text( + 0.05, + 0.05, + f"Prompt:\n{prompt_A[:200]}{'...' if len(prompt_A) > 200 else ''}", + ha="left", + va="bottom", + fontsize=8, + wrap=True, + bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue", alpha=0.7), + ) img_A_np = np.array(image_A) if img_A_np.ndim == 3 and img_A_np.shape[0] == 3: # CHW format img_A_np = np.transpose(img_A_np, (1, 2, 0)) img_A_np = np.clip(img_A_np, 0, 1) # Ensure values are in [0,1] axes[0].imshow(img_A_np) - axes[0].set_title(f'Image A - Score: {score_A_val:.3f}') - axes[0].axis('off') + axes[0].set_title(f"Image A - Score: {score_A_val:.3f}") + axes[0].axis("off") img_B_np = np.array(image_B) if img_B_np.ndim == 3 and img_B_np.shape[0] == 3: # CHW format @@ -679,18 +713,19 @@ def _log_training_visualization(self, inputs, rewards_A, rewards_B): img_B_np = np.clip(img_B_np, 0, 1) # Ensure values are in [0,1] axes[1].imshow(img_B_np) - axes[1].set_title(f'Image B - Score: {score_B_val:.3f}') - axes[1].axis('off') + axes[1].set_title(f"Image B - Score: {score_B_val:.3f}") + axes[1].axis("off") # Add prediction info winner = "A" if score_diff > 0 else "B" plt.suptitle( - f'Step {step} - Sample {i} | Predicted Winner: Image {winner} | Diff: {score_diff:.3f}', - fontsize=14) + f"Step {step} - Sample {i} | Predicted Winner: Image {winner} | Diff: {score_diff:.3f}", + fontsize=14, + ) plt.tight_layout() # Log figure to tensorboard - writer.add_figure(f'train_viz/sample_{i}_comparison', fig, step) + writer.add_figure(f"train_viz/sample_{i}_comparison", fig, step) plt.close(fig) except Exception as viz_error: print(f"Warning: Could not extract images for visualization: {viz_error}") @@ -700,17 +735,20 @@ def _log_training_visualization(self, inputs, rewards_A, rewards_B): all_scores_A = rewards_A.float().detach().cpu().numpy() all_scores_B = rewards_B.float().detach().cpu().numpy() - writer.add_histogram('train_viz/all_scores_A', all_scores_A, step) - writer.add_histogram('train_viz/all_scores_B', all_scores_B, step) - writer.add_scalar('train_viz/mean_score_A', float(all_scores_A.mean()), step) - writer.add_scalar('train_viz/mean_score_B', float(all_scores_B.mean()), step) - writer.add_scalar('train_viz/mean_score_diff', float((all_scores_A - all_scores_B).mean()), step) + writer.add_histogram("train_viz/all_scores_A", all_scores_A, step) + writer.add_histogram("train_viz/all_scores_B", all_scores_B, step) + writer.add_scalar("train_viz/mean_score_A", float(all_scores_A.mean()), step) + writer.add_scalar("train_viz/mean_score_B", float(all_scores_B.mean()), step) + writer.add_scalar( + "train_viz/mean_score_diff", + float((all_scores_A - all_scores_B).mean()), + step, + ) except Exception as e: print(f"Error in training visualization: {e}") def _save_checkpoint(self, model, trial, metrics=None): - if isinstance(self.model, PeftModel): checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" @@ -726,8 +764,9 @@ def _save_checkpoint(self, model, trial, metrics=None): # pdb.set_trace() if not self.args.save_full_model: - non_lora_weights = get_peft_state_non_lora_maybe_zero_3(self.model.named_parameters(), - require_grad_only=True) + non_lora_weights = get_peft_state_non_lora_maybe_zero_3( + self.model.named_parameters(), require_grad_only=True + ) torch.save( non_lora_weights, os.path.join(output_dir, "non_lora_state_dict.pth"), @@ -750,7 +789,7 @@ def _save(self, output_dir: str | None = None, state_dict=None): logger.info(f"Saving model checkpoint to {output_dir}") # pdb.set_trace() - supported_classes = ((PreTrainedModel, ) if not is_peft_available() else (PreTrainedModel, PeftModel)) + supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel) # Save a trained model and configuration using `save_pretrained()`. # They can then be reloaded using `from_pretrained()` if not isinstance(self.model, supported_classes): @@ -773,16 +812,15 @@ def _save(self, output_dir: str | None = None, state_dict=None): ) else: torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + elif not self.args.save_full_model: + state_dict = {k: v for k, v in state_dict.items() if "wte" not in k} + self.model.save_pretrained( + output_dir, + state_dict=state_dict, + safe_serialization=self.args.save_safetensors, + ) else: - if not self.args.save_full_model: - state_dict = {k: v for k, v in state_dict.items() if "wte" not in k} - self.model.save_pretrained( - output_dir, - state_dict=state_dict, - safe_serialization=self.args.save_safetensors, - ) - else: - torch.save(state_dict, os.path.join(output_dir, "model.pth")) + torch.save(state_dict, os.path.join(output_dir, "model.pth")) if self.tokenizer is not None: os.makedirs(os.path.join(output_dir, "tokenizer"), exist_ok=True) @@ -809,9 +847,11 @@ def compute_multi_attr_accuracy(eval_pred, metainfo_idxs=None) -> dict[str, floa accuracy = np.sum(rewards_chosen > rewards_rejected) / total_count - metrics.update({ - "Acc": accuracy, - "R_chosen_avg": rewards_chosen_avg, - "R_rejected_avg": rewards_rejected_avg, - }) + metrics.update( + { + "Acc": accuracy, + "R_chosen_avg": rewards_chosen_avg, + "R_rejected_avg": rewards_rejected_avg, + } + ) return metrics diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/train.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/train.py index 37be76916a..0d69bdba9d 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/train.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/train.py @@ -1,24 +1,29 @@ import json import os -import fire from dataclasses import asdict from functools import partial + +import fire import torch +from peft import LoraConfig, get_peft_model +from transformers import AutoProcessor +from trl import get_kbit_device_map, get_quantization_config + +from .dataset.data_collator_qwen import QWen2VLDataCollator +from .dataset.pairwise_dataset import PairwiseOriginalDataset +from .model.differentiable_image_processor import Qwen2VLImageProcessor from .model.qwen2vl_trainer import ( + PartialEmbeddingUpdateCallback, Qwen2VLRewardModelBT, VLMRewardTrainer, compute_multi_attr_accuracy, - PartialEmbeddingUpdateCallback, ) -from .dataset.pairwise_dataset import PairwiseOriginalDataset -from .dataset.data_collator_qwen import QWen2VLDataCollator -from .utils.parser import ModelConfig, PEFTLoraConfig, TrainingConfig, DataConfig -from .utils.training_utils import load_model_from_checkpoint, find_target_linear_names -from .utils.parser import parse_args_with_yaml -from transformers import AutoProcessor -from peft import LoraConfig, get_peft_model -from trl import get_kbit_device_map, get_quantization_config -from .model.differentiable_image_processor import Qwen2VLImageProcessor +from .utils.parser import DataConfig, ModelConfig, PEFTLoraConfig, TrainingConfig, parse_args_with_yaml +from .utils.training_utils import ( + find_target_linear_names, + load_model_from_checkpoint, +) + try: import flash_attn except ImportError: @@ -34,19 +39,24 @@ def create_model_and_processor( differentiable=False, ): # create model - torch_dtype = (model_config.torch_dtype if model_config.torch_dtype in ["auto", None] else getattr( - torch, model_config.torch_dtype)) + torch_dtype = ( + model_config.torch_dtype + if model_config.torch_dtype in ["auto", None] + else getattr(torch, model_config.torch_dtype) + ) quantization_config = get_quantization_config(model_config) - model_kwargs = dict(revision=model_config.model_revision, - device_map=get_kbit_device_map() if quantization_config is not None else None, - quantization_config=quantization_config, - use_cache=False) + model_kwargs = dict( + revision=model_config.model_revision, + device_map=get_kbit_device_map() if quantization_config is not None else None, + quantization_config=quantization_config, + use_cache=False, + ) # create processor and set padding - processor = AutoProcessor.from_pretrained(model_config.model_name_or_path, - padding_side="right", - cache_dir=cache_dir) + processor = AutoProcessor.from_pretrained( + model_config.model_name_or_path, padding_side="right", cache_dir=cache_dir + ) if differentiable: processor.image_processor = Qwen2VLImageProcessor() @@ -63,8 +73,9 @@ def create_model_and_processor( reward_token=model_config.reward_token, special_token_ids=special_token_ids, torch_dtype=torch_dtype, - attn_implementation=("flash_attention_2" - if not training_args.disable_flash_attn2 and flash_attn is not None else "sdpa"), + attn_implementation=( + "flash_attention_2" if not training_args.disable_flash_attn2 and flash_attn is not None else "sdpa" + ), cache_dir=cache_dir, rm_head_type=model_config.rm_head_type, rm_head_kwargs=model_config.rm_head_kwargs, @@ -137,18 +148,26 @@ def set_requires_grad(parameters, requires_grad): def train(config, local_rank=0, debug=False): - ## ===> Step 1: Parse arguments - (data_config, training_args, model_config, peft_lora_config), config_path = (parse_args_with_yaml( - (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig), config, is_train=True)) + ( + ( + data_config, + training_args, + model_config, + peft_lora_config, + ), + config_path, + ) = parse_args_with_yaml((DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig), config, is_train=True) training_args.output_dir = os.path.join(training_args.output_dir, config.split("/")[-1].split(".")[0]) training_args.logging_dir = training_args.output_dir # check valid (lora config) - assert not (peft_lora_config.lora_enable and model_config.freeze_llm - ), "When using LoRA, the LLM should not be frozen. If you want to freeze the LLM, please disable LoRA." + assert not (peft_lora_config.lora_enable and model_config.freeze_llm), ( + "When using LoRA, the LLM should not be frozen. If you want to freeze the LLM, please disable LoRA." + ) if not peft_lora_config.lora_enable: - assert (not peft_lora_config.vision_lora - ), "Error: model_config.lora_enable is not enabled, but model_config.vision_lora is enabled." + assert not peft_lora_config.vision_lora, ( + "Error: model_config.lora_enable is not enabled, but model_config.vision_lora is enabled." + ) else: if peft_lora_config.lora_namespan_exclude is None: peft_lora_config.lora_namespan_exclude = [] @@ -183,13 +202,17 @@ def train(config, local_rank=0, debug=False): set_requires_grad(model_to_configure.visual.parameters(), not model_config.freeze_vision_tower) set_requires_grad(model_to_configure.visual.merger.parameters(), model_config.tune_merger) - if model_config.trainable_visual_layers: # This is inverse order to index of model.visual.blocks, set -1 to unfreeze all layers - assert model_config.trainable_visual_layers <= len( - model_to_configure.visual.blocks - ), "trainable_visual_layers should be less than or equal to the number of visual blocks" - freeze_layer_num = len( - model_to_configure.visual.blocks - ) - model_config.trainable_visual_layers if model_config.trainable_visual_layers > 0 else 0 + if ( + model_config.trainable_visual_layers + ): # This is inverse order to index of model.visual.blocks, set -1 to unfreeze all layers + assert model_config.trainable_visual_layers <= len(model_to_configure.visual.blocks), ( + "trainable_visual_layers should be less than or equal to the number of visual blocks" + ) + freeze_layer_num = ( + len(model_to_configure.visual.blocks) - model_config.trainable_visual_layers + if model_config.trainable_visual_layers > 0 + else 0 + ) for index, layer in enumerate(model_to_configure.visual.blocks): if index < freeze_layer_num: set_requires_grad(layer.parameters(), False) @@ -227,9 +250,8 @@ def train(config, local_rank=0, debug=False): ) compute_metrics = partial(compute_multi_attr_accuracy) - actual_batch_size = (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * - num_gpu) - total_steps = (training_args.num_train_epochs * len(train_dataset) // actual_batch_size) + actual_batch_size = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * num_gpu + total_steps = training_args.num_train_epochs * len(train_dataset) // actual_batch_size if training_args.save_epochs is not None: training_args.save_steps = round(training_args.save_epochs * len(train_dataset) / actual_batch_size) if training_args.eval_epochs is not None: diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/parser.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/parser.py index b806310e31..cdc0783e7d 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/parser.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/parser.py @@ -1,8 +1,9 @@ +import contextlib +from dataclasses import dataclass, field from typing import Any, Literal + from omegaconf import OmegaConf -from transformers import HfArgumentParser -from dataclasses import dataclass, field -from transformers import TrainingArguments +from transformers import HfArgumentParser, TrainingArguments @dataclass @@ -61,10 +62,10 @@ class PEFTLoraConfig: num_lora_modules: int = -1 def __post_init__(self): - if (isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1): + if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1: self.lora_target_modules = self.lora_target_modules[0] - if (isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1): + if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1: self.lora_namespan_exclude = self.lora_namespan_exclude[0] @@ -91,8 +92,8 @@ class ModelConfig: bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4" use_bnb_nested_quant: bool = False reward_token: Literal["last", "mean", "special"] = "last" - loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled"] = ("regular") - loss_hyperparameters: dict = field(default_factory=lambda: {}) + loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled"] = "regular" + loss_hyperparameters: dict = field(default_factory=dict) checkpoint_path: str | None = None def __post_init__(self): @@ -117,12 +118,12 @@ def parse_args_with_yaml( ) -> tuple[Any, ...]: """ Parse arguments using HfArgumentParser with OmegaConf for YAML support. - + Args: dataclass_types: Tuple of dataclass types for HfArgumentParser args: Optional arguments (if None, will read from sys.argv) allow_extra_keys: Whether to allow extra keys in config - + Returns: Tuple of parsed dataclass instances """ @@ -130,13 +131,40 @@ def parse_args_with_yaml( # Load YAML config and merge with command line overrides args = OmegaConf.to_container(OmegaConf.load(config_path)) if not is_train: - args.pop('deepspeed', None) + args.pop("deepspeed", None) + + @contextlib.contextmanager + def _disable_accelerate_state_reset(enabled: bool): + if not enabled: + yield + return + try: + from accelerate.state import AcceleratorState, PartialState + except Exception: + # If accelerate is unavailable, just continue. + yield + return + orig_acc_reset = AcceleratorState._reset_state + orig_partial_reset = PartialState._reset_state + + def _no_reset_state(*_args, **_kwargs): + return None + + AcceleratorState._reset_state = staticmethod(_no_reset_state) + PartialState._reset_state = staticmethod(_no_reset_state) + try: + yield + finally: + AcceleratorState._reset_state = orig_acc_reset + PartialState._reset_state = orig_partial_reset # Parse with HfArgumentParser parser = HfArgumentParser(dataclass_types) - return parser.parse_dict(args, allow_extra_keys=allow_extra_keys), config_path + with _disable_accelerate_state_reset(enabled=not is_train): + return parser.parse_dict(args, allow_extra_keys=allow_extra_keys), config_path if __name__ == "__main__": data_config, training_args, model_config, peft_lora_config = parse_args_with_yaml( - (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig)) + (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig) + ) diff --git a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/training_utils.py b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/training_utils.py index c8e1de4c7b..69d3af487d 100755 --- a/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/training_utils.py +++ b/fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/training_utils.py @@ -1,7 +1,8 @@ -import torch -import os import glob +import os + import safetensors +import torch def maybe_zero_3(param, ignore_status=False, name=None): @@ -9,8 +10,9 @@ def maybe_zero_3(param, ignore_status=False, name=None): from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus if hasattr(param, "ds_id"): - if param.ds_status == ZeroParamStatus.NOT_AVAILABLE and not ignore_status: - print(f"Parameter {name} is not available in ZeRO-3, please check the ZeRO-3 status.") + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if not ignore_status: + print(f"Parameter {name} is not available in ZeRO-3, please check the ZeRO-3 status.") with zero.GatheredParameters([param]): param = param.data.detach().cpu().clone() else: @@ -52,8 +54,9 @@ def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True): return to_return -def _insert_adapter_name_into_state_dict(state_dict: dict[str, torch.Tensor], adapter_name: str, - parameter_prefix: str) -> dict[str, torch.Tensor]: +def _insert_adapter_name_into_state_dict( + state_dict: dict[str, torch.Tensor], adapter_name: str, parameter_prefix: str +) -> dict[str, torch.Tensor]: """Utility function to remap the state_dict keys to fit the PEFT model by inserting the adapter name.""" peft_model_state_dict = {} for key, val in state_dict.items(): @@ -107,9 +110,9 @@ def load_model_from_checkpoint(model, checkpoint_dir, checkpoint_step): lora_state_dict = safetensors.torch.load_file(lora_ckpt) non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu") - lora_state_dict = _insert_adapter_name_into_state_dict(lora_state_dict, - adapter_name="default", - parameter_prefix="lora_") + lora_state_dict = _insert_adapter_name_into_state_dict( + lora_state_dict, adapter_name="default", parameter_prefix="lora_" + ) model_state_dict = model.state_dict() model_state_dict.update(non_lora_state_dict) @@ -119,14 +122,12 @@ def load_model_from_checkpoint(model, checkpoint_dir, checkpoint_step): return model, checkpoint_step -def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=None, verbose=False): +def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=[], verbose=False): """ Find the target linear modules for LoRA. """ linear_cls = torch.nn.Linear embedding_cls = torch.nn.Embedding - if lora_namespan_exclude is None: - lora_namespan_exclude = [] lora_module_names = [] for name, module in model.named_modules(): @@ -134,7 +135,7 @@ def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=N # print(f"Excluding module: {name}") continue - if isinstance(module, linear_cls | embedding_cls): + if isinstance(module, (linear_cls, embedding_cls)): lora_module_names.append(name) if num_lora_modules > 0: diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/__init__.py b/fastvideo/train/methods/rl/reward/VideoAlign/__init__.py index 640f0c63ce..bcda1c90bc 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/__init__.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/__init__.py @@ -1,7 +1,7 @@ """Vendored runtime subset of VideoAlign. Source: https://github.com/KlingAIResearch/VideoAlign -Commit: 219ab9db64c045e5181a2202d11f686439351292 +Commit: aba26b658fec7d9fd30c295187b548ea673c8769 Purpose: Runtime reward inference integration for FastVideo GenRL. This is temporary minimal vendoring for PR integration. It is expected to be diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/checkpoints/README.md b/fastvideo/train/methods/rl/reward/VideoAlign/checkpoints/README.md index 457a3b6b74..a0b6e0af27 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/checkpoints/README.md +++ b/fastvideo/train/methods/rl/reward/VideoAlign/checkpoints/README.md @@ -4,5 +4,10 @@ Please download our checkpoints from [Huggingface](https://huggingface.co/KwaiVG cd checkpoints git lfs install git clone https://huggingface.co/KwaiVGI/VideoReward +# Move all files from VideoReward to checkpoints directory +mv VideoReward/* . +mv VideoReward/.* . 2>/dev/null || true # Move hidden files, ignore errors if none exist +# Remove the empty VideoReward directory +rmdir VideoReward cd .. -``` +``` \ No newline at end of file diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/data.py b/fastvideo/train/methods/rl/reward/VideoAlign/data.py index 9a021daabe..e86edcb0bd 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/data.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/data.py @@ -2,6 +2,7 @@ import torch from .prompt_template import build_prompt + # from qwen_vl_utils import process_vision_info from .vision_process import process_vision_info @@ -23,14 +24,16 @@ class DataConfig: use_tied_data: bool = True -def convert_GSB_csv_to_reward_data(example, - data_dir, - eval_dims=None, - max_pixels=448 * 448, - fps=2.0, - num_frames=None, - prompt_template_type="none", - sample_type="uniform"): +def convert_GSB_csv_to_reward_data( + example, + data_dir, + eval_dims=["VQ"], + max_pixels=448 * 448, + fps=2.0, + num_frames=None, + prompt_template_type="none", + sample_type="uniform", +): """ Convert Good/Same/Bad csv data to reward data. @@ -45,45 +48,45 @@ def convert_GSB_csv_to_reward_data(example, Returns: dict: A dictionary containing the reward data. """ - if eval_dims is None: - eval_dims = ["VQ"] - A_data = [{ - "role": - "user", - "content": [ - { - "type": "video", - "video": f"file://{data_dir}/{example['path_A']}", - "max_pixels": max_pixels, - "fps": fps if num_frames is None else None, - "nframes": min(num_frames, example["num_frames_A"]) if num_frames is not None else None, - "sample_type": sample_type, - }, - { - "type": "text", - "text": build_prompt(example["prompt"], eval_dims, prompt_template_type) - }, - ], - }] - B_data = [{ - "role": - "user", - "content": [ - { - "type": "video", - "video": f"file://{data_dir}/{example['path_B']}", - "max_pixels": max_pixels, - "fps": fps if num_frames is None else None, - "nframes": min(num_frames, example["num_frames_B"]) if num_frames is not None else None, - "sample_type": sample_type, - }, - { - "type": "text", - "text": build_prompt(example["prompt"], eval_dims, prompt_template_type) - }, - ], - }] + A_data = [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": f"file://{data_dir}/{example['path_A']}", + "max_pixels": max_pixels, + "fps": fps if num_frames is None else None, + "nframes": (min(num_frames, example["num_frames_A"]) if num_frames is not None else None), + "sample_type": sample_type, + }, + { + "type": "text", + "text": build_prompt(example["prompt"], eval_dims, prompt_template_type), + }, + ], + } + ] + B_data = [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": f"file://{data_dir}/{example['path_B']}", + "max_pixels": max_pixels, + "fps": fps if num_frames is None else None, + "nframes": (min(num_frames, example["num_frames_B"]) if num_frames is not None else None), + "sample_type": sample_type, + }, + { + "type": "text", + "text": build_prompt(example["prompt"], eval_dims, prompt_template_type), + }, + ], + } + ] chosen_labels = [] A_scores = [] @@ -127,8 +130,8 @@ def convert_GSB_csv_to_reward_data(example, A_scores = torch.tensor(A_scores, dtype=torch.float) B_scores = torch.tensor(B_scores, dtype=torch.float) metainfo_idx = None - if 'metainfo_idx' in example: - metainfo_idx = example['metainfo_idx'] + if "metainfo_idx" in example: + metainfo_idx = example["metainfo_idx"] return { "A_data": A_data, @@ -141,7 +144,6 @@ def convert_GSB_csv_to_reward_data(example, class QWen2VLDataCollator: - def __init__(self, processor, add_noise=False, p_shuffle_frames=0.0, p_color_jitter=0.0): self.processor = processor self.add_noise = add_noise @@ -156,25 +158,28 @@ def _clean_message(self, message): """ remove unnecessary keys from message(very very necessary) """ - message_content = message[0]["content"][0] - out_message = [{ - "role": - "user", - "content": [ - { - "type": "video", - "video": message_content["video"], - "max_pixels": message_content["max_pixels"], - "fps": message_content.get("fps", None), - "nframes": message_content.get("nframes", None), - "sample_type": message_content.get("sample_type", "uniform"), - }, - { - "type": "text", - "text": message[0]["content"][1]["text"] - }, - ], - }] + out_message = [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": message[0]["content"][0]["video"], + "max_pixels": message[0]["content"][0]["max_pixels"], + "fps": (message[0]["content"][0]["fps"] if "fps" in message[0]["content"][0] else None), + "nframes": ( + message[0]["content"][0]["nframes"] if "nframes" in message[0]["content"][0] else None + ), + "sample_type": ( + message[0]["content"][0]["sample_type"] + if "sample_type" in message[0]["content"][0] + else "uniform" + ), + }, + {"type": "text", "text": message[0]["content"][1]["text"]}, + ], + } + ] if out_message[0]["content"][0]["fps"] is None: out_message[0]["content"][0].pop("fps") @@ -183,20 +188,21 @@ def _clean_message(self, message): return out_message - def _pad_sequence(self, sequences, attention_mask, max_len, padding_side='right'): + def _pad_sequence(self, sequences, attention_mask, max_len, padding_side="right"): """ Pad the sequences to the maximum length. """ - assert padding_side in ['right', 'left'] + assert padding_side in ["right", "left"] if sequences.shape[1] >= max_len: return sequences, attention_mask pad_len = max_len - sequences.shape[1] - padding = (0, pad_len) if padding_side == 'right' else (pad_len, 0) + padding = (0, pad_len) if padding_side == "right" else (pad_len, 0) - sequences_padded = torch.nn.functional.pad(sequences, padding, 'constant', - self.processor.tokenizer.pad_token_id) - attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, 'constant', 0) + sequences_padded = torch.nn.functional.pad( + sequences, padding, "constant", self.processor.tokenizer.pad_token_id + ) + attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, "constant", 0) return sequences_padded, attention_mask_padded @@ -245,12 +251,12 @@ def __call__(self, features, enable_noise=True): # pdb.set_trace() max_len = max(batch_A["input_ids"].shape[1], batch_B["input_ids"].shape[1]) - batch_A["input_ids"], batch_A["attention_mask"] = self._pad_sequence(batch_A["input_ids"], - batch_A["attention_mask"], max_len, - "right") - batch_B["input_ids"], batch_B["attention_mask"] = self._pad_sequence(batch_B["input_ids"], - batch_B["attention_mask"], max_len, - "right") + batch_A["input_ids"], batch_A["attention_mask"] = self._pad_sequence( + batch_A["input_ids"], batch_A["attention_mask"], max_len, "right" + ) + batch_B["input_ids"], batch_B["attention_mask"] = self._pad_sequence( + batch_B["input_ids"], batch_B["attention_mask"], max_len, "right" + ) # print(f"Batch A: {batch_A['input_ids'].shape}, Batch B: {batch_B['input_ids'].shape}") chosen_label = torch.stack([torch.tensor(feature["chosen_label"]) for feature in features]) diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/inference.py b/fastvideo/train/methods/rl/reward/VideoAlign/inference.py index b3231348e3..f34cc7d360 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/inference.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/inference.py @@ -3,13 +3,11 @@ from collections.abc import Mapping import torch -from .vision_process import process_vision_info - from .data import DataConfig -from .utils import ModelConfig, PEFTLoraConfig, TrainingConfig -from .utils import load_model_from_checkpoint -from .train_reward import create_model_and_processor from .prompt_template import build_prompt +from .train_reward import create_model_and_processor +from .utils import ModelConfig, PEFTLoraConfig, TrainingConfig, load_model_from_checkpoint +from .vision_process import process_vision_info def load_configs_from_json(config_path): @@ -20,15 +18,31 @@ def load_configs_from_json(config_path): del config_dict["data_config"]["meta_data"] del config_dict["data_config"]["data_dir"] - return config_dict["data_config"], None, config_dict["model_config"], config_dict["peft_lora_config"], \ - config_dict.get("inference_config", None) + return ( + config_dict["data_config"], + None, + config_dict["model_config"], + config_dict["peft_lora_config"], + config_dict["inference_config"] if "inference_config" in config_dict else None, + ) class VideoVLMRewardInference: - - def __init__(self, load_from_pretrained, load_from_pretrained_step=-1, device='cuda', dtype=torch.bfloat16): + def __init__( + self, + load_from_pretrained, + load_from_pretrained_step=-1, + device="cuda", + dtype=torch.bfloat16, + ): config_path = os.path.join(load_from_pretrained, "model_config.json") - data_config, _, model_config, peft_lora_config, inference_config = load_configs_from_json(config_path) + ( + data_config, + _, + model_config, + peft_lora_config, + inference_config, + ) = load_configs_from_json(config_path) data_config = DataConfig(**data_config) model_config = ModelConfig(**model_config) peft_lora_config = PEFTLoraConfig(**peft_lora_config) @@ -38,8 +52,8 @@ def __init__(self, load_from_pretrained, load_from_pretrained_step=-1, device='c load_from_pretrained_step=load_from_pretrained_step, gradient_checkpointing=False, disable_flash_attn2=False, - bf16=dtype == torch.bfloat16, - fp16=dtype == torch.float16, + bf16=True if dtype == torch.bfloat16 else False, + fp16=True if dtype == torch.float16 else False, output_dir="", ) @@ -66,26 +80,26 @@ def __init__(self, load_from_pretrained, load_from_pretrained_step=-1, device='c def _norm(self, reward): if self.inference_config is None: return reward - else: - reward['VQ'] = (reward['VQ'] - self.inference_config['VQ_mean']) / self.inference_config['VQ_std'] - reward['MQ'] = (reward['MQ'] - self.inference_config['MQ_mean']) / self.inference_config['MQ_std'] - reward['TA'] = (reward['TA'] - self.inference_config['TA_mean']) / self.inference_config['TA_std'] - return reward + reward["VQ"] = (reward["VQ"] - self.inference_config["VQ_mean"]) / self.inference_config["VQ_std"] + reward["MQ"] = (reward["MQ"] - self.inference_config["MQ_mean"]) / self.inference_config["MQ_std"] + reward["TA"] = (reward["TA"] - self.inference_config["TA_mean"]) / self.inference_config["TA_std"] + return reward - def _pad_sequence(self, sequences, attention_mask, max_len, padding_side='right'): + def _pad_sequence(self, sequences, attention_mask, max_len, padding_side="right"): """ Pad the sequences to the maximum length. """ - assert padding_side in ['right', 'left'] + assert padding_side in ["right", "left"] if sequences.shape[1] >= max_len: return sequences, attention_mask pad_len = max_len - sequences.shape[1] - padding = (0, pad_len) if padding_side == 'right' else (pad_len, 0) + padding = (0, pad_len) if padding_side == "right" else (pad_len, 0) - sequences_padded = torch.nn.functional.pad(sequences, padding, 'constant', - self.processor.tokenizer.pad_token_id) - attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, 'constant', 0) + sequences_padded = torch.nn.functional.pad( + sequences, padding, "constant", self.processor.tokenizer.pad_token_id + ) + attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, "constant", 0) return sequences_padded, attention_mask_padded @@ -96,9 +110,9 @@ def _prepare_input(self, data): """ if isinstance(data, Mapping): return type(data)({k: self._prepare_input(v) for k, v in data.items()}) - elif isinstance(data, tuple | list): + if isinstance(data, (tuple, list)): return type(data)(self._prepare_input(v) for v in data) - elif isinstance(data, torch.Tensor): + if isinstance(data, torch.Tensor): kwargs = {"device": self.device} ## TODO: Maybe need to add dtype # if self.is_deepspeed_enabled and (torch.is_floating_point(data) or torch.is_complex(data)): @@ -132,47 +146,57 @@ def prepare_batch( max_pixels = self.data_config.max_frame_pixels if max_pixels is None else max_pixels if num_frames is None: - chat_data = [[ - { - "role": - "user", - "content": [ - { - "type": "video", - "video": f"file://{video_path}", - "max_pixels": max_pixels, - "fps": fps, - "sample_type": self.data_config.sample_type, - }, - { - "type": "text", - "text": build_prompt(prompt, self.data_config.eval_dim, - self.data_config.prompt_template_type) - }, - ], - }, - ] for video_path, prompt in zip(video_paths, prompts, strict=False)] + chat_data = [ + [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": f"file://{video_path}", + "max_pixels": max_pixels, + "fps": fps, + "sample_type": self.data_config.sample_type, + }, + { + "type": "text", + "text": build_prompt( + prompt, + self.data_config.eval_dim, + self.data_config.prompt_template_type, + ), + }, + ], + }, + ] + for video_path, prompt in zip(video_paths, prompts) + ] else: - chat_data = [[ - { - "role": - "user", - "content": [ - { - "type": "video", - "video": f"file://{video_path}", - "max_pixels": max_pixels, - "nframes": num_frames, - "sample_type": self.data_config.sample_type, - }, - { - "type": "text", - "text": build_prompt(prompt, self.data_config.eval_dim, - self.data_config.prompt_template_type) - }, - ], - }, - ] for video_path, prompt in zip(video_paths, prompts, strict=False)] + chat_data = [ + [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": f"file://{video_path}", + "max_pixels": max_pixels, + "nframes": num_frames, + "sample_type": self.data_config.sample_type, + }, + { + "type": "text", + "text": build_prompt( + prompt, + self.data_config.eval_dim, + self.data_config.prompt_template_type, + ), + }, + ], + }, + ] + for video_path, prompt in zip(video_paths, prompts) + ] image_inputs, video_inputs = process_vision_info(chat_data) batch = self.processor( @@ -186,7 +210,15 @@ def prepare_batch( batch = self._prepare_inputs(batch) return batch - def reward(self, video_paths, prompts, fps=None, num_frames=None, max_pixels=None, use_norm=True): + def reward( + self, + video_paths, + prompts, + fps=None, + num_frames=None, + max_pixels=None, + use_norm=True, + ): """ Inputs: video_paths: List[str], B paths of the videos. @@ -204,11 +236,11 @@ def reward(self, video_paths, prompts, fps=None, num_frames=None, max_pixels=Non batch = self.prepare_batch(video_paths, prompts, fps, num_frames, max_pixels) rewards = self.model(return_dict=True, **batch)["logits"] - rewards = [{'VQ': reward[0].item(), 'MQ': reward[1].item(), 'TA': reward[2].item()} for reward in rewards] + rewards = [{"VQ": reward[0].item(), "MQ": reward[1].item(), "TA": reward[2].item()} for reward in rewards] for i in range(len(rewards)): if use_norm: rewards[i] = self._norm(rewards[i]) - rewards[i]['Overall'] = rewards[i]['VQ'] + rewards[i]['MQ'] + rewards[i]['TA'] + rewards[i]["Overall"] = rewards[i]["VQ"] + rewards[i]["MQ"] + rewards[i]["TA"] return rewards diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/prompt_template.py b/fastvideo/train/methods/rl/reward/VideoAlign/prompt_template.py index be54ce3baf..c8940f02e2 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/prompt_template.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/prompt_template.py @@ -9,12 +9,21 @@ """ DIMENSION_DESCRIPTIONS = { - 'VQ': ['visual quality', 'the quality of the video in terms of clearness, resolution, brightness, and color'], - 'TA': ['text-to-video alignment', 'the alignment between the text prompt and the video content and motion'], - 'MQ': ['motion quality', 'the quality of the motion in terms of consistency, smoothness, and completeness'], - 'Overall': [ - 'Overall Performance', - 'the overall performance of the video in terms of visual quality, text-to-video alignment, and motion quality' + "VQ": [ + "visual quality", + "the quality of the video in terms of clearness, resolution, brightness, and color", + ], + "TA": [ + "text-to-video alignment", + "the alignment between the text prompt and the video content and motion", + ], + "MQ": [ + "motion quality", + "the quality of the motion in terms of consistency, smoothness, and completeness", + ], + "Overall": [ + "Overall Performance", + "the overall performance of the video in terms of visual quality, text-to-video alignment, and motion quality", ], } @@ -103,7 +112,7 @@ def build_prompt(prompt, dimension, template_type): if isinstance(dimension, list) and len(dimension) > 1: dimension_name = ", ".join([DIMENSION_DESCRIPTIONS[d][0] for d in dimension]) - dimension_name = f'overall performance({dimension_name})' + dimension_name = f"overall performance({dimension_name})" dimension_description = "the overall performance of the video" else: if isinstance(dimension, list): @@ -113,17 +122,20 @@ def build_prompt(prompt, dimension, template_type): if template_type == "none": return prompt - elif template_type == "simple": - return SIMPLE_PROMPT.format(dimension_name=dimension_name, - dimension_description=dimension_description, - text_prompt=prompt) - elif template_type == "video_score": - return VIDEOSCORE_QUERY_PROMPT.format(dimension_name=dimension_name, - dimension_description=dimension_description, - text_prompt=prompt) - elif template_type == "detailed_special": + if template_type == "simple": + return SIMPLE_PROMPT.format( + dimension_name=dimension_name, + dimension_description=dimension_description, + text_prompt=prompt, + ) + if template_type == "video_score": + return VIDEOSCORE_QUERY_PROMPT.format( + dimension_name=dimension_name, + dimension_description=dimension_description, + text_prompt=prompt, + ) + if template_type == "detailed_special": return DETAILED_PROMPT_WITH_SPECIAL_TOKEN.format(text_prompt=prompt) - elif template_type == "detailed": + if template_type == "detailed": return DETAILED_PROMPT.format(text_prompt=prompt) - else: - raise ValueError("Invalid template type") + raise ValueError("Invalid template type") diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/train_reward.py b/fastvideo/train/methods/rl/reward/VideoAlign/train_reward.py index f4a79dc7a2..10501330d0 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/train_reward.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/train_reward.py @@ -6,15 +6,19 @@ from functools import partial import torch -from datasets import load_dataset +from .data import DataConfig, QWen2VLDataCollator, convert_GSB_csv_to_reward_data from peft import LoraConfig, get_peft_model +from .trainer import ( + PartialEmbeddingUpdateCallback, + Qwen2VLRewardModelBT, + VideoVLMRewardTrainer, + compute_multi_attr_accuracy, +) from transformers import AutoProcessor, HfArgumentParser from trl import get_kbit_device_map, get_quantization_config +from .utils import ModelConfig, PEFTLoraConfig, TrainingConfig, load_model_from_checkpoint -from .trainer import Qwen2VLRewardModelBT, VideoVLMRewardTrainer, compute_multi_attr_accuracy, PartialEmbeddingUpdateCallback -from .data import DataConfig, QWen2VLDataCollator, convert_GSB_csv_to_reward_data -from .utils import ModelConfig, PEFTLoraConfig, TrainingConfig -from .utils import load_model_from_checkpoint +from datasets import load_dataset def save_configs_to_json(data_config, training_args, model_config, peft_lora_config): @@ -40,14 +44,12 @@ def save_configs_to_json(data_config, training_args, model_config, peft_lora_con json.dump(config_dict, f, indent=4) -def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=None, verbose=False): +def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=[], verbose=False): """ Find the target linear modules for LoRA. """ linear_cls = torch.nn.Linear embedding_cls = torch.nn.Embedding - if lora_namespan_exclude is None: - lora_namespan_exclude = [] lora_module_names = [] for name, module in model.named_modules(): @@ -55,7 +57,7 @@ def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=N # print(f"Excluding module: {name}") continue - if isinstance(module, linear_cls | embedding_cls): + if isinstance(module, (linear_cls, embedding_cls)): lora_module_names.append(name) if num_lora_modules > 0: @@ -77,21 +79,24 @@ def create_model_and_processor( cache_dir=None, ): # create model - torch_dtype = (model_config.torch_dtype if model_config.torch_dtype in ["auto", None] else getattr( - torch, model_config.torch_dtype)) + torch_dtype = ( + model_config.torch_dtype + if model_config.torch_dtype in ["auto", None] + else getattr(torch, model_config.torch_dtype) + ) quantization_config = get_quantization_config(model_config) model_kwargs = dict( revision=model_config.model_revision, device_map=get_kbit_device_map() if quantization_config is not None else None, quantization_config=quantization_config, - use_cache=bool(training_args.gradient_checkpointing), + use_cache=True if training_args.gradient_checkpointing else False, ) # pdb.set_trace() # create processor and set padding - processor = AutoProcessor.from_pretrained(model_config.model_name_or_path, - padding_side="right", - cache_dir=cache_dir) + processor = AutoProcessor.from_pretrained( + model_config.model_name_or_path, padding_side="right", cache_dir=cache_dir + ) special_token_ids = None if model_config.use_special_tokens: @@ -105,9 +110,10 @@ def create_model_and_processor( reward_token=model_config.reward_token, special_token_ids=special_token_ids, torch_dtype=torch_dtype, - attn_implementation="flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa", + attn_implementation=("flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa"), cache_dir=cache_dir, - **model_kwargs) + **model_kwargs, + ) if model_config.use_special_tokens: model.resize_token_embeddings(len(processor.tokenizer)) @@ -118,9 +124,11 @@ def create_model_and_processor( # create lora and peft model if peft_lora_config.lora_enable: - target_modules = find_target_linear_names(model, - num_lora_modules=peft_lora_config.num_lora_modules, - lora_namespan_exclude=peft_lora_config.lora_namespan_exclude) + target_modules = find_target_linear_names( + model, + num_lora_modules=peft_lora_config.num_lora_modules, + lora_namespan_exclude=peft_lora_config.lora_namespan_exclude, + ) peft_config = LoraConfig( target_modules=target_modules, r=peft_lora_config.lora_r, @@ -144,13 +152,13 @@ def create_model_and_processor( def create_dataset(data_config, meta_file=None): if meta_file is None: meta_file = data_config.meta_data - dataset = load_dataset('csv', data_files=meta_file) + dataset = load_dataset("csv", data_files=meta_file) def add_idx(example, idx): - example['metainfo_idx'] = idx + example["metainfo_idx"] = idx return example - dataset['train'] = dataset['train'].map(lambda example, idx: add_idx(example, idx), with_indices=True) + dataset["train"] = dataset["train"].map(lambda example, idx: add_idx(example, idx), with_indices=True) if not data_config.use_tied_data: filter_func = lambda example: any(example[f"{dim}"] != "same" for dim in data_config.eval_dim) @@ -167,8 +175,12 @@ def add_idx(example, idx): data_config.prompt_template_type, sample_type=data_config.sample_type, ) - dataset = dataset.map(convert_func, remove_columns=dataset['train'].column_names, load_from_cache_file=False) - dataset = dataset['train'] + dataset = dataset.map( + convert_func, + remove_columns=dataset["train"].column_names, + load_from_cache_file=False, + ) + dataset = dataset["train"] # pdb.set_trace() return dataset @@ -176,15 +188,22 @@ def add_idx(example, idx): def train(): ## ===> Step 1: Parse arguments parser = HfArgumentParser((DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig)) - data_config, training_args, model_config, peft_lora_config = parser.parse_args_into_dataclasses() + ( + data_config, + training_args, + model_config, + peft_lora_config, + ) = parser.parse_args_into_dataclasses() # pdb.set_trace() # check valid (lora config) - assert not (peft_lora_config.lora_enable and model_config.freeze_llm - ), 'When using LoRA, the LLM should not be frozen. If you want to freeze the LLM, please disable LoRA.' + assert not (peft_lora_config.lora_enable and model_config.freeze_llm), ( + "When using LoRA, the LLM should not be frozen. If you want to freeze the LLM, please disable LoRA." + ) if not peft_lora_config.lora_enable: - assert not peft_lora_config.vision_lora, \ + assert not peft_lora_config.vision_lora, ( "Error: model_config.lora_enable is not enabled, but model_config.vision_lora is enabled." + ) else: if peft_lora_config.lora_namespan_exclude is not None: peft_lora_config.lora_namespan_exclude = ast.literal_eval(peft_lora_config.lora_namespan_exclude) @@ -204,8 +223,11 @@ def train(): ## load model if training_args.load_from_pretrained is not None: - model, checkpoint_step = load_model_from_checkpoint(model, training_args.load_from_pretrained, - training_args.load_from_pretrained_step) + model, checkpoint_step = load_model_from_checkpoint( + model, + training_args.load_from_pretrained, + training_args.load_from_pretrained_step, + ) model.train() if peft_lora_config.lora_enable: @@ -239,8 +261,8 @@ def train(): # valid_dataset = valid_dataset.select(indices) else: dataset = train_dataset.train_test_split(test_size=0.02) - train_dataset = dataset['train'] - valid_dataset = dataset['test'] + train_dataset = dataset["train"] + valid_dataset = dataset["test"] else: valid_dataset = None @@ -305,7 +327,7 @@ def train(): if training_args.local_rank == -1 or training_args.local_rank == 0: model_state_dict = model.state_dict() - torch.save(model_state_dict, os.path.join(training_args.output_dir, 'final_model.pth')) + torch.save(model_state_dict, os.path.join(training_args.output_dir, "final_model.pth")) model.config.save_pretrained(training_args.output_dir) diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/trainer.py b/fastvideo/train/methods/rl/reward/VideoAlign/trainer.py index c6039be763..74cee0ff71 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/trainer.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/trainer.py @@ -1,34 +1,34 @@ -import os import math -# from training.train_utils import get_peft_state_maybe_zero_3, get_peft_state_non_lora_maybe_zero_3 +import os +# from training.train_utils import get_peft_state_maybe_zero_3, get_peft_state_non_lora_maybe_zero_3 +import numpy as np import pandas as pd import safetensors -import numpy as np import torch -import torch.nn as nn -import datasets -from torch.utils.data import Dataset, DataLoader from peft import PeftModel +from torch import nn +from torch.utils.data import DataLoader, Dataset from transformers import Qwen2VLForConditionalGeneration from transformers.modeling_utils import PreTrainedModel -from transformers.trainer import TrainerCallback from transformers.trainer import ( - is_sagemaker_mp_enabled, - is_peft_available, - is_datasets_available, - WEIGHTS_NAME, - TRAINING_ARGS_NAME, - SAFE_WEIGHTS_NAME, PREFIX_CHECKPOINT_DIR, - logger, + SAFE_WEIGHTS_NAME, + TRAINING_ARGS_NAME, + WEIGHTS_NAME, + TrainerCallback, + is_datasets_available, + is_peft_available, + is_sagemaker_mp_enabled, is_torch_xla_available, + logger, ) - from transformers.trainer_pt_utils import nested_detach from trl import RewardTrainer from .utils import get_peft_state_non_lora_maybe_zero_3 +import datasets + if is_torch_xla_available(): pass else: @@ -36,7 +36,6 @@ class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration): - def __init__(self, config, output_dim=4, reward_token="last", special_token_ids=None): super().__init__(config) # pdb.set_trace() @@ -68,8 +67,9 @@ def forward( ): ## modified from the origin class Qwen2VLForConditionalGeneration output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = (output_hidden_states - if output_hidden_states is not None else self.config.output_hidden_states) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # pdb.set_trace() if inputs_embeds is None: @@ -107,21 +107,23 @@ def forward( logits = self.rm_head(hidden_states) # [B, L, N] - batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] ## get sequence length if self.config.pad_token_id is None and batch_size != 1: raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") if self.config.pad_token_id is None: sequence_lengths = -1 + elif input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) else: - if input_ids is not None: - # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility - sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 - sequence_lengths = sequence_lengths % input_ids.shape[-1] - sequence_lengths = sequence_lengths.to(logits.device) - else: - sequence_lengths = -1 + sequence_lengths = -1 ## get the last token's logits if self.reward_token == "last": @@ -129,7 +131,7 @@ def forward( elif self.reward_token == "mean": ## get the mean of all valid tokens' logits valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1) - pooled_logits = torch.stack([logits[i, :valid_lengths[i]].mean(dim=0) for i in range(batch_size)]) + pooled_logits = torch.stack([logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)]) elif self.reward_token == "special": # special_token_ids = self.tokenizer.convert_tokens_to_ids(self.special_tokens) # create a mask for special tokens @@ -165,9 +167,9 @@ def _convert_A_B_to_chosen_rejected(rewards_A, rewards_B, scores_A, scores_B, ch nontied_mask: [B, N] (preference labels that is not tied) valid_mask: [B, N] (all valid labels) """ - chosen_mask = (chosen_label == 1) + chosen_mask = chosen_label == 1 # rejected_mask = (chosen_label == -1) - rejected_mask = (chosen_label != 1) + rejected_mask = chosen_label != 1 if label_dim is not None: N = chosen_label.size(1) chosen_mask = chosen_mask[:, label_dim].unsqueeze(1).expand(-1, N) @@ -188,7 +190,14 @@ def _convert_A_B_to_chosen_rejected(rewards_A, rewards_B, scores_A, scores_B, ch # rewards_chosen = rewards_chosen * valid_mask # rewards_rejected = rewards_rejected * valid_mask - return rewards_chosen, rewards_rejected, scores_chosen, scores_rejected, nontied_mask, valid_mask + return ( + rewards_chosen, + rewards_rejected, + scores_chosen, + scores_rejected, + nontied_mask, + valid_mask, + ) class PartialEmbeddingUpdateCallback(TrainerCallback): @@ -211,14 +220,13 @@ def on_step_end(self, args, state, control, **kwargs): model = kwargs.get("model") tokenizer = kwargs.get("tokenizer") - index_no_updates = torch.ones((len(tokenizer), ), dtype=torch.bool) + index_no_updates = torch.ones((len(tokenizer),), dtype=torch.bool) index_no_updates[self.special_token_ids] = False with torch.no_grad(): model.get_input_embeddings().weight[index_no_updates] = self.orig_embeds_params[index_no_updates] class VideoVLMRewardTrainer(RewardTrainer): - def __init__(self, loss_type="regular", enable_noise_in_eval=False, *args, **kwargs): super().__init__(*args, **kwargs) @@ -246,12 +254,20 @@ def get_eval_dataloader(self, eval_dataset: str | Dataset | None = None) -> Data # If we have persistent workers, don't do a fork bomb especially as eval datasets # don't change during training dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" - if (hasattr(self, "_eval_dataloaders") and dataloader_key in self._eval_dataloaders - and self.args.dataloader_persistent_workers): + if ( + hasattr(self, "_eval_dataloaders") + and dataloader_key in self._eval_dataloaders + and self.args.dataloader_persistent_workers + ): return self.accelerator.prepare(self._eval_dataloaders[dataloader_key]) - eval_dataset = (self.eval_dataset[eval_dataset] if isinstance(eval_dataset, str) else - eval_dataset if eval_dataset is not None else self.eval_dataset) + eval_dataset = ( + self.eval_dataset[eval_dataset] + if isinstance(eval_dataset, str) + else eval_dataset + if eval_dataset is not None + else self.eval_dataset + ) data_collator = lambda features: self.data_collator(features, enable_noise=self.enable_noise_in_eval) if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset): @@ -316,82 +332,84 @@ def create_optimizer(self): optimizer_grouped_parameters = [ { "params": [ - p for n, p in opt_model.named_parameters() + p + for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in special_lr_parameters and p.requires_grad) ], - "weight_decay": - self.args.weight_decay, + "weight_decay": self.args.weight_decay, }, { "params": [ - p for n, p in opt_model.named_parameters() + p + for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in special_lr_parameters and p.requires_grad) ], - "weight_decay": - 0.0, + "weight_decay": 0.0, }, ] if visual_parameters: - optimizer_grouped_parameters.extend([ - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n in decay_parameters and n in visual_parameters and p.requires_grad) - ], - "weight_decay": - self.args.weight_decay, - "lr": - self.args.vision_lr, - }, - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n not in decay_parameters and n in visual_parameters and p.requires_grad) - ], - "weight_decay": - 0.0, - "lr": - self.args.vision_lr, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and n in visual_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + "lr": self.args.vision_lr, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and n in visual_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + "lr": self.args.vision_lr, + }, + ] + ) if merger_parameters: - optimizer_grouped_parameters.extend([ - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n in decay_parameters and n in merger_parameters and p.requires_grad) - ], - "weight_decay": - self.args.weight_decay, - "lr": - self.args.merger_lr, - }, - { - "params": [ - p for n, p in opt_model.named_parameters() - if (n not in decay_parameters and n in merger_parameters and p.requires_grad) - ], - "weight_decay": - 0.0, - "lr": - self.args.merger_lr, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and n in merger_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + "lr": self.args.merger_lr, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and n in merger_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + "lr": self.args.merger_lr, + }, + ] + ) else: optimizer_grouped_parameters = [ { - "params": - [p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)], - "weight_decay": - self.args.weight_decay, + "params": [ + p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, }, { - "params": - [p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)], - "weight_decay": - 0.0, + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, }, ] @@ -400,14 +418,16 @@ def create_optimizer(self): special_token_embeddings.requires_grad = True - optimizer_grouped_parameters.extend([ - { - # "params": [p for n, p in opt_model.get_input_embeddings().named_parameters() if (p.requires_grad)], - "params": [special_token_embeddings], - "lr": self.args.special_token_lr, - "weight_decay": 0.0, - }, - ]) + optimizer_grouped_parameters.extend( + [ + { + # "params": [p for n, p in opt_model.get_input_embeddings().named_parameters() if (p.requires_grad)], + "params": [special_token_embeddings], + "lr": self.args.special_token_lr, + "weight_decay": 0.0, + }, + ] + ) optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model) @@ -425,12 +445,24 @@ def compute_loss( inputs, return_outputs=False, ): - rewards_A = model(return_dict=True, **inputs['A'])["logits"] - rewards_B = model(return_dict=True, **inputs['B'])["logits"] + rewards_A = model(return_dict=True, **inputs["A"])["logits"] + rewards_B = model(return_dict=True, **inputs["B"])["logits"] # calculate loss, optionally modulate with margin # get chosen and rejected rewards from the chosen label - rewards_chosen, rewards_rejected, scores_chosen, scores_rejected, nontied_mask, valid_mask = _convert_A_B_to_chosen_rejected( - rewards_A, rewards_B, inputs["A_scores"], inputs["B_scores"], inputs["chosen_label"]) + ( + rewards_chosen, + rewards_rejected, + scores_chosen, + scores_rejected, + nontied_mask, + valid_mask, + ) = _convert_A_B_to_chosen_rejected( + rewards_A, + rewards_B, + inputs["A_scores"], + inputs["B_scores"], + inputs["chosen_label"], + ) # pdb.set_trace() inputs["margin"] = scores_chosen - scores_rejected @@ -450,14 +482,14 @@ def compute_loss( out_mask = nontied_mask elif self.loss_type == "scaled": # Bradley-Terry model with scaled margin - loss = (-(inputs["margin"] + 0.0) * nn.functional.logsigmoid(rewards_chosen - rewards_rejected)) + loss = -(inputs["margin"] + 0.0) * nn.functional.logsigmoid(rewards_chosen - rewards_rejected) out_mask = nontied_mask elif self.loss_type == "reg": # regression loss rewards = torch.stack([rewards_A, rewards_B], dim=1) scores = torch.stack([inputs["A_scores"], inputs["B_scores"]], dim=1) out_mask = scores != 0.0 - scores = (scores - 3.0) # rescale + scores = scores - 3.0 # rescale # pdb.set_trace() loss = nn.functional.mse_loss(rewards, scores, reduction="none") elif self.loss_type == "btt": @@ -466,9 +498,11 @@ def compute_loss( log_k = math.log(k) log_k2_sub_1 = math.log(k**2 - 1) bt_loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) - same_loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) \ - -nn.functional.logsigmoid(rewards_rejected - rewards_chosen - log_k) \ - -log_k2_sub_1 + same_loss = ( + -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) + - nn.functional.logsigmoid(rewards_rejected - rewards_chosen - log_k) + - log_k2_sub_1 + ) loss = bt_loss * nontied_mask + same_loss * (1 - nontied_mask) out_mask = valid_mask else: @@ -517,7 +551,6 @@ def prediction_step( return loss, logits, labels def _save_checkpoint(self, model, trial, metrics=None): - if isinstance(self.model, PeftModel): checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" @@ -533,9 +566,13 @@ def _save_checkpoint(self, model, trial, metrics=None): # pdb.set_trace() if not self.args.save_full_model: - non_lora_weights = get_peft_state_non_lora_maybe_zero_3(self.model.named_parameters(), - require_grad_only=True) - torch.save(non_lora_weights, os.path.join(output_dir, "non_lora_state_dict.pth")) + non_lora_weights = get_peft_state_non_lora_maybe_zero_3( + self.model.named_parameters(), require_grad_only=True + ) + torch.save( + non_lora_weights, + os.path.join(output_dir, "non_lora_state_dict.pth"), + ) # safetensors.torch.save(non_lora_weights, os.path.join(output_dir, "non_lora_model.safetensors")) if not self.args.save_only_model: @@ -554,7 +591,7 @@ def _save(self, output_dir: str | None = None, state_dict=None): logger.info(f"Saving model checkpoint to {output_dir}") # pdb.set_trace() - supported_classes = (PreTrainedModel, ) if not is_peft_available() else (PreTrainedModel, PeftModel) + supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel) # Save a trained model and configuration using `save_pretrained()`. # They can then be reloaded using `from_pretrained()` if not isinstance(self.model, supported_classes): @@ -562,25 +599,30 @@ def _save(self, output_dir: str | None = None, state_dict=None): state_dict = self.model.state_dict() if isinstance(self.accelerator.unwrap_model(self.model), supported_classes): - self.accelerator.unwrap_model(self.model).save_pretrained(output_dir, - state_dict=state_dict, - safe_serialization=self.args.save_safetensors) + self.accelerator.unwrap_model(self.model).save_pretrained( + output_dir, + state_dict=state_dict, + safe_serialization=self.args.save_safetensors, + ) else: logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") if self.args.save_safetensors: - safetensors.torch.save_file(state_dict, - os.path.join(output_dir, SAFE_WEIGHTS_NAME), - metadata={"format": "pt"}) + safetensors.torch.save_file( + state_dict, + os.path.join(output_dir, SAFE_WEIGHTS_NAME), + metadata={"format": "pt"}, + ) else: torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + elif not self.args.save_full_model: + state_dict = {k: v for k, v in state_dict.items() if "wte" not in k} + self.model.save_pretrained( + output_dir, + state_dict=state_dict, + safe_serialization=self.args.save_safetensors, + ) else: - if not self.args.save_full_model: - state_dict = {k: v for k, v in state_dict.items() if "wte" not in k} - self.model.save_pretrained(output_dir, - state_dict=state_dict, - safe_serialization=self.args.save_safetensors) - else: - torch.save(state_dict, os.path.join(output_dir, 'model.pth')) + torch.save(state_dict, os.path.join(output_dir, "model.pth")) if self.tokenizer is not None: os.makedirs(os.path.join(output_dir, "tokenizer"), exist_ok=True) @@ -599,7 +641,7 @@ def compute_multi_attr_accuracy(eval_pred, metainfo_idxs=None, eval_dims=None, s label_curr = labels[:, idx] # pdb.set_trace() ## calculate the average scores of rewards_chosen and rewards_rejected - valid_mask = (label_curr != 0) + valid_mask = label_curr != 0 rewards_chosen = np.where(label_curr == 1, pred_curr[:, 0], pred_curr[:, 1]) rewards_rejected = np.where(label_curr == -1, pred_curr[:, 0], pred_curr[:, 1]) @@ -612,11 +654,13 @@ def compute_multi_attr_accuracy(eval_pred, metainfo_idxs=None, eval_dims=None, s accuracy = np.array(pred_curr == label_curr, dtype=float) accuracy = np.sum(accuracy * valid_mask) / np.sum(valid_mask) - metrics.update({ - f"accuracy_{eval_dim}": accuracy, - f"rewards_chosen_avg_{eval_dim}": rewards_chosen_avg, - f"rewards_rejected_avg_{eval_dim}": rewards_rejected_avg, - }) + metrics.update( + { + f"accuracy_{eval_dim}": accuracy, + f"rewards_chosen_avg_{eval_dim}": rewards_chosen_avg, + f"rewards_rejected_avg_{eval_dim}": rewards_rejected_avg, + } + ) if save_path is not None and metainfo_idxs is not None: df = pd.DataFrame(metainfo_idxs, columns=["metainfo_idx"]) diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/utils.py b/fastvideo/train/methods/rl/reward/VideoAlign/utils.py index c6f0a9fc53..5af88f599f 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/utils.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/utils.py @@ -1,6 +1,5 @@ -import os import glob -import logging +import os from dataclasses import dataclass, field from typing import Literal @@ -95,9 +94,11 @@ def __post_init__(self): def maybe_zero_3(param, ignore_status=False, name=None): from deepspeed import zero from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + if hasattr(param, "ds_id"): - if param.ds_status == ZeroParamStatus.NOT_AVAILABLE and not ignore_status: - logging.warning("%s: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: %s", name, param.ds_status) + if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: + if not ignore_status: + logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}") with zero.GatheredParameters([param]): param = param.data.detach().cpu().clone() else: @@ -142,8 +143,9 @@ def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True): ########## Load Models From Folder ########## -def _insert_adapter_name_into_state_dict(state_dict: dict[str, torch.Tensor], adapter_name: str, - parameter_prefix: str) -> dict[str, torch.Tensor]: +def _insert_adapter_name_into_state_dict( + state_dict: dict[str, torch.Tensor], adapter_name: str, parameter_prefix: str +) -> dict[str, torch.Tensor]: """Utility function to remap the state_dict keys to fit the PEFT model by inserting the adapter name.""" peft_model_state_dict = {} for key, val in state_dict.items(): @@ -162,15 +164,26 @@ def _insert_adapter_name_into_state_dict(state_dict: dict[str, torch.Tensor], ad def save_video(tensor, path): from torchvision.io import write_video + tensor = tensor * 255.0 tensor = tensor.permute(0, 2, 3, 1) tensor = tensor.clamp(0, 255).byte() - write_video(path, tensor, 4, video_codec='h264') + write_video(path, tensor, 4, video_codec="h264") def load_model_from_checkpoint(model, checkpoint_dir, checkpoint_step): checkpoint_paths = glob.glob(os.path.join(checkpoint_dir, "checkpoint-*")) checkpoint_paths.sort(key=lambda x: int(x.split("-")[-1]), reverse=True) + if not checkpoint_paths: + raise FileNotFoundError( + f"No VideoAlign checkpoint-* directories found under {checkpoint_dir}." + ) + + before_rm_head = { + key: value.detach().cpu().clone() + for key, value in model.state_dict().items() + if "rm_head" in key + } if checkpoint_step is None or checkpoint_step == -1: # get the latest checkpoint @@ -193,16 +206,42 @@ def load_model_from_checkpoint(model, checkpoint_dir, checkpoint_step): model_state_dict = torch.load(full_ckpt, map_location="cpu") model.load_state_dict(model_state_dict) else: + if not os.path.exists(lora_ckpt) or not os.path.exists(non_lora_ckpt): + raise FileNotFoundError( + "VideoAlign checkpoint must contain either model.pth or both " + f"adapter_model.safetensors and non_lora_state_dict.pth. Got {checkpoint_path}." + ) lora_state_dict = safetensors.torch.load_file(lora_ckpt) non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu") + if not any("rm_head" in key for key in non_lora_state_dict): + raise RuntimeError( + f"{non_lora_ckpt} does not contain rm_head weights. Refusing " + "to use a randomly initialized VideoAlign reward head." + ) - lora_state_dict = _insert_adapter_name_into_state_dict(lora_state_dict, - adapter_name="default", - parameter_prefix="lora_") + lora_state_dict = _insert_adapter_name_into_state_dict( + lora_state_dict, adapter_name="default", parameter_prefix="lora_" + ) model_state_dict = model.state_dict() model_state_dict.update(non_lora_state_dict) model_state_dict.update(lora_state_dict) model.load_state_dict(model_state_dict) + after_rm_head = { + key: value.detach().cpu() + for key, value in model.state_dict().items() + if "rm_head" in key + } + unchanged = ( + before_rm_head + and set(before_rm_head) == set(after_rm_head) + and all(torch.equal(before_rm_head[key], after_rm_head[key]) for key in before_rm_head) + ) + if unchanged: + raise RuntimeError( + f"VideoAlign checkpoint {checkpoint_path} did not overwrite rm_head " + "weights. Refusing to score rewards with a randomly initialized head." + ) + return model, checkpoint_step diff --git a/fastvideo/train/methods/rl/reward/VideoAlign/vision_process.py b/fastvideo/train/methods/rl/reward/VideoAlign/vision_process.py index 19401f7a04..3b7387a0b1 100644 --- a/fastvideo/train/methods/rl/reward/VideoAlign/vision_process.py +++ b/fastvideo/train/methods/rl/reward/VideoAlign/vision_process.py @@ -7,6 +7,7 @@ import math import os import sys +import time import warnings from functools import lru_cache from io import BytesIO @@ -50,11 +51,13 @@ def floor_by_factor(number: int, factor: int) -> int: return math.floor(number / factor) * factor -def smart_resize(height: int, - width: int, - factor: int = IMAGE_FACTOR, - min_pixels: int = MIN_PIXELS, - max_pixels: int = MAX_PIXELS) -> tuple[int, int]: +def smart_resize( + height: int, + width: int, + factor: int = IMAGE_FACTOR, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS, +) -> tuple[int, int]: """ Rescales the image so that the following conditions are met: @@ -66,7 +69,8 @@ def smart_resize(height: int, """ if max(height, width) / min(height, width) > MAX_RATIO: raise ValueError( - f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}") + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) h_bar = max(factor, round_by_factor(height, factor)) w_bar = max(factor, round_by_factor(width, factor)) if h_bar * w_bar > max_pixels: @@ -81,7 +85,10 @@ def smart_resize(height: int, def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image: - image = ele["image"] if "image" in ele else ele["image_url"] + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] image_obj = None if isinstance(image, Image.Image): image_obj = image @@ -125,7 +132,7 @@ def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACT def smart_nframes( ele: dict, total_frames: int, - video_fps: int | float, + video_fps: float, ) -> int: """calculate the number of frames for video used for model inputs. @@ -155,14 +162,47 @@ def smart_nframes( nframes = total_frames / video_fps * fps nframes = min(max(nframes, min_frames), max_frames) nframes = round_by_factor(nframes, FRAME_FACTOR) - if nframes > total_frames: - nframes = total_frames + nframes = min(nframes, total_frames) if not (nframes >= FRAME_FACTOR and nframes <= total_frames): raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") return nframes -def _read_video_torchvision(ele: dict, ) -> torch.Tensor: +def _get_video_fps_fallback(video_path: str) -> float: + """Get video fps using PyAV or OpenCV as fallback when torchvision info doesn't have it.""" + try: + # Try PyAV first (since torchvision uses pyav backend) + import av + + container = av.open(video_path) + video_stream = container.streams.video[0] + fps = float(video_stream.average_rate) + container.close() + if fps > 0: + return fps + except Exception: + pass + + try: + # Try OpenCV as backup + import cv2 + + cap = cv2.VideoCapture(video_path) + fps = cap.get(cv2.CAP_PROP_FPS) + cap.release() + if fps > 0: + return float(fps) + except Exception: + pass + logger.error("Error getting video fps using PyAV or OpenCV, using default fallback 30.0 fps.") + + # Default fallback + return 30.0 + + +def _read_video_torchvision( + ele: dict, +) -> torch.Tensor: """read video using torchvision.io.read_video Args: @@ -175,14 +215,13 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor: torch.Tensor: the video tensor with shape (T, C, H, W). """ video_path = ele["video"] + # Remove file:// prefix - torchvision doesn't support it (especially for relative paths) + if video_path.startswith("file://"): + video_path = video_path[7:] if version.parse(torchvision.__version__) < version.parse("0.19.0"): if "http://" in video_path or "https://" in video_path: - warnings.warn( - "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.", - stacklevel=2, - ) - if "file://" in video_path: - video_path = video_path[7:] + warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.") + st = time.time() video, audio, info = io.read_video( video_path, start_pts=ele.get("video_start", 0.0), @@ -191,12 +230,24 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor: output_format="TCHW", ) - total_frames, video_fps = video.size(0), info["video_fps"] + total_frames = video.size(0) + if total_frames == 0: + raise ValueError( + f"No frames were read from video: {video_path}. " + f"This may be caused by invalid video_start ({ele.get('video_start', 0.0)}) " + f"or video_end ({ele.get('video_end')}) parameters, or the video file may be corrupted." + ) + # Try to get video_fps from info, use fallback methods if not available + if "video_fps" in info: + video_fps = info["video_fps"] + else: + # Fallback: use PyAV or OpenCV to get real fps + video_fps = _get_video_fps_fallback(video_path) # logger.info(f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") - if ele['sample_type'] == 'uniform': + if ele["sample_type"] == "uniform": nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() - elif ele['sample_type'] == 'multi_pts': + elif ele["sample_type"] == "multi_pts": frames_each_pts = 6 num_pts = 4 fps = 8 @@ -208,7 +259,7 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor: pts = torch.linspace(start_pt, end_pt, num_pts).round().long().tolist() idx = [] for pt in pts: - idx.extend(frames_idx[pt - frames_each_pts // 2:pt + frames_each_pts // 2]) + idx.extend(frames_idx[pt - frames_each_pts // 2 : pt + frames_each_pts // 2]) video = video[idx] return video @@ -220,7 +271,9 @@ def is_decord_available() -> bool: return importlib.util.find_spec("decord") is not None -def _read_video_decord(ele: dict, ) -> torch.Tensor: +def _read_video_decord( + ele: dict, +) -> torch.Tensor: """read video using decord.VideoReader Args: @@ -233,19 +286,21 @@ def _read_video_decord(ele: dict, ) -> torch.Tensor: torch.Tensor: the video tensor with shape (T, C, H, W). """ import decord + video_path = ele["video"] + st = time.time() vr = decord.VideoReader(video_path) # TODO: support start_pts and end_pts - if 'video_start' in ele or 'video_end' in ele: + if "video_start" in ele or "video_end" in ele: raise NotImplementedError("not support start_pts and end_pts in decord for now.") total_frames, video_fps = len(vr), vr.get_avg_fps() # logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") - if ele['sample_type'] == 'uniform': + if ele["sample_type"] == "uniform": nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) # nframes = max(nframes, 8) # import pdb; pdb.set_trace() idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() - elif ele['sample_type'] == 'multi_pts': + elif ele["sample_type"] == "multi_pts": frames_each_pts = 6 num_pts = 4 fps = 8 @@ -257,7 +312,7 @@ def _read_video_decord(ele: dict, ) -> torch.Tensor: pts = torch.linspace(start_pt, end_pt, num_pts).round().long().tolist() idx = [] for pt in pts: - idx.extend(frames_idx[pt - frames_each_pts // 2:pt + frames_each_pts // 2]) + idx.extend(frames_idx[pt - frames_each_pts // 2 : pt + frames_each_pts // 2]) video = vr.get_batch(idx).asnumpy() video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format return video @@ -292,7 +347,10 @@ def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | l min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS) total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS) - max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05)) + max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05), + ) max_pixels = ele.get("max_pixels", max_pixels) if "resized_height" in ele and "resized_width" in ele: resized_height, resized_width = smart_resize( @@ -315,21 +373,18 @@ def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | l antialias=True, ).float() return video - else: - assert isinstance(ele["video"], list | tuple) - process_info = ele.copy() - process_info.pop("type", None) - process_info.pop("video", None) - images = [ - fetch_image({ - "image": video_element, - **process_info - }, size_factor=image_factor) for video_element in ele["video"] - ] - nframes = ceil_by_factor(len(images), FRAME_FACTOR) - if len(images) < nframes: - images.extend([images[-1]] * (nframes - len(images))) - return images + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + images = [ + fetch_image({"image": video_element, **process_info}, size_factor=image_factor) + for video_element in ele["video"] + ] + nframes = ceil_by_factor(len(images), FRAME_FACTOR) + if len(images) < nframes: + images.extend([images[-1]] * (nframes - len(images))) + return images def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]: @@ -340,8 +395,12 @@ def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[di for message in conversation: if isinstance(message["content"], list): for ele in message["content"]: - if ("image" in ele or "image_url" in ele or "video" in ele - or ele["type"] in ("image", "image_url", "video")): + if ( + "image" in ele + or "image_url" in ele + or "video" in ele + or ele["type"] in ("image", "image_url", "video") + ): vision_infos.append(ele) return vision_infos diff --git a/fastvideo/train/methods/rl/reward/hpsv3.py b/fastvideo/train/methods/rl/reward/hpsv3.py index 11c01f0d45..ede491d603 100644 --- a/fastvideo/train/methods/rl/reward/hpsv3.py +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -98,6 +98,11 @@ def load_state_dict_with_key_remap( assign=False, ): state_dict = _remap_hpsv3_state_dict(state_dict) + if not assign: + try: + assign = any(getattr(param, "is_meta", False) for param in self.parameters()) + except Exception: + assign = False return original_load_state_dict( self, state_dict, @@ -180,9 +185,10 @@ def _get_hpsv3_inferencer(device: torch.device | str) -> Any: _patch_hpsv3_state_dict_loader() except ImportError as exc: - msg = ("Failed to import HPSv3. Ensure the HPSv3 submodule is " - "checked out under fastvideo/train/methods/rl/reward/HPSv3 " - "and that its transformers dependencies are compatible.") + msg = ("Failed to import the vendored HPSv3 reward runtime under " + "fastvideo/train/methods/rl/reward/HPSv3. Install " + "FastVideo with the GenRL extra and verify the vendored " + "runtime files are present.") raise ImportError(msg) from exc inf = HPSv3RewardInferencer(device=device) _patch_hpsv3_runtime_model(inf.model) diff --git a/fastvideo/train/methods/rl/reward/videoalign.py b/fastvideo/train/methods/rl/reward/videoalign.py index 9cf1b84f42..6e10aef6aa 100644 --- a/fastvideo/train/methods/rl/reward/videoalign.py +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -329,8 +329,11 @@ def _resolve_videoalign_checkpoint_path(checkpoint_path: str | None) -> str: """Return a local VideoAlign checkpoint directory.""" if checkpoint_path is not None: return os.path.abspath(checkpoint_path) + env_checkpoint_path = os.environ.get("VIDEOALIGN_CHECKPOINT_PATH") + if env_checkpoint_path: + return os.path.abspath(env_checkpoint_path) return snapshot_download( - repo_id="KlingTeam/VideoReward", + repo_id="KwaiVGI/VideoReward", repo_type="model", allow_patterns=( "model_config.json", @@ -353,10 +356,10 @@ def _get_inferencer( inference_mod = _patch_videoalign_modules() VideoVLMRewardInference = inference_mod.VideoVLMRewardInference except ImportError as exc: - msg = ("VideoAlign not found. Ensure the " - "VideoAlign submodule is checked out " - "under fastvideo/train/methods/rl/" - "reward/VideoAlign") + msg = ("Failed to import the vendored VideoAlign reward runtime " + "under fastvideo/train/methods/rl/reward/VideoAlign. " + "Install FastVideo with the GenRL extra and verify the " + "vendored runtime files are present.") raise ImportError(msg) from exc inf = VideoVLMRewardInference( diff --git a/fastvideo/train/methods/rl/utils/data.py b/fastvideo/train/methods/rl/utils/data.py index fc44a4c170..c2b29acf28 100644 --- a/fastvideo/train/methods/rl/utils/data.py +++ b/fastvideo/train/methods/rl/utils/data.py @@ -54,11 +54,22 @@ def __init__(self, dataset: str, split: str = "train"): self._load_all_prompts() def _load_all_prompts(self) -> None: + if not os.path.exists(self.file_path): + raise FileNotFoundError("GenRL prompt file not found: " + f"{self.file_path}. Expected train.json/test.json " + "from examples/train/prepare_genrl_assets.py.") + + saw_content = False with open(self.file_path, encoding="utf-8") as f: - for raw_line in f: + for line_no, raw_line in enumerate(f, start=1): line = raw_line.strip() if not line: continue + if not saw_content and line.startswith("version https://git-lfs.github.com"): + raise RuntimeError(f"{self.file_path} is a Git LFS pointer, not the " + "real prompt JSON. Rerun " + "`python examples/train/prepare_genrl_assets.py`.") + saw_content = True try: item = json.loads(line) prompt = item.get("prompt", "") @@ -68,9 +79,15 @@ def _load_all_prompts(self) -> None: self._metadatas.append(metadata) except json.JSONDecodeError as e: logger.warning( - "Skipping invalid JSON line: %s", + "Skipping invalid JSON line %d in %s: %s", + line_no, + self.file_path, e, ) + if not self._prompts: + raise ValueError("No usable prompts found in " + f"{self.file_path}. Expected JSONL rows with a " + "`prompt` field.") def __len__(self) -> int: return len(self._prompts)