feat(lora): add native Qwen3 dense and MoE policy runtime - #488
Draft
indevn wants to merge 22 commits into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds native LoRA support for Qwen3 dense and MoE policy training with colocated rollout, fully compatible with our existing tensor parallelism (TP) infrastructure.
The key design choice is to keep LoRA A/B matrices as parameters within AReno's native projection modules, rather than wrapping them externally. This lets us reuse the existing TP sharding, grouped-expert GEMMs, optimizer, offload, and checkpoint logic without special-casing LoRA everywhere. We also support standard PEFT adapter import/export so you can warm-start from or serve with existing adapters.
A new E2E test covers the full TP2/DP2 colocated rollout → train → next-rollout cycle for both dense and MoE variants.
Supported scope
qwen3andqwen3_moe.q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj(any non-empty subset).gate_proj,up_proj, anddown_projusing AReno grouped GEMMs. On replicated-KV topologies wheretp_size > num_key_value_heads,k_proj/v_projLoRA is rejected.adapter_config.jsonandadapter_model.safetensors; rank, alpha, dropout, and targets are read from artifact metadata on import.save_checkpointwrites an adapter-only PEFT warm-start artifact. It does not provide exact resume of optimizer state, trainer cursor, RNG, or adapter version.0only. RS-LoRA, DoRA, bias training, rank/alpha patterns, andmodules_to_saveare not supported.Not included in this PR: independent rollout-device LoRA synchronization, PPO/DPO reference or critic lifecycle, exact resume, non-Qwen model families, or a generic LoRA kernel framework.
Implementation notes
Tests