Skip to content

feat(lora): add native Qwen3 dense and MoE policy runtime - #488

Draft
indevn wants to merge 22 commits into
inclusionAI:mainfrom
indevn:codex/qwen-lora-pr-replay-v1
Draft

feat(lora): add native Qwen3 dense and MoE policy runtime#488
indevn wants to merge 22 commits into
inclusionAI:mainfrom
indevn:codex/qwen-lora-pr-replay-v1

Conversation

@indevn

@indevn indevn commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

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

  • Model families: qwen3 and qwen3_moe.
  • Algorithms: policy-only GRPO/GSPO-style training through the existing policy trainer. PPO and DPO are rejected at the config boundary because reference/critic lifecycle support is not included.
  • Topologies covered by the E2E contract: TP2/DP2, colocated train and rollout engine.
  • Dense targets: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj (any non-empty subset).
  • MoE targets: attention projections plus routed expert gate_proj, up_proj, and down_proj using AReno grouped GEMMs. On replicated-KV topologies where tp_size > num_key_value_heads, k_proj/v_proj LoRA is rejected.
  • Adapter format: standard PEFT adapter_config.json and adapter_model.safetensors; rank, alpha, dropout, and targets are read from artifact metadata on import.
  • Checkpoint semantics: LoRA save_checkpoint writes an adapter-only PEFT warm-start artifact. It does not provide exact resume of optimizer state, trainer cursor, RNG, or adapter version.
  • Dropout: 0 only. RS-LoRA, DoRA, bias training, rank/alpha patterns, and modules_to_save are 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

  • Dense column/row-parallel projections preserve native AReno sharding and gradient collective contracts.
  • Routed experts attach canonical per-expert A/B parameters and execute their two low-rank products through the existing grouped GEMM facility.
  • FP32 route accumulation is scoped to active routed-expert LoRA. Base/fullweight Qwen3-MoE retains the upstream BF16 unpermute path while the broader accuracy/memory/performance trade-off remains under study.
  • Base parameters are frozen and the optimizer receives only adapter parameters.
  • Periodic LoRA checkpoint saves are routed once at the engine API boundary to PEFT export; lower layers do not repeat the policy decision.

Tests

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant