[feat] GenRL: add Wan LoRA adapter support - #1404
Conversation
Extracted from hao-ai-lab#1391. Note: this PR should explicitly discuss LoRA insertion relative to FSDP/HSDP wrapping. GenRL-Stack: 5/6
|
There was a problem hiding this comment.
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.
| return any( | ||
| module_name == target | ||
| or module_name.endswith(f".{target}") | ||
| or target in module_name | ||
| for target in target_modules | ||
| ) |
There was a problem hiding this comment.
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
)
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
use_lora, rank, alpha, init, path, and target-module options toGenRLWanModel.disable_adapter()compatibility so LoRA training can use the frozen base model for reference KL.Test Plan
Test Results
Test output