Skip to content

[feat] GenRL: add Wan LoRA adapter support - #1404

Merged
Davids048 merged 2 commits into
hao-ai-lab:py/add_rlfrom
Abecid:abecid/genrl-lora-support
Jun 4, 2026
Merged

[feat] GenRL: add Wan LoRA adapter support#1404
Davids048 merged 2 commits into
hao-ai-lab:py/add_rlfrom
Abecid:abecid/genrl-lora-support

Conversation

@Abecid

@Abecid Abecid commented May 27, 2026

Copy link
Copy Markdown

Extracted from #1391.

GenRL-Stack: 5/6

Purpose

Add FastVideo LoRA adapter support for GenRL Wan training, isolated from PPO-loop and reward-model changes.

This PR should be reviewed carefully for FSDP/HSDP interaction because review feedback noted that LoRA insertion after transformer loading may affect sharding/checkpoint assumptions.

Fixes #

Changes

  • Add use_lora, rank, alpha, init, path, and target-module options to GenRLWanModel.
  • Convert matching Wan transformer linear modules to FastVideo LoRA layers.
  • Add disable_adapter() compatibility so LoRA training can use the frozen base model for reference KL.
  • Implement LoRA initialization modes for GenRL Wan.
  • Fix RowParallel LoRA forward path.
  • Replace missing LoRA weight assertion with a runtime error.

Test Plan

python -m py_compile \
  fastvideo/train/models/wan/wan_genrl.py \
  fastvideo/layers/lora/linear.py

Test Results

Test output
py_compile passed locally.

Extracted from hao-ai-lab#1391.
Note: this PR should explicitly discuss LoRA insertion relative to FSDP/HSDP wrapping.

GenRL-Stack: 5/6
@mergify mergify Bot added scope: training Training pipeline, methods, configs scope: model Model architecture (DiTs, encoders, VAEs) labels May 27, 2026
@mergify

mergify Bot commented May 27, 2026

Copy link
Copy Markdown
Contributor

⚠️ PR title format required

Your PR title must start with a type tag in brackets. Examples:

  • [feat] Add new model support
  • [bugfix] Fix VAE tiling corruption
  • [refactor] Restructure training pipeline
  • [perf] Optimize attention kernel
  • [ci] Update test infrastructure
  • [infra] Add activation trace hooks
  • [docs] Add inference guide
  • [misc] Clean up configs
  • [new-model] Port Flux2 to FastVideo
  • [skill] Add add-model agent skill

Valid tags: feat, feature, bugfix, fix, refactor, perf, ci, infra, doc, docs, misc, chore, kernel, new-model, skill, skills

Please update your PR title and the merge protection check will pass automatically.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for FastVideo LoRA training in the Wan GenRL model, including inline forward logic for LoRA layers and helper functions for module conversion, weight initialization, and disabling adapters. The review feedback highlights a critical bug where slicing lora_A during training mode can cause an out-of-bounds IndexError on ranks > 0, suggests more robustly checking both lora_A and lora_B individually for DTensor instances, and recommends tightening the LoRA target module matching logic to avoid loose substring matches that could convert unintended modules.

Comment thread fastvideo/layers/lora/linear.py
Comment on lines +41 to +46
return any(
module_name == target
or module_name.endswith(f".{target}")
or target in module_name
for target in target_modules
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Using target in module_name for LoRA target matching is extremely loose and can lead to unexpected modules being matched and converted (e.g., if target is 'q', it will match any module name containing the letter 'q', such as 'sequence'). It is much safer to stick to exact matches or suffix matches (endswith), which is also the standard behavior in PEFT.

    return any(
        module_name == target
        or module_name.endswith(f".{target}")
        for target in target_modules
    )

Comment thread fastvideo/layers/lora/linear.py Outdated
@Davids048 Davids048 changed the title [genrl]: add Wan LoRA adapter support [feat] GenRL: add Wan LoRA adapter support May 27, 2026
@mergify mergify Bot added the type: feat New feature or capability label May 27, 2026
@Davids048
Davids048 merged commit 571c8ad into hao-ai-lab:py/add_rl Jun 4, 2026
3 of 5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

scope: model Model architecture (DiTs, encoders, VAEs) scope: training Training pipeline, methods, configs type: feat New feature or capability

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants