[feat] GenRL: stabilize reward model compatibility

[Abecid]

Extracted from #1391.

GenRL-Stack: 1/6

Purpose

Stabilize GenRL reward model loading and device handling for HPSv3 and VideoAlign rewards.

This splits the reward compatibility portion out of #1391 so reward-model import/runtime fixes can be reviewed independently from PPO loop, LoRA, and config changes.

Fixes #

Changes

• Patch HPSv3 runtime compatibility with newer transformers Qwen2-VL naming/import behavior.
• Patch VideoAlign runtime compatibility with newer Qwen2-VL model/key layouts.
• Improve reward image/video layout handling for one-frame and multi-frame inputs.
• Add reward model cache/device lifecycle helpers used by GenRL GPU reward mode.

Test Plan

python -m py_compile \
  fastvideo/train/methods/rl/reward/hpsv3.py \
  fastvideo/train/methods/rl/reward/videoalign.py \
  fastvideo/train/methods/rl/reward/utils.py \
  fastvideo/train/methods/rl/utils/rewards.py

Test Results

Test output

py_compile passed locally.

[gemini-code-assist]

Code Review

This pull request introduces patches and utilities to maintain compatibility with newer transformers releases and older Qwen2-VL checkpoint keys for the HPSv3 and VideoAlign reward models. It also adds robust device-moving capabilities, OpenCV-based video reading fallbacks, and memory cleanup utilities before PPO training. The review feedback identifies several critical issues: potential descending/negative indices in VideoAlign's frame selection when nframes is exactly 6, duplicate submodule registration when assigning embed_tokens directly to PyTorch modules, layout transposition bugs in prepare_images for single-frame video batches, and a missing meta-device parameter assignment fallback in the HPSv3 state dict loader.

[gemini-code-assist]

In _select_videoalign_frame_indices, when sample_type == "multi_pts", end_pt is calculated as int(nframes - frames_each_pts // 2 - 1). Since nframes is guaranteed to be at least frames_each_pts (6), when nframes is exactly 6, end_pt becomes 2. Since start_pt is 3, this results in start_pt > end_pt, causing torch.linspace to generate descending values. This leads to negative slice start indices (e.g., pt - 3 = -1), which in Python returns an empty slice, resulting in missing frames.

Removing the - 1 from end_pt ensures that end_pt is at least start_pt (3) and the slice indices remain valid and within bounds.

Suggested change

	start_pt = int(frames_each_pts // 2)
	end_pt = int(nframes - frames_each_pts // 2 - 1)
	start_pt = int(frames_each_pts // 2)
	end_pt = int(nframes - frames_each_pts // 2)

[gemini-code-assist]

Assigning candidate.embed_tokens = language_model.embed_tokens on a PyTorch nn.Module triggers nn.Module.__setattr__, which registers embed_tokens as a duplicate submodule of candidate. This causes model.parameters() and model.named_parameters() to yield duplicate parameter tensors under different names, which can lead to issues during optimization (e.g., duplicate optimizer states, double-counted gradients) or checkpoint saving.

Bypassing nn.Module.__setattr__ by writing directly to __dict__ provides the expected attribute alias without registering it as a duplicate submodule.

Suggested change

	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.embed_tokens = language_model.embed_tokens
	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

[gemini-code-assist]

Assigning candidate.embed_tokens = language_model.embed_tokens on a PyTorch nn.Module triggers nn.Module.__setattr__, which registers embed_tokens as a duplicate submodule of candidate. This causes model.parameters() and model.named_parameters() to yield duplicate parameter tensors under different names, which can lead to issues during optimization (e.g., duplicate optimizer states, double-counted gradients) or checkpoint saving.

Bypassing nn.Module.__setattr__ by writing directly to __dict__ provides the expected attribute alias without registering it as a duplicate submodule.

Suggested change

	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.embed_tokens = language_model.embed_tokens
	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

[gemini-code-assist]

In prepare_images, when images.ndim == 5 and the input layout is (N, C, F, H, W) with C=3 and F=1 (e.g., shape (N, 3, 1, H, W)), images.shape[2] is 1 (which is in (1, 3)). Because images.shape[2] in (1, 3) is checked first, it incorrectly matches and transposes the array as if it were (N, F, C, H, W), resulting in an incorrect shape of (N, 3, H, W, 1) instead of (N, 1, H, W, 3).

Adding an explicit check for images.shape[1] == 3 and images.shape[2] == 1 resolves this ambiguity and ensures the layout is correctly transposed.

Suggested change

	if images.shape[-1] in (1, 3):
	pass
	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):
	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):

[gemini-code-assist]

In _patch_load_state_dict, the load_state_dict_with_key_remap wrapper does not handle the case where parameters are on the meta device. If the model is loaded on a meta device (e.g., during DeepSpeed or Accelerate initialization), load_state_dict will fail unless assign=True is used.

Adding the same assign fallback logic used in videoalign.py makes the HPSv3 loader much more robust.

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