Skip to content

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

Merged
xsuler merged 29 commits into
inclusionAI:mainfrom
indevn:codex/qwen-lora-pr-replay-v1
Aug 27, 2026
Merged

feat(lora): add native Qwen3 dense and MoE policy runtime#488
xsuler merged 29 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

@indevn
indevn force-pushed the codex/qwen-lora-pr-replay-v1 branch from 49403a2 to 96b592f Compare August 26, 2026 13:05
@xsuler
xsuler marked this pull request as ready for review August 26, 2026 22:04
@xsuler
xsuler self-requested a review August 26, 2026 22:04

@xsuler xsuler left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

lgtm

@xsuler
xsuler force-pushed the codex/qwen-lora-pr-replay-v1 branch from 3891f6b to bbc59ef Compare August 26, 2026 23:48
@xsuler
xsuler merged commit dcc3507 into inclusionAI:main Aug 27, 2026
5 checks passed
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.

2 participants