[feat] GenRL: stabilize reward model compatibility (#1400)
Co-authored-by: Davids048 <jundasu@ucsd.edu>
This commit is contained in:
Adam
committed by
Davids048
Davids048
co-authored by
Davids048
Davids048
parent
5ae633a551
commit
2c19cbdebd
-1
Submodule GenRL deleted from 82a867ed24
@@ -28,6 +28,129 @@ if _HPSV3_ROOT.exists():
|
||||
|
||||
# Global cache of HPSv3 inferencers keyed by device.
|
||||
_HPSV3_INFERENCERS: dict[str, Any] = {}
|
||||
_HPSV3_LOAD_PATCHED = False
|
||||
|
||||
|
||||
def _patch_transformers_video_input_alias() -> None:
|
||||
"""Keep HPSv3 compatible with newer transformers releases.
|
||||
|
||||
HPSv3 imports ``VideoInput`` from ``transformers.image_utils`` for type
|
||||
annotations. Some transformers versions used by FastVideo no longer
|
||||
export that alias, even though the runtime image utilities HPSv3 needs are
|
||||
still present.
|
||||
"""
|
||||
from transformers import image_utils
|
||||
|
||||
if not hasattr(image_utils, "VideoInput"):
|
||||
image_utils.VideoInput = image_utils.ImageInput
|
||||
|
||||
|
||||
def _remap_hpsv3_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Adapt HPSv3 checkpoints saved with older Qwen2-VL key names."""
|
||||
remapped = {}
|
||||
for key, value in state_dict.items():
|
||||
if key.startswith("visual."):
|
||||
key = f"model.{key}"
|
||||
elif key.startswith("model.layers."):
|
||||
key = f"model.language_model.{key[len('model.'):]}"
|
||||
elif key.startswith("model.embed_tokens."):
|
||||
key = f"model.language_model.{key[len('model.'):]}"
|
||||
elif key.startswith("model.norm."):
|
||||
key = f"model.language_model.{key[len('model.'):]}"
|
||||
|
||||
key = key.replace(
|
||||
"base_model.model.visual.",
|
||||
"base_model.model.model.visual.",
|
||||
1,
|
||||
)
|
||||
key = key.replace(
|
||||
"base_model.model.model.layers.",
|
||||
"base_model.model.model.language_model.layers.",
|
||||
1,
|
||||
)
|
||||
key = key.replace(
|
||||
"base_model.model.model.embed_tokens.",
|
||||
"base_model.model.model.language_model.embed_tokens.",
|
||||
1,
|
||||
)
|
||||
key = key.replace(
|
||||
"base_model.model.model.norm.",
|
||||
"base_model.model.model.language_model.norm.",
|
||||
1,
|
||||
)
|
||||
remapped[key] = value
|
||||
return remapped
|
||||
|
||||
|
||||
def _walk_model_graph(model: Any):
|
||||
"""Yield common wrapper/base model objects without importing PEFT."""
|
||||
stack = [model]
|
||||
seen = set()
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
if current is None or id(current) in seen:
|
||||
continue
|
||||
seen.add(id(current))
|
||||
yield current
|
||||
for attr in ("base_model", "model"):
|
||||
child = getattr(current, attr, None)
|
||||
if child is not None:
|
||||
stack.append(child)
|
||||
|
||||
|
||||
def _patch_load_state_dict(cls: Any) -> None:
|
||||
"""Patch a model class to accept old Qwen2-VL checkpoint keys."""
|
||||
if getattr(cls, "_fastvideo_qwen2vl_key_remap", False):
|
||||
return
|
||||
|
||||
original_load_state_dict = cls.load_state_dict
|
||||
|
||||
def load_state_dict_with_key_remap(
|
||||
self,
|
||||
state_dict,
|
||||
strict=True,
|
||||
assign=False,
|
||||
):
|
||||
state_dict = _remap_hpsv3_state_dict(state_dict)
|
||||
return original_load_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
strict=strict,
|
||||
assign=assign,
|
||||
)
|
||||
|
||||
cls.load_state_dict = load_state_dict_with_key_remap
|
||||
cls._fastvideo_qwen2vl_key_remap = True
|
||||
|
||||
|
||||
def _patch_hpsv3_state_dict_loader() -> None:
|
||||
"""Patch HPSv3 reward model loading for transformers key drift."""
|
||||
global _HPSV3_LOAD_PATCHED
|
||||
if _HPSV3_LOAD_PATCHED:
|
||||
return
|
||||
|
||||
from hpsv3.model.qwen2vl_trainer import Qwen2VLRewardModelBT
|
||||
|
||||
_patch_load_state_dict(Qwen2VLRewardModelBT)
|
||||
try:
|
||||
from peft import PeftModel
|
||||
except ImportError:
|
||||
PeftModel = None
|
||||
if PeftModel is not None:
|
||||
_patch_load_state_dict(PeftModel)
|
||||
_HPSV3_LOAD_PATCHED = True
|
||||
|
||||
|
||||
def _patch_hpsv3_runtime_model(model: Any) -> None:
|
||||
"""Add aliases expected by HPSv3's older Qwen2-VL forward."""
|
||||
for candidate in _walk_model_graph(model):
|
||||
language_model = getattr(candidate, "language_model", None)
|
||||
if (
|
||||
language_model is not None
|
||||
and not hasattr(candidate, "embed_tokens")
|
||||
and hasattr(language_model, "embed_tokens")
|
||||
):
|
||||
candidate.__dict__["embed_tokens"] = language_model.embed_tokens
|
||||
|
||||
|
||||
def _normalize_device(device) -> str:
|
||||
@@ -36,6 +159,19 @@ def _normalize_device(device) -> str:
|
||||
return str(torch.device(device))
|
||||
|
||||
|
||||
def _move_hpsv3_inferencer(inferencer: Any, device) -> None:
|
||||
"""Move an HPSv3 inferencer across devices.
|
||||
|
||||
HPSv3RewardInferencer does not expose ``.to()``, but it stores its torch
|
||||
module on ``.model`` and reads ``.device`` when preparing batches.
|
||||
"""
|
||||
device_str = _normalize_device(device)
|
||||
model = getattr(inferencer, "model", None)
|
||||
if model is not None and hasattr(model, "to"):
|
||||
model.to(device)
|
||||
inferencer.device = device_str
|
||||
|
||||
|
||||
def set_hpsv3_device(device) -> None:
|
||||
"""Move cached HPSv3 inferencer to given device."""
|
||||
key = _normalize_device(device)
|
||||
@@ -44,7 +180,7 @@ def set_hpsv3_device(device) -> None:
|
||||
# Move from any existing device.
|
||||
for old_key, inf in list(_HPSV3_INFERENCERS.items()):
|
||||
if old_key != key:
|
||||
inf.to(device)
|
||||
_move_hpsv3_inferencer(inf, device)
|
||||
_HPSV3_INFERENCERS[key] = inf
|
||||
del _HPSV3_INFERENCERS[old_key]
|
||||
return
|
||||
@@ -55,15 +191,18 @@ def _get_hpsv3_inferencer(device):
|
||||
key = _normalize_device(device)
|
||||
if key not in _HPSV3_INFERENCERS:
|
||||
try:
|
||||
_patch_transformers_video_input_alias()
|
||||
from hpsv3 import HPSv3RewardInferencer
|
||||
_patch_hpsv3_state_dict_loader()
|
||||
except ImportError as exc:
|
||||
msg = (
|
||||
"hpsv3 package not found. Ensure the "
|
||||
"HPSv3 submodule is checked out under "
|
||||
"fastvideo/train/methods/rl/reward/HPSv3"
|
||||
"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."
|
||||
)
|
||||
raise ImportError(msg) from exc
|
||||
inf = HPSv3RewardInferencer(device=device)
|
||||
_patch_hpsv3_runtime_model(inf.model)
|
||||
_HPSV3_INFERENCERS[key] = inf
|
||||
return _HPSV3_INFERENCERS[key]
|
||||
|
||||
|
||||
@@ -24,11 +24,20 @@ def prepare_images(
|
||||
|
||||
if images.ndim == 4:
|
||||
# Image batch: (N, C, H, W) or (N, H, W, C)
|
||||
if images.shape[1] in (1, 3):
|
||||
if images.shape[-1] in (1, 3):
|
||||
pass
|
||||
elif images.shape[1] in (1, 3):
|
||||
images = images.transpose(0, 2, 3, 1)
|
||||
elif images.ndim == 5:
|
||||
# Video batch: (N, F, C, H, W) or (N, C, F, H, W)
|
||||
if images.shape[2] in (1, 3):
|
||||
# Video batch: (N, F, H, W, C), (N, F, C, H, W),
|
||||
# or (N, C, F, H, W). Check channel-last first because
|
||||
# one-frame videos have shape[1] == 1.
|
||||
if images.shape[-1] in (1, 3):
|
||||
pass
|
||||
elif images.shape[1] == 3 and images.shape[2] == 1:
|
||||
# (N, C=3, F=1, H, W) -> (N, F, H, W, C)
|
||||
images = images.transpose(0, 2, 3, 4, 1)
|
||||
elif images.shape[2] in (1, 3):
|
||||
# (N, F, C, H, W) -> (N, F, H, W, C)
|
||||
images = images.transpose(0, 1, 3, 4, 2)
|
||||
elif images.shape[1] in (1, 3):
|
||||
|
||||
@@ -7,6 +7,7 @@ from __future__ import annotations
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from importlib import import_module, util
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
@@ -29,6 +30,7 @@ if os.path.isdir(_VIDEOALIGN_ROOT):
|
||||
|
||||
# Global cache of VideoAlign inferencers.
|
||||
_VIDEOALIGN_INFERENCERS: dict[str, Any] = {}
|
||||
_VIDEOALIGN_PATCHED = False
|
||||
|
||||
|
||||
def _normalize_device_str(device) -> str:
|
||||
@@ -37,6 +39,283 @@ def _normalize_device_str(device) -> str:
|
||||
return str(torch.device(device))
|
||||
|
||||
|
||||
def _move_videoalign_inferencer(inferencer: Any, device) -> None:
|
||||
"""Move a VideoAlign inferencer across devices."""
|
||||
device_str = _normalize_device_str(device)
|
||||
model = getattr(inferencer, "model", None)
|
||||
if model is not None and hasattr(model, "to"):
|
||||
model.to(device)
|
||||
inferencer.device = device_str
|
||||
|
||||
|
||||
def _remap_qwen2vl_state_dict_keys(
|
||||
state_dict: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Adapt checkpoints saved with older Qwen2-VL key names."""
|
||||
remapped = {}
|
||||
for key, value in state_dict.items():
|
||||
if key.startswith("visual."):
|
||||
key = f"model.{key}"
|
||||
elif key.startswith("model.layers."):
|
||||
key = f"model.language_model.{key[len('model.'):]}"
|
||||
elif key.startswith("model.embed_tokens."):
|
||||
key = f"model.language_model.{key[len('model.'):]}"
|
||||
elif key.startswith("model.norm."):
|
||||
key = f"model.language_model.{key[len('model.'):]}"
|
||||
|
||||
key = key.replace(
|
||||
"base_model.model.visual.",
|
||||
"base_model.model.model.visual.",
|
||||
1,
|
||||
)
|
||||
key = key.replace(
|
||||
"base_model.model.model.layers.",
|
||||
"base_model.model.model.language_model.layers.",
|
||||
1,
|
||||
)
|
||||
key = key.replace(
|
||||
"base_model.model.model.embed_tokens.",
|
||||
"base_model.model.model.language_model.embed_tokens.",
|
||||
1,
|
||||
)
|
||||
key = key.replace(
|
||||
"base_model.model.model.norm.",
|
||||
"base_model.model.model.language_model.norm.",
|
||||
1,
|
||||
)
|
||||
remapped[key] = value
|
||||
return remapped
|
||||
|
||||
|
||||
def _walk_model_graph(model: Any):
|
||||
"""Yield common wrapper/base model objects without importing PEFT."""
|
||||
stack = [model]
|
||||
seen = set()
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
if current is None or id(current) in seen:
|
||||
continue
|
||||
seen.add(id(current))
|
||||
yield current
|
||||
for attr in ("base_model", "model"):
|
||||
child = getattr(current, attr, None)
|
||||
if child is not None:
|
||||
stack.append(child)
|
||||
|
||||
|
||||
def _patch_load_state_dict(cls: Any) -> None:
|
||||
"""Patch a model class to accept old VideoAlign checkpoint keys."""
|
||||
if getattr(cls, "_fastvideo_qwen2vl_key_remap", False):
|
||||
return
|
||||
|
||||
original_load_state_dict = cls.load_state_dict
|
||||
|
||||
def load_state_dict_with_key_remap(
|
||||
self,
|
||||
state_dict,
|
||||
strict=True,
|
||||
assign=False,
|
||||
):
|
||||
state_dict = _remap_qwen2vl_state_dict_keys(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,
|
||||
strict=strict,
|
||||
assign=assign,
|
||||
)
|
||||
|
||||
cls.load_state_dict = load_state_dict_with_key_remap
|
||||
cls._fastvideo_qwen2vl_key_remap = True
|
||||
|
||||
|
||||
def _select_videoalign_frame_indices(
|
||||
vision_mod: Any,
|
||||
ele: dict[str, Any],
|
||||
total_frames: int,
|
||||
video_fps: float,
|
||||
) -> list[int]:
|
||||
sample_type = ele.get("sample_type", "uniform")
|
||||
if sample_type == "uniform":
|
||||
nframes = vision_mod.smart_nframes(
|
||||
ele,
|
||||
total_frames=total_frames,
|
||||
video_fps=video_fps,
|
||||
)
|
||||
return torch.linspace(
|
||||
0,
|
||||
total_frames - 1,
|
||||
nframes,
|
||||
).round().long().tolist()
|
||||
if sample_type == "multi_pts":
|
||||
frames_each_pts = 6
|
||||
num_pts = 4
|
||||
fps = 8
|
||||
nframes = max(
|
||||
frames_each_pts,
|
||||
int(total_frames * fps // video_fps),
|
||||
)
|
||||
frame_idx = torch.linspace(
|
||||
0,
|
||||
total_frames - 1,
|
||||
nframes,
|
||||
).round().long().tolist()
|
||||
start_pt = int(frames_each_pts // 2)
|
||||
end_pt = int(nframes - frames_each_pts // 2)
|
||||
pts = torch.linspace(
|
||||
start_pt,
|
||||
end_pt,
|
||||
num_pts,
|
||||
).round().long().tolist()
|
||||
idx = []
|
||||
for pt in pts:
|
||||
idx.extend(
|
||||
frame_idx[
|
||||
pt - frames_each_pts // 2:
|
||||
pt + frames_each_pts // 2
|
||||
]
|
||||
)
|
||||
return idx
|
||||
raise ValueError(f"Unsupported VideoAlign sample_type: {sample_type}")
|
||||
|
||||
|
||||
def _read_video_opencv(
|
||||
vision_mod: Any,
|
||||
ele: dict[str, Any],
|
||||
) -> torch.Tensor:
|
||||
"""Read local MP4s without relying on torchvision.io.read_video."""
|
||||
import cv2
|
||||
|
||||
video_path = ele["video"]
|
||||
if video_path.startswith("file://"):
|
||||
video_path = video_path[7:]
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
raise ValueError(f"Could not open video: {video_path}")
|
||||
|
||||
video_fps = float(cap.get(cv2.CAP_PROP_FPS) or 30.0)
|
||||
total_file_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
|
||||
start_frame = max(
|
||||
0,
|
||||
int(round(float(ele.get("video_start", 0.0) or 0.0) * video_fps)),
|
||||
)
|
||||
end_sec = ele.get("video_end")
|
||||
if end_sec is None:
|
||||
end_frame = total_file_frames if total_file_frames > 0 else None
|
||||
else:
|
||||
end_frame = int(round(float(end_sec) * video_fps))
|
||||
if total_file_frames > 0:
|
||||
end_frame = min(end_frame, total_file_frames)
|
||||
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
|
||||
frames = []
|
||||
current_frame = start_frame
|
||||
while end_frame is None or current_frame < end_frame:
|
||||
ok, frame = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
current_frame += 1
|
||||
cap.release()
|
||||
|
||||
if not frames:
|
||||
raise ValueError(f"No frames were read from video: {video_path}.")
|
||||
|
||||
idx = _select_videoalign_frame_indices(
|
||||
vision_mod,
|
||||
ele,
|
||||
total_frames=len(frames),
|
||||
video_fps=video_fps,
|
||||
)
|
||||
video = np.stack([frames[i] for i in idx], axis=0)
|
||||
return torch.from_numpy(video).permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
def _torchvision_read_video_available() -> bool:
|
||||
try:
|
||||
torchvision_io = import_module("torchvision.io")
|
||||
except Exception:
|
||||
return False
|
||||
return hasattr(torchvision_io, "read_video")
|
||||
|
||||
|
||||
def _patch_videoalign_video_reader() -> None:
|
||||
"""Register an OpenCV reader for torchvision builds without read_video."""
|
||||
vision_mod = import_module("vision_process")
|
||||
if "opencv" not in vision_mod.VIDEO_READER_BACKENDS:
|
||||
|
||||
def read_video_opencv(ele):
|
||||
return _read_video_opencv(vision_mod, ele)
|
||||
|
||||
vision_mod.VIDEO_READER_BACKENDS["opencv"] = read_video_opencv
|
||||
|
||||
if _torchvision_read_video_available():
|
||||
return
|
||||
|
||||
vision_mod.FORCE_QWENVL_VIDEO_READER = "opencv"
|
||||
if hasattr(vision_mod.get_video_reader_backend, "cache_clear"):
|
||||
vision_mod.get_video_reader_backend.cache_clear()
|
||||
|
||||
|
||||
def _patch_videoalign_modules() -> Any:
|
||||
"""Patch VideoAlign for the FastVideo dependency set."""
|
||||
global _VIDEOALIGN_PATCHED
|
||||
|
||||
inference_mod = import_module("inference")
|
||||
if _VIDEOALIGN_PATCHED:
|
||||
return inference_mod
|
||||
|
||||
train_reward_mod = import_module("train_reward")
|
||||
trainer_mod = import_module("trainer")
|
||||
_patch_videoalign_video_reader()
|
||||
|
||||
if util.find_spec("flash_attn") is None:
|
||||
for mod in (train_reward_mod, inference_mod):
|
||||
original_create = mod.create_model_and_processor
|
||||
|
||||
def create_model_and_processor_sdpa(
|
||||
*args,
|
||||
_original_create=original_create,
|
||||
**kwargs,
|
||||
):
|
||||
training_args = kwargs.get("training_args")
|
||||
if training_args is not None:
|
||||
training_args.disable_flash_attn2 = True
|
||||
return _original_create(*args, **kwargs)
|
||||
|
||||
mod.create_model_and_processor = create_model_and_processor_sdpa
|
||||
|
||||
_patch_load_state_dict(trainer_mod.Qwen2VLRewardModelBT)
|
||||
try:
|
||||
peft_mod = import_module("peft")
|
||||
except ImportError:
|
||||
peft_mod = None
|
||||
if peft_mod is not None:
|
||||
_patch_load_state_dict(peft_mod.PeftModel)
|
||||
|
||||
_VIDEOALIGN_PATCHED = True
|
||||
return inference_mod
|
||||
|
||||
|
||||
def _patch_videoalign_runtime_model(model: Any) -> None:
|
||||
"""Add aliases expected by VideoAlign's older Qwen2-VL forward."""
|
||||
for candidate in _walk_model_graph(model):
|
||||
language_model = getattr(candidate, "language_model", None)
|
||||
if (
|
||||
language_model is not None
|
||||
and not hasattr(candidate, "embed_tokens")
|
||||
and hasattr(language_model, "embed_tokens")
|
||||
):
|
||||
candidate.__dict__["embed_tokens"] = language_model.embed_tokens
|
||||
|
||||
|
||||
def set_videoalign_device(device) -> None:
|
||||
"""Move cached VideoAlign inferencers to device."""
|
||||
key = _normalize_device_str(device)
|
||||
@@ -45,7 +324,7 @@ def set_videoalign_device(device) -> None:
|
||||
):
|
||||
if old_key != key and old_key.split(":")[0] != key:
|
||||
new_key = inf._key_prefix + ":" + key
|
||||
inf.to(device)
|
||||
_move_videoalign_inferencer(inf, device)
|
||||
_VIDEOALIGN_INFERENCERS[new_key] = inf
|
||||
del _VIDEOALIGN_INFERENCERS[old_key]
|
||||
|
||||
@@ -56,10 +335,13 @@ def _get_inferencer(
|
||||
):
|
||||
"""Get or create VideoAlign inferencer."""
|
||||
if checkpoint_path is None:
|
||||
checkpoint_path = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..", "..", "..", "..", "..",
|
||||
"data", "VideoReward",
|
||||
checkpoint_path = os.environ.get(
|
||||
"VIDEOALIGN_CHECKPOINT_PATH",
|
||||
os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"VideoAlign",
|
||||
"checkpoints",
|
||||
),
|
||||
)
|
||||
checkpoint_path = os.path.abspath(checkpoint_path)
|
||||
|
||||
@@ -67,7 +349,8 @@ def _get_inferencer(
|
||||
cache_key = f"{checkpoint_path}:{key}"
|
||||
if cache_key not in _VIDEOALIGN_INFERENCERS:
|
||||
try:
|
||||
from inference import VideoVLMRewardInference
|
||||
inference_mod = _patch_videoalign_modules()
|
||||
VideoVLMRewardInference = inference_mod.VideoVLMRewardInference
|
||||
except ImportError as exc:
|
||||
msg = (
|
||||
"VideoAlign not found. Ensure the "
|
||||
@@ -81,6 +364,7 @@ def _get_inferencer(
|
||||
load_from_pretrained=checkpoint_path,
|
||||
device=device,
|
||||
)
|
||||
_patch_videoalign_runtime_model(inf.model)
|
||||
inf._key_prefix = checkpoint_path or "default"
|
||||
_VIDEOALIGN_INFERENCERS[cache_key] = inf
|
||||
return _VIDEOALIGN_INFERENCERS[cache_key]
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import importlib
|
||||
import inspect
|
||||
import time
|
||||
@@ -20,9 +21,11 @@ from fastvideo.train.methods.rl.reward import (
|
||||
videoalign_ta_score,
|
||||
)
|
||||
from fastvideo.train.methods.rl.reward.hpsv3 import (
|
||||
_HPSV3_INFERENCERS,
|
||||
set_hpsv3_device,
|
||||
)
|
||||
from fastvideo.train.methods.rl.reward.videoalign import (
|
||||
_VIDEOALIGN_INFERENCERS,
|
||||
set_videoalign_device,
|
||||
)
|
||||
|
||||
@@ -154,6 +157,28 @@ def move_reward_models(reward_cfg, device) -> None:
|
||||
set_videoalign_device(device)
|
||||
|
||||
|
||||
def clear_reward_models(reward_cfg) -> None:
|
||||
"""Drop cached GPU-backed reward models before PPO training."""
|
||||
cleared = False
|
||||
if _has_reward(
|
||||
reward_cfg,
|
||||
{"hpsv3_general", "hpsv3_percentile"},
|
||||
):
|
||||
_HPSV3_INFERENCERS.clear()
|
||||
cleared = True
|
||||
if _has_reward(
|
||||
reward_cfg,
|
||||
{"videoalign_mq", "videoalign_ta"},
|
||||
):
|
||||
_VIDEOALIGN_INFERENCERS.clear()
|
||||
cleared = True
|
||||
if cleared:
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def reward_models_on_device(reward_cfg, device):
|
||||
"""Temporarily move reward models to device."""
|
||||
@@ -173,10 +198,9 @@ def reward_models_on_device(reward_cfg, device):
|
||||
_t2 = time.perf_counter()
|
||||
move_reward_models(reward_cfg, "cpu")
|
||||
if use_cuda:
|
||||
import gc
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
_t3 = time.perf_counter()
|
||||
logger.info(
|
||||
"[rewards] move_to_cpu+gc=%.1fs",
|
||||
|
||||
Reference in New Issue
Block a user