From 3e36909eab29d105ebd3d7e0cd9cd89016b44c95 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Mon, 17 Aug 2026 20:49:36 +0800 Subject: [PATCH 01/29] feat(lora): add native Qwen3 policy runtime --- areno/__init__.py | 5 + areno/accel/csrc/extension.cpp | 13 +- areno/accel/csrc/linear.cu | 176 ++++++++++--------- areno/accel/linear.py | 26 ++- areno/adapters/__init__.py | 6 + areno/adapters/config.py | 72 ++++++++ areno/adapters/lora.py | 273 ++++++++++++++++++++++++++++++ areno/adapters/peft.py | 124 ++++++++++++++ areno/api/__init__.py | 3 +- areno/api/backend/base.py | 5 + areno/api/backend/common.py | 10 +- areno/api/backend/cuda/backend.py | 23 ++- areno/api/config.py | 2 + areno/api/models.py | 1 + areno/api/trainer.py | 7 +- areno/api/trainer_config.py | 8 + areno/cli/serve.py | 34 ++++ areno/cli/train.py | 51 +++++- areno/engine/api.py | 34 +++- areno/engine/config.py | 30 ++++ areno/engine/data/batch.py | 2 + areno/engine/layers/linear.py | 27 ++- areno/engine/protocol.py | 8 + areno/engine/runtime/common.py | 5 +- areno/engine/runtime/rollout.py | 11 ++ areno/engine/training.py | 3 + areno/engine/worker.py | 45 ++++- areno/models/qwen3/model.py | 44 ++++- 28 files changed, 939 insertions(+), 109 deletions(-) create mode 100644 areno/adapters/__init__.py create mode 100644 areno/adapters/config.py create mode 100644 areno/adapters/lora.py create mode 100644 areno/adapters/peft.py diff --git a/areno/__init__.py b/areno/__init__.py index 5284a29c..e707bb5c 100644 --- a/areno/__init__.py +++ b/areno/__init__.py @@ -48,6 +48,10 @@ def __getattr__(name: str): from areno.engine import config return getattr(config, name) + if name == "LoraConfig": + from areno.adapters import LoraConfig + + return LoraConfig if name in {"RolloutOutput", "SamplingParams", "TrainStats"}: from areno.engine import data @@ -62,6 +66,7 @@ def __getattr__(name: str): __all__ = [ "ArenoEngine", "EngineConfig", + "LoraConfig", "ModelConfig", "OptimizerConfig", "RolloutOutput", diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index 149180c8..6cee02e9 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -15,7 +15,10 @@ std::vector areno_linear_backward_cuda( torch::Tensor grad_output, torch::Tensor input, torch::Tensor weight, - bool use_bias); + bool use_bias, + bool need_grad_input, + bool need_grad_weight, + bool need_grad_bias); torch::Tensor areno_causal_attention_forward_cuda( torch::Tensor q, torch::Tensor k, @@ -71,12 +74,16 @@ std::vector areno_grouped_linear_backward_cuda( torch::Tensor grad_output, torch::Tensor input, torch::Tensor weight, - std::vector tokens_per_expert); + std::vector tokens_per_expert, + bool need_grad_input, + bool need_grad_weight); std::vector areno_grouped_linear_backward_counts_cuda( torch::Tensor grad_output, torch::Tensor input, torch::Tensor weight, - torch::Tensor tokens_per_expert); + torch::Tensor tokens_per_expert, + bool need_grad_input, + bool need_grad_weight); std::vector areno_depthwise_causal_conv1d_silu_forward_cuda(torch::Tensor input, torch::Tensor weight); std::vector areno_depthwise_causal_conv1d_silu_decode_cuda( torch::Tensor current, diff --git a/areno/accel/csrc/linear.cu b/areno/accel/csrc/linear.cu index 243e8c00..89d69f01 100644 --- a/areno/accel/csrc/linear.cu +++ b/areno/accel/csrc/linear.cu @@ -193,57 +193,64 @@ std::vector areno_linear_backward_cuda( torch::Tensor grad_output, torch::Tensor input, torch::Tensor weight, - bool use_bias) { + bool use_bias, + bool need_grad_input, + bool need_grad_weight, + bool need_grad_bias) { TORCH_CHECK(grad_output.is_cuda(), "areno_linear grad_output must be CUDA"); TORCH_CHECK(input.is_cuda(), "areno_linear input must be CUDA"); TORCH_CHECK(weight.is_cuda(), "areno_linear weight must be CUDA"); TORCH_CHECK(input.scalar_type() == weight.scalar_type(), "areno_linear input and weight dtype must match"); TORCH_CHECK(grad_output.scalar_type() == input.scalar_type(), "areno_linear grad dtype must match input"); - auto grad_input = torch::empty_like(input); - auto grad_weight = torch::empty_like(weight); - auto grad_bias = use_bias ? areno_accel::reduce_bias_grad(grad_output) : torch::empty({0}, grad_output.options()); + auto grad_input = need_grad_input ? torch::empty_like(input) : torch::empty({0}, input.options()); + auto grad_weight = need_grad_weight ? torch::empty_like(weight) : torch::empty({0}, weight.options()); + auto grad_bias = need_grad_bias ? areno_accel::reduce_bias_grad(grad_output) : torch::empty({0}, grad_output.options()); int64_t k = input.size(-1); int64_t m = input.numel() / k; int64_t n = weight.size(0); const at::cuda::OptionalCUDAGuard guard(device_of(input)); - cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); - auto dtype = areno_accel::cuda_type(input.scalar_type()); - - areno_accel::gemm_row_major( - handle, - CUBLAS_OP_N, - CUBLAS_OP_N, - k, - m, - n, - weight.data_ptr(), - dtype, - k, - grad_output.data_ptr(), - dtype, - n, - grad_input.data_ptr(), - dtype, - k); - - areno_accel::gemm_row_major( - handle, - CUBLAS_OP_N, - CUBLAS_OP_T, - k, - n, - m, - input.data_ptr(), - dtype, - k, - grad_output.data_ptr(), - dtype, - n, - grad_weight.data_ptr(), - dtype, - k); + if (need_grad_input || need_grad_weight) { + cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); + auto dtype = areno_accel::cuda_type(input.scalar_type()); + if (need_grad_input) { + areno_accel::gemm_row_major( + handle, + CUBLAS_OP_N, + CUBLAS_OP_N, + k, + m, + n, + weight.data_ptr(), + dtype, + k, + grad_output.data_ptr(), + dtype, + n, + grad_input.data_ptr(), + dtype, + k); + } + if (need_grad_weight) { + areno_accel::gemm_row_major( + handle, + CUBLAS_OP_N, + CUBLAS_OP_T, + k, + n, + m, + input.data_ptr(), + dtype, + k, + grad_output.data_ptr(), + dtype, + n, + grad_weight.data_ptr(), + dtype, + k); + } + } return {grad_input, grad_weight, grad_bias}; } @@ -336,7 +343,9 @@ std::vector areno_grouped_linear_backward_cuda( torch::Tensor grad_output, torch::Tensor input, torch::Tensor weight, - std::vector tokens_per_expert) { + std::vector tokens_per_expert, + bool need_grad_input, + bool need_grad_weight) { TORCH_CHECK(grad_output.is_cuda(), "areno_grouped_linear grad_output must be CUDA"); TORCH_CHECK(input.is_cuda(), "areno_grouped_linear input must be CUDA"); TORCH_CHECK(weight.is_cuda(), "areno_grouped_linear weight must be CUDA"); @@ -360,8 +369,8 @@ std::vector areno_grouped_linear_backward_cuda( } TORCH_CHECK(total_tokens == input.size(0), "tokens_per_expert sum must match input rows"); - auto grad_input = torch::empty_like(input); - auto grad_weight = torch::zeros_like(weight); + auto grad_input = need_grad_input ? torch::empty_like(input) : torch::empty({0}, input.options()); + auto grad_weight = need_grad_weight ? torch::zeros_like(weight) : torch::empty({0}, weight.options()); const at::cuda::OptionalCUDAGuard guard(device_of(input)); cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); auto dtype = areno_accel::cuda_type(input.scalar_type()); @@ -379,42 +388,44 @@ std::vector areno_grouped_linear_backward_cuda( const void* expert_weight = weight_base + expert * n * k * elem_size; const void* expert_input = input_base + offset * k * elem_size; const void* expert_grad_output = grad_output_base + offset * n * elem_size; - void* expert_grad_input = grad_input_base + offset * k * elem_size; - void* expert_grad_weight = grad_weight_base + expert * n * k * elem_size; - - areno_accel::gemm_row_major( - handle, - CUBLAS_OP_N, - CUBLAS_OP_N, - k, - m, - n, - expert_weight, - dtype, - k, - expert_grad_output, - dtype, - n, - expert_grad_input, - dtype, - k); - - areno_accel::gemm_row_major( - handle, - CUBLAS_OP_N, - CUBLAS_OP_T, - k, - n, - m, - expert_input, - dtype, - k, - expert_grad_output, - dtype, - n, - expert_grad_weight, - dtype, - k); + if (need_grad_input) { + void* expert_grad_input = grad_input_base + offset * k * elem_size; + areno_accel::gemm_row_major( + handle, + CUBLAS_OP_N, + CUBLAS_OP_N, + k, + m, + n, + expert_weight, + dtype, + k, + expert_grad_output, + dtype, + n, + expert_grad_input, + dtype, + k); + } + if (need_grad_weight) { + void* expert_grad_weight = grad_weight_base + expert * n * k * elem_size; + areno_accel::gemm_row_major( + handle, + CUBLAS_OP_N, + CUBLAS_OP_T, + k, + n, + m, + expert_input, + dtype, + k, + expert_grad_output, + dtype, + n, + expert_grad_weight, + dtype, + k); + } } offset += m; } @@ -425,7 +436,9 @@ std::vector areno_grouped_linear_backward_counts_cuda( torch::Tensor grad_output, torch::Tensor input, torch::Tensor weight, - torch::Tensor tokens_per_expert) { + torch::Tensor tokens_per_expert, + bool need_grad_input, + bool need_grad_weight) { TORCH_CHECK(tokens_per_expert.is_cuda(), "areno_grouped_linear tokens_per_expert must be CUDA"); TORCH_CHECK(tokens_per_expert.dim() == 1, "areno_grouped_linear tokens_per_expert must be 1D"); TORCH_CHECK(tokens_per_expert.scalar_type() == at::kLong || tokens_per_expert.scalar_type() == at::kInt, "areno_grouped_linear tokens_per_expert must be int32 or int64"); @@ -443,5 +456,6 @@ std::vector areno_grouped_linear_backward_counts_cuda( counts[static_cast(i)] = static_cast(ptr[i]); } } - return areno_grouped_linear_backward_cuda(grad_output, input, weight, counts); + return areno_grouped_linear_backward_cuda( + grad_output, input, weight, counts, need_grad_input, need_grad_weight); } diff --git a/areno/accel/linear.py b/areno/accel/linear.py index a62a0ae1..35f7667e 100644 --- a/areno/accel/linear.py +++ b/areno/accel/linear.py @@ -29,15 +29,27 @@ def forward(ctx, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | Non return out @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + def backward( + ctx, grad_output: torch.Tensor + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: x, weight = ctx.saved_tensors + need_grad_input = bool(ctx.needs_input_grad[0]) + need_grad_weight = bool(ctx.needs_input_grad[1]) + need_grad_bias = ctx.use_bias and bool(ctx.needs_input_grad[2]) grad_input, grad_weight, grad_bias = _extension().areno_linear_backward( grad_output.contiguous(), x.contiguous(), weight.contiguous(), ctx.use_bias, + need_grad_input, + need_grad_weight, + need_grad_bias, + ) + return ( + grad_input if need_grad_input else None, + grad_weight if need_grad_weight else None, + grad_bias if need_grad_bias else None, ) - return grad_input, grad_weight, grad_bias if ctx.use_bias else None @torch._dynamo.disable @@ -68,13 +80,16 @@ def forward(ctx, x: torch.Tensor, weight: torch.Tensor, tokens_per_expert: list[ @staticmethod def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]: x, weight = ctx.saved_tensors + need_grad_input, need_grad_weight, _ = ctx.needs_input_grad grad_input, grad_weight = _extension().areno_grouped_linear_backward( grad_output.contiguous(), x.contiguous(), weight.contiguous(), ctx.tokens_per_expert, + need_grad_input, + need_grad_weight, ) - return grad_input, grad_weight, None + return grad_input if need_grad_input else None, grad_weight if need_grad_weight else None, None class _GroupedLinearCounts(torch.autograd.Function): @@ -91,13 +106,16 @@ def forward(ctx, x: torch.Tensor, weight: torch.Tensor, tokens_per_expert: torch @staticmethod def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]: x, weight, tokens_per_expert = ctx.saved_tensors + need_grad_input, need_grad_weight, _ = ctx.needs_input_grad grad_input, grad_weight = _extension().areno_grouped_linear_backward_counts( grad_output.contiguous(), x.contiguous(), weight.contiguous(), tokens_per_expert.contiguous(), + need_grad_input, + need_grad_weight, ) - return grad_input, grad_weight, None + return grad_input if need_grad_input else None, grad_weight if need_grad_weight else None, None @torch._dynamo.disable diff --git a/areno/adapters/__init__.py b/areno/adapters/__init__.py new file mode 100644 index 00000000..2c05280e --- /dev/null +++ b/areno/adapters/__init__.py @@ -0,0 +1,6 @@ +"""Native adapter runtime exposed by AReno.""" + +from areno.adapters.config import LoraConfig +from areno.adapters.lora import AdapterRegistry, LoraSlot, initialize_lora + +__all__ = ["AdapterRegistry", "LoraConfig", "LoraSlot", "initialize_lora"] diff --git a/areno/adapters/config.py b/areno/adapters/config.py new file mode 100644 index 00000000..e2eb452a --- /dev/null +++ b/areno/adapters/config.py @@ -0,0 +1,72 @@ +"""Public configuration for the Qwen3 dense and MoE LoRA runtime.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +QWEN3_DENSE_TARGETS = ( + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +) + + +@dataclass(frozen=True, slots=True) +class LoraConfig: + """Supported PEFT-compatible LoRA subset for Qwen3 dense and MoE models. + + When ``adapter_path`` is set, its standard PEFT metadata is authoritative + for rank, alpha, dropout, and targets. + """ + + rank: int = 8 + alpha: float = 16.0 + dropout: float = 0.0 + target_modules: tuple[str, ...] = QWEN3_DENSE_TARGETS + adapter_path: str | None = None + + def __post_init__(self) -> None: + if self.adapter_path is not None: + adapter_config = _read_adapter_config(self.adapter_path) + object.__setattr__(self, "rank", int(adapter_config["r"])) + object.__setattr__(self, "alpha", float(adapter_config["lora_alpha"])) + object.__setattr__(self, "dropout", float(adapter_config.get("lora_dropout", 0.0))) + object.__setattr__(self, "target_modules", tuple(adapter_config["target_modules"])) + object.__setattr__(self, "target_modules", tuple(self.target_modules)) + if self.rank < 1: + raise ValueError("lora rank must be >= 1") + if self.alpha <= 0: + raise ValueError("lora alpha must be > 0") + if self.dropout != 0.0: + raise ValueError("native Qwen3 LoRA currently requires dropout=0") + requested = set(self.target_modules) + supported = set(QWEN3_DENSE_TARGETS) + if not requested or not requested <= supported: + raise ValueError(f"target_modules must be a non-empty subset of {QWEN3_DENSE_TARGETS}") + + @property + def scale(self) -> float: + return float(self.alpha) / float(self.rank) + + +def _read_adapter_config(path: str) -> dict: + adapter_config = json.loads((Path(path) / "adapter_config.json").read_text(encoding="utf-8")) + if str(adapter_config.get("peft_type", "")).upper() != "LORA": + raise ValueError("adapter_path must contain a PEFT LoRA artifact") + unsupported = [] + if adapter_config.get("bias", "none") != "none" or bool(adapter_config.get("lora_bias", False)): + unsupported.append("bias") + if bool(adapter_config.get("fan_in_fan_out", False)): + unsupported.append("fan_in_fan_out") + for option in ("use_rslora", "use_dora", "rank_pattern", "alpha_pattern", "modules_to_save"): + if adapter_config.get(option): + unsupported.append(option) + if unsupported: + raise ValueError(f"unsupported PEFT LoRA options: {', '.join(unsupported)}") + return adapter_config diff --git a/areno/adapters/lora.py b/areno/adapters/lora.py new file mode 100644 index 00000000..2a6019e1 --- /dev/null +++ b/areno/adapters/lora.py @@ -0,0 +1,273 @@ +"""TP-aware native LoRA slots for Qwen3 dense and routed-expert projections.""" + +from __future__ import annotations + +import hashlib +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from areno.accel import areno_grouped_linear +from areno.adapters.config import LoraConfig +from areno.engine.layers.linear import RowParallelLinear, mark_tensor_parallel_parameter +from areno.engine.parallel.context import get_tp_context + + +class LoraSlot(nn.Module): + """One canonical LoRA A/B pair owned by its native projection module.""" + + def __init__( + self, + *, + logical_name: str, + base_weight: nn.Parameter, + global_in_features: int, + global_out_features: int, + local_in_features: int, + local_out_features: int, + row_parallel: bool, + config: LoraConfig, + seed: int, + ) -> None: + super().__init__() + ctx = get_tp_context() + self.logical_name = logical_name + self.rank = int(config.rank) + self.global_in_features = int(global_in_features) + self.global_out_features = int(global_out_features) + self.local_in_features = int(local_in_features) + self.local_out_features = int(local_out_features) + self.row_parallel = bool(row_parallel) + self.lora_A = nn.Parameter( + torch.empty(self.rank, self.local_in_features, device=base_weight.device, dtype=base_weight.dtype) + ) + self.lora_B = nn.Parameter( + torch.empty(self.local_out_features, self.rank, device=base_weight.device, dtype=base_weight.dtype) + ) + self.register_buffer("scale", torch.tensor(config.scale, device=base_weight.device, dtype=torch.float32)) + if row_parallel: + mark_tensor_parallel_parameter(self.lora_A, True, sequence_parallel=True) + mark_tensor_parallel_parameter(self.lora_B, False, sequence_parallel=True, tp_grad_allreduce=True) + else: + mark_tensor_parallel_parameter(self.lora_A, False, sequence_parallel=True, tp_grad_allreduce=True) + mark_tensor_parallel_parameter(self.lora_B, True, sequence_parallel=True) + self._reset_parameters(seed, ctx.rank, ctx.world_size) + + @torch.no_grad() + def _reset_parameters(self, seed: int, tp_rank: int, tp_size: int) -> None: + material = f"{int(seed)}:{self.logical_name}".encode() + target_seed = int.from_bytes(hashlib.sha256(material).digest()[:8], "big") % (2**63) + generator = torch.Generator(device="cpu") + generator.manual_seed(target_seed) + canonical_A = torch.empty(self.rank, self.global_in_features, dtype=torch.float32) + nn.init.kaiming_uniform_(canonical_A, a=math.sqrt(5), generator=generator) + if self.row_parallel: + shard = self.global_in_features // tp_size + canonical_A = canonical_A[:, tp_rank * shard : (tp_rank + 1) * shard] + self.lora_A.copy_(canonical_A.to(device=self.lora_A.device, dtype=self.lora_A.dtype)) + self.lora_B.zero_() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scale + + +class RoutedExpertLoraSlot(nn.Module): + """One expert-sharded canonical LoRA A/B pair for grouped Qwen3-MoE GEMMs.""" + + def __init__( + self, + *, + logical_name: str, + base_weight: nn.Parameter, + local_num_experts: int, + local_expert_start: int, + in_features: int, + out_features: int, + config: LoraConfig, + seed: int, + ) -> None: + super().__init__() + self.logical_name = logical_name + self.rank = int(config.rank) + self.local_num_experts = int(local_num_experts) + self.local_expert_start = int(local_expert_start) + self.in_features = int(in_features) + self.out_features = int(out_features) + self.lora_A = nn.Parameter( + torch.empty( + self.local_num_experts, + self.rank, + self.in_features, + device=base_weight.device, + dtype=base_weight.dtype, + ) + ) + self.lora_B = nn.Parameter( + torch.empty( + self.local_num_experts, + self.out_features, + self.rank, + device=base_weight.device, + dtype=base_weight.dtype, + ) + ) + self.register_buffer("scale", torch.tensor(config.scale, device=base_weight.device, dtype=torch.float32)) + mark_tensor_parallel_parameter(self.lora_A, True, sequence_parallel=False, tp_grad_allreduce=False) + mark_tensor_parallel_parameter(self.lora_B, True, sequence_parallel=False, tp_grad_allreduce=False) + self._reset_parameters(seed) + + @torch.no_grad() + def _reset_parameters(self, seed: int) -> None: + for local_expert_id in range(self.local_num_experts): + expert_id = self.local_expert_start + local_expert_id + material = f"{int(seed)}:{self.logical_name}:expert={expert_id}".encode() + target_seed = int.from_bytes(hashlib.sha256(material).digest()[:8], "big") % (2**63) + generator = torch.Generator(device="cpu") + generator.manual_seed(target_seed) + initial_A = torch.empty(self.rank, self.in_features, dtype=torch.float32) + nn.init.kaiming_uniform_(initial_A, a=math.sqrt(5), generator=generator) + self.lora_A[local_expert_id].copy_(initial_A.to(device=self.lora_A.device, dtype=self.lora_A.dtype)) + self.lora_B.zero_() + + def forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + hidden = areno_grouped_linear(x.contiguous(), self.lora_A, tokens_per_expert) + return areno_grouped_linear(hidden, self.lora_B, tokens_per_expert) * self.scale + + +class AdapterRegistry: + """Non-owning index over LoRA slots; projection modules remain sole owners.""" + + def __init__(self, slots: dict[str, LoraSlot | RoutedExpertLoraSlot], config: LoraConfig) -> None: + self.slots = slots + self.config = config + self.version = 0 + + def named_parameters(self): + for name, slot in self.slots.items(): + yield f"{name}.lora_A.weight", slot.lora_A + yield f"{name}.lora_B.weight", slot.lora_B + + def parameters(self) -> tuple[nn.Parameter, ...]: + return tuple(parameter for _, parameter in self.named_parameters()) + + def increment_version(self) -> int: + self.version += 1 + return self.version + + +def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> AdapterRegistry: + """Freeze a Qwen3 dense or MoE base and attach canonical targets.""" + + model_config = getattr(model, "config", None) + if getattr(model_config, "model_type", None) not in {"qwen3", "qwen3_moe"}: + raise ValueError("native LoRA currently supports Qwen3 models only") + for parameter in model.parameters(): + parameter.requires_grad_(False) + + requested = set(config.target_modules) + slots: dict[str, LoraSlot | RoutedExpertLoraSlot] = {} + for layer_index, layer in enumerate(model.layers): + prefix = f"layers.{layer_index}" + qkv = layer.self_attn.qkv_proj + for component_index, component in enumerate(("q_proj", "k_proj", "v_proj")): + if component not in requested: + continue + logical_name = f"{prefix}.self_attn.{component}" + slot = LoraSlot( + logical_name=logical_name, + base_weight=qkv.weight, + global_in_features=qkv.in_features, + global_out_features=qkv.out_features[component_index], + local_in_features=qkv.in_features, + local_out_features=qkv.local_out_features[component_index], + row_parallel=False, + config=config, + seed=seed, + ) + qkv.install_lora_component(component, component_index, slot) + slots[logical_name] = slot + + if "o_proj" in requested: + owner = layer.self_attn.o_proj + logical_name = f"{prefix}.self_attn.o_proj" + slot = _row_slot(logical_name, owner, config, seed) + owner.install_lora(slot) + slots[logical_name] = slot + + if getattr(model_config, "enable_moe_block", False): + _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, slots) + else: + gate_up = layer.mlp.gate_up_proj + for component_index, component in enumerate(("gate_proj", "up_proj")): + if component not in requested: + continue + logical_name = f"{prefix}.mlp.{component}" + slot = LoraSlot( + logical_name=logical_name, + base_weight=gate_up.weight, + global_in_features=gate_up.in_features, + global_out_features=gate_up.out_features[component_index], + local_in_features=gate_up.in_features, + local_out_features=gate_up.local_out_features[component_index], + row_parallel=False, + config=config, + seed=seed, + ) + gate_up.install_lora_component(component, component_index, slot) + slots[logical_name] = slot + + if "down_proj" in requested: + owner = layer.mlp.down_proj + logical_name = f"{prefix}.mlp.down_proj" + slot = _row_slot(logical_name, owner, config, seed) + owner.install_lora(slot) + slots[logical_name] = slot + + return AdapterRegistry(slots, config) + + +def _row_slot(logical_name: str, owner: RowParallelLinear, config: LoraConfig, seed: int) -> LoraSlot: + return LoraSlot( + logical_name=logical_name, + base_weight=owner.weight, + global_in_features=owner.in_features, + global_out_features=owner.out_features, + local_in_features=owner.local_in_features, + local_out_features=owner.out_features, + row_parallel=True, + config=config, + seed=seed, + ) + + +def _install_moe_slots( + experts: nn.Module, + prefix: str, + requested: set[str], + config: LoraConfig, + seed: int, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], +) -> None: + components = ( + ("gate_proj", experts.hidden_size, experts.intermediate_size, experts.gate_up_weight), + ("up_proj", experts.hidden_size, experts.intermediate_size, experts.gate_up_weight), + ("down_proj", experts.intermediate_size, experts.hidden_size, experts.down_weight), + ) + for component, in_features, out_features, base_weight in components: + if component not in requested: + continue + logical_name = f"{prefix}.mlp.experts.{{expert}}.{component}" + slot = RoutedExpertLoraSlot( + logical_name=logical_name, + base_weight=base_weight, + local_num_experts=experts.local_num_experts, + local_expert_start=experts.local_expert_start, + in_features=in_features, + out_features=out_features, + config=config, + seed=seed, + ) + experts.install_lora_component(component, slot) + slots[logical_name] = slot diff --git a/areno/adapters/peft.py b/areno/adapters/peft.py new file mode 100644 index 00000000..addd8d72 --- /dev/null +++ b/areno/adapters/peft.py @@ -0,0 +1,124 @@ +"""Standard PEFT safetensors import/export for native TP LoRA slots.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import torch +import torch.distributed as dist +from safetensors.torch import load_file, save_file + +from areno.adapters.lora import AdapterRegistry, RoutedExpertLoraSlot +from areno.engine.parallel.context import get_tp_context + +_PREFIX = "base_model.model.model." + + +@torch.no_grad() +def load_peft_adapter(registry: AdapterRegistry, path: str | Path) -> None: + """Copy one supported PEFT adapter into the registry's stable A/B storage.""" + + input_path = Path(path) + tensors = load_file(input_path / "adapter_model.safetensors", device="cpu") + ctx = get_tp_context() + for logical_name, slot in registry.slots.items(): + if isinstance(slot, RoutedExpertLoraSlot): + for local_expert_id in range(slot.local_num_experts): + expert_id = slot.local_expert_start + local_expert_id + expert_name = logical_name.format(expert=expert_id) + slot.lora_A[local_expert_id].copy_( + tensors[_key(expert_name, "A")].to(device=slot.lora_A.device, dtype=slot.lora_A.dtype) + ) + slot.lora_B[local_expert_id].copy_( + tensors[_key(expert_name, "B")].to(device=slot.lora_B.device, dtype=slot.lora_B.dtype) + ) + continue + canonical_A = tensors[_key(logical_name, "A")] + canonical_B = tensors[_key(logical_name, "B")] + if slot.row_parallel: + width = slot.local_in_features + local_A = canonical_A[:, ctx.rank * width : (ctx.rank + 1) * width] + local_B = canonical_B + else: + height = slot.local_out_features + local_A = canonical_A + local_B = canonical_B[ctx.rank * height : (ctx.rank + 1) * height] + slot.lora_A.copy_(local_A.to(device=slot.lora_A.device, dtype=slot.lora_A.dtype)) + slot.lora_B.copy_(local_B.to(device=slot.lora_B.device, dtype=slot.lora_B.dtype)) + + +@torch.no_grad() +def export_peft_adapter( + registry: AdapterRegistry, + path: str | Path, + *, + base_model_name_or_path: str | None, +) -> str | None: + """Gather the authoritative DP0 TP shards and write a PEFT adapter.""" + + ctx = get_tp_context() + if ctx.dp_rank != 0: + return None + state: dict[str, torch.Tensor] = {} + for logical_name, slot in registry.slots.items(): + if isinstance(slot, RoutedExpertLoraSlot): + gathered_A = _all_gather(slot.lora_A.detach(), ctx.world_size, ctx.group) + gathered_B = _all_gather(slot.lora_B.detach(), ctx.world_size, ctx.group) + if ctx.rank != 0: + continue + canonical_A = torch.cat(gathered_A, dim=0) + canonical_B = torch.cat(gathered_B, dim=0) + for expert_id in range(canonical_A.shape[0]): + expert_name = logical_name.format(expert=expert_id) + state[_key(expert_name, "A")] = canonical_A[expert_id].float().cpu().contiguous() + state[_key(expert_name, "B")] = canonical_B[expert_id].float().cpu().contiguous() + continue + if slot.row_parallel: + gathered_A = _all_gather(slot.lora_A.detach(), ctx.world_size, ctx.group) + if ctx.rank != 0: + continue + canonical_A = torch.cat(gathered_A, dim=1) + canonical_B = slot.lora_B.detach() + else: + gathered_B = _all_gather(slot.lora_B.detach(), ctx.world_size, ctx.group) + if ctx.rank != 0: + continue + canonical_A = slot.lora_A.detach() + canonical_B = torch.cat(gathered_B, dim=0) + state[_key(logical_name, "A")] = canonical_A.float().cpu().contiguous() + state[_key(logical_name, "B")] = canonical_B.float().cpu().contiguous() + if ctx.rank != 0: + return None + + output_path = Path(path) + output_path.mkdir(parents=True, exist_ok=True) + config = { + "base_model_name_or_path": base_model_name_or_path, + "bias": "none", + "fan_in_fan_out": False, + "inference_mode": True, + "lora_alpha": registry.config.alpha, + "lora_dropout": registry.config.dropout, + "peft_type": "LORA", + "r": registry.config.rank, + "target_modules": list(registry.config.target_modules), + "task_type": "CAUSAL_LM", + } + (output_path / "adapter_config.json").write_text( + json.dumps(config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + save_file(state, output_path / "adapter_model.safetensors") + return str(output_path) + + +def _all_gather(tensor: torch.Tensor, world_size: int, group) -> list[torch.Tensor]: + if world_size == 1: + return [tensor] + gathered = [torch.empty_like(tensor) for _ in range(world_size)] + dist.all_gather(gathered, tensor, group=group) + return gathered + + +def _key(logical_name: str, component: str) -> str: + return f"{_PREFIX}{logical_name}.lora_{component}.weight" diff --git a/areno/api/__init__.py b/areno/api/__init__.py index cc8f9812..5230f2c0 100644 --- a/areno/api/__init__.py +++ b/areno/api/__init__.py @@ -26,7 +26,7 @@ register_algorithm, sft_loss_fn, ) -from areno.api.config import CudaConfig, MlxConfig, default_backend_type +from areno.api.config import CudaConfig, LoraConfig, MlxConfig, default_backend_type from areno.api.data import PromptBatch, PromptItem from areno.api.models import ( BackendType, @@ -49,6 +49,7 @@ "AlgorithmSpec", "CudaConfig", "MlxConfig", + "LoraConfig", "PromptBatch", "PromptItem", "AgentBatch", diff --git a/areno/api/backend/base.py b/areno/api/backend/base.py index a37da9e8..9b487665 100644 --- a/areno/api/backend/base.py +++ b/areno/api/backend/base.py @@ -133,6 +133,11 @@ def save_checkpoint(self, ctx: Context, path: str) -> str: raise NotImplementedError(f"{type(self).__name__} does not support checkpoint saving") + def export_adapter(self, ctx: Context, path: str) -> str: + """Persist a standard adapter artifact, or raise when unsupported.""" + + raise NotImplementedError(f"{type(self).__name__} does not support adapter export") + def ensure_roles(self, ctx: Context, roles: dict[str, ModelRole]) -> None: """Prepare backend-owned auxiliary model roles, or raise if unsupported.""" diff --git a/areno/api/backend/common.py b/areno/api/backend/common.py index f9ab445e..3306d57b 100644 --- a/areno/api/backend/common.py +++ b/areno/api/backend/common.py @@ -122,13 +122,21 @@ def group_rollout_sequences( sequences: list[RolloutSequence], prompt_count: int, n_samples: int, + *, + adapter_version: int | None = None, ) -> list[RolloutResult]: """Restore a flat prompt-major sequence list to the public result layout.""" expected = prompt_count * n_samples if len(sequences) != expected: raise ValueError(f"backend returned {len(sequences)} sequences; expected {expected}") - return [RolloutResult(sequences=sequences[start : start + n_samples]) for start in range(0, expected, n_samples)] + return [ + RolloutResult( + sequences=sequences[start : start + n_samples], + adapter_version=adapter_version, + ) + for start in range(0, expected, n_samples) + ] __all__ = [ diff --git a/areno/api/backend/cuda/backend.py b/areno/api/backend/cuda/backend.py index 5c4ee6fc..9252f028 100644 --- a/areno/api/backend/cuda/backend.py +++ b/areno/api/backend/cuda/backend.py @@ -156,6 +156,8 @@ def initialize(self, ctx: Context): raise ValueError(f"training device count must equal world_size={world_size}") if cfg.rollout_tp_size is not None and cfg.rollout_devices is None: raise ValueError("rollout_tp_size requires rollout_devices") + if cfg.lora is not None and cfg.uses_separate_rollout_engine(): + raise ValueError("native LoRA currently supports colocated rollout only") if not cfg.uses_separate_rollout_engine(): self._train_engine = ArenoEngine.from_pretrained( @@ -169,6 +171,7 @@ def initialize(self, ctx: Context): runtime_config=RuntimeConfig(**cfg.runtime), loss_fn=dispatch_loss, policy_sync_bucket_mb=cfg.policy_sync_bucket_mb, + lora_config=cfg.lora, ) return self._policy_sync_bucket_bytes = cfg.policy_sync_bucket_mb * 1024 * 1024 @@ -355,7 +358,12 @@ def rollout_batch( RolloutSequence(resp_tokens=tokens, resp_logprobs=rollout.logprobs[index, : len(tokens)].tolist()) for index, tokens in enumerate(rollout.response_ids) ] - return group_rollout_sequences(sequences, len(prompt_tokens), n_samples) + return group_rollout_sequences( + sequences, + len(prompt_tokens), + n_samples, + adapter_version=rollout.adapter_version, + ) def begin_rollout_session(self, ctx: Context) -> None: """Prepare colocated actor state before rollout requests are issued.""" @@ -459,7 +467,12 @@ async def rollout_batch_async( RolloutSequence(resp_tokens=tokens, resp_logprobs=rollout.logprobs[index, : len(tokens)].tolist()) for index, tokens in enumerate(rollout.response_ids) ] - return group_rollout_sequences(sequences, len(prompt_tokens), n_samples) + return group_rollout_sequences( + sequences, + len(prompt_tokens), + n_samples, + adapter_version=rollout.adapter_version, + ) def train( self, @@ -510,6 +523,8 @@ def train( metric_rows.append({str(key): float(value) for key, value in stats.metrics.items()}) metrics = reduce_microbatch_metrics(metric_rows) result = {"loss": sum(losses) / max(len(losses), 1)} + if stats_list and stats_list[-1].adapter_version is not None: + result["adapter_version"] = stats_list[-1].adapter_version result.update(metrics) result.update(self._pending_policy_sync_metrics) self._pending_policy_sync_metrics = {} @@ -531,6 +546,10 @@ def save_checkpoint(self, ctx: Context, path: str) -> str: engine = self._require_train_engine() return save_checkpoint(engine, path) + def export_adapter(self, ctx: Context, path: str) -> str: + del ctx + return self._require_train_engine().export_adapter(path) + def ensure_roles(self, ctx: Context, roles: dict[str, ModelRole]) -> None: engine = self._require_train_engine() engine.ensure_roles(roles) diff --git a/areno/api/config.py b/areno/api/config.py index effd5aa5..169bb773 100644 --- a/areno/api/config.py +++ b/areno/api/config.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from typing import Any +from areno.adapters.config import LoraConfig from areno.api.models import BackendType @@ -35,6 +36,7 @@ class CudaConfig: runtime: dict[str, Any] = field(default_factory=dict) max_running_prompts: int = 64 decode_progress_interval_s: float = 10.0 + lora: LoraConfig | None = None def uses_separate_rollout_engine(self) -> bool: """Return whether rollout runs on its own CUDA device partition.""" diff --git a/areno/api/models.py b/areno/api/models.py index fa737d4a..b6d8d7a9 100644 --- a/areno/api/models.py +++ b/areno/api/models.py @@ -64,6 +64,7 @@ class RolloutResult(BaseModel): """All sampled completions for one prompt.""" sequences: list[RolloutSequence] = Field(default_factory=list) + adapter_version: int | None = Field(default=None) class TrainSequence(BaseModel): diff --git a/areno/api/trainer.py b/areno/api/trainer.py index 4ae21f28..f1a53945 100644 --- a/areno/api/trainer.py +++ b/areno/api/trainer.py @@ -536,10 +536,15 @@ def train_values( ) def save_checkpoint(self, path: str) -> str: - """Save a checkpoint in the selected backend's native format.""" + """Save a native backend checkpoint, or a PEFT artifact for native LoRA.""" return self._backend.save_checkpoint(self._ctx, path) + def export_adapter(self, path: str) -> str: + """Export the live native LoRA weights as a standard PEFT adapter.""" + + return self._backend.export_adapter(self._ctx, path) + def close(self) -> None: """Release backend workers and local resources such as metric writers.""" diff --git a/areno/api/trainer_config.py b/areno/api/trainer_config.py index c39b15ee..f1fedcb3 100644 --- a/areno/api/trainer_config.py +++ b/areno/api/trainer_config.py @@ -13,6 +13,7 @@ from dataclasses import dataclass +from areno.adapters.config import LoraConfig from areno.api.defaults import DEFAULT_METRICS_LOG_DIR @@ -76,6 +77,7 @@ class TrainerConfig: agent_timeout_s: float = 300.0 train_tool_results: bool = False chat_template_enable_thinking: bool | None = None + lora: LoraConfig | None = None def __post_init__(self) -> None: if self.backend is None: @@ -116,6 +118,10 @@ def __post_init__(self) -> None: self.multimodal_projector_lr_decay_steps, self.multimodal_projector_lr_decay_style, ) + if self.lora is not None and self.backend != "cuda": + raise ValueError("native LoRA is only supported by the CUDA backend") + if self.lora is not None and self.algo.lower() in {"ppo", "dpo"}: + raise ValueError("native LoRA does not yet support PPO/DPO reference and critic roles") @staticmethod def _validate_multimodal_optimizer_group( @@ -210,6 +216,7 @@ def cuda_config(self): "eager_decode": self.eager_decode, "attn_backend": self.attn_backend, }, + lora=self.lora, ) @@ -257,6 +264,7 @@ def cuda_config(self): "eager_decode": self.eager_decode, "attn_backend": self.attn_backend, }, + lora=self.lora, ) def mlx_config(self): diff --git a/areno/cli/serve.py b/areno/cli/serve.py index b8d5666b..64a0fc19 100644 --- a/areno/cli/serve.py +++ b/areno/cli/serve.py @@ -21,6 +21,7 @@ from pydantic import BaseModel, Field from areno.api import MLX, BackendType, MlxConfig, SamplingParams, Trainer, default_backend_type +from areno.adapters import LoraConfig from areno.api.multimodal import ( expand_image_tokens, image_token_counts_from_features, @@ -204,6 +205,7 @@ def __init__( world_size: int, eager_decode: bool, attn_backend: str, + lora: LoraConfig | None, ) -> None: from areno.engine.config import RuntimeConfig @@ -218,6 +220,7 @@ def __init__( devices=list(range(world_size)), runtime_config=RuntimeConfig(eager_decode=bool(eager_decode), attn_backend=attn_backend), loss_fn=_serve_loss_fn, + lora_config=lora, ) self.max_model_len = int(self._engine.config.model.max_position_embeddings) @@ -306,8 +309,11 @@ def _create_serve_runtime( decode_progress_interval_s: float, eager_decode: bool, attn_backend: str, + lora: LoraConfig | None, ) -> _CudaServeRuntime | _MlxServeRuntime: if backend_type == MLX: + if lora is not None: + raise ValueError("native LoRA serving is only supported by the CUDA backend") if world_size != 1 or tp_size != 1: raise ValueError("MLX serving requires --world-size 1 and --tp-size 1") return _MlxServeRuntime( @@ -322,6 +328,7 @@ def _create_serve_runtime( world_size=world_size, eager_decode=eager_decode, attn_backend=attn_backend, + lora=lora, ) @@ -336,6 +343,7 @@ def create_app( eager_decode: bool = False, attn_backend: Literal["flash", "native"] = "flash", chat_template_enable_thinking: bool | None = None, + lora: LoraConfig | None = None, ) -> FastAPI: """Construct the FastAPI app: load tokenizer/engine, install routes and lifecycle hooks.""" if world_size < 1: @@ -365,6 +373,7 @@ def create_app( decode_progress_interval_s=decode_progress_interval_s, eager_decode=eager_decode, attn_backend=attn_backend, + lora=lora, ) if backend_type == MLX: tokenizer = engine.tokenizer @@ -980,6 +989,16 @@ def _normalize_stop(stop: str | list[str] | None) -> list[str]: is_flag=True, help="Pass enable_thinking=False to tokenizer chat templates when supported.", ) +@click.option("--lora-rank", type=int, default=None, help="Enable native LoRA with this rank.") +@click.option("--lora-alpha", type=float, default=16.0, show_default=True, help="Native LoRA alpha.") +@click.option("--lora-dropout", type=float, default=0.0, show_default=True, help="Native LoRA dropout (must be 0).") +@click.option( + "--lora-target-modules", + default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", + show_default=True, + help="Comma-separated dense Qwen3 projection targets.", +) +@click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter to serve.") def serve_command( model_path: str, model_hub: Literal["hf", "modelscope"], @@ -993,11 +1012,25 @@ def serve_command( eager_decode: bool, attn_backend: Literal["flash", "native"], disable_thinking: bool, + lora_rank: int | None, + lora_alpha: float, + lora_dropout: float, + lora_target_modules: str, + lora_adapter_path: str | None, ) -> None: """Click entry point: build the app and hand it to uvicorn.""" import uvicorn model_path = resolve_model_ref(model_path, model_hub=model_hub) + lora = None + if lora_rank is not None or lora_adapter_path is not None: + lora = LoraConfig( + rank=8 if lora_rank is None else lora_rank, + alpha=lora_alpha, + dropout=lora_dropout, + target_modules=tuple(item.strip() for item in lora_target_modules.split(",") if item.strip()), + adapter_path=lora_adapter_path, + ) from areno.cli.dashboard_registry import register_dashboard_job register_dashboard_job( @@ -1027,6 +1060,7 @@ def serve_command( eager_decode=eager_decode, attn_backend=attn_backend, chat_template_enable_thinking=False if disable_thinking else None, + lora=lora, ) uvicorn.run(app, host=host, port=port) diff --git a/areno/cli/train.py b/areno/cli/train.py index e8f5ff3c..a3ecf74f 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -19,7 +19,7 @@ import logging import shutil import textwrap -from dataclasses import fields +from dataclasses import asdict, fields, is_dataclass from pathlib import Path from types import SimpleNamespace from typing import TYPE_CHECKING @@ -132,6 +132,11 @@ def flash_attention_unsupported_model_reason(model_config): "score_micro_bs", "gradient_accumulation_steps", "activation_checkpointing", + "lora_rank", + "lora_alpha", + "lora_dropout", + "lora_target_modules", + "lora_adapter_path", "lr", "min_lr", "lr_decay_steps", @@ -244,6 +249,7 @@ def _trainer_config_from_options(**options) -> TrainerConfig: args.multimodal_projector_min_lr = getattr(args, "multimodal_projector_min_lr", None) args.multimodal_projector_lr_decay_steps = getattr(args, "multimodal_projector_lr_decay_steps", None) args.multimodal_projector_lr_decay_style = getattr(args, "multimodal_projector_lr_decay_style", None) + args.lora = _lora_config_from_options(args) if args.backend == "mlx": if args.train_devices is not None or args.rollout_devices is not None or args.rollout_tp_size is not None: raise click.UsageError("MLX does not use CUDA device or rollout TP options") @@ -400,6 +406,26 @@ def _require_positive_float(value: float, option_name: str) -> None: raise click.UsageError(f"{option_name} must be positive") +def _lora_config_from_options(args): + rank = getattr(args, "lora_rank", None) + adapter_path = getattr(args, "lora_adapter_path", None) + if rank is None and adapter_path is None: + return None + from areno.adapters import LoraConfig + + targets = tuple(item.strip() for item in args.lora_target_modules.split(",") if item.strip()) + try: + return LoraConfig( + rank=8 if rank is None else rank, + alpha=args.lora_alpha, + dropout=args.lora_dropout, + target_modules=targets, + adapter_path=adapter_path, + ) + except ValueError as exc: + raise click.UsageError(str(exc)) from exc + + def _parse_cuda_devices(value: str | None, option_name: str) -> list[int] | None: """Parse CUDA indices and inclusive ranges such as ``0..3,8``.""" @@ -815,6 +841,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: args.multimodal_projector_min_lr = getattr(args, "multimodal_projector_min_lr", None) args.multimodal_projector_lr_decay_steps = getattr(args, "multimodal_projector_lr_decay_steps", None) args.multimodal_projector_lr_decay_style = getattr(args, "multimodal_projector_lr_decay_style", None) + lora = getattr(args, "lora", None) algorithm = get_algorithm(args.algo) chat_template_enable_thinking = False if args.disable_thinking else None if algorithm.name == "dpo": @@ -873,6 +900,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: chat_template_enable_thinking=chat_template_enable_thinking, ref_ckpt=args.ref_ckpt, dpo_beta=args.dpo_beta, + lora=lora, ) if algorithm.name == "sft": return TrainerConfig( @@ -928,6 +956,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: agent_timeout_s=args.agent_timeout_s, train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, + lora=lora, ) if algorithm.name != "ppo": return PolicyTrainerConfig( @@ -995,6 +1024,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: agent_timeout_s=args.agent_timeout_s, train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, + lora=lora, ) return PPOTrainerConfig( algo=algorithm.name, @@ -1075,6 +1105,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: agent_timeout_s=args.agent_timeout_s, train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, + lora=lora, ) @@ -1140,13 +1171,17 @@ def _write_dashboard_run_config(config: TrainerConfig) -> None: def _training_config_settings(config: TrainerConfig) -> dict: used: set[str] = set() + def value(name: str): + item = getattr(config, name) + return asdict(item) if is_dataclass(item) else item + def section(title: str, names: list[str]) -> dict: items = [] for name in names: if not hasattr(config, name): continue used.add(name) - items.append({"key": name, "value": getattr(config, name)}) + items.append({"key": name, "value": value(name)}) return {"title": title, "items": items} sections = [ @@ -1248,7 +1283,7 @@ def section(title: str, names: list[str]) -> dict: extras = [] for field in fields(config): if field.name not in used: - extras.append({"key": field.name, "value": getattr(config, field.name)}) + extras.append({"key": field.name, "value": value(field.name)}) if extras: sections.append({"title": "Other", "items": extras}) if isinstance(config, RolloutTrainerConfig): @@ -1593,6 +1628,16 @@ def _dataset_builder_for_suffix(suffix: str) -> str: help="Override global concurrent rollout prompts; defaults to batch-size * n-samples.", ) @click.option("--lr", type=float, default=1.0e-6, show_default=True, help="Policy optimizer learning rate.") +@click.option("--lora-rank", type=int, default=None, help="Enable native LoRA with this rank.") +@click.option("--lora-alpha", type=float, default=16.0, show_default=True, help="Native LoRA alpha.") +@click.option("--lora-dropout", type=float, default=0.0, show_default=True, help="Native LoRA dropout (must be 0).") +@click.option( + "--lora-target-modules", + default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", + show_default=True, + help="Comma-separated Qwen3 projection targets (MoE MLP targets apply to each routed expert).", +) +@click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter used to initialize native LoRA.") @click.option("--min-lr", type=float, default=1.0e-7, show_default=True, help="Policy optimizer minimum learning rate.") @click.option("--lr-decay-steps", type=int, default=1000, show_default=True, help="Policy LR decay steps.") @click.option("--lr-decay-style", default="cosine", show_default=True, help="Policy LR decay style.") diff --git a/areno/engine/api.py b/areno/engine/api.py index 77d68e14..4faf9e27 100644 --- a/areno/engine/api.py +++ b/areno/engine/api.py @@ -26,11 +26,13 @@ import torch +from areno.adapters.config import LoraConfig from areno.engine.checkpoints.io import resolve_model_path from areno.engine.config import EngineConfig, OptimizerConfig, RuntimeConfig from areno.engine.data import RolloutOutput, SamplingParams, TrainStats, to_cpu from areno.engine.protocol import ( EnsureRolesPayload, + ExportAdapterPayload, Op, RoleSpecPayload, RolloutCacheProbePayload, @@ -110,9 +112,19 @@ def _merge_dp_rollouts_by_prompt_indices( [row[2] for row in materialized], [row[3] for row in materialized], metrics=None, + adapter_version=_rollout_version(non_empty=[output for output in outputs if output is not None]), ) +def _rollout_version(*, non_empty: list[RolloutOutput]) -> int | None: + versions = {output.adapter_version for output in non_empty} + if not versions: + return None + if len(versions) != 1: + raise RuntimeError(f"rollout DP replicas reported different adapter versions: {versions}") + return versions.pop() + + class ArenoEngine: """User-facing coordinator for one local training/inference engine. @@ -198,6 +210,7 @@ def from_pretrained( start: bool = True, cluster_kwargs: dict[str, Any] | None = None, policy_sync_bucket_mb: int = 64, + lora_config: LoraConfig | None = None, ) -> ArenoEngine: """Build an engine by reading model config from a checkpoint path. @@ -228,6 +241,8 @@ def from_pretrained( runtime=runtime_config or RuntimeConfig(), role=role, policy_sync_bucket_mb=policy_sync_bucket_mb, + lora=lora_config, + lora_seed=torch.initial_seed(), ) return cls(cfg, start=start, cluster_kwargs=cluster_kwargs) @@ -687,16 +702,27 @@ def train_values( return merge_metric_dicts(rank0_results) or {} def save_checkpoint(self, path: str) -> str: - """Ask workers to write a HuggingFace-compatible checkpoint. + """Save base weights, or the standard PEFT artifact in native LoRA mode. - Dispatches ``Op.SAVE_CHECKPOINT`` (blocking). Workers cooperatively - write shards to ``path``; only rank 0's returned path is propagated - back to the caller. + Fullweight workers cooperatively write HuggingFace shards to ``path``. + Native LoRA keeps the base frozen, so its checkpoint is the adapter-only + PEFT artifact consumed by training and serving. """ + if self.config.lora is not None: + return self.export_adapter(path) results = self.cluster.call(Op.SAVE_CHECKPOINT, SaveCheckpointPayload(path=path)) return results[0]["path"] + def export_adapter(self, path: str) -> str: + """Export the live native LoRA weights as a standard PEFT adapter.""" + + results = self.cluster.call(Op.EXPORT_ADAPTER, ExportAdapterPayload(path=path)) + result = next((result for result in results if result is not None), None) + if result is None: + raise RuntimeError("native LoRA export did not produce an artifact") + return result["path"] + def _transport_payload(self, payload: Any) -> Any: """Move tensors to CPU shared memory for zero-copy IPC to workers.""" diff --git a/areno/engine/config.py b/areno/engine/config.py index d52f08dc..55c6340e 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -16,6 +16,8 @@ import torch +from areno.adapters.config import LoraConfig + # AReno's flash path uses flash-attn features beyond the Turing-compatible # forward kernels, including paged KV/cache and training paths, so require # Ampere+ even though flash-attn 2.x has partial sm75 forward support. @@ -119,6 +121,19 @@ def resolve_compile_model(self, *, model: ModelConfig, devices: list[int]) -> No ) self.compile_model = False + def resolve_eager_decode(self, *, model: ModelConfig, lora: LoraConfig | None) -> None: + """Use eager decode when routed-expert adapters need grouped execution.""" + + if self.eager_decode or lora is None: + return + if model.model_type == "qwen3_moe" and {"gate_proj", "up_proj", "down_proj"} & set(lora.target_modules): + warnings.warn( + "Qwen3-MoE expert LoRA uses grouped execution during rollout; falling back to eager decode.", + RuntimeWarning, + stacklevel=2, + ) + self.eager_decode = True + @dataclass(slots=True) class ModelConfig: @@ -263,6 +278,8 @@ class EngineConfig: dummy_load: bool = False role: Literal["train", "rollout"] = "train" policy_sync_bucket_mb: int = 64 + lora: LoraConfig | None = None + lora_seed: int = 0 def __post_init__(self) -> None: """Infer DP/devices and validate the distributed layout.""" @@ -270,6 +287,18 @@ def __post_init__(self) -> None: if self.sequence_parallel is not None: self.model.sequence_parallel = bool(self.sequence_parallel) self.model.validate_tp(self.tp_size) + if self.lora is not None: + replicated_kv_targets = {"k_proj", "v_proj"} & set(self.lora.target_modules) + if ( + self.model.model_type in {"qwen3", "qwen3_moe"} + and self.tp_size > self.model.num_key_value_heads + and replicated_kv_targets + ): + targets = ", ".join(sorted(replicated_kv_targets)) + raise ValueError( + f"Qwen3-MoE replicated-KV topology does not support LoRA targets {targets}; " + "omit k_proj/v_proj or use tp_size <= num_key_value_heads" + ) if self.devices is None: if torch.cuda.is_available(): device_count = torch.cuda.device_count() @@ -305,6 +334,7 @@ def __post_init__(self) -> None: raise ValueError("runtime.kv_block_size must be a multiple of 256 for FlashAttention paged KV") self.runtime.resolve_attn_backend(model=self.model, devices=self.devices) self.runtime.resolve_compile_model(model=self.model, devices=self.devices) + self.runtime.resolve_eager_decode(model=self.model, lora=self.lora) self.model.attn_backend = self.runtime.attn_backend @property diff --git a/areno/engine/data/batch.py b/areno/engine/data/batch.py index bfdd3536..6f195fa9 100644 --- a/areno/engine/data/batch.py +++ b/areno/engine/data/batch.py @@ -22,6 +22,7 @@ class TrainStats: loss: float stepped: bool = True metrics: dict[str, float] | None = None + adapter_version: int | None = None @dataclass(slots=True) @@ -50,6 +51,7 @@ class RolloutOutput: logprobs: torch.Tensor finish_reason: list[str] metrics: dict[str, float] | None = None + adapter_version: int | None = None def to_device(obj: Any, device: torch.device, _memo: dict[int, torch.Tensor] | None = None) -> Any: diff --git a/areno/engine/layers/linear.py b/areno/engine/layers/linear.py index e7271b93..ae9033ee 100644 --- a/areno/engine/layers/linear.py +++ b/areno/engine/layers/linear.py @@ -135,6 +135,8 @@ class MergedColumnParallelLinear(nn.Module): def __init__(self, in_features: int, out_features: list[int] | tuple[int, ...], bias: bool = False): super().__init__() + self.lora_slots = nn.ModuleDict() + self._lora_component_indices: dict[str, int] = {} ctx = get_tp_context() if not out_features: raise ValueError("out_features must not be empty") @@ -152,6 +154,12 @@ def __init__(self, in_features: int, out_features: list[int] | tuple[int, ...], mark_tensor_parallel_parameter(self.bias, True, sequence_parallel=True) self.reset_parameters() + def install_lora_component(self, component: str, component_index: int, slot: nn.Module) -> None: + """Attach one canonical adapter to a fused output component.""" + + self.lora_slots[component] = slot + self._lora_component_indices[component] = component_index + def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: @@ -165,7 +173,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if is_sequence_parallel_active() else copy_to_tensor_parallel_region(x) ) - return _areno_linear_forward(x, self.weight, self.bias) + out = _areno_linear_forward(x, self.weight, self.bias) + if not self.lora_slots: + return out + parts = list(out.split(self.local_out_features, dim=-1)) + for component, slot in self.lora_slots.items(): + index = self._lora_component_indices[component] + parts[index] = parts[index] + slot(x) + return torch.cat(parts, dim=-1) class QKVParallelLinear(MergedColumnParallelLinear): @@ -186,6 +201,8 @@ def __init__( bias: bool = False, ): nn.Module.__init__(self) + self.lora_slots = nn.ModuleDict() + self._lora_component_indices: dict[str, int] = {} self.head_dim = head_dim self.num_heads = num_heads self.num_kv_heads = num_kv_heads @@ -224,6 +241,7 @@ def __init__(self, in_features: int, out_features: int, bias: bool = False, inpu self.out_features = out_features self.local_in_features = end - start self.input_is_parallel = input_is_parallel + self.lora_slot: nn.Module | None = None self.weight = nn.Parameter(torch.empty(out_features, self.local_in_features)) # Bias lives on each rank as a replica (not TP-sharded) and is added # post-reduction so it is not summed `world_size` times. @@ -232,6 +250,11 @@ def __init__(self, in_features: int, out_features: int, bias: bool = False, inpu mark_tensor_parallel_parameter(self.bias, False, sequence_parallel=True) self.reset_parameters() + def install_lora(self, slot: nn.Module) -> None: + """Attach the adapter before compilation and optimizer construction.""" + + self.lora_slot = slot + def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: @@ -246,6 +269,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: start, end = _shard_range(self.in_features, ctx.rank, ctx.world_size) x = x[..., start:end] out = _areno_linear_forward(x, self.weight, None) + if self.lora_slot is not None: + out = out + self.lora_slot(x) # Partial sum -> cross-rank reduction. SP mode also re-shards along # the sequence dim via reduce-scatter, saving activation memory. out = reduce_scatter_to_sequence_parallel_region(out) if is_sequence_parallel_active() else all_reduce(out) diff --git a/areno/engine/protocol.py b/areno/engine/protocol.py index abed8c4b..cf96bc6f 100644 --- a/areno/engine/protocol.py +++ b/areno/engine/protocol.py @@ -43,6 +43,7 @@ class Op(Enum): ROLLOUT_SESSION_SYNC = auto() ROLLOUT_SESSION_END = auto() SAVE_CHECKPOINT = auto() + EXPORT_ADAPTER = auto() POLICY_SYNC_PLAN = auto() POLICY_SYNC_PUBLISH = auto() POLICY_SYNC_RECEIVE = auto() @@ -156,6 +157,13 @@ class SaveCheckpointPayload: path: str +@dataclass(slots=True) +class ExportAdapterPayload: + """Typed payload for Op.EXPORT_ADAPTER.""" + + path: str + + @dataclass(slots=True, frozen=True) class PolicySyncPayload: """Version and fixed GPU bucket bound for one policy transfer.""" diff --git a/areno/engine/runtime/common.py b/areno/engine/runtime/common.py index 7031e8b8..028ad70a 100644 --- a/areno/engine/runtime/common.py +++ b/areno/engine/runtime/common.py @@ -115,7 +115,10 @@ def merge_train_stats(results: list[dict[str, Any]]) -> TrainStats: loss = sum(float(result["loss"]) for result in results) / len(results) stepped = all(bool(result["stepped"]) for result in results) metrics = merge_metric_dicts([result.get("metrics") for result in results]) - return TrainStats(loss=loss, stepped=stepped, metrics=metrics) + versions = {result.get("adapter_version") for result in results} + if len(versions) != 1: + raise RuntimeError(f"DP replicas reported different adapter versions: {versions}") + return TrainStats(loss=loss, stepped=stepped, metrics=metrics, adapter_version=versions.pop()) def merge_metric_dicts(metrics_list: list[dict[str, Any] | None]) -> dict[str, float] | None: diff --git a/areno/engine/runtime/rollout.py b/areno/engine/runtime/rollout.py index 36c1da93..eac6fa3a 100644 --- a/areno/engine/runtime/rollout.py +++ b/areno/engine/runtime/rollout.py @@ -38,6 +38,7 @@ def _merge_rollouts(outputs: list[RolloutOutput]) -> RolloutOutput: logprobs=logprobs, finish_reason=finish_reason, metrics=_merge_rollout_metrics(outputs), + adapter_version=_merge_adapter_version(outputs), ) @@ -73,6 +74,7 @@ def _merge_dp_rollouts_in_input_order(outputs: list[RolloutOutput | None], total finish_reason, logprob_rows, metrics=_merge_rollout_metrics(outputs), + adapter_version=_merge_adapter_version(outputs), ) @@ -83,6 +85,7 @@ def _build_rollout_from_rows( logprob_rows: list[torch.Tensor], *, metrics: dict[str, float] | None, + adapter_version: int | None = None, ) -> RolloutOutput: """Build padded rollout tensors from variable-length Python rows.""" if not prompt_ids: @@ -97,6 +100,7 @@ def _build_rollout_from_rows( logprobs=logprobs, finish_reason=finish_reason, metrics=metrics, + adapter_version=adapter_version, ) @@ -125,3 +129,10 @@ def _merge_rollout_metrics(outputs: list[RolloutOutput]) -> dict[str, float] | N for key, value in output.metrics.items(): merged[key] = merged.get(key, 0.0) + float(value) return merged or None + + +def _merge_adapter_version(outputs: list[RolloutOutput]) -> int | None: + versions = {output.adapter_version for output in outputs} + if len(versions) != 1: + raise RuntimeError(f"rollout shards reported different adapter versions: {versions}") + return versions.pop() diff --git a/areno/engine/training.py b/areno/engine/training.py index 7b89efe2..cbd4b0ba 100644 --- a/areno/engine/training.py +++ b/areno/engine/training.py @@ -184,6 +184,8 @@ def _train_step( worker.optimizer.step() worker.optimizer.zero_grad(set_to_none=True) worker._global_step += 1 + if worker.adapter_registry is not None: + worker.adapter_registry.increment_version() if worker.device.type == "cuda": torch.cuda.empty_cache() else: @@ -201,6 +203,7 @@ def _train_step( "loss": float(loss.detach().cpu()), "stepped": stepped, "global_step": worker._global_step, + "adapter_version": (worker.adapter_registry.version if worker.adapter_registry is not None else None), "metrics": _merge_metrics( metrics, None, diff --git a/areno/engine/worker.py b/areno/engine/worker.py index 29cb4ac2..649d7d2a 100644 --- a/areno/engine/worker.py +++ b/areno/engine/worker.py @@ -21,6 +21,8 @@ import torch.distributed as dist from areno.api.backend.cuda.roles import RoleManager, WorkerRole +from areno.adapters import initialize_lora +from areno.adapters.peft import export_peft_adapter, load_peft_adapter from areno.engine.config import EngineConfig from areno.engine.data import RolloutOutput from areno.engine.data.sampling import _truncate_generated @@ -30,6 +32,7 @@ from areno.engine.policy_sync import policy_plan_metadata, transfer_policy_weights from areno.engine.protocol import ( Command, + ExportAdapterPayload, Op, PolicySyncPayload, RolloutCacheProbePayload, @@ -63,10 +66,18 @@ def __init__(self, config: EngineConfig): if config.model_path is not None and not config.dummy_load: load_model_weights(self.model, config.model, config.model_path) configure_multimodal_training(self.model, config.optimizer, trainable=config.role == "train") + self.adapter_registry = ( + initialize_lora(self.model, config.lora, seed=config.lora_seed) if config.lora is not None else None + ) if config.runtime.compile_model: self.model = torch.compile(self.model) + if self.adapter_registry is not None and config.lora.adapter_path is not None: + load_peft_adapter(self.adapter_registry, config.lora.adapter_path) opt = config.optimizer - self.optimizer = build_optimizer(self.model.parameters(), opt, ctx) if config.role == "train" else None + optimizer_parameters = ( + self.adapter_registry.parameters() if self.adapter_registry is not None else self.model.parameters() + ) + self.optimizer = build_optimizer(optimizer_parameters, opt, ctx) if config.role == "train" else None self.grad_clip_norm = opt.grad_clip_norm self.base_lr = opt.lr self.min_lr = opt.min_lr @@ -165,6 +176,8 @@ def handle(self, cmd: Command): return self.train_values(cmd.payload) if cmd.op is Op.SAVE_CHECKPOINT: return self.save_checkpoint(cmd.payload) + if cmd.op is Op.EXPORT_ADAPTER: + return self.export_adapter(cmd.payload) if cmd.op is Op.POLICY_SYNC_PLAN: return self.policy_sync_plan(cmd.payload) if cmd.op is Op.POLICY_SYNC_PUBLISH: @@ -219,11 +232,14 @@ def ensure_roles(self, payload: dict) -> None: def infer_rollout(self, payload: dict, finished_callback=None, refill_callback=None) -> RolloutOutput | None: """Delegate rollout generation to `InferenceManager`.""" - return self.inference.infer_rollout( + output = self.inference.infer_rollout( payload, finished_callback=finished_callback, refill_callback=refill_callback, ) + if output is not None and self.adapter_registry is not None: + output.adapter_version = self.adapter_registry.version + return output def probe_rollout_cache(self, payload: RolloutCacheProbePayload) -> float: """Allocate rollout KV cache and capture decode graphs without decoding.""" @@ -279,7 +295,9 @@ def send_empty_requests() -> None: ( self._rank, WorkerResult( - ok=True, payload=_empty_rollout() if ctx.is_rank0 else None, request_id=request_id + ok=True, + payload=self._stamp_adapter_version(_empty_rollout()) if ctx.is_rank0 else None, + request_id=request_id, ), ) ) @@ -315,6 +333,7 @@ def send_finished( row_ids, truncate_stop_token_ids, ) + result_payload = self._stamp_adapter_version(result_payload) request_id = request_ids[request_idx] self._result_queue.put( (self._rank, WorkerResult(ok=True, payload=result_payload, request_id=request_id)) @@ -374,6 +393,12 @@ def refill_waiting(state) -> list[int]: if idx not in sent ] + def _stamp_adapter_version(self, output: RolloutOutput) -> RolloutOutput: + adapter_registry = getattr(self, "adapter_registry", None) + if adapter_registry is not None: + output.adapter_version = adapter_registry.version + return output + def _next_refill_command(self) -> Command | None: """Fetch the next queued command consistently across TP ranks.""" @@ -627,6 +652,18 @@ def save_checkpoint(self, payload: SaveCheckpointPayload) -> dict | None: path = save_model_weights(self.model, self.config.model, payload.path, self.config.model_path) return {"path": path} if path is not None else None + def export_adapter(self, payload: ExportAdapterPayload) -> dict | None: + """Write the native adapter in standard PEFT format.""" + + if self.adapter_registry is None: + raise RuntimeError("export_adapter requires native LoRA") + self._prepare_actor_onloaded() + path = export_peft_adapter( + self.adapter_registry, + payload.path, + base_model_name_or_path=self.config.model_path, + ) + return {"path": path} if path is not None else None def _rollout_payloads_compatible(first: RolloutPayload, other: RolloutPayload) -> bool: """Return whether two rollout payloads can share one InferenceBatchState.""" @@ -793,6 +830,7 @@ def _slice_rollout_output(output: RolloutOutput, start: int, end: int) -> Rollou logprobs=logprobs, finish_reason=finish_reason, metrics=output.metrics, + adapter_version=output.adapter_version, ) @@ -815,4 +853,5 @@ def _slice_rollout_output_rows(output: RolloutOutput, rows: list[int]) -> Rollou logprobs=logprobs, finish_reason=finish_reason, metrics=output.metrics, + adapter_version=output.adapter_version, ) diff --git a/areno/models/qwen3/model.py b/areno/models/qwen3/model.py index 9693858f..cdd16727 100644 --- a/areno/models/qwen3/model.py +++ b/areno/models/qwen3/model.py @@ -103,9 +103,35 @@ def __init__(self, config: ModelConfig): self.down_weight = nn.Parameter( torch.empty(self.local_num_experts, self.hidden_size, self.intermediate_size, dtype=config.dtype) ) + self.lora_slots = nn.ModuleDict() mark_tensor_parallel_parameter(self.gate_up_weight, True, sequence_parallel=False, tp_grad_allreduce=False) mark_tensor_parallel_parameter(self.down_weight, True, sequence_parallel=False, tp_grad_allreduce=False) + def install_lora_component(self, component: str, slot: nn.Module) -> None: + """Attach one grouped canonical adapter before optimizer construction.""" + + self.lora_slots[component] = slot + + def has_active_lora(self) -> bool: + return bool(self.lora_slots) + + def _gate_up_forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + base = _areno_grouped_linear_no_compile(x.contiguous(), self.gate_up_weight, tokens_per_expert) + if not self.lora_slots: + return base + gate, up = base.chunk(2, dim=-1) + if "gate_proj" in self.lora_slots: + gate = gate + self.lora_slots["gate_proj"](x, tokens_per_expert) + if "up_proj" in self.lora_slots: + up = up + self.lora_slots["up_proj"](x, tokens_per_expert) + return torch.cat((gate, up), dim=-1) + + def _down_forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: + out = _areno_grouped_linear_no_compile(x, self.down_weight, tokens_per_expert) + if "down_proj" in self.lora_slots: + out = out + self.lora_slots["down_proj"](x, tokens_per_expert) + return out + def forward(self, flat: torch.Tensor, topk_idx: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: x, route_weight, token_idx, tokens_per_expert = _areno_moe_topk_permute_no_compile( flat, @@ -124,14 +150,21 @@ def forward(self, flat: torch.Tensor, topk_idx: torch.Tensor, topk_weight: torch + self.down_weight.reshape(-1)[0] * 0 + topk_weight.sum().to(dtype=self.gate_up_weight.dtype) * 0 ) + for slot in self.lora_slots.values(): + zero = zero + slot.lora_A.reshape(-1)[0] * 0 + slot.lora_B.reshape(-1)[0] * 0 return all_reduce(flat.new_zeros(flat.shape) + zero) - hidden = _areno_grouped_linear_no_compile(x.contiguous(), self.gate_up_weight, tokens_per_expert) + hidden = self._gate_up_forward(x, tokens_per_expert) log_once("qwen3_moe_silu_and_mul", "using ARENO fused silu_and_mul kernel for Qwen3-MoE experts") hidden = ( _areno_silu_and_mul_no_compile(hidden) * route_weight.unsqueeze(-1).to(dtype=hidden.dtype) ).contiguous() - out = _areno_grouped_linear_no_compile(hidden, self.down_weight, tokens_per_expert) - out = _areno_moe_unpermute_no_compile(out, token_idx, flat.shape) + out = self._down_forward(hidden, tokens_per_expert) + if self.has_active_lora(): + # Stabilize routed-expert LoRA without changing the base/fullweight MoE path. + out = _areno_moe_unpermute_no_compile(out.float(), token_idx, flat.shape) + out = out.to(dtype=flat.dtype) + else: + out = _areno_moe_unpermute_no_compile(out, token_idx, flat.shape) return all_reduce(out) def local_routes(self, topk_idx: torch.Tensor, topk_weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -212,7 +245,7 @@ def forward_with_routes( batch, seqlen, hidden = hidden_states.shape flat = hidden_states.reshape(-1, hidden) with sequence_parallel_region(False): - if self.training: + if self.training or self.experts.has_active_lora(): out = self.experts(flat, topk_idx.to(torch.long), topk_weight) else: out = self._forward_fused_moe(flat, topk_idx, topk_weight) @@ -225,6 +258,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @torch.no_grad() def prepare_infer_weights(self) -> None: + if self.experts.has_active_lora(): + self.clear_infer_weights() + return self._infer_w1_weight = self._updated_infer_weight( self._infer_w1_weight, self.experts.gate_up_weight.detach().to(dtype=self.experts.gate_up_weight.dtype).contiguous(), From 74636b9c1486acc8b285889aec434a733f71e6a3 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Mon, 17 Aug 2026 20:49:44 +0800 Subject: [PATCH 02/29] test(lora): cover Qwen3 train-rollout PEFT loop --- tests/test_config_data_cpu.py | 53 +++++++ tests/test_qwen3_lora_e2e.py | 240 +++++++++++++++++++++++++++++ tests/test_train_cli_config_cpu.py | 27 ++++ 3 files changed, 320 insertions(+) create mode 100644 tests/test_qwen3_lora_e2e.py diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 91eeb213..1589bf3a 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import sys import tempfile import types @@ -12,6 +13,7 @@ import torch from click.testing import CliRunner +from areno.adapters import LoraConfig from areno.api.data import PromptBatch, PromptItem from areno.api.trainer_config import RolloutTrainerConfig, TrainerConfig from areno.cli import train as train_cli @@ -199,6 +201,57 @@ def test_engine_config_resolves_sequence_parallel_override_before_model_config(s tp1 = EngineConfig(model=model, tp_size=1, devices=[0], sequence_parallel=True) self.assertFalse(tp1.effective_sequence_parallel) + def test_engine_config_rejects_replicated_kv_lora_targets(self): + """Replicated Qwen3 KV requires range-aware LoRA support.""" + for model_type in ("qwen3", "qwen3_moe"): + with self.subTest(model_type=model_type): + model = ModelConfig( + model_type=model_type, + num_attention_heads=8, + num_key_value_heads=2, + intermediate_size=16, + vocab_size=32, + ) + + with self.assertRaisesRegex(ValueError, "replicated-KV.*k_proj"): + EngineConfig(model=model, tp_size=4, devices=[0, 1, 2, 3], lora=LoraConfig()) + + EngineConfig( + model=model, + tp_size=4, + devices=[0, 1, 2, 3], + lora=LoraConfig(target_modules=("q_proj", "o_proj")), + ) + + def test_trainer_config_rejects_lora_ppo_and_dpo(self): + """Reference and critic roles are outside the initial native-LoRA scope.""" + for algo in ("ppo", "dpo"): + with self.subTest(algo=algo), self.assertRaisesRegex(ValueError, "PPO/DPO"): + TrainerConfig(algo=algo, ckpt="actor", dataset_path="dataset", lora=LoraConfig()) + + def test_adapter_path_uses_peft_metadata(self): + """A PEFT artifact should configure non-default native slots itself.""" + with tempfile.TemporaryDirectory() as adapter_path: + Path(adapter_path, "adapter_config.json").write_text( + json.dumps( + { + "peft_type": "LORA", + "r": 4, + "lora_alpha": 8, + "lora_dropout": 0, + "bias": "none", + "target_modules": ["q_proj", "o_proj"], + } + ), + encoding="utf-8", + ) + + config = LoraConfig(adapter_path=adapter_path) + + self.assertEqual(config.rank, 4) + self.assertEqual(config.alpha, 8) + self.assertEqual(config.target_modules, ("q_proj", "o_proj")) + def test_runtime_config_attn_backend_propagates_to_model_config(self): """The runtime attention backend should reach model layer construction.""" model = ModelConfig(num_attention_heads=4, num_key_value_heads=4, intermediate_size=16, vocab_size=32) diff --git a/tests/test_qwen3_lora_e2e.py b/tests/test_qwen3_lora_e2e.py new file mode 100644 index 00000000..dd287a8f --- /dev/null +++ b/tests/test_qwen3_lora_e2e.py @@ -0,0 +1,240 @@ +"""Qwen3 dense and MoE TP2/DP2 native-LoRA rollout/train/PEFT E2E.""" + +from __future__ import annotations + +import os +import sys +from importlib import util as importlib_util +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file + +from areno.adapters import LoraConfig +from areno.api import ArenoConfig, Trainer +from areno.api.algorithms import get_algorithm +from areno.api.trainer_config import PolicyTrainerConfig +from areno.api.trainers.policy_only import PolicyOnlyTrainer + + +class _ObservedTrainer: + def __init__(self, inner: Trainer) -> None: + self.inner = inner + self.rollout_versions: list[int | None] = [] + self.train_versions: list[int | None] = [] + + def __getattr__(self, name: str): + return getattr(self.inner, name) + + async def rollout_token_batch_async(self, prompt_tokens, n_samples, sampling_params, *, prompt_features=None): + results = await self.inner.rollout_token_batch_async( + prompt_tokens, + n_samples, + sampling_params, + prompt_features=prompt_features, + ) + self.rollout_versions.extend(result.adapter_version for result in results) + return results + + def train(self, batch_data, loss_fn, mini_bs=8, gradient_accumulation_steps=None): + result = self.inner.train(batch_data, loss_fn, mini_bs, gradient_accumulation_steps) + self.train_versions.append(result.get("adapter_version")) + return result + + +@pytest.mark.parametrize( + ("model_env", "model_kind"), + (("ARENO_E2E_QWEN3_MODEL", "dense"), ("ARENO_E2E_QWEN3_MOE_MODEL", "moe")), +) +def test_qwen3_lora_tp2_dp2_rollout_train_peft(tmp_path: Path, model_env: str, model_kind: str) -> None: + model_path_value = os.getenv(model_env) + if not model_path_value: + pytest.skip(f"set {model_env} to run the 4-GPU Qwen3 {model_kind} LoRA E2E") + model_path = Path(model_path_value) + initial_path = tmp_path / "adapter-initial" + checkpoint_path = tmp_path / "checkpoints" + final_path = checkpoint_path / "step_000002" + reexported_path = tmp_path / "adapter-reexported" + lora = LoraConfig(rank=8, alpha=16.0) + backend_config = ArenoConfig( + tp_size=2, + dp_size=2, + devices=[0, 1, 2, 3], + lora=lora, + max_running_prompts=4, + optimizer={ + "lr": 1.0e-4, + "min_lr": 1.0e-4, + "lr_decay_style": "constant", + "weight_decay": 0.0, + "grad_clip_norm": 1.0, + }, + runtime={ + "compile_model": False, + "activation_checkpointing": False, + "keep_rollout_state": False, + }, + ) + inner = Trainer(4, os.fspath(model_path), custom_config=backend_config) + observed = _ObservedTrainer(inner) + config = PolicyTrainerConfig( + algo="grpo", + ckpt=os.fspath(model_path), + dataset_path="e2e://in-memory", + save_path=os.fspath(checkpoint_path), + save_interval=2, + epochs=1, + max_steps=2, + world_size=4, + tp_size=2, + train_devices=[0, 1, 2, 3], + batch_size=1, + mini_bs=4, + n_samples=4, + max_running_prompts=4, + max_prompt_tokens=64, + max_new_tokens=16, + optimizer_lr=1.0e-4, + optimizer_min_lr=1.0e-4, + lr_decay_style="constant", + weight_decay=0.0, + activation_checkpointing=False, + keep_rollout_state=False, + metrics_log_dir=None, + lora=lora, + ) + dataset = [ + {"prompt": "Write one uncommon English noun. Output only the noun."}, + {"prompt": "Invent one short fictional name. Output only the name."}, + ] + + def reward_fn(record) -> float: + return float(record.metadata["sample_index"]) + + policy = PolicyOnlyTrainer( + config, + instance=observed, + dataset=dataset, + reward_fn=reward_fn, + loss_fn=get_algorithm("grpo").make_loss_fn(config), + ) + + observed.init() + try: + parity_tokens = observed.get_tokenizer().encode("A short adapter parity check.", add_special_tokens=True) + observed.export_adapter(os.fspath(initial_path)) + initial_native_logprobs = observed.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + policy._fit_initialized() + trained_logprobs = observed.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + finally: + observed.close() + + assert observed.rollout_versions == [0, 1] + assert observed.train_versions == [1, 2] + assert (final_path / "adapter_config.json").is_file() + assert (final_path / "adapter_model.safetensors").is_file() + initial = load_file(initial_path / "adapter_model.safetensors") + final = load_file(final_path / "adapter_model.safetensors") + changed = {name for name in initial if not torch.equal(initial[name], final[name])} + assert any(".self_attn." in name for name in changed) + if model_kind == "moe": + assert any(".experts." in name for name in changed) + else: + assert any(".mlp." in name for name in changed) + + if model_kind == "dense": + initial_peft_logprobs = _peft_logprobs(model_path, initial_path, parity_tokens) + peft_logprobs = _peft_logprobs(model_path, final_path, parity_tokens) + else: + expert_key = next(name for name in changed if ".experts." in name) + peft_logprobs = _peft_logprobs( + model_path, + final_path, + parity_tokens, + expected_state=final, + representative_key=expert_key, + ) + imported = Trainer( + 4, + os.fspath(model_path), + custom_config=ArenoConfig( + tp_size=2, + dp_size=2, + devices=[0, 1, 2, 3], + lora=LoraConfig(adapter_path=os.fspath(final_path)), + runtime={"compile_model": False, "activation_checkpointing": False}, + ), + ) + imported.init() + try: + imported.export_adapter(os.fspath(reexported_path)) + areno_logprobs = imported.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + repeated_logprobs = imported.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + finally: + imported.close() + reexported = load_file(reexported_path / "adapter_model.safetensors") + assert reexported.keys() == final.keys() + assert all(torch.equal(reexported[name], final[name]) for name in final) + torch.testing.assert_close(torch.tensor(repeated_logprobs), torch.tensor(areno_logprobs), rtol=0.0, atol=1.0e-5) + torch.testing.assert_close( + torch.tensor(areno_logprobs), + torch.tensor(trained_logprobs), + rtol=0.0, + atol=1.0e-5, + ) + if model_kind == "dense": + native_delta = torch.tensor(areno_logprobs[1:]) - torch.tensor(initial_native_logprobs[1:]) + peft_delta = torch.tensor(peft_logprobs) - torch.tensor(initial_peft_logprobs) + torch.testing.assert_close(native_delta, peft_delta, rtol=0.0, atol=1.5e-1) + else: + assert torch.isfinite(torch.tensor(peft_logprobs)).all() + + +def _peft_logprobs( + model_path: Path, + adapter_path: Path, + token_ids: list[int], + *, + expected_state: dict[str, torch.Tensor] | None = None, + representative_key: str | None = None, +) -> list[float]: + peft_source = os.getenv("ARENO_E2E_PEFT_SOURCE") + if peft_source: + sys.path.insert(0, peft_source) + original_find_spec = importlib_util.find_spec + + def find_spec_without_torchao(name, *args, **kwargs): + if name == "torchao": + return None + return original_find_spec(name, *args, **kwargs) + + importlib_util.find_spec = find_spec_without_torchao + try: + from peft import PeftModel, get_peft_model_state_dict + from transformers import AutoModelForCausalLM + + base = AutoModelForCausalLM.from_pretrained(model_path, dtype=torch.bfloat16).to("cuda:0") + model = PeftModel.from_pretrained(base, os.fspath(adapter_path), autocast_adapter_dtype=False).eval() + if expected_state is not None: + loaded = get_peft_model_state_dict(model, save_embedding_layers=False) + assert loaded.keys() == expected_state.keys() + assert representative_key is not None + torch.testing.assert_close( + loaded[representative_key].detach().cpu().float(), + expected_state[representative_key].float(), + rtol=0.0, + atol=0.0, + ) + tokens = torch.tensor([token_ids], device="cuda:0", dtype=torch.long) + with torch.inference_mode(): + logits = model(input_ids=tokens).logits[0, :-1].float() + selected = logits.log_softmax(dim=-1).gather(-1, tokens[0, 1:].unsqueeze(-1)).squeeze(-1) + result = selected.cpu().tolist() + del model, base, tokens, logits, selected + torch.cuda.empty_cache() + return result + finally: + importlib_util.find_spec = original_find_spec + if peft_source: + sys.path.remove(peft_source) diff --git a/tests/test_train_cli_config_cpu.py b/tests/test_train_cli_config_cpu.py index c0d6eaa2..02b768cb 100644 --- a/tests/test_train_cli_config_cpu.py +++ b/tests/test_train_cli_config_cpu.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import re from dataclasses import replace from types import SimpleNamespace @@ -385,6 +386,27 @@ def test_training_config_summary_can_colorize_output(): assert "AReno training config" in summary +def test_dashboard_run_config_serializes_lora(tmp_path): + cfg = _trainer_config_from_options(**_options(lora_rank=8, lora_alpha=16.0, metrics_log_dir=str(tmp_path))) + + train_cli._write_dashboard_run_config(cfg) + + payload = json.loads(next(tmp_path.glob("areno_run_config.*.json")).read_text()) + settings = payload["settings"]["sections"] + other = next(section for section in settings if section["title"] == "Other") + lora = next(item["value"] for item in other["items"] if item["key"] == "lora") + assert lora["rank"] == 8 + assert lora["target_modules"] == [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + + def test_training_config_summary_wraps_for_narrow_terminals(monkeypatch): monkeypatch.setattr( train_cli.shutil, "get_terminal_size", lambda fallback: train_cli.shutil.os.terminal_size((48, 24)) @@ -916,6 +938,11 @@ def _options(**overrides): top_k=-1, top_p=1.0, max_running_prompts=None, + lora_rank=None, + lora_alpha=16.0, + lora_dropout=0.0, + lora_target_modules="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", + lora_adapter_path=None, lr=1e-6, min_lr=1e-7, lr_decay_steps=100, From 3c18f9319a94cdb00350503f5a0e91bcd2d4df44 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Tue, 18 Aug 2026 21:14:37 +0800 Subject: [PATCH 03/29] feat(lora): reuse actor base for reference scoring --- areno/adapters/lora.py | 65 +++++++++++++++++++++++++++--- areno/api/backend/cuda/backend.py | 1 + areno/api/backend/cuda/roles.py | 58 ++++++++++++++------------ areno/api/config.py | 3 +- areno/api/trainer_config.py | 6 ++- areno/cli/train.py | 13 ++++++ areno/engine/api.py | 10 +++++ areno/engine/config.py | 7 +++- areno/engine/layers/linear.py | 4 +- areno/engine/protocol.py | 1 + areno/models/qwen3/model.py | 20 +++++---- tests/test_config_data_cpu.py | 31 +++++++++++--- tests/test_qwen3_lora_e2e.py | 11 +++++ tests/test_train_cli_config_cpu.py | 10 +++++ 14 files changed, 190 insertions(+), 50 deletions(-) diff --git a/areno/adapters/lora.py b/areno/adapters/lora.py index 2a6019e1..a4d3d85f 100644 --- a/areno/adapters/lora.py +++ b/areno/adapters/lora.py @@ -4,6 +4,8 @@ import hashlib import math +from collections.abc import Iterator +from contextlib import contextmanager import torch import torch.nn.functional as F @@ -15,6 +17,17 @@ from areno.engine.parallel.context import get_tp_context +class _AdapterRuntimeState: + """Shared control state for one model's adapter view.""" + + def __init__(self) -> None: + self.base_only_depth = 0 + + @property + def enabled(self) -> bool: + return self.base_only_depth == 0 + + class LoraSlot(nn.Module): """One canonical LoRA A/B pair owned by its native projection module.""" @@ -30,6 +43,7 @@ def __init__( row_parallel: bool, config: LoraConfig, seed: int, + runtime_state: _AdapterRuntimeState, ) -> None: super().__init__() ctx = get_tp_context() @@ -40,6 +54,7 @@ def __init__( self.local_in_features = int(local_in_features) self.local_out_features = int(local_out_features) self.row_parallel = bool(row_parallel) + self._runtime_state = runtime_state self.lora_A = nn.Parameter( torch.empty(self.rank, self.local_in_features, device=base_weight.device, dtype=base_weight.dtype) ) @@ -72,6 +87,10 @@ def _reset_parameters(self, seed: int, tp_rank: int, tp_size: int) -> None: def forward(self, x: torch.Tensor) -> torch.Tensor: return F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scale + @property + def enabled(self) -> bool: + return self._runtime_state.enabled + class RoutedExpertLoraSlot(nn.Module): """One expert-sharded canonical LoRA A/B pair for grouped Qwen3-MoE GEMMs.""" @@ -87,6 +106,7 @@ def __init__( out_features: int, config: LoraConfig, seed: int, + runtime_state: _AdapterRuntimeState, ) -> None: super().__init__() self.logical_name = logical_name @@ -95,6 +115,7 @@ def __init__( self.local_expert_start = int(local_expert_start) self.in_features = int(in_features) self.out_features = int(out_features) + self._runtime_state = runtime_state self.lora_A = nn.Parameter( torch.empty( self.local_num_experts, @@ -135,13 +156,23 @@ def forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Ten hidden = areno_grouped_linear(x.contiguous(), self.lora_A, tokens_per_expert) return areno_grouped_linear(hidden, self.lora_B, tokens_per_expert) * self.scale + @property + def enabled(self) -> bool: + return self._runtime_state.enabled + class AdapterRegistry: """Non-owning index over LoRA slots; projection modules remain sole owners.""" - def __init__(self, slots: dict[str, LoraSlot | RoutedExpertLoraSlot], config: LoraConfig) -> None: + def __init__( + self, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], + config: LoraConfig, + runtime_state: _AdapterRuntimeState, + ) -> None: self.slots = slots self.config = config + self._runtime_state = runtime_state self.version = 0 def named_parameters(self): @@ -156,6 +187,16 @@ def increment_version(self) -> int: self.version += 1 return self.version + @contextmanager + def base_only(self) -> Iterator[None]: + """Temporarily expose the frozen base policy without evaluating A/B.""" + + self._runtime_state.base_only_depth += 1 + try: + yield + finally: + self._runtime_state.base_only_depth -= 1 + def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> AdapterRegistry: """Freeze a Qwen3 dense or MoE base and attach canonical targets.""" @@ -167,6 +208,7 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt parameter.requires_grad_(False) requested = set(config.target_modules) + runtime_state = _AdapterRuntimeState() slots: dict[str, LoraSlot | RoutedExpertLoraSlot] = {} for layer_index, layer in enumerate(model.layers): prefix = f"layers.{layer_index}" @@ -185,6 +227,7 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt row_parallel=False, config=config, seed=seed, + runtime_state=runtime_state, ) qkv.install_lora_component(component, component_index, slot) slots[logical_name] = slot @@ -192,12 +235,12 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt if "o_proj" in requested: owner = layer.self_attn.o_proj logical_name = f"{prefix}.self_attn.o_proj" - slot = _row_slot(logical_name, owner, config, seed) + slot = _row_slot(logical_name, owner, config, seed, runtime_state) owner.install_lora(slot) slots[logical_name] = slot if getattr(model_config, "enable_moe_block", False): - _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, slots) + _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots) else: gate_up = layer.mlp.gate_up_proj for component_index, component in enumerate(("gate_proj", "up_proj")): @@ -214,6 +257,7 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt row_parallel=False, config=config, seed=seed, + runtime_state=runtime_state, ) gate_up.install_lora_component(component, component_index, slot) slots[logical_name] = slot @@ -221,14 +265,20 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt if "down_proj" in requested: owner = layer.mlp.down_proj logical_name = f"{prefix}.mlp.down_proj" - slot = _row_slot(logical_name, owner, config, seed) + slot = _row_slot(logical_name, owner, config, seed, runtime_state) owner.install_lora(slot) slots[logical_name] = slot - return AdapterRegistry(slots, config) + return AdapterRegistry(slots, config, runtime_state) -def _row_slot(logical_name: str, owner: RowParallelLinear, config: LoraConfig, seed: int) -> LoraSlot: +def _row_slot( + logical_name: str, + owner: RowParallelLinear, + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, +) -> LoraSlot: return LoraSlot( logical_name=logical_name, base_weight=owner.weight, @@ -239,6 +289,7 @@ def _row_slot(logical_name: str, owner: RowParallelLinear, config: LoraConfig, s row_parallel=True, config=config, seed=seed, + runtime_state=runtime_state, ) @@ -248,6 +299,7 @@ def _install_moe_slots( requested: set[str], config: LoraConfig, seed: int, + runtime_state: _AdapterRuntimeState, slots: dict[str, LoraSlot | RoutedExpertLoraSlot], ) -> None: components = ( @@ -268,6 +320,7 @@ def _install_moe_slots( out_features=out_features, config=config, seed=seed, + runtime_state=runtime_state, ) experts.install_lora_component(component, slot) slots[logical_name] = slot diff --git a/areno/api/backend/cuda/backend.py b/areno/api/backend/cuda/backend.py index 9252f028..fb6e74dc 100644 --- a/areno/api/backend/cuda/backend.py +++ b/areno/api/backend/cuda/backend.py @@ -172,6 +172,7 @@ def initialize(self, ctx: Context): loss_fn=dispatch_loss, policy_sync_bucket_mb=cfg.policy_sync_bucket_mb, lora_config=cfg.lora, + reference_mode=cfg.reference_mode, ) return self._policy_sync_bucket_bytes = cfg.policy_sync_bucket_mb * 1024 * 1024 diff --git a/areno/api/backend/cuda/roles.py b/areno/api/backend/cuda/roles.py index 8a901aa3..ec0f7b61 100644 --- a/areno/api/backend/cuda/roles.py +++ b/areno/api/backend/cuda/roles.py @@ -2,6 +2,8 @@ from __future__ import annotations +from contextlib import nullcontext + import torch from areno.engine.checkpoints.io import SafetensorsIndex @@ -322,6 +324,7 @@ class RoleManager: def __init__(self, worker): self.worker = worker self.roles: dict[str, WorkerRole] = {} + self.actor_base_roles: set[str] = set() def ensure_roles(self, payload: EnsureRolesPayload) -> None: """Lazily instantiate non-actor roles.""" @@ -330,14 +333,17 @@ def ensure_roles(self, payload: EnsureRolesPayload) -> None: worker._prepare_actor_offloaded() model_sources: dict[str, torch.nn.Module] = {} actor_path = canonical_model_path(worker.config.model_path) - if actor_path is not None: + if actor_path is not None and worker.adapter_registry is None: model_sources[actor_path] = unwrap_model(worker.model) for role in self.roles.values(): role_path = canonical_model_path(role.path) if role_path is not None: model_sources.setdefault(role_path, role.model) for name, spec in payload.roles.items(): - if name == "actor" or name in self.roles: + if name == "actor" or name in self.roles or name in self.actor_base_roles: + continue + if spec.reference_mode == "reuse_actor_base": + self.actor_base_roles.add(name) continue path = spec.path cache_key = canonical_model_path(path) @@ -369,36 +375,36 @@ def score_logprobs(self, payload: ScorePayload) -> list[list[float]] | None: worker = self.worker ctx = get_tp_context() role_name = payload.role - if role_name == "actor": - worker._prepare_actor_for_inference() - model = worker.model - offload_role = None - sequence_parallel = worker.config.effective_sequence_parallel - else: - offload_role = self.roles[role_name] - worker._prepare_actor_offloaded() - offload_role.onload_for_inference(worker.device) - model = offload_role.model - sequence_parallel = offload_role.sequence_parallel - model.eval() - try: - token_rows = payload.token_rows_by_dp[ctx.dp_rank] - features = payload.features_by_dp[ctx.dp_rank] if payload.features_by_dp is not None else None - local = ( - [] - if not token_rows - else self._score_logprob_rows( + actor_view = role_name == "actor" or role_name in self.actor_base_roles + base_only = role_name in self.actor_base_roles + view = worker.adapter_registry.base_only() if base_only else nullcontext() + with view: + if actor_view: + worker._prepare_actor_for_inference() + model = worker.model + offload_role = None + sequence_parallel = worker.config.effective_sequence_parallel + else: + offload_role = self.roles[role_name] + worker._prepare_actor_offloaded() + offload_role.onload_for_inference(worker.device) + model = offload_role.model + sequence_parallel = offload_role.sequence_parallel + model.eval() + try: + token_rows = payload.token_rows_by_dp[ctx.dp_rank] + features = payload.features_by_dp[ctx.dp_rank] if payload.features_by_dp is not None else None + local = [] if not token_rows else self._score_logprob_rows( model, token_rows, payload, features=features, sequence_parallel=sequence_parallel, ) - ) - return local if ctx.rank == 0 else None - finally: - if offload_role is not None: - offload_role.offload() + return local if ctx.rank == 0 else None + finally: + if offload_role is not None: + offload_role.offload() def _score_logprob_rows( self, diff --git a/areno/api/config.py b/areno/api/config.py index 169bb773..3f37485b 100644 --- a/areno/api/config.py +++ b/areno/api/config.py @@ -4,7 +4,7 @@ import platform from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal from areno.adapters.config import LoraConfig from areno.api.models import BackendType @@ -37,6 +37,7 @@ class CudaConfig: max_running_prompts: int = 64 decode_progress_interval_s: float = 10.0 lora: LoraConfig | None = None + reference_mode: Literal["independent", "reuse_actor_base"] = "independent" def uses_separate_rollout_engine(self) -> bool: """Return whether rollout runs on its own CUDA device partition.""" diff --git a/areno/api/trainer_config.py b/areno/api/trainer_config.py index f1fedcb3..940b227d 100644 --- a/areno/api/trainer_config.py +++ b/areno/api/trainer_config.py @@ -12,6 +12,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Literal from areno.adapters.config import LoraConfig from areno.api.defaults import DEFAULT_METRICS_LOG_DIR @@ -78,6 +79,7 @@ class TrainerConfig: train_tool_results: bool = False chat_template_enable_thinking: bool | None = None lora: LoraConfig | None = None + reference_mode: Literal["independent", "reuse_actor_base"] = "independent" def __post_init__(self) -> None: if self.backend is None: @@ -120,8 +122,6 @@ def __post_init__(self) -> None: ) if self.lora is not None and self.backend != "cuda": raise ValueError("native LoRA is only supported by the CUDA backend") - if self.lora is not None and self.algo.lower() in {"ppo", "dpo"}: - raise ValueError("native LoRA does not yet support PPO/DPO reference and critic roles") @staticmethod def _validate_multimodal_optimizer_group( @@ -217,6 +217,7 @@ def cuda_config(self): "attn_backend": self.attn_backend, }, lora=self.lora, + reference_mode=self.reference_mode, ) @@ -265,6 +266,7 @@ def cuda_config(self): "attn_backend": self.attn_backend, }, lora=self.lora, + reference_mode=self.reference_mode, ) def mlx_config(self): diff --git a/areno/cli/train.py b/areno/cli/train.py index a3ecf74f..84a478ca 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -137,6 +137,7 @@ def flash_attention_unsupported_model_reason(model_config): "lora_dropout", "lora_target_modules", "lora_adapter_path", + "reference_mode", "lr", "min_lr", "lr_decay_steps", @@ -249,6 +250,7 @@ def _trainer_config_from_options(**options) -> TrainerConfig: args.multimodal_projector_min_lr = getattr(args, "multimodal_projector_min_lr", None) args.multimodal_projector_lr_decay_steps = getattr(args, "multimodal_projector_lr_decay_steps", None) args.multimodal_projector_lr_decay_style = getattr(args, "multimodal_projector_lr_decay_style", None) + args.reference_mode = getattr(args, "reference_mode", "independent") args.lora = _lora_config_from_options(args) if args.backend == "mlx": if args.train_devices is not None or args.rollout_devices is not None or args.rollout_tp_size is not None: @@ -901,6 +903,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: ref_ckpt=args.ref_ckpt, dpo_beta=args.dpo_beta, lora=lora, + reference_mode=args.reference_mode, ) if algorithm.name == "sft": return TrainerConfig( @@ -957,6 +960,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, lora=lora, + reference_mode=args.reference_mode, ) if algorithm.name != "ppo": return PolicyTrainerConfig( @@ -1025,6 +1029,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, lora=lora, + reference_mode=args.reference_mode, ) return PPOTrainerConfig( algo=algorithm.name, @@ -1106,6 +1111,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, lora=lora, + reference_mode=args.reference_mode, ) @@ -1638,6 +1644,13 @@ def _dataset_builder_for_suffix(suffix: str) -> str: help="Comma-separated Qwen3 projection targets (MoE MLP targets apply to each routed expert).", ) @click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter used to initialize native LoRA.") +@click.option( + "--reference-mode", + type=click.Choice(["independent", "reuse_actor_base"]), + default="independent", + show_default=True, + help="Use the frozen actor base as the PPO/DPO reference instead of loading a second reference model.", +) @click.option("--min-lr", type=float, default=1.0e-7, show_default=True, help="Policy optimizer minimum learning rate.") @click.option("--lr-decay-steps", type=int, default=1000, show_default=True, help="Policy LR decay steps.") @click.option("--lr-decay-style", default="cosine", show_default=True, help="Policy LR decay style.") diff --git a/areno/engine/api.py b/areno/engine/api.py index 4faf9e27..01a18619 100644 --- a/areno/engine/api.py +++ b/areno/engine/api.py @@ -30,6 +30,7 @@ from areno.engine.checkpoints.io import resolve_model_path from areno.engine.config import EngineConfig, OptimizerConfig, RuntimeConfig from areno.engine.data import RolloutOutput, SamplingParams, TrainStats, to_cpu +from areno.engine.modeling import canonical_model_path from areno.engine.protocol import ( EnsureRolesPayload, ExportAdapterPayload, @@ -211,6 +212,7 @@ def from_pretrained( cluster_kwargs: dict[str, Any] | None = None, policy_sync_bucket_mb: int = 64, lora_config: LoraConfig | None = None, + reference_mode: str = "independent", ) -> ArenoEngine: """Build an engine by reading model config from a checkpoint path. @@ -243,6 +245,7 @@ def from_pretrained( policy_sync_bucket_mb=policy_sync_bucket_mb, lora=lora_config, lora_seed=torch.initial_seed(), + reference_mode=reference_mode, ) return cls(cfg, start=start, cluster_kwargs=cluster_kwargs) @@ -562,12 +565,19 @@ def ensure_roles(self, roles: dict[str, Any]) -> None: path=str(spec.path), trainable=bool(spec.trainable), optimizer_lr=getattr(spec, "optimizer_lr", None), + reference_mode=self._role_reference_mode(name, str(spec.path)), ) for name, spec in roles.items() } ) self.cluster.call(Op.ENSURE_ROLES, payload) + def _role_reference_mode(self, name: str, path: str) -> str: + mode = self.config.reference_mode if name == "ref" else "independent" + if mode == "reuse_actor_base" and canonical_model_path(path) != canonical_model_path(self.config.model_path): + raise ValueError("reuse_actor_base requires the reference checkpoint to match the actor base checkpoint") + return mode + def score_logprobs( self, role: str, diff --git a/areno/engine/config.py b/areno/engine/config.py index 55c6340e..1575d590 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -280,6 +280,7 @@ class EngineConfig: policy_sync_bucket_mb: int = 64 lora: LoraConfig | None = None lora_seed: int = 0 + reference_mode: Literal["independent", "reuse_actor_base"] = "independent" def __post_init__(self) -> None: """Infer DP/devices and validate the distributed layout.""" @@ -287,6 +288,10 @@ def __post_init__(self) -> None: if self.sequence_parallel is not None: self.model.sequence_parallel = bool(self.sequence_parallel) self.model.validate_tp(self.tp_size) + if self.reference_mode not in {"independent", "reuse_actor_base"}: + raise ValueError("reference_mode must be one of: independent, reuse_actor_base") + if self.reference_mode == "reuse_actor_base" and self.lora is None: + raise ValueError("reference_mode='reuse_actor_base' requires native LoRA") if self.lora is not None: replicated_kv_targets = {"k_proj", "v_proj"} & set(self.lora.target_modules) if ( @@ -296,7 +301,7 @@ def __post_init__(self) -> None: ): targets = ", ".join(sorted(replicated_kv_targets)) raise ValueError( - f"Qwen3-MoE replicated-KV topology does not support LoRA targets {targets}; " + f"Qwen3 replicated-KV topology does not support LoRA targets {targets}; " "omit k_proj/v_proj or use tp_size <= num_key_value_heads" ) if self.devices is None: diff --git a/areno/engine/layers/linear.py b/areno/engine/layers/linear.py index ae9033ee..e6bd62d6 100644 --- a/areno/engine/layers/linear.py +++ b/areno/engine/layers/linear.py @@ -178,6 +178,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return out parts = list(out.split(self.local_out_features, dim=-1)) for component, slot in self.lora_slots.items(): + if not slot.enabled: + continue index = self._lora_component_indices[component] parts[index] = parts[index] + slot(x) return torch.cat(parts, dim=-1) @@ -269,7 +271,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: start, end = _shard_range(self.in_features, ctx.rank, ctx.world_size) x = x[..., start:end] out = _areno_linear_forward(x, self.weight, None) - if self.lora_slot is not None: + if self.lora_slot is not None and self.lora_slot.enabled: out = out + self.lora_slot(x) # Partial sum -> cross-rank reduction. SP mode also re-shards along # the sequence dim via reduce-scatter, saving activation memory. diff --git a/areno/engine/protocol.py b/areno/engine/protocol.py index cf96bc6f..cb96806e 100644 --- a/areno/engine/protocol.py +++ b/areno/engine/protocol.py @@ -119,6 +119,7 @@ class RoleSpecPayload: path: str trainable: bool optimizer_lr: float | None = None + reference_mode: str = "independent" @dataclass(slots=True) diff --git a/areno/models/qwen3/model.py b/areno/models/qwen3/model.py index cdd16727..67dbd5b0 100644 --- a/areno/models/qwen3/model.py +++ b/areno/models/qwen3/model.py @@ -112,12 +112,15 @@ def install_lora_component(self, component: str, slot: nn.Module) -> None: self.lora_slots[component] = slot - def has_active_lora(self) -> bool: + def has_lora(self) -> bool: return bool(self.lora_slots) + def has_active_lora(self) -> bool: + return self.has_lora() and next(iter(self.lora_slots.values())).enabled + def _gate_up_forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: base = _areno_grouped_linear_no_compile(x.contiguous(), self.gate_up_weight, tokens_per_expert) - if not self.lora_slots: + if not self.has_active_lora(): return base gate, up = base.chunk(2, dim=-1) if "gate_proj" in self.lora_slots: @@ -128,7 +131,7 @@ def _gate_up_forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> def _down_forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor: out = _areno_grouped_linear_no_compile(x, self.down_weight, tokens_per_expert) - if "down_proj" in self.lora_slots: + if self.has_active_lora() and "down_proj" in self.lora_slots: out = out + self.lora_slots["down_proj"](x, tokens_per_expert) return out @@ -150,8 +153,9 @@ def forward(self, flat: torch.Tensor, topk_idx: torch.Tensor, topk_weight: torch + self.down_weight.reshape(-1)[0] * 0 + topk_weight.sum().to(dtype=self.gate_up_weight.dtype) * 0 ) - for slot in self.lora_slots.values(): - zero = zero + slot.lora_A.reshape(-1)[0] * 0 + slot.lora_B.reshape(-1)[0] * 0 + if self.has_active_lora(): + for slot in self.lora_slots.values(): + zero = zero + slot.lora_A.reshape(-1)[0] * 0 + slot.lora_B.reshape(-1)[0] * 0 return all_reduce(flat.new_zeros(flat.shape) + zero) hidden = self._gate_up_forward(x, tokens_per_expert) log_once("qwen3_moe_silu_and_mul", "using ARENO fused silu_and_mul kernel for Qwen3-MoE experts") @@ -159,7 +163,7 @@ def forward(self, flat: torch.Tensor, topk_idx: torch.Tensor, topk_weight: torch _areno_silu_and_mul_no_compile(hidden) * route_weight.unsqueeze(-1).to(dtype=hidden.dtype) ).contiguous() out = self._down_forward(hidden, tokens_per_expert) - if self.has_active_lora(): + if self.has_lora(): # Stabilize routed-expert LoRA without changing the base/fullweight MoE path. out = _areno_moe_unpermute_no_compile(out.float(), token_idx, flat.shape) out = out.to(dtype=flat.dtype) @@ -245,7 +249,7 @@ def forward_with_routes( batch, seqlen, hidden = hidden_states.shape flat = hidden_states.reshape(-1, hidden) with sequence_parallel_region(False): - if self.training or self.experts.has_active_lora(): + if self.training or self.experts.has_lora(): out = self.experts(flat, topk_idx.to(torch.long), topk_weight) else: out = self._forward_fused_moe(flat, topk_idx, topk_weight) @@ -258,7 +262,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @torch.no_grad() def prepare_infer_weights(self) -> None: - if self.experts.has_active_lora(): + if self.experts.has_lora(): self.clear_infer_weights() return self._infer_w1_weight = self._updated_infer_weight( diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 1589bf3a..b4054f73 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -223,11 +223,32 @@ def test_engine_config_rejects_replicated_kv_lora_targets(self): lora=LoraConfig(target_modules=("q_proj", "o_proj")), ) - def test_trainer_config_rejects_lora_ppo_and_dpo(self): - """Reference and critic roles are outside the initial native-LoRA scope.""" - for algo in ("ppo", "dpo"): - with self.subTest(algo=algo), self.assertRaisesRegex(ValueError, "PPO/DPO"): - TrainerConfig(algo=algo, ckpt="actor", dataset_path="dataset", lora=LoraConfig()) + def test_reference_view_requires_lora_at_engine_boundary(self): + """The actor base can be reused only when the actor owns a native adapter.""" + model = ModelConfig(num_attention_heads=4, num_key_value_heads=4, intermediate_size=16, vocab_size=32) + + with self.assertRaisesRegex(ValueError, "requires native LoRA"): + EngineConfig(model=model, tp_size=1, devices=[0], reference_mode="reuse_actor_base") + + config = EngineConfig( + model=model, + tp_size=1, + devices=[0], + lora=LoraConfig(), + reference_mode="reuse_actor_base", + ) + self.assertEqual(config.reference_mode, "reuse_actor_base") + + def test_trainer_config_propagates_lora_reference_view(self): + config = TrainerConfig( + algo="dpo", + ckpt="actor", + dataset_path="dataset", + lora=LoraConfig(), + reference_mode="reuse_actor_base", + ) + + self.assertEqual(config.cuda_config().reference_mode, "reuse_actor_base") def test_adapter_path_uses_peft_metadata(self): """A PEFT artifact should configure non-default native slots itself.""" diff --git a/tests/test_qwen3_lora_e2e.py b/tests/test_qwen3_lora_e2e.py index dd287a8f..ee673b2e 100644 --- a/tests/test_qwen3_lora_e2e.py +++ b/tests/test_qwen3_lora_e2e.py @@ -14,6 +14,7 @@ from areno.adapters import LoraConfig from areno.api import ArenoConfig, Trainer from areno.api.algorithms import get_algorithm +from areno.api.roles import ModelRole from areno.api.trainer_config import PolicyTrainerConfig from areno.api.trainers.policy_only import PolicyOnlyTrainer @@ -62,6 +63,7 @@ def test_qwen3_lora_tp2_dp2_rollout_train_peft(tmp_path: Path, model_env: str, m dp_size=2, devices=[0, 1, 2, 3], lora=lora, + reference_mode="reuse_actor_base", max_running_prompts=4, optimizer={ "lr": 1.0e-4, @@ -127,11 +129,20 @@ def reward_fn(record) -> float: initial_native_logprobs = observed.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] policy._fit_initialized() trained_logprobs = observed.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + observed.ensure_roles({"ref": ModelRole("ref", os.fspath(model_path), trainable=False)}) + reference_logprobs = observed.score_logprobs("ref", [parity_tokens], microbatch_size=1)[0] + restored_actor_logprobs = observed.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] finally: observed.close() assert observed.rollout_versions == [0, 1] assert observed.train_versions == [1, 2] + torch.testing.assert_close( + torch.tensor(reference_logprobs), torch.tensor(initial_native_logprobs), rtol=0.0, atol=1.0e-5 + ) + torch.testing.assert_close( + torch.tensor(restored_actor_logprobs), torch.tensor(trained_logprobs), rtol=0.0, atol=1.0e-5 + ) assert (final_path / "adapter_config.json").is_file() assert (final_path / "adapter_model.safetensors").is_file() initial = load_file(initial_path / "adapter_model.safetensors") diff --git a/tests/test_train_cli_config_cpu.py b/tests/test_train_cli_config_cpu.py index 02b768cb..4fc9e69c 100644 --- a/tests/test_train_cli_config_cpu.py +++ b/tests/test_train_cli_config_cpu.py @@ -407,6 +407,15 @@ def test_dashboard_run_config_serializes_lora(tmp_path): ] +def test_train_config_propagates_reference_view(): + config = _trainer_config_from_options( + **_options(algo="dpo", lora_rank=8, reference_mode="reuse_actor_base", reward_ckpt=None) + ) + + assert config.reference_mode == "reuse_actor_base" + assert config.cuda_config().reference_mode == "reuse_actor_base" + + def test_training_config_summary_wraps_for_narrow_terminals(monkeypatch): monkeypatch.setattr( train_cli.shutil, "get_terminal_size", lambda fallback: train_cli.shutil.os.terminal_size((48, 24)) @@ -943,6 +952,7 @@ def _options(**overrides): lora_dropout=0.0, lora_target_modules="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", lora_adapter_path=None, + reference_mode="independent", lr=1e-6, min_lr=1e-7, lr_decay_steps=100, From 914edf44a615cedb07f75877e5d9e999d6894124 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Tue, 18 Aug 2026 21:24:21 +0800 Subject: [PATCH 04/29] fix(lora): tighten reference replay boundaries --- areno/cli/train.py | 1 + areno/engine/config.py | 4 ++-- tests/test_config_data_cpu.py | 32 +++++++++++++++----------------- 3 files changed, 18 insertions(+), 19 deletions(-) diff --git a/areno/cli/train.py b/areno/cli/train.py index 84a478ca..375a47ae 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -843,6 +843,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: args.multimodal_projector_min_lr = getattr(args, "multimodal_projector_min_lr", None) args.multimodal_projector_lr_decay_steps = getattr(args, "multimodal_projector_lr_decay_steps", None) args.multimodal_projector_lr_decay_style = getattr(args, "multimodal_projector_lr_decay_style", None) + args.reference_mode = getattr(args, "reference_mode", "independent") lora = getattr(args, "lora", None) algorithm = get_algorithm(args.algo) chat_template_enable_thinking = False if args.disable_thinking else None diff --git a/areno/engine/config.py b/areno/engine/config.py index 1575d590..c13a387b 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -295,13 +295,13 @@ def __post_init__(self) -> None: if self.lora is not None: replicated_kv_targets = {"k_proj", "v_proj"} & set(self.lora.target_modules) if ( - self.model.model_type in {"qwen3", "qwen3_moe"} + self.model.model_type == "qwen3_moe" and self.tp_size > self.model.num_key_value_heads and replicated_kv_targets ): targets = ", ".join(sorted(replicated_kv_targets)) raise ValueError( - f"Qwen3 replicated-KV topology does not support LoRA targets {targets}; " + f"Qwen3-MoE replicated-KV topology does not support LoRA targets {targets}; " "omit k_proj/v_proj or use tp_size <= num_key_value_heads" ) if self.devices is None: diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index b4054f73..aa10d423 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -203,25 +203,23 @@ def test_engine_config_resolves_sequence_parallel_override_before_model_config(s def test_engine_config_rejects_replicated_kv_lora_targets(self): """Replicated Qwen3 KV requires range-aware LoRA support.""" - for model_type in ("qwen3", "qwen3_moe"): - with self.subTest(model_type=model_type): - model = ModelConfig( - model_type=model_type, - num_attention_heads=8, - num_key_value_heads=2, - intermediate_size=16, - vocab_size=32, - ) + model = ModelConfig( + model_type="qwen3_moe", + num_attention_heads=8, + num_key_value_heads=2, + intermediate_size=16, + vocab_size=32, + ) - with self.assertRaisesRegex(ValueError, "replicated-KV.*k_proj"): - EngineConfig(model=model, tp_size=4, devices=[0, 1, 2, 3], lora=LoraConfig()) + with self.assertRaisesRegex(ValueError, "Qwen3-MoE replicated-KV.*k_proj"): + EngineConfig(model=model, tp_size=4, devices=[0, 1, 2, 3], lora=LoraConfig()) - EngineConfig( - model=model, - tp_size=4, - devices=[0, 1, 2, 3], - lora=LoraConfig(target_modules=("q_proj", "o_proj")), - ) + EngineConfig( + model=model, + tp_size=4, + devices=[0, 1, 2, 3], + lora=LoraConfig(target_modules=("q_proj", "o_proj")), + ) def test_reference_view_requires_lora_at_engine_boundary(self): """The actor base can be reused only when the actor owns a native adapter.""" From 53df820eaef5dc809b31c876bf2e386a85405fb4 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 10:30:15 +0800 Subject: [PATCH 05/29] fix(runtime): preserve autograd across role onload --- areno/api/backend/cuda/roles.py | 38 +++++++++++++++++---------------- areno/engine/worker.py | 29 +++++++++++++------------ 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/areno/api/backend/cuda/roles.py b/areno/api/backend/cuda/roles.py index ec0f7b61..8768eb8b 100644 --- a/areno/api/backend/cuda/roles.py +++ b/areno/api/backend/cuda/roles.py @@ -274,30 +274,32 @@ def parameters(self): def onload(self, device: torch.device) -> None: """Move this role's model, value head, and optimizer state to `device`.""" - self.model.to(device) - self.model.onload_train_weights(device) - if self.value_head is not None: - self.value_head.to(device) - if self.optimizer is not None: - if self.optimizer_offload_mode == "disk": - self.optimizer.configure_state_offload( - mode="disk", - directory=self.optimizer_offload_dir, - batch_size=self.optimizer_offload_batch_size, - ) - self.optimizer.prefetch_state() - else: - self.optimizer.onload_state(device) + with torch.inference_mode(False), torch.no_grad(): + self.model.to(device) + self.model.onload_train_weights(device) + if self.value_head is not None: + self.value_head.to(device) + if self.optimizer is not None: + if self.optimizer_offload_mode == "disk": + self.optimizer.configure_state_offload( + mode="disk", + directory=self.optimizer_offload_dir, + batch_size=self.optimizer_offload_batch_size, + ) + self.optimizer.prefetch_state() + else: + self.optimizer.onload_state(device) def onload_for_inference(self, device: torch.device) -> None: """Move this role to `device` and materialize derived inference weights.""" - self.model.to(device) - self.model.onload_train_weights(device) + with torch.inference_mode(False), torch.no_grad(): + self.model.to(device) + self.model.onload_train_weights(device) + if self.value_head is not None: + self.value_head.to(device) self.model.prepare_infer_weights() self.model.offload_train_weights() - if self.value_head is not None: - self.value_head.to(device) def offload(self) -> None: """Free all HBM held by this role.""" diff --git a/areno/engine/worker.py b/areno/engine/worker.py index 649d7d2a..eccd73e6 100644 --- a/areno/engine/worker.py +++ b/areno/engine/worker.py @@ -510,20 +510,21 @@ def _prepare_actor_onloaded(self) -> None: """Move the actor model + optimizer state back to `device` if offloaded.""" if self._actor_on_device: return - self.model.to(self.device) - self.model.onload_train_weights(self.device) - if self.optimizer is not None: - mode, directory, batch_size = self._optimizer_offload_options() - if mode == "disk": - # Keep mmap-backed state out of HBM. The optimizer step loads - # only its current bucket after TrainingManager starts prefetch. - self.optimizer.configure_state_offload( - mode=mode, - directory=directory, - batch_size=batch_size, - ) - else: - self.optimizer.onload_state(self.device) + with torch.inference_mode(False), torch.no_grad(): + self.model.to(self.device) + self.model.onload_train_weights(self.device) + if self.optimizer is not None: + mode, directory, batch_size = self._optimizer_offload_options() + if mode == "disk": + # Keep mmap-backed state out of HBM. The optimizer step loads + # only its current bucket after TrainingManager starts prefetch. + self.optimizer.configure_state_offload( + mode=mode, + directory=directory, + batch_size=batch_size, + ) + else: + self.optimizer.onload_state(self.device) self._actor_on_device = True def _prepare_actor_for_inference(self) -> None: From 4dc1cab1dbdfa17b2521cde5ff59b3dd4f352aec Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 10:30:28 +0800 Subject: [PATCH 06/29] test(lora): cover two-step PPO DPO reference view --- tests/test_qwen3_lora_e2e.py | 181 ++++++++++++++++++++++++++++++++++- 1 file changed, 180 insertions(+), 1 deletion(-) diff --git a/tests/test_qwen3_lora_e2e.py b/tests/test_qwen3_lora_e2e.py index ee673b2e..191c5999 100644 --- a/tests/test_qwen3_lora_e2e.py +++ b/tests/test_qwen3_lora_e2e.py @@ -15,7 +15,8 @@ from areno.api import ArenoConfig, Trainer from areno.api.algorithms import get_algorithm from areno.api.roles import ModelRole -from areno.api.trainer_config import PolicyTrainerConfig +from areno.api.trainer_config import DPOTrainerConfig, PolicyTrainerConfig, PPOTrainerConfig +from areno.api.trainer_factory import build_trainer from areno.api.trainers.policy_only import PolicyOnlyTrainer @@ -44,6 +45,184 @@ def train(self, batch_data, loss_fn, mini_bs=8, gradient_accumulation_steps=None return result +class _ObservedReferenceTrainer(_ObservedTrainer): + """Observe the public PPO/DPO lifecycle without inspecting adapter slots.""" + + def __init__(self, inner: Trainer) -> None: + super().__init__(inner) + self.reference_versions: list[int] = [] + self.critic_train_count = 0 + self.roles: set[str] = set() + self.parity_scores: dict[str, list[float]] = {} + self.parity_tokens: list[int] | None = None + + def ensure_roles(self, roles: dict[str, ModelRole]) -> None: + self.roles = set(roles) + self.inner.ensure_roles(roles) + self.parity_tokens = self.inner.get_tokenizer().encode( + "A fixed actor-base reference check.", add_special_tokens=True + ) + self.parity_scores["initial_actor"] = self.inner.score_logprobs( + "actor", [self.parity_tokens], microbatch_size=1 + )[0] + self.parity_scores["initial_reference"] = self.inner.score_logprobs( + "ref", [self.parity_tokens], microbatch_size=1 + )[0] + + def score_logprobs(self, role, token_rows, *, features=None, microbatch_size=None): + if role == "ref": + self.reference_versions.append(len(self.train_versions)) + return self.inner.score_logprobs( + role, + token_rows, + features=features, + microbatch_size=microbatch_size, + ) + + def train_values( + self, + role, + batch_data, + mini_bs, + gradient_accumulation_steps=None, + *, + cliprange_value=0.5, + value_loss_coef=0.5, + ): + result = self.inner.train_values( + role, + batch_data, + mini_bs, + gradient_accumulation_steps, + cliprange_value=cliprange_value, + value_loss_coef=value_loss_coef, + ) + self.critic_train_count += 1 + return result + + def close(self) -> None: + if self.parity_tokens is not None and self.train_versions == [1, 2]: + self.parity_scores["final_actor_before_reference"] = self.inner.score_logprobs( + "actor", [self.parity_tokens], microbatch_size=1 + )[0] + self.parity_scores["final_reference"] = self.inner.score_logprobs( + "ref", [self.parity_tokens], microbatch_size=1 + )[0] + self.parity_scores["final_actor_after_reference"] = self.inner.score_logprobs( + "actor", [self.parity_tokens], microbatch_size=1 + )[0] + self.inner.close() + + +@pytest.mark.parametrize("algorithm", ("ppo", "dpo")) +def test_qwen3_lora_tp2_dp2_reference_two_step(algorithm: str) -> None: + model_path_value = os.getenv("ARENO_E2E_QWEN3_MODEL") + if not model_path_value: + pytest.skip("set ARENO_E2E_QWEN3_MODEL to run the 4-GPU Qwen3 LoRA PPO/DPO E2E") + model_path = Path(model_path_value) + common = { + "algo": algorithm, + "ckpt": os.fspath(model_path), + "dataset_path": f"e2e://{algorithm}-in-memory", + "epochs": 1, + "max_steps": 2, + "world_size": 4, + "tp_size": 2, + "train_devices": [0, 1, 2, 3], + "batch_size": 2, + "score_micro_bs": 2, + "gradient_accumulation_steps": 1, + "max_prompt_tokens": 64, + "optimizer_lr": 1.0e-4, + "optimizer_min_lr": 1.0e-4, + "lr_decay_style": "constant", + "weight_decay": 0.0, + "activation_checkpointing": False, + "keep_rollout_state": False, + "eager_decode": True, + "metrics_log_dir": None, + "lora": LoraConfig(rank=8, alpha=16.0), + "reference_mode": "reuse_actor_base", + "ref_ckpt": os.fspath(model_path), + } + if algorithm == "ppo": + config = PPOTrainerConfig( + **common, + mini_bs=2, + n_samples=1, + greedy=True, + max_new_tokens=3, + max_running_prompts=2, + critic_ckpt=os.fspath(model_path), + critic_lr=1.0e-4, + critic_warmup_steps=0, + ) + dataset = [ + {"prompt": "Reply with the English word for the number one."}, + {"prompt": "Reply with the English word for the number two."}, + {"prompt": "Reply with the English word for the number three."}, + {"prompt": "Reply with the English word for the number four."}, + ] + + def reward_fn(record) -> float: + return float(record.metadata["prompt_index"]) + + else: + config = DPOTrainerConfig(**common, mini_bs=4, max_new_tokens=8, dpo_beta=0.1) + dataset = [ + {"prompt": "What is 1 + 1?", "chosen": "2", "rejected": "3"}, + {"prompt": "What is 2 + 2?", "chosen": "4", "rejected": "5"}, + {"prompt": "What is 3 + 3?", "chosen": "6", "rejected": "7"}, + {"prompt": "What is 4 + 4?", "chosen": "8", "rejected": "9"}, + ] + reward_fn = None + + backend_config = config.areno_config() + backend_config.dp_size = 2 + backend_config.runtime["compile_model"] = False + observed = _ObservedReferenceTrainer( + Trainer( + config.world_size, + config.ckpt, + custom_config=backend_config, + metrics_log_dir=None, + score_micro_bs=config.score_micro_bs, + ) + ) + trainer = build_trainer( + config, + instance=observed, + dataset=dataset, + reward_fn=reward_fn, + loss_fn=get_algorithm(algorithm).make_loss_fn(config), + ) + trainer.fit() + + assert observed.reference_versions == [0, 1] + assert observed.train_versions == [1, 2] + assert observed.roles == ({"actor", "ref", "critic"} if algorithm == "ppo" else {"ref"}) + assert observed.rollout_versions == ([0, 0, 1, 1] if algorithm == "ppo" else []) + assert observed.critic_train_count == (2 if algorithm == "ppo" else 0) + torch.testing.assert_close( + torch.tensor(observed.parity_scores["initial_reference"]), + torch.tensor(observed.parity_scores["initial_actor"]), + rtol=0.0, + atol=1.0e-5, + ) + torch.testing.assert_close( + torch.tensor(observed.parity_scores["final_reference"]), + torch.tensor(observed.parity_scores["initial_reference"]), + rtol=0.0, + atol=1.0e-5, + ) + torch.testing.assert_close( + torch.tensor(observed.parity_scores["final_actor_after_reference"]), + torch.tensor(observed.parity_scores["final_actor_before_reference"]), + rtol=0.0, + atol=1.0e-5, + ) + + @pytest.mark.parametrize( ("model_env", "model_kind"), (("ARENO_E2E_QWEN3_MODEL", "dense"), ("ARENO_E2E_QWEN3_MOE_MODEL", "moe")), From 7c1048fb51e8ade7bfd0e082a56c7f406fd4cc2f Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 17:21:51 +0800 Subject: [PATCH 07/29] feat(lora): sync adapters to independent rollout engines --- areno/api/backend/cuda/backend.py | 15 ++-- areno/engine/policy_sync.py | 89 +++++++++++++++++++++- areno/engine/worker.py | 2 + tests/test_policy_tensor_sync_cpu.py | 107 ++++++++++++++++++++++++++- 4 files changed, 203 insertions(+), 10 deletions(-) diff --git a/areno/api/backend/cuda/backend.py b/areno/api/backend/cuda/backend.py index fb6e74dc..80a4bfca 100644 --- a/areno/api/backend/cuda/backend.py +++ b/areno/api/backend/cuda/backend.py @@ -1,7 +1,7 @@ """CUDA adapter from the public `Trainer` API onto `ArenoEngine`. -areno runs a co-located train + rollout engine in the same process group. -This file is the thin glue that: +areno can run colocated or independent train and rollout engines. This file is +the thin glue that: - starts the engine with the dataclass-validated `CudaConfig`, - forwards rollout requests through `generate_rollout` while translating @@ -156,9 +156,6 @@ def initialize(self, ctx: Context): raise ValueError(f"training device count must equal world_size={world_size}") if cfg.rollout_tp_size is not None and cfg.rollout_devices is None: raise ValueError("rollout_tp_size requires rollout_devices") - if cfg.lora is not None and cfg.uses_separate_rollout_engine(): - raise ValueError("native LoRA currently supports colocated rollout only") - if not cfg.uses_separate_rollout_engine(): self._train_engine = ArenoEngine.from_pretrained( cfg.model_path or ctx.model_path, @@ -226,6 +223,8 @@ def initialize(self, ctx: Context): optimizer_config=OptimizerConfig(**cfg.optimizer), loss_fn=dispatch_loss, role="train", + lora_config=cfg.lora, + reference_mode=cfg.reference_mode, cluster_kwargs={"world_spec": world_spec, "partition": train_partition}, **common, ) @@ -239,6 +238,7 @@ def initialize(self, ctx: Context): runtime_config=rollout_runtime, loss_fn=None, role="rollout", + lora_config=cfg.lora, policy_sync_bucket_mb=cfg.policy_sync_bucket_mb, start=False, cluster_kwargs={"world_spec": world_spec, "partition": rollout_partition}, @@ -515,7 +515,10 @@ def train( ) stats_list = engine.step(packs, gradient_accumulation_steps=gradient_accumulation_steps) if self._separate_rollout and any(bool(stats.stepped) for stats in stats_list): - self._train_policy_version += 1 + adapter_versions = [stats.adapter_version for stats in stats_list if stats.adapter_version is not None] + self._train_policy_version = ( + int(adapter_versions[-1]) if adapter_versions else self._train_policy_version + 1 + ) train_time_s = time.perf_counter() - train_start metric_rows: list[dict[str, float]] = [] for stats in stats_list: diff --git a/areno/engine/policy_sync.py b/areno/engine/policy_sync.py index 275ec6e8..2ace948a 100644 --- a/areno/engine/policy_sync.py +++ b/areno/engine/policy_sync.py @@ -8,7 +8,8 @@ import torch import torch.distributed as dist -from areno.engine.checkpoints.io import PolicyTensorLayout +from areno.adapters.lora import AdapterRegistry, RoutedExpertLoraSlot +from areno.engine.checkpoints.io import PolicyTensorLayout, PolicyTensorPiece, PolicyTensorStore from areno.engine.parallel.context import get_tp_context from areno.engine.protocol import PolicySyncPayload from areno.models.registry import build_policy_weight_plan @@ -25,9 +26,13 @@ class PolicyTensorMeta: def build_policy_plan(worker) -> tuple[dict[str, object], tuple[PolicyTensorMeta, ...]]: - """Build and cache live adapter tasks plus transport metadata.""" + """Build and cache live policy tasks plus transport metadata.""" - plan = build_policy_weight_plan(worker.model, worker.config.model) + plan = ( + build_adapter_policy_plan(worker.adapter_registry) + if worker.adapter_registry is not None + else build_policy_weight_plan(worker.model, worker.config.model) + ) metadata = [] for key, task in plan.items(): layout = task.policy_layout() @@ -44,6 +49,84 @@ def build_policy_plan(worker) -> tuple[dict[str, object], tuple[PolicyTensorMeta return plan, tuple(metadata) +def build_adapter_policy_plan(registry: AdapterRegistry) -> PolicyTensorStore: + """Describe canonical A/B tensors without materializing frozen base weights.""" + + ctx = get_tp_context() + plan = PolicyTensorStore() + for logical_name, slot in registry.slots.items(): + key_a = f"{logical_name}.lora_A.weight" + key_b = f"{logical_name}.lora_B.weight" + if isinstance(slot, RoutedExpertLoraSlot): + global_experts = slot.local_num_experts * ctx.world_size + plan.add_layout( + key_a, + _sharded_layout( + slot.lora_A, + (global_experts, slot.rank, slot.in_features), + dim=0, + start=slot.local_expert_start, + ), + ) + plan.add_layout( + key_b, + _sharded_layout( + slot.lora_B, + (global_experts, slot.out_features, slot.rank), + dim=0, + start=slot.local_expert_start, + ), + ) + elif slot.row_parallel: + plan.add_layout( + key_a, + _sharded_layout( + slot.lora_A, + (slot.rank, slot.global_in_features), + dim=1, + start=ctx.rank * slot.local_in_features, + ), + ) + plan.add_layout(key_b, _replicated_layout(slot.lora_B)) + else: + plan.add_layout(key_a, _replicated_layout(slot.lora_A)) + plan.add_layout( + key_b, + _sharded_layout( + slot.lora_B, + (slot.global_out_features, slot.rank), + dim=0, + start=ctx.rank * slot.local_out_features, + ), + ) + return plan + + +def _sharded_layout( + tensor: torch.Tensor, + shape: tuple[int, ...], + *, + dim: int, + start: int, +) -> PolicyTensorLayout: + local_size = tensor.shape[dim] + return PolicyTensorLayout( + shape=shape, + dtype=tensor.dtype, + pieces=(PolicyTensorPiece(tensor.detach(), shape, dim, start, start + local_size),), + ) + + +def _replicated_layout(tensor: torch.Tensor) -> PolicyTensorLayout: + shape = tuple(tensor.shape) + return PolicyTensorLayout( + shape=shape, + dtype=tensor.dtype, + pieces=(PolicyTensorPiece(tensor.detach(), shape, 0, 0, shape[0]),), + replicated=True, + ) + + def policy_plan_metadata(worker) -> tuple[PolicyTensorMeta, ...]: """Return canonical metadata without transferring weights.""" diff --git a/areno/engine/worker.py b/areno/engine/worker.py index eccd73e6..ac7ad63b 100644 --- a/areno/engine/worker.py +++ b/areno/engine/worker.py @@ -214,6 +214,8 @@ def receive_policy(self, payload: PolicySyncPayload) -> dict[str, object]: self.model.offload_train_weights() self._train_state_ready = False self._loaded_policy_version = payload.version + if self.adapter_registry is not None: + self.adapter_registry.version = payload.version return result def _prepare_policy_receive(self) -> None: diff --git a/tests/test_policy_tensor_sync_cpu.py b/tests/test_policy_tensor_sync_cpu.py index 23a1739b..2beb4a87 100644 --- a/tests/test_policy_tensor_sync_cpu.py +++ b/tests/test_policy_tensor_sync_cpu.py @@ -3,7 +3,10 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp +from torch import nn +from areno.adapters import LoraConfig +from areno.adapters.lora import AdapterRegistry, LoraSlot, RoutedExpertLoraSlot, _AdapterRuntimeState from areno.engine.checkpoints.io import ( PolicyFlatPiece, PolicyTensorLayout, @@ -13,7 +16,12 @@ _TensorParallelGatherTask, ) from areno.engine.parallel.context import TPContext, destroy_process_group, init_process_group, set_tp_context -from areno.engine.policy_sync import PolicyTensorMeta, assign_policy_owners, transfer_policy_weights +from areno.engine.policy_sync import ( + PolicyTensorMeta, + assign_policy_owners, + build_adapter_policy_plan, + transfer_policy_weights, +) from areno.engine.protocol import PolicySyncPayload, find_free_port @@ -48,6 +56,103 @@ def _write(layouts: list[PolicyTensorLayout], canonical: torch.Tensor, *, chunk_ layout.write_chunk(offset, chunk) +def _adapter_registry(rank: int, world_size: int) -> AdapterRegistry: + _set_rank(rank, world_size) + state = _AdapterRuntimeState() + config = LoraConfig(rank=2, alpha=4.0) + base = nn.Parameter(torch.zeros(1)) + column = LoraSlot( + logical_name="column", + base_weight=base, + global_in_features=4, + global_out_features=6, + local_in_features=4, + local_out_features=6 // world_size, + row_parallel=False, + config=config, + seed=1, + runtime_state=state, + ) + row = LoraSlot( + logical_name="row", + base_weight=base, + global_in_features=6, + global_out_features=4, + local_in_features=6 // world_size, + local_out_features=4, + row_parallel=True, + config=config, + seed=1, + runtime_state=state, + ) + expert = RoutedExpertLoraSlot( + logical_name="experts.{expert}.proj", + base_weight=base, + local_num_experts=4 // world_size, + local_expert_start=rank * (4 // world_size), + in_features=3, + out_features=5, + config=config, + seed=1, + runtime_state=state, + ) + return AdapterRegistry( + {"column": column, "row": row, "experts.{expert}.proj": expert}, + config, + state, + ) + + +def test_adapter_plan_maps_row_column_and_expert_factors_tp2_to_tp1() -> None: + expected = { + "column.lora_A.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4), + "column.lora_B.weight": torch.arange(12, dtype=torch.float32).reshape(6, 2), + "row.lora_A.weight": torch.arange(12, dtype=torch.float32).reshape(2, 6), + "row.lora_B.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2), + "experts.{expert}.proj.lora_A.weight": torch.arange(24, dtype=torch.float32).reshape(4, 2, 3), + "experts.{expert}.proj.lora_B.weight": torch.arange(40, dtype=torch.float32).reshape(4, 5, 2), + } + train_plans = [] + for rank in range(2): + registry = _adapter_registry(rank, 2) + registry.slots["column"].lora_A.data.copy_(expected["column.lora_A.weight"]) + registry.slots["column"].lora_B.data.copy_(expected["column.lora_B.weight"].chunk(2, dim=0)[rank]) + registry.slots["row"].lora_A.data.copy_(expected["row.lora_A.weight"].chunk(2, dim=1)[rank]) + registry.slots["row"].lora_B.data.copy_(expected["row.lora_B.weight"]) + registry.slots["experts.{expert}.proj"].lora_A.data.copy_( + expected["experts.{expert}.proj.lora_A.weight"].chunk(2, dim=0)[rank] + ) + registry.slots["experts.{expert}.proj"].lora_B.data.copy_( + expected["experts.{expert}.proj.lora_B.weight"].chunk(2, dim=0)[rank] + ) + train_plans.append(build_adapter_policy_plan(registry)) + + rollout_registry = _adapter_registry(0, 1) + for parameter in rollout_registry.parameters(): + parameter.data.zero_() + rollout_plan = build_adapter_policy_plan(rollout_registry) + + assert train_plans[0]["column.lora_A.weight"].policy_layout().replicated + assert train_plans[0]["row.lora_B.weight"].policy_layout().replicated + for key, expected_tensor in expected.items(): + canonical = _canonical([plan[key].policy_layout() for plan in train_plans]) + torch.testing.assert_close(canonical, expected_tensor) + _write([rollout_plan[key].policy_layout()], canonical) + + torch.testing.assert_close(rollout_registry.slots["column"].lora_A, expected["column.lora_A.weight"]) + torch.testing.assert_close(rollout_registry.slots["column"].lora_B, expected["column.lora_B.weight"]) + torch.testing.assert_close(rollout_registry.slots["row"].lora_A, expected["row.lora_A.weight"]) + torch.testing.assert_close(rollout_registry.slots["row"].lora_B, expected["row.lora_B.weight"]) + torch.testing.assert_close( + rollout_registry.slots["experts.{expert}.proj"].lora_A, + expected["experts.{expert}.proj.lora_A.weight"], + ) + torch.testing.assert_close( + rollout_registry.slots["experts.{expert}.proj"].lora_B, + expected["experts.{expert}.proj.lora_B.weight"], + ) + + def test_column_parallel_layout_reshards_tp2_to_tp1() -> None: full = torch.arange(24, dtype=torch.float32).reshape(6, 4) train_layouts = [] From f0c3fd5758dc353704a1b940ae43dc1ae3fef912 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 17:27:32 +0800 Subject: [PATCH 08/29] test(lora): model authoritative replicated publisher --- tests/test_policy_tensor_sync_cpu.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_policy_tensor_sync_cpu.py b/tests/test_policy_tensor_sync_cpu.py index 2beb4a87..5ce76868 100644 --- a/tests/test_policy_tensor_sync_cpu.py +++ b/tests/test_policy_tensor_sync_cpu.py @@ -40,9 +40,13 @@ def _canonical(layouts: list[PolicyTensorLayout], *, chunk_size: int = 5) -> tor output = torch.empty(layouts[0].numel, dtype=layouts[0].dtype) for offset in range(0, output.numel(), chunk_size): chunk = torch.zeros(min(chunk_size, output.numel() - offset), dtype=output.dtype) - for layout in layouts: + for rank, layout in enumerate(layouts): contribution = torch.empty_like(chunk) - layout.read_chunk(offset, contribution) + layout.read_chunk( + offset, + contribution, + include_replicated=not layout.replicated or rank == 0, + ) chunk.add_(contribution) output[offset : offset + chunk.numel()].copy_(chunk) return output.reshape(layouts[0].shape) From e408343fb31878a96c9d244024f420074e3f348d Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 17:32:14 +0800 Subject: [PATCH 09/29] test(lora): cover independent rollout adapter sync --- tests/test_qwen3_lora_e2e.py | 128 ++++++++++++++++++++++++++++++++++- 1 file changed, 126 insertions(+), 2 deletions(-) diff --git a/tests/test_qwen3_lora_e2e.py b/tests/test_qwen3_lora_e2e.py index 191c5999..905f3bad 100644 --- a/tests/test_qwen3_lora_e2e.py +++ b/tests/test_qwen3_lora_e2e.py @@ -1,4 +1,4 @@ -"""Qwen3 dense and MoE TP2/DP2 native-LoRA rollout/train/PEFT E2E.""" +"""Qwen3 dense and MoE native-LoRA rollout/train/PEFT E2E.""" from __future__ import annotations @@ -12,7 +12,7 @@ from safetensors.torch import load_file from areno.adapters import LoraConfig -from areno.api import ArenoConfig, Trainer +from areno.api import ArenoConfig, SamplingParams, Trainer from areno.api.algorithms import get_algorithm from areno.api.roles import ModelRole from areno.api.trainer_config import DPOTrainerConfig, PolicyTrainerConfig, PPOTrainerConfig @@ -25,6 +25,7 @@ def __init__(self, inner: Trainer) -> None: self.inner = inner self.rollout_versions: list[int | None] = [] self.train_versions: list[int | None] = [] + self.train_results: list[dict[str, float]] = [] def __getattr__(self, name: str): return getattr(self.inner, name) @@ -42,6 +43,7 @@ async def rollout_token_batch_async(self, prompt_tokens, n_samples, sampling_par def train(self, batch_data, loss_fn, mini_bs=8, gradient_accumulation_steps=None): result = self.inner.train(batch_data, loss_fn, mini_bs, gradient_accumulation_steps) self.train_versions.append(result.get("adapter_version")) + self.train_results.append(result) return result @@ -381,6 +383,128 @@ def reward_fn(record) -> float: assert torch.isfinite(torch.tensor(peft_logprobs)).all() +@pytest.mark.parametrize( + ("model_env", "model_kind", "rollout_tp_size", "rollout_devices"), + ( + ("ARENO_E2E_QWEN3_MODEL", "dense", 1, [2]), + ("ARENO_E2E_QWEN3_MODEL", "dense", 2, [2, 3]), + ("ARENO_E2E_QWEN3_MOE_MODEL", "moe", 2, [2, 3]), + ), +) +def test_qwen3_lora_independent_rollout_two_step( + tmp_path: Path, + model_env: str, + model_kind: str, + rollout_tp_size: int, + rollout_devices: list[int], +) -> None: + model_path_value = os.getenv(model_env) + if not model_path_value: + pytest.skip(f"set {model_env} to run the Qwen3 {model_kind} independent-rollout LoRA E2E") + model_path = Path(model_path_value) + initial_path = tmp_path / "adapter-initial" + trained_path = tmp_path / "adapter-trained" + rollout_path = tmp_path / "adapter-rollout" + lora = LoraConfig(rank=8, alpha=16.0) + config = PolicyTrainerConfig( + algo="grpo", + ckpt=os.fspath(model_path), + dataset_path="e2e://independent-rollout", + epochs=1, + max_steps=2, + world_size=2, + tp_size=2, + train_devices=[0, 1], + rollout_tp_size=rollout_tp_size, + rollout_devices=rollout_devices, + batch_size=1, + mini_bs=4, + n_samples=4, + greedy=True, + max_running_prompts=4, + max_prompt_tokens=64, + max_new_tokens=4, + optimizer_lr=1.0e-4, + optimizer_min_lr=1.0e-4, + lr_decay_style="constant", + weight_decay=0.0, + activation_checkpointing=False, + keep_rollout_state=False, + eager_decode=True, + metrics_log_dir=None, + lora=lora, + ) + backend_config = config.areno_config() + backend_config.dp_size = 1 + backend_config.runtime["compile_model"] = False + inner = Trainer(2, os.fspath(model_path), custom_config=backend_config) + observed = _ObservedTrainer(inner) + dataset = [ + {"prompt": "Write one uncommon English noun. Output only the noun."}, + {"prompt": "Write one uncommon English noun. Output only the noun."}, + ] + + def reward_fn(record) -> float: + return float(record.metadata["sample_index"]) + + policy = PolicyOnlyTrainer( + config, + instance=observed, + dataset=dataset, + reward_fn=reward_fn, + loss_fn=get_algorithm("grpo").make_loss_fn(config), + ) + + observed.init() + try: + observed.export_adapter(os.fspath(initial_path)) + policy._fit_initialized() + observed.export_adapter(os.fspath(trained_path)) + + sampling_params = SamplingParams(greedy=True, max_new_tokens=2, max_prompt_len=64) + prompt_tokens = observed.get_tokenizer().encode( + "A fixed independent-rollout adapter check.", add_special_tokens=True + ) + observed.begin_rollout_session() + try: + final_rollout = observed.rollout_token_batch([prompt_tokens], 1, sampling_params) + finally: + observed.end_rollout_session() + observed.finish_step() + + backend = inner._backend + assert backend is not None + backend._require_rollout_engine().export_adapter(os.fspath(rollout_path)) + finally: + observed.close() + + assert observed.rollout_versions == [0, 1] + assert observed.train_versions == [1, 2] + assert final_rollout[0].adapter_version == 2 + assert observed.train_results[1]["policy_sync_tensors"] > 0 + assert observed.train_results[1]["policy_sync_bytes"] > 0 + + initial = load_file(initial_path / "adapter_model.safetensors") + trained = load_file(trained_path / "adapter_model.safetensors") + rollout = load_file(rollout_path / "adapter_model.safetensors") + changed = {name for name in initial if not torch.equal(initial[name], trained[name])} + assert any(".self_attn." in name for name in changed) + if model_kind == "moe": + assert any(".experts." in name for name in changed) + representative_keys = [ + next(name for name in changed if ".self_attn." in name), + next(name for name in changed if ".experts." in name), + ] + else: + assert any(".mlp." in name for name in changed) + representative_keys = [ + next(name for name in changed if ".self_attn." in name), + next(name for name in changed if ".mlp." in name), + ] + assert rollout.keys() == trained.keys() + assert all(torch.equal(rollout[name], trained[name]) for name in representative_keys) + + def _peft_logprobs( model_path: Path, adapter_path: Path, From 39c97d2f1bc6a435b1e16cad43f590abfafee379 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 18:20:01 +0800 Subject: [PATCH 10/29] feat(lora): support replicated KV adapters --- areno/adapters/lora.py | 18 +++++ areno/adapters/peft.py | 28 ++++++- areno/engine/config.py | 12 --- areno/engine/policy_sync.py | 17 ++++- areno/engine/training.py | 25 +++++- tests/test_config_data_cpu.py | 9 +-- tests/test_policy_tensor_sync_cpu.py | 55 ++++++++++++++ tests/test_qwen3_lora_e2e.py | 110 +++++++++++++++++++++++++++ 8 files changed, 247 insertions(+), 27 deletions(-) diff --git a/areno/adapters/lora.py b/areno/adapters/lora.py index a4d3d85f..bd8f8e84 100644 --- a/areno/adapters/lora.py +++ b/areno/adapters/lora.py @@ -44,6 +44,7 @@ def __init__( config: LoraConfig, seed: int, runtime_state: _AdapterRuntimeState, + output_range: tuple[int, int] | None = None, ) -> None: super().__init__() ctx = get_tp_context() @@ -54,6 +55,16 @@ def __init__( self.local_in_features = int(local_in_features) self.local_out_features = int(local_out_features) self.row_parallel = bool(row_parallel) + if output_range is None: + output_range = ( + (0, self.global_out_features) + if self.row_parallel + else (ctx.rank * self.local_out_features, (ctx.rank + 1) * self.local_out_features) + ) + self.output_start, self.output_end = (int(value) for value in output_range) + self.output_replicated = ( + not self.row_parallel and self.local_out_features * ctx.world_size > self.global_out_features + ) self._runtime_state = runtime_state self.lora_A = nn.Parameter( torch.empty(self.rank, self.local_in_features, device=base_weight.device, dtype=base_weight.dtype) @@ -68,6 +79,12 @@ def __init__( else: mark_tensor_parallel_parameter(self.lora_A, False, sequence_parallel=True, tp_grad_allreduce=True) mark_tensor_parallel_parameter(self.lora_B, True, sequence_parallel=True) + if self.output_replicated: + setattr( + self.lora_B, + "tp_replicated_output_range", + (self.output_start, self.output_end, self.global_out_features), + ) self._reset_parameters(seed, ctx.rank, ctx.world_size) @torch.no_grad() @@ -225,6 +242,7 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt local_in_features=qkv.in_features, local_out_features=qkv.local_out_features[component_index], row_parallel=False, + output_range=qkv.shard_ranges[component_index], config=config, seed=seed, runtime_state=runtime_state, diff --git a/areno/adapters/peft.py b/areno/adapters/peft.py index addd8d72..17f1b85c 100644 --- a/areno/adapters/peft.py +++ b/areno/adapters/peft.py @@ -9,7 +9,7 @@ import torch.distributed as dist from safetensors.torch import load_file, save_file -from areno.adapters.lora import AdapterRegistry, RoutedExpertLoraSlot +from areno.adapters.lora import AdapterRegistry, LoraSlot, RoutedExpertLoraSlot from areno.engine.parallel.context import get_tp_context _PREFIX = "base_model.model.model." @@ -41,9 +41,8 @@ def load_peft_adapter(registry: AdapterRegistry, path: str | Path) -> None: local_A = canonical_A[:, ctx.rank * width : (ctx.rank + 1) * width] local_B = canonical_B else: - height = slot.local_out_features local_A = canonical_A - local_B = canonical_B[ctx.rank * height : (ctx.rank + 1) * height] + local_B = canonical_B[slot.output_start : slot.output_end] slot.lora_A.copy_(local_A.to(device=slot.lora_A.device, dtype=slot.lora_A.dtype)) slot.lora_B.copy_(local_B.to(device=slot.lora_B.device, dtype=slot.lora_B.dtype)) @@ -85,7 +84,11 @@ def export_peft_adapter( if ctx.rank != 0: continue canonical_A = slot.lora_A.detach() - canonical_B = torch.cat(gathered_B, dim=0) + canonical_B = ( + _gather_replicated_column(slot, gathered_B, ctx.world_size) + if slot.output_replicated + else torch.cat(gathered_B, dim=0) + ) state[_key(logical_name, "A")] = canonical_A.float().cpu().contiguous() state[_key(logical_name, "B")] = canonical_B.float().cpu().contiguous() if ctx.rank != 0: @@ -120,5 +123,22 @@ def _all_gather(tensor: torch.Tensor, world_size: int, group) -> list[torch.Tens return gathered +def _gather_replicated_column(slot: LoraSlot, gathered: list[torch.Tensor], world_size: int) -> torch.Tensor: + """Keep one authoritative copy of every unique replicated output range.""" + + local_rows = slot.local_out_features + unique_shards = slot.global_out_features // local_rows + ranks_per_shard = world_size // unique_shards + output = torch.empty( + (slot.global_out_features, slot.rank), + device=gathered[0].device, + dtype=gathered[0].dtype, + ) + for shard_index in range(unique_shards): + start = shard_index * local_rows + output[start : start + local_rows].copy_(gathered[shard_index * ranks_per_shard]) + return output + + def _key(logical_name: str, component: str) -> str: return f"{_PREFIX}{logical_name}.lora_{component}.weight" diff --git a/areno/engine/config.py b/areno/engine/config.py index c13a387b..38359855 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -292,18 +292,6 @@ def __post_init__(self) -> None: raise ValueError("reference_mode must be one of: independent, reuse_actor_base") if self.reference_mode == "reuse_actor_base" and self.lora is None: raise ValueError("reference_mode='reuse_actor_base' requires native LoRA") - if self.lora is not None: - replicated_kv_targets = {"k_proj", "v_proj"} & set(self.lora.target_modules) - if ( - self.model.model_type == "qwen3_moe" - and self.tp_size > self.model.num_key_value_heads - and replicated_kv_targets - ): - targets = ", ".join(sorted(replicated_kv_targets)) - raise ValueError( - f"Qwen3-MoE replicated-KV topology does not support LoRA targets {targets}; " - "omit k_proj/v_proj or use tp_size <= num_key_value_heads" - ) if self.devices is None: if torch.cuda.is_available(): device_count = torch.cuda.device_count() diff --git a/areno/engine/policy_sync.py b/areno/engine/policy_sync.py index 2ace948a..bd2e33e3 100644 --- a/areno/engine/policy_sync.py +++ b/areno/engine/policy_sync.py @@ -8,7 +8,7 @@ import torch import torch.distributed as dist -from areno.adapters.lora import AdapterRegistry, RoutedExpertLoraSlot +from areno.adapters.lora import AdapterRegistry, LoraSlot, RoutedExpertLoraSlot from areno.engine.checkpoints.io import PolicyTensorLayout, PolicyTensorPiece, PolicyTensorStore from areno.engine.parallel.context import get_tp_context from areno.engine.protocol import PolicySyncPayload @@ -96,7 +96,8 @@ def build_adapter_policy_plan(registry: AdapterRegistry) -> PolicyTensorStore: slot.lora_B, (slot.global_out_features, slot.rank), dim=0, - start=ctx.rank * slot.local_out_features, + start=slot.output_start, + publish=_column_range_publisher(slot, ctx.rank, ctx.world_size), ), ) return plan @@ -108,15 +109,25 @@ def _sharded_layout( *, dim: int, start: int, + publish: bool = True, ) -> PolicyTensorLayout: local_size = tensor.shape[dim] return PolicyTensorLayout( shape=shape, dtype=tensor.dtype, - pieces=(PolicyTensorPiece(tensor.detach(), shape, dim, start, start + local_size),), + pieces=(PolicyTensorPiece(tensor.detach(), shape, dim, start, start + local_size, publish=publish),), ) +def _column_range_publisher(slot: LoraSlot, tp_rank: int, tp_size: int) -> bool: + if not slot.output_replicated: + return True + unique_shards = slot.global_out_features // slot.local_out_features + ranks_per_shard = tp_size // unique_shards + owner_rank = (slot.output_start // slot.local_out_features) * ranks_per_shard + return tp_rank == owner_rank + + def _replicated_layout(tensor: torch.Tensor) -> PolicyTensorLayout: shape = tuple(tensor.shape) return PolicyTensorLayout( diff --git a/areno/engine/training.py b/areno/engine/training.py index cbd4b0ba..13cfd371 100644 --- a/areno/engine/training.py +++ b/areno/engine/training.py @@ -282,11 +282,32 @@ def _sync_tensor_parallel_replicated_gradients(self) -> None: ctx = get_tp_context() if ctx.world_size == 1: return + ranged = [] for param in worker.model.parameters(): grad = param_grad(param) - if grad is None or not bool(getattr(param, "tp_grad_allreduce", False)): + if grad is None: continue - dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=ctx.group) + output_range = getattr(param, "tp_replicated_output_range", None) + if output_range is not None: + start, end, global_size = output_range + canonical_numel = global_size * grad[0].numel() + ranged.append((grad, start, end, global_size, canonical_numel)) + elif bool(getattr(param, "tp_grad_allreduce", False)): + dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=ctx.group) + if not ranged: + return + packed = ranged[0][0].new_zeros(sum(item[-1] for item in ranged)) + offset = 0 + for grad, start, end, global_size, canonical_numel in ranged: + canonical = packed.narrow(0, offset, canonical_numel).view(global_size, *grad.shape[1:]) + canonical[start:end].copy_(grad) + offset += canonical_numel + dist.all_reduce(packed, op=dist.ReduceOp.SUM, group=ctx.group) + offset = 0 + for grad, start, end, global_size, canonical_numel in ranged: + canonical = packed.narrow(0, offset, canonical_numel).view(global_size, *grad.shape[1:]) + grad.copy_(canonical[start:end]) + offset += canonical_numel def _sync_data_parallel_gradients(self) -> None: """Average resident full gradients across data-parallel replicas.""" diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index aa10d423..94affcdb 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -201,8 +201,8 @@ def test_engine_config_resolves_sequence_parallel_override_before_model_config(s tp1 = EngineConfig(model=model, tp_size=1, devices=[0], sequence_parallel=True) self.assertFalse(tp1.effective_sequence_parallel) - def test_engine_config_rejects_replicated_kv_lora_targets(self): - """Replicated Qwen3 KV requires range-aware LoRA support.""" + def test_engine_config_allows_replicated_kv_lora_targets(self): + """Range-aware LoRA supports Qwen3 KV replication across wider TP.""" model = ModelConfig( model_type="qwen3_moe", num_attention_heads=8, @@ -211,14 +211,11 @@ def test_engine_config_rejects_replicated_kv_lora_targets(self): vocab_size=32, ) - with self.assertRaisesRegex(ValueError, "Qwen3-MoE replicated-KV.*k_proj"): - EngineConfig(model=model, tp_size=4, devices=[0, 1, 2, 3], lora=LoraConfig()) - EngineConfig( model=model, tp_size=4, devices=[0, 1, 2, 3], - lora=LoraConfig(target_modules=("q_proj", "o_proj")), + lora=LoraConfig(), ) def test_reference_view_requires_lora_at_engine_boundary(self): diff --git a/tests/test_policy_tensor_sync_cpu.py b/tests/test_policy_tensor_sync_cpu.py index 5ce76868..12b26abb 100644 --- a/tests/test_policy_tensor_sync_cpu.py +++ b/tests/test_policy_tensor_sync_cpu.py @@ -157,6 +157,61 @@ def test_adapter_plan_maps_row_column_and_expert_factors_tp2_to_tp1() -> None: ) +def test_adapter_plan_publishes_replicated_column_range_once() -> None: + config = LoraConfig(rank=2, alpha=4.0) + plans = [] + for rank in range(4): + _set_rank(rank, 4) + state = _AdapterRuntimeState() + start = 0 if rank < 2 else 1 + slot = LoraSlot( + logical_name="kv", + base_weight=nn.Parameter(torch.zeros(1)), + global_in_features=4, + global_out_features=2, + local_in_features=4, + local_out_features=1, + row_parallel=False, + config=config, + seed=1, + runtime_state=state, + output_range=(start, start + 1), + ) + slot.lora_B.data.fill_(float(start + 1)) + plans.append(build_adapter_policy_plan(AdapterRegistry({"kv": slot}, config, state))) + + canonical = _canonical([plan["kv.lora_B.weight"].policy_layout() for plan in plans]) + torch.testing.assert_close(canonical, torch.tensor([[1.0, 1.0], [2.0, 2.0]])) + + rollout_slots = [] + for rank in range(4): + _set_rank(rank, 4) + state = _AdapterRuntimeState() + start = 0 if rank < 2 else 1 + slot = LoraSlot( + logical_name="kv", + base_weight=nn.Parameter(torch.zeros(1)), + global_in_features=4, + global_out_features=2, + local_in_features=4, + local_out_features=1, + row_parallel=False, + config=config, + seed=1, + runtime_state=state, + output_range=(start, start + 1), + ) + slot.lora_B.data.zero_() + rollout_slots.append(slot) + plan = build_adapter_policy_plan(AdapterRegistry({"kv": slot}, config, state)) + _write([plan["kv.lora_B.weight"].policy_layout()], canonical) + + torch.testing.assert_close(rollout_slots[0].lora_B, rollout_slots[1].lora_B) + torch.testing.assert_close(rollout_slots[2].lora_B, rollout_slots[3].lora_B) + torch.testing.assert_close(rollout_slots[0].lora_B, torch.ones(1, 2)) + torch.testing.assert_close(rollout_slots[2].lora_B, torch.full((1, 2), 2.0)) + + def test_column_parallel_layout_reshards_tp2_to_tp1() -> None: full = torch.arange(24, dtype=torch.float32).reshape(6, 4) train_layouts = [] diff --git a/tests/test_qwen3_lora_e2e.py b/tests/test_qwen3_lora_e2e.py index 905f3bad..392833d7 100644 --- a/tests/test_qwen3_lora_e2e.py +++ b/tests/test_qwen3_lora_e2e.py @@ -383,6 +383,116 @@ def reward_fn(record) -> float: assert torch.isfinite(torch.tensor(peft_logprobs)).all() +def test_qwen3_moe_lora_tp8_replicated_kv_roundtrip(tmp_path: Path) -> None: + model_path_value = os.getenv("ARENO_E2E_QWEN3_MOE_MODEL") + if not model_path_value: + pytest.skip("set ARENO_E2E_QWEN3_MOE_MODEL to run the 8-GPU replicated-KV LoRA E2E") + model_path = Path(model_path_value) + initial_path = tmp_path / "adapter-initial" + checkpoint_path = tmp_path / "checkpoints" + final_path = checkpoint_path / "step_000001" + reexported_path = tmp_path / "adapter-reexported" + lora = LoraConfig(rank=8, alpha=16.0) + backend_config = ArenoConfig( + tp_size=8, + dp_size=1, + devices=list(range(8)), + lora=lora, + max_running_prompts=2, + optimizer={ + "lr": 1.0e-4, + "min_lr": 1.0e-4, + "lr_decay_style": "constant", + "weight_decay": 0.0, + }, + runtime={ + "compile_model": False, + "activation_checkpointing": False, + "keep_rollout_state": False, + "eager_decode": True, + }, + ) + observed = _ObservedTrainer(Trainer(8, os.fspath(model_path), custom_config=backend_config)) + config = PolicyTrainerConfig( + algo="grpo", + ckpt=os.fspath(model_path), + dataset_path="e2e://replicated-kv", + save_path=os.fspath(checkpoint_path), + save_interval=1, + epochs=1, + max_steps=1, + world_size=8, + tp_size=8, + train_devices=list(range(8)), + batch_size=1, + mini_bs=2, + n_samples=2, + greedy=True, + max_running_prompts=2, + max_prompt_tokens=64, + max_new_tokens=4, + optimizer_lr=1.0e-4, + optimizer_min_lr=1.0e-4, + lr_decay_style="constant", + weight_decay=0.0, + activation_checkpointing=False, + keep_rollout_state=False, + eager_decode=True, + metrics_log_dir=None, + lora=lora, + ) + + def reward_fn(record) -> float: + return float(record.metadata["sample_index"]) + + policy = PolicyOnlyTrainer( + config, + instance=observed, + dataset=[{"prompt": "Write one uncommon English noun. Output only the noun."}], + reward_fn=reward_fn, + loss_fn=get_algorithm("grpo").make_loss_fn(config), + ) + observed.init() + try: + parity_tokens = observed.get_tokenizer().encode("A replicated KV adapter check.", add_special_tokens=True) + observed.export_adapter(os.fspath(initial_path)) + policy._fit_initialized() + trained_logprobs = observed.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + finally: + observed.close() + + initial = load_file(initial_path / "adapter_model.safetensors") + final = load_file(final_path / "adapter_model.safetensors") + kv_keys = [name for name in final if ".self_attn.k_proj." in name or ".self_attn.v_proj." in name] + assert kv_keys + assert any(not torch.equal(initial[name], final[name]) for name in kv_keys) + assert all(final[name].shape[0] == 512 for name in kv_keys if name.endswith("lora_B.weight")) + + imported = Trainer( + 8, + os.fspath(model_path), + custom_config=ArenoConfig( + tp_size=8, + dp_size=1, + devices=list(range(8)), + lora=LoraConfig(adapter_path=os.fspath(final_path)), + runtime={"compile_model": False, "activation_checkpointing": False, "eager_decode": True}, + ), + ) + imported.init() + try: + imported.export_adapter(os.fspath(reexported_path)) + imported_logprobs = imported.score_logprobs("actor", [parity_tokens], microbatch_size=1)[0] + finally: + imported.close() + reexported = load_file(reexported_path / "adapter_model.safetensors") + assert reexported.keys() == final.keys() + assert all(torch.equal(reexported[name], final[name]) for name in final) + torch.testing.assert_close( + torch.tensor(imported_logprobs), torch.tensor(trained_logprobs), rtol=0.0, atol=1.0e-5 + ) + + @pytest.mark.parametrize( ("model_env", "model_kind", "rollout_tp_size", "rollout_devices"), ( From 83a135cdc79b5bd94b23e6b36ba5716f3b4b0b38 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 18:42:00 +0800 Subject: [PATCH 11/29] test(runtime): report sequence parallel train metric --- areno/engine/training.py | 1 + 1 file changed, 1 insertion(+) diff --git a/areno/engine/training.py b/areno/engine/training.py index 13cfd371..f04b12bc 100644 --- a/areno/engine/training.py +++ b/areno/engine/training.py @@ -209,6 +209,7 @@ def _train_step( None, {"lr": current_lr}, multimodal_lrs, + {"sequence_parallel": float(model_kwargs["train_meta"].sequence_parallel)}, {"grad_norm": grad_norm} if grad_norm is not None else None, multimodal_grad_metrics, {"clipped_grad_norm": clipped_grad_norm} if clipped_grad_norm is not None else None, From 6d625e34ef5abbafc70c8a8903da8bb64323ad66 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 22:01:14 +0800 Subject: [PATCH 12/29] feat(lora): add Bailing Tiny V3 native targets --- areno/adapters/config.py | 27 +++- areno/adapters/lora.py | 240 ++++++++++++++++++++++++++++-- areno/engine/config.py | 8 +- areno/engine/layers/linear.py | 8 + areno/models/bailing_v3/model.py | 54 ++++++- tests/test_bailing_v3_lora_cpu.py | 134 +++++++++++++++++ 6 files changed, 445 insertions(+), 26 deletions(-) create mode 100644 tests/test_bailing_v3_lora_cpu.py diff --git a/areno/adapters/config.py b/areno/adapters/config.py index e2eb452a..fb114892 100644 --- a/areno/adapters/config.py +++ b/areno/adapters/config.py @@ -1,4 +1,4 @@ -"""Public configuration for the Qwen3 dense and MoE LoRA runtime.""" +"""Public configuration for the native LoRA runtime.""" from __future__ import annotations @@ -16,10 +16,27 @@ "down_proj", ) +BAILING_V3_TARGETS = ( + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "dense", + "gate_proj", + "up_proj", + "down_proj", +) + +NATIVE_LORA_TARGETS = tuple(dict.fromkeys((*QWEN3_DENSE_TARGETS, *BAILING_V3_TARGETS))) + @dataclass(frozen=True, slots=True) class LoraConfig: - """Supported PEFT-compatible LoRA subset for Qwen3 dense and MoE models. + """Supported PEFT-compatible subset for native LoRA model families. When ``adapter_path`` is set, its standard PEFT metadata is authoritative for rank, alpha, dropout, and targets. @@ -44,11 +61,11 @@ def __post_init__(self) -> None: if self.alpha <= 0: raise ValueError("lora alpha must be > 0") if self.dropout != 0.0: - raise ValueError("native Qwen3 LoRA currently requires dropout=0") + raise ValueError("native LoRA currently requires dropout=0") requested = set(self.target_modules) - supported = set(QWEN3_DENSE_TARGETS) + supported = set(NATIVE_LORA_TARGETS) if not requested or not requested <= supported: - raise ValueError(f"target_modules must be a non-empty subset of {QWEN3_DENSE_TARGETS}") + raise ValueError(f"target_modules must be a non-empty subset of {NATIVE_LORA_TARGETS}") @property def scale(self) -> float: diff --git a/areno/adapters/lora.py b/areno/adapters/lora.py index bd8f8e84..f17db3ec 100644 --- a/areno/adapters/lora.py +++ b/areno/adapters/lora.py @@ -1,4 +1,4 @@ -"""TP-aware native LoRA slots for Qwen3 dense and routed-expert projections.""" +"""TP-aware native LoRA slots for dense and routed-expert projections.""" from __future__ import annotations @@ -13,7 +13,7 @@ from areno.accel import areno_grouped_linear from areno.adapters.config import LoraConfig -from areno.engine.layers.linear import RowParallelLinear, mark_tensor_parallel_parameter +from areno.engine.layers.linear import ColumnParallelLinear, RowParallelLinear, mark_tensor_parallel_parameter from areno.engine.parallel.context import get_tp_context @@ -110,7 +110,7 @@ def enabled(self) -> bool: class RoutedExpertLoraSlot(nn.Module): - """One expert-sharded canonical LoRA A/B pair for grouped Qwen3-MoE GEMMs.""" + """One expert-sharded canonical LoRA A/B pair for grouped MoE GEMMs.""" def __init__( self, @@ -216,23 +216,47 @@ def base_only(self) -> Iterator[None]: def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> AdapterRegistry: - """Freeze a Qwen3 dense or MoE base and attach canonical targets.""" + """Freeze one supported native base and attach its canonical targets.""" model_config = getattr(model, "config", None) - if getattr(model_config, "model_type", None) not in {"qwen3", "qwen3_moe"}: - raise ValueError("native LoRA currently supports Qwen3 models only") + model_type = getattr(model_config, "model_type", None) + if model_type not in {"qwen3", "qwen3_moe", "bailing_moe_v3"}: + raise ValueError("native LoRA currently supports Qwen3 and Bailing-MoE V3 models only") + if model_type == "bailing_moe_v3" and not bool(getattr(model_config, "no_kda_lora", False)): + raise ValueError("Bailing-MoE V3 native LoRA currently requires no_kda_lora=true") for parameter in model.parameters(): parameter.requires_grad_(False) requested = set(config.target_modules) runtime_state = _AdapterRuntimeState() slots: dict[str, LoraSlot | RoutedExpertLoraSlot] = {} + if model_type == "bailing_moe_v3": + matched = _initialize_bailing_v3_lora(model, requested, config, seed, runtime_state, slots) + else: + matched = _initialize_qwen3_lora(model, requested, config, seed, runtime_state, slots) + missing = requested - matched + if missing: + raise ValueError(f"target_modules are not present in {model_type}: {', '.join(sorted(missing))}") + return AdapterRegistry(slots, config, runtime_state) + + +def _initialize_qwen3_lora( + model: nn.Module, + requested: set[str], + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], +) -> set[str]: + matched: set[str] = set() + model_config = model.config for layer_index, layer in enumerate(model.layers): prefix = f"layers.{layer_index}" qkv = layer.self_attn.qkv_proj for component_index, component in enumerate(("q_proj", "k_proj", "v_proj")): if component not in requested: continue + matched.add(component) logical_name = f"{prefix}.self_attn.{component}" slot = LoraSlot( logical_name=logical_name, @@ -251,6 +275,7 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt slots[logical_name] = slot if "o_proj" in requested: + matched.add("o_proj") owner = layer.self_attn.o_proj logical_name = f"{prefix}.self_attn.o_proj" slot = _row_slot(logical_name, owner, config, seed, runtime_state) @@ -258,12 +283,15 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt slots[logical_name] = slot if getattr(model_config, "enable_moe_block", False): - _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots) + matched.update( + _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots) + ) else: gate_up = layer.mlp.gate_up_proj for component_index, component in enumerate(("gate_proj", "up_proj")): if component not in requested: continue + matched.add(component) logical_name = f"{prefix}.mlp.{component}" slot = LoraSlot( logical_name=logical_name, @@ -281,13 +309,196 @@ def initialize_lora(model: nn.Module, config: LoraConfig, *, seed: int) -> Adapt slots[logical_name] = slot if "down_proj" in requested: + matched.add("down_proj") owner = layer.mlp.down_proj logical_name = f"{prefix}.mlp.down_proj" slot = _row_slot(logical_name, owner, config, seed, runtime_state) owner.install_lora(slot) slots[logical_name] = slot - return AdapterRegistry(slots, config, runtime_state) + return matched + + +def _initialize_bailing_v3_lora( + model: nn.Module, + requested: set[str], + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], +) -> set[str]: + matched: set[str] = set() + for layer_index, layer in enumerate(model.layers): + prefix = f"layers.{layer_index}" + attention = layer.attention + attention_prefix = f"{prefix}.attention" + if hasattr(attention, "q_conv1d_weight"): + for component in ("q_proj", "k_proj", "v_proj"): + if component in requested: + matched.add(component) + _install_column_slot( + f"{attention_prefix}.{component}", + getattr(attention, component), + config, + seed, + runtime_state, + slots, + ) + if "o_proj" in requested: + matched.add("o_proj") + _install_row_slot( + f"{attention_prefix}.o_proj", + attention.o_proj, + config, + seed, + runtime_state, + slots, + ) + else: + if "q_proj" in requested and attention.q_proj is not None: + matched.add("q_proj") + _install_column_slot( + f"{attention_prefix}.q_proj", + attention.q_proj, + config, + seed, + runtime_state, + slots, + ) + for component in ("q_a_proj", "kv_a_proj_with_mqa"): + owner = getattr(attention, component, None) + if component in requested and owner is not None: + matched.add(component) + slot = _replicated_slot( + f"{attention_prefix}.{component}", owner, config, seed, runtime_state + ) + attention.install_lora_component(component, slot) + slots[slot.logical_name] = slot + for component in ("q_b_proj", "kv_b_proj"): + owner = getattr(attention, component, None) + if component in requested and owner is not None: + matched.add(component) + _install_column_slot( + f"{attention_prefix}.{component}", owner, config, seed, runtime_state, slots + ) + if "dense" in requested: + matched.add("dense") + _install_row_slot( + f"{attention_prefix}.dense", + attention.dense, + config, + seed, + runtime_state, + slots, + ) + + mlp_prefix = f"{prefix}.mlp" + if hasattr(layer.mlp, "experts"): + matched.update( + _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots) + ) + if layer.mlp.shared_experts is not None: + matched.update( + _install_dense_mlp_slots( + layer.mlp.shared_experts, + f"{mlp_prefix}.shared_experts", + requested, + config, + seed, + runtime_state, + slots, + ) + ) + else: + matched.update( + _install_dense_mlp_slots( + layer.mlp, mlp_prefix, requested, config, seed, runtime_state, slots + ) + ) + return matched + + +def _install_column_slot( + logical_name: str, + owner: ColumnParallelLinear, + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], +) -> None: + slot = LoraSlot( + logical_name=logical_name, + base_weight=owner.weight, + global_in_features=owner.in_features, + global_out_features=owner.out_features, + local_in_features=owner.in_features, + local_out_features=owner.local_out_features, + row_parallel=False, + config=config, + seed=seed, + runtime_state=runtime_state, + ) + owner.install_lora(slot) + slots[logical_name] = slot + + +def _install_row_slot( + logical_name: str, + owner: RowParallelLinear, + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], +) -> None: + slot = _row_slot(logical_name, owner, config, seed, runtime_state) + owner.install_lora(slot) + slots[logical_name] = slot + + +def _replicated_slot( + logical_name: str, + owner: nn.Linear, + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, +) -> LoraSlot: + return LoraSlot( + logical_name=logical_name, + base_weight=owner.weight, + global_in_features=owner.in_features, + global_out_features=owner.out_features, + local_in_features=owner.in_features, + local_out_features=owner.out_features, + row_parallel=False, + output_range=(0, owner.out_features), + config=config, + seed=seed, + runtime_state=runtime_state, + ) + + +def _install_dense_mlp_slots( + mlp: nn.Module, + prefix: str, + requested: set[str], + config: LoraConfig, + seed: int, + runtime_state: _AdapterRuntimeState, + slots: dict[str, LoraSlot | RoutedExpertLoraSlot], +) -> set[str]: + matched: set[str] = set() + for component in ("gate_proj", "up_proj"): + if component in requested: + matched.add(component) + _install_column_slot( + f"{prefix}.{component}", getattr(mlp, component), config, seed, runtime_state, slots + ) + if "down_proj" in requested: + matched.add("down_proj") + _install_row_slot( + f"{prefix}.down_proj", mlp.down_proj, config, seed, runtime_state, slots + ) + return matched def _row_slot( @@ -319,15 +530,19 @@ def _install_moe_slots( seed: int, runtime_state: _AdapterRuntimeState, slots: dict[str, LoraSlot | RoutedExpertLoraSlot], -) -> None: +) -> set[str]: + gate_up_weight = experts.gate_up_weight if hasattr(experts, "gate_up_weight") else experts.linear_fc1.weight + down_weight = experts.down_weight if hasattr(experts, "down_weight") else experts.linear_fc2.weight components = ( - ("gate_proj", experts.hidden_size, experts.intermediate_size, experts.gate_up_weight), - ("up_proj", experts.hidden_size, experts.intermediate_size, experts.gate_up_weight), - ("down_proj", experts.intermediate_size, experts.hidden_size, experts.down_weight), + ("gate_proj", experts.hidden_size, experts.intermediate_size, gate_up_weight), + ("up_proj", experts.hidden_size, experts.intermediate_size, gate_up_weight), + ("down_proj", experts.intermediate_size, experts.hidden_size, down_weight), ) + matched: set[str] = set() for component, in_features, out_features, base_weight in components: if component not in requested: continue + matched.add(component) logical_name = f"{prefix}.mlp.experts.{{expert}}.{component}" slot = RoutedExpertLoraSlot( logical_name=logical_name, @@ -342,3 +557,4 @@ def _install_moe_slots( ) experts.install_lora_component(component, slot) slots[logical_name] = slot + return matched diff --git a/areno/engine/config.py b/areno/engine/config.py index 38359855..0d043973 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -126,9 +126,13 @@ def resolve_eager_decode(self, *, model: ModelConfig, lora: LoraConfig | None) - if self.eager_decode or lora is None: return - if model.model_type == "qwen3_moe" and {"gate_proj", "up_proj", "down_proj"} & set(lora.target_modules): + if model.model_type in {"qwen3_moe", "bailing_moe_v3"} and { + "gate_proj", + "up_proj", + "down_proj", + } & set(lora.target_modules): warnings.warn( - "Qwen3-MoE expert LoRA uses grouped execution during rollout; falling back to eager decode.", + "routed-expert LoRA uses grouped execution during rollout; falling back to eager decode.", RuntimeWarning, stacklevel=2, ) diff --git a/areno/engine/layers/linear.py b/areno/engine/layers/linear.py index e6bd62d6..b95c3200 100644 --- a/areno/engine/layers/linear.py +++ b/areno/engine/layers/linear.py @@ -86,6 +86,7 @@ def __init__( input_grad_allreduce: bool = True, ): super().__init__() + self.lora_slot: nn.Module | None = None ctx = get_tp_context() start, end = _shard_range(out_features, ctx.rank, ctx.world_size) self.in_features = in_features @@ -99,6 +100,11 @@ def __init__( mark_tensor_parallel_parameter(self.bias, True, sequence_parallel=True) self.reset_parameters() + def install_lora(self, slot: nn.Module) -> None: + """Attach one adapter before compilation and optimizer construction.""" + + self.lora_slot = slot + def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: @@ -116,6 +122,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: elif self.input_grad_allreduce: x = copy_to_tensor_parallel_region(x) out = _areno_linear_forward(x, self.weight, self.bias) + if self.lora_slot is not None and self.lora_slot.enabled: + out = out + self.lora_slot(x) if self.gather_output: # Concatenate column-shards along the last dim to recover the # full output (only used when downstream code needs it dense). diff --git a/areno/models/bailing_v3/model.py b/areno/models/bailing_v3/model.py index 176679ef..41cb5b26 100644 --- a/areno/models/bailing_v3/model.py +++ b/areno/models/bailing_v3/model.py @@ -289,7 +289,7 @@ def forward(self, hidden_states: torch.Tensor, num_padding_tokens: int = 0) -> t with sequence_parallel_region(False): topk_idx, topk_weight, _ = self.gate(hidden_states, num_padding_tokens) flat = expert_input.view(-1, hidden) - if self.training: + if self.training or self.experts.has_lora(): # Permute/unpermute path is autograd-friendly. out = self.experts(flat, topk_idx, topk_weight).view(bsz, seqlen, hidden) else: @@ -311,6 +311,9 @@ def prepare_infer_weights(self) -> None: down projection. Buffers are reused across calls if the shape/device already match to avoid reallocating on every weight refresh. """ + if self.experts.has_lora(): + self.clear_infer_weights() + return gate_weights, up_weights, down_weights = self.experts.expert_weights() self._infer_gate_weight = self._updated_infer_weight( self._infer_gate_weight, @@ -411,11 +414,21 @@ def __init__(self, config: ModelConfig): self.hidden_size, dtype=config.dtype, ) + self.lora_slots = nn.ModuleDict() # Expert weights are sharded by EP (collapsed into TP); flag them as # not-TP/not-SP so the standard TP collectives leave them alone. for param in self.parameters(): mark_tensor_parallel_parameter(param, False, sequence_parallel=False) + def install_lora_component(self, component: str, slot: nn.Module) -> None: + self.lora_slots[component] = slot + + def has_lora(self) -> bool: + return bool(self.lora_slots) + + def has_active_lora(self) -> bool: + return self.has_lora() and next(iter(self.lora_slots.values())).enabled + def forward(self, flat: torch.Tensor, topk_idx: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: return self._forward_fused_permute(flat, topk_idx, topk_weight) @@ -438,15 +451,29 @@ def _forward_fused_permute( # collective sync with peers. return all_reduce(flat.new_zeros(flat.shape)) hidden, _ = _grouped_linear_forward(self.linear_fc1, x.contiguous(), tokens_per_expert) + if self.has_active_lora(): + gate, up = hidden.chunk(2, dim=-1) + if "gate_proj" in self.lora_slots: + gate = gate + self.lora_slots["gate_proj"](x, tokens_per_expert) + if "up_proj" in self.lora_slots: + up = up + self.lora_slots["up_proj"](x, tokens_per_expert) + hidden = torch.cat((gate, up), dim=-1) # Apply routing weight before fc2 so it stays inside the fp32 reduction. hidden = ( _areno_silu_and_mul_no_compile(hidden) * sorted_route_weight.unsqueeze(-1).to(dtype=hidden.dtype) ).contiguous() expert_out, _ = _grouped_linear_forward(self.linear_fc2, hidden, tokens_per_expert) + if self.has_active_lora() and "down_proj" in self.lora_slots: + expert_out = expert_out + self.lora_slots["down_proj"](hidden, tokens_per_expert) # Unpermute back to original (batch, seq) order, then scale and reduce. - out = _areno_moe_unpermute_no_compile( - expert_out, sorted_token_idx, merging_probs=None, restore_shape=flat.shape - ) + if self.has_lora(): + out = _areno_moe_unpermute_no_compile( + expert_out.float(), sorted_token_idx, merging_probs=None, restore_shape=flat.shape + ).to(dtype=flat.dtype) + else: + out = _areno_moe_unpermute_no_compile( + expert_out, sorted_token_idx, merging_probs=None, restore_shape=flat.shape + ) return all_reduce(out * self.config.routed_scaling_factor) def local_routes(self, topk_idx: torch.Tensor, topk_weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -676,6 +703,7 @@ class BailingSoftmaxAttention(nn.Module): def __init__(self, config: ModelConfig, layer_idx: int): super().__init__() + self.lora_slots = nn.ModuleDict() ctx = get_tp_context() self.layer_idx = layer_idx # Head-dim split: rope vs non-rope channels on Q/K, plus separate V dim. @@ -774,6 +802,15 @@ def __init__(self, config: ModelConfig, layer_idx: int): self.k_cache = torch.tensor([]) self.v_cache = torch.tensor([]) + def install_lora_component(self, component: str, slot: nn.Module) -> None: + """Attach an adapter to one replicated MLA projection.""" + + self.lora_slots[component] = slot + + def _with_lora(self, component: str, x: torch.Tensor, output: torch.Tensor) -> torch.Tensor: + slot = self.lora_slots[component] if component in self.lora_slots else None + return output + slot(x) if slot is not None and slot.enabled else output + def forward( self, hidden_states: torch.Tensor, @@ -827,11 +864,14 @@ def _project( q = self.q_proj(mla_input) else: assert self.q_a_proj is not None and self.q_a_layernorm is not None and self.q_b_proj is not None - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(mla_input))) + q_a = self._with_lora("q_a_proj", mla_input, self.q_a_proj(mla_input)) + q = self.q_b_proj(self.q_a_layernorm(q_a)) bsz, seqlen = q.shape[:2] q = q.view(bsz, seqlen, self.local_heads, self.head_dim) q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - kv_a = self.kv_a_proj_with_mqa(mla_input) + kv_a = self._with_lora( + "kv_a_proj_with_mqa", mla_input, self.kv_a_proj_with_mqa(mla_input) + ) compressed_kv, k_rope = kv_a.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) if is_sequence_parallel_active(): k_rope = gather_from_sequence_parallel_region(k_rope) @@ -1514,7 +1554,7 @@ def clear_infer_weights(self) -> None: @torch.no_grad() def offload_train_weights(self) -> None: for layer in self.layers: - if isinstance(layer.mlp, BailingSparseMoeBlock): + if isinstance(layer.mlp, BailingSparseMoeBlock) and not layer.mlp.experts.has_lora(): layer.mlp.experts.offload_to_cpu() @torch.no_grad() diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py new file mode 100644 index 00000000..b0b2a62d --- /dev/null +++ b/tests/test_bailing_v3_lora_cpu.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch +import pytest +from torch import nn + +from areno.adapters.config import BAILING_V3_TARGETS, LoraConfig +from areno.adapters.lora import initialize_lora +from areno.engine.layers import linear +from areno.engine.layers.linear import ColumnParallelLinear, RowParallelLinear + + +class _ReplicatedAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.q_proj = None + self.q_a_proj = nn.Linear(8, 4, bias=False) + self.q_b_proj = ColumnParallelLinear(4, 12, bias=False) + self.kv_a_proj_with_mqa = nn.Linear(8, 6, bias=False) + self.kv_b_proj = ColumnParallelLinear(4, 12, bias=False) + self.dense = RowParallelLinear(8, 8, bias=False) + self.lora_slots = nn.ModuleDict() + + def install_lora_component(self, component: str, slot: nn.Module) -> None: + self.lora_slots[component] = slot + + +class _KDAAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.q_conv1d_weight = nn.Parameter(torch.empty(8, 1, 2)) + self.q_proj = ColumnParallelLinear(8, 8, bias=False) + self.k_proj = ColumnParallelLinear(8, 8, bias=False) + self.v_proj = ColumnParallelLinear(8, 8, bias=False) + self.o_proj = RowParallelLinear(8, 8, bias=False) + + +class _DenseMLP(nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear(8, 12, bias=False) + self.up_proj = ColumnParallelLinear(8, 12, bias=False) + self.down_proj = RowParallelLinear(12, 8, bias=False) + + +class _GroupedExperts(nn.Module): + def __init__(self) -> None: + super().__init__() + self.hidden_size = 8 + self.intermediate_size = 4 + self.local_num_experts = 2 + self.local_expert_start = 0 + self.linear_fc1 = nn.Linear(8, 8, bias=False) + self.linear_fc1.weight = nn.Parameter(torch.empty(2, 8, 8)) + self.linear_fc2 = nn.Linear(4, 8, bias=False) + self.linear_fc2.weight = nn.Parameter(torch.empty(2, 8, 4)) + self.lora_slots = nn.ModuleDict() + + def install_lora_component(self, component: str, slot: nn.Module) -> None: + self.lora_slots[component] = slot + + +class _SparseMLP(nn.Module): + def __init__(self) -> None: + super().__init__() + self.experts = _GroupedExperts() + self.shared_experts = _DenseMLP() + + +class _Layer(nn.Module): + def __init__(self, attention: nn.Module, mlp: nn.Module) -> None: + super().__init__() + self.attention = attention + self.mlp = mlp + + +class _BailingModel(nn.Module): + def __init__(self, *, no_kda_lora: bool = True) -> None: + super().__init__() + self.config = SimpleNamespace(model_type="bailing_moe_v3", no_kda_lora=no_kda_lora) + self.layers = nn.ModuleList( + ( + _Layer(_KDAAttention(), _DenseMLP()), + _Layer(_ReplicatedAttention(), _SparseMLP()), + ) + ) + + +def _single_tp() -> SimpleNamespace: + return SimpleNamespace(rank=0, world_size=1, group=None) + + +def test_bailing_v3_full_profile_attaches_native_slots(monkeypatch) -> None: + monkeypatch.setattr(linear, "get_tp_context", _single_tp) + monkeypatch.setattr("areno.adapters.lora.get_tp_context", _single_tp) + model = _BailingModel() + profile = ( + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "dense", + "gate_proj", + "up_proj", + "down_proj", + ) + + registry = initialize_lora(model, LoraConfig(rank=4, alpha=4, target_modules=profile), seed=42) + + assert set(profile) <= set(BAILING_V3_TARGETS) + assert "layers.0.attention.q_proj" in registry.slots + assert "layers.1.attention.q_a_proj" in registry.slots + assert "layers.1.mlp.shared_experts.gate_proj" in registry.slots + assert "layers.1.mlp.experts.{expert}.down_proj" in registry.slots + assert all(not parameter.requires_grad for name, parameter in model.named_parameters() if "lora_" not in name) + assert all(parameter.requires_grad for parameter in registry.parameters()) + + +def test_bailing_v3_requires_non_factorized_kda(monkeypatch) -> None: + monkeypatch.setattr(linear, "get_tp_context", _single_tp) + monkeypatch.setattr("areno.adapters.lora.get_tp_context", _single_tp) + + with pytest.raises(ValueError, match="no_kda_lora=true"): + initialize_lora( + _BailingModel(no_kda_lora=False), + LoraConfig(target_modules=("q_proj",)), + seed=42, + ) From 2865aad5c4b82d32a0578d63db4585cc99fb5935 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 22:02:58 +0800 Subject: [PATCH 13/29] style(lora): order Bailing test imports --- tests/test_bailing_v3_lora_cpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index b0b2a62d..f80fbe65 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -2,8 +2,8 @@ from types import SimpleNamespace -import torch import pytest +import torch from torch import nn from areno.adapters.config import BAILING_V3_TARGETS, LoraConfig From 1367b9c23d3ac787995bf351f7f6cbfad2bdc98e Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 22:09:59 +0800 Subject: [PATCH 14/29] docs(lora): generalize target option help --- areno/cli/serve.py | 2 +- areno/cli/train.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/areno/cli/serve.py b/areno/cli/serve.py index 64a0fc19..7962108e 100644 --- a/areno/cli/serve.py +++ b/areno/cli/serve.py @@ -996,7 +996,7 @@ def _normalize_stop(stop: str | list[str] | None) -> list[str]: "--lora-target-modules", default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", show_default=True, - help="Comma-separated dense Qwen3 projection targets.", + help="Comma-separated native projection targets.", ) @click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter to serve.") def serve_command( diff --git a/areno/cli/train.py b/areno/cli/train.py index 375a47ae..7b89a27e 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -1642,7 +1642,7 @@ def _dataset_builder_for_suffix(suffix: str) -> str: "--lora-target-modules", default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", show_default=True, - help="Comma-separated Qwen3 projection targets (MoE MLP targets apply to each routed expert).", + help="Comma-separated native projection targets (MoE MLP targets apply to each routed expert).", ) @click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter used to initialize native LoRA.") @click.option( From 234f9acb8ecf2ab32d73eb3d88407c1957af827a Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 19 Aug 2026 22:42:01 +0800 Subject: [PATCH 15/29] fix(attention): trim padded native decode values --- .../engine/layers/attention_backend/infer.py | 1 + tests/test_config_data_cpu.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/areno/engine/layers/attention_backend/infer.py b/areno/engine/layers/attention_backend/infer.py index 7551c56c..f1179d39 100644 --- a/areno/engine/layers/attention_backend/infer.py +++ b/areno/engine/layers/attention_backend/infer.py @@ -128,6 +128,7 @@ def forward( window_size=call.window_size, softmax_scale=call.softmax_scale, ) + out = call.trim_value_dim(out) return out.view(q.shape[0], q.shape[1], q.shape[2], call.value_dim) require_flash_attention_supported(call, mode="decode attention") # When value head dim < cache head dim we pad to match the cache diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 94affcdb..60440327 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -517,6 +517,31 @@ def fake_native_prefill(q_arg, k_arg, v_arg, meta_arg, window_size, softmax_scal self.assertIs(captured["meta"], meta) self.assertEqual(tuple(out.shape), (1, 2, 2, 4)) + def test_native_decode_pads_value_dim_and_trims_output(self): + """Native decode should return the original V dim after using a QK-sized cache.""" + backend = FlashAttnInferBackend("native") + q = torch.zeros(1, 1, 2, 6) + k = torch.zeros(1, 1, 2, 6) + v = torch.zeros(1, 1, 2, 4) + k_cache = torch.zeros(1, 2, 2, 6) + v_cache = torch.zeros(1, 2, 2, 6) + meta = InferMeta( + mode="decode", + cache_seqlens=torch.tensor([1], dtype=torch.int32), + block_table=torch.zeros(1, 1, dtype=torch.int32), + ) + captured = {} + + def fake_native_decode(**kwargs): + captured["v_update_shape"] = tuple(kwargs["v_update"].shape) + return torch.ones_like(kwargs["q"]) + + with patch("areno.engine.layers.attention_backend.infer._native_decode", fake_native_decode): + out = backend(q, k, v, k_cache, v_cache, meta) + + self.assertEqual(captured["v_update_shape"], (1, 2, 6)) + self.assertEqual(tuple(out.shape), (1, 1, 2, 4)) + def test_native_attention_backend_does_not_require_flash_attn_import(self): """Native train/infer backends should construct without flash-attn installed.""" from areno.engine.layers.attention_backend.infer import build_infer_attention_backend From c6327966f0bd661e5b54974e6c3132b936f36483 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Thu, 20 Aug 2026 16:54:39 +0800 Subject: [PATCH 16/29] feat(lora): complete Bailing native adapter contracts --- areno/adapters/config.py | 15 ++- areno/adapters/lora.py | 2 +- areno/adapters/peft.py | 31 +++++ areno/cli/serve.py | 5 +- areno/cli/train.py | 5 +- areno/engine/api.py | 1 + areno/engine/config.py | 7 ++ areno/engine/worker.py | 3 +- areno/models/bailing_v3/model.py | 14 ++- tests/test_bailing_v3_lora_cpu.py | 65 ++++++++++- tests/test_config_data_cpu.py | 16 +++ tests/test_peft_adapter_cpu.py | 180 ++++++++++++++++++++++++++++++ 12 files changed, 335 insertions(+), 9 deletions(-) create mode 100644 tests/test_peft_adapter_cpu.py diff --git a/areno/adapters/config.py b/areno/adapters/config.py index fb114892..885c9da7 100644 --- a/areno/adapters/config.py +++ b/areno/adapters/config.py @@ -20,6 +20,8 @@ "q_proj", "k_proj", "v_proj", + "f_proj", + "g_proj", "o_proj", "q_a_proj", "q_b_proj", @@ -81,7 +83,18 @@ def _read_adapter_config(path: str) -> dict: unsupported.append("bias") if bool(adapter_config.get("fan_in_fan_out", False)): unsupported.append("fan_in_fan_out") - for option in ("use_rslora", "use_dora", "rank_pattern", "alpha_pattern", "modules_to_save"): + for option in ( + "use_rslora", + "use_dora", + "rank_pattern", + "alpha_pattern", + "modules_to_save", + "alora_invocation_tokens", + "layer_replication", + "trainable_token_indices", + "target_parameters", + "use_qalora", + ): if adapter_config.get(option): unsupported.append(option) if unsupported: diff --git a/areno/adapters/lora.py b/areno/adapters/lora.py index f17db3ec..8c91b5fc 100644 --- a/areno/adapters/lora.py +++ b/areno/adapters/lora.py @@ -333,7 +333,7 @@ def _initialize_bailing_v3_lora( attention = layer.attention attention_prefix = f"{prefix}.attention" if hasattr(attention, "q_conv1d_weight"): - for component in ("q_proj", "k_proj", "v_proj"): + for component in ("q_proj", "k_proj", "v_proj", "f_proj", "g_proj"): if component in requested: matched.add(component) _install_column_slot( diff --git a/areno/adapters/peft.py b/areno/adapters/peft.py index 17f1b85c..b5e2b1b9 100644 --- a/areno/adapters/peft.py +++ b/areno/adapters/peft.py @@ -22,6 +22,21 @@ def load_peft_adapter(registry: AdapterRegistry, path: str | Path) -> None: input_path = Path(path) tensors = load_file(input_path / "adapter_model.safetensors", device="cpu") ctx = get_tp_context() + expected_shapes = _expected_peft_shapes(registry, ctx.world_size) + actual_keys = set(tensors) + expected_keys = set(expected_shapes) + if actual_keys != expected_keys: + missing = sorted(expected_keys - actual_keys) + unexpected = sorted(actual_keys - expected_keys) + raise ValueError( + "PEFT adapter tensor keys do not match the native LoRA registry: " + f"missing={missing[:3]}, unexpected={unexpected[:3]}" + ) + for key, expected_shape in expected_shapes.items(): + actual_shape = tuple(tensors[key].shape) + if actual_shape != expected_shape: + raise ValueError(f"PEFT adapter tensor {key!r} has shape {actual_shape}, expected {expected_shape}") + for logical_name, slot in registry.slots.items(): if isinstance(slot, RoutedExpertLoraSlot): for local_expert_id in range(slot.local_num_experts): @@ -142,3 +157,19 @@ def _gather_replicated_column(slot: LoraSlot, gathered: list[torch.Tensor], worl def _key(logical_name: str, component: str) -> str: return f"{_PREFIX}{logical_name}.lora_{component}.weight" + + +def _expected_peft_shapes(registry: AdapterRegistry, tp_size: int) -> dict[str, tuple[int, ...]]: + """Return the one canonical PEFT tensor contract represented by a registry.""" + + shapes: dict[str, tuple[int, ...]] = {} + for logical_name, slot in registry.slots.items(): + if isinstance(slot, RoutedExpertLoraSlot): + for expert_id in range(slot.local_num_experts * tp_size): + expert_name = logical_name.format(expert=expert_id) + shapes[_key(expert_name, "A")] = (slot.rank, slot.in_features) + shapes[_key(expert_name, "B")] = (slot.out_features, slot.rank) + continue + shapes[_key(logical_name, "A")] = (slot.rank, slot.global_in_features) + shapes[_key(logical_name, "B")] = (slot.global_out_features, slot.rank) + return shapes diff --git a/areno/cli/serve.py b/areno/cli/serve.py index 7962108e..aa9aeb9d 100644 --- a/areno/cli/serve.py +++ b/areno/cli/serve.py @@ -996,7 +996,10 @@ def _normalize_stop(stop: str | list[str] | None) -> list[str]: "--lora-target-modules", default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", show_default=True, - help="Comma-separated native projection targets.", + help=( + "Comma-separated native projection targets (selected Bailing V3 KDA q/k/v/f/g projections " + "use independent canonical adapters)." + ), ) @click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter to serve.") def serve_command( diff --git a/areno/cli/train.py b/areno/cli/train.py index 7b89a27e..9ba0b3ef 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -1642,7 +1642,10 @@ def _dataset_builder_for_suffix(suffix: str) -> str: "--lora-target-modules", default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj", show_default=True, - help="Comma-separated native projection targets (MoE MLP targets apply to each routed expert).", + help=( + "Comma-separated native projection targets (MoE MLP targets apply to each routed expert; " + "selected Bailing V3 KDA q/k/v/f/g projections use independent canonical adapters)." + ), ) @click.option("--lora-adapter-path", default=None, help="Standard PEFT adapter used to initialize native LoRA.") @click.option( diff --git a/areno/engine/api.py b/areno/engine/api.py index 01a18619..71b52187 100644 --- a/areno/engine/api.py +++ b/areno/engine/api.py @@ -233,6 +233,7 @@ def from_pretrained( cfg = EngineConfig( model=model_config, model_path=model_path, + base_model_name_or_path=model, train_loss_fn=loss_fn, tp_size=tp_size, sequence_parallel=sequence_parallel, diff --git a/areno/engine/config.py b/areno/engine/config.py index 0d043973..f6d6d5eb 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -285,6 +285,7 @@ class EngineConfig: lora: LoraConfig | None = None lora_seed: int = 0 reference_mode: Literal["independent", "reuse_actor_base"] = "independent" + base_model_name_or_path: str | None = None def __post_init__(self) -> None: """Infer DP/devices and validate the distributed layout.""" @@ -296,6 +297,12 @@ def __post_init__(self) -> None: raise ValueError("reference_mode must be one of: independent, reuse_actor_base") if self.reference_mode == "reuse_actor_base" and self.lora is None: raise ValueError("reference_mode='reuse_actor_base' requires native LoRA") + if ( + self.lora is not None + and self.model.model_type == "bailing_moe_v3" + and self.model.moe_router_bias_update_rate != 0.0 + ): + raise ValueError("native LoRA requires moe_router_bias_update_rate=0 to keep the base policy frozen") if self.devices is None: if torch.cuda.is_available(): device_count = torch.cuda.device_count() diff --git a/areno/engine/worker.py b/areno/engine/worker.py index ac7ad63b..5e33c715 100644 --- a/areno/engine/worker.py +++ b/areno/engine/worker.py @@ -664,10 +664,11 @@ def export_adapter(self, payload: ExportAdapterPayload) -> dict | None: path = export_peft_adapter( self.adapter_registry, payload.path, - base_model_name_or_path=self.config.model_path, + base_model_name_or_path=(self.config.base_model_name_or_path or self.config.model_path), ) return {"path": path} if path is not None else None + def _rollout_payloads_compatible(first: RolloutPayload, other: RolloutPayload) -> bool: """Return whether two rollout payloads can share one InferenceBatchState.""" diff --git a/areno/models/bailing_v3/model.py b/areno/models/bailing_v3/model.py index 41cb5b26..269fa06d 100644 --- a/areno/models/bailing_v3/model.py +++ b/areno/models/bailing_v3/model.py @@ -447,9 +447,17 @@ def _forward_fused_permute( self.local_num_experts, ) if x.shape[0] == 0: - # No tokens routed to this rank — still need to all_reduce to keep - # collective sync with peers. - return all_reduce(flat.new_zeros(flat.shape)) + # Every TP/DP replica must produce gradients for the same parameter + # set even when this rank owns no active routes. + zero = ( + self.linear_fc1.weight.reshape(-1)[0] * 0 + + self.linear_fc2.weight.reshape(-1)[0] * 0 + + topk_weight.sum().to(dtype=self.linear_fc1.weight.dtype) * 0 + ) + if self.has_active_lora(): + for slot in self.lora_slots.values(): + zero = zero + slot.lora_A.reshape(-1)[0] * 0 + slot.lora_B.reshape(-1)[0] * 0 + return all_reduce(flat.new_zeros(flat.shape) + zero) hidden, _ = _grouped_linear_forward(self.linear_fc1, x.contiguous(), tokens_per_expert) if self.has_active_lora(): gate, up = hidden.chunk(2, dim=-1) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index f80fbe65..3997d22a 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -34,6 +34,9 @@ def __init__(self) -> None: self.q_proj = ColumnParallelLinear(8, 8, bias=False) self.k_proj = ColumnParallelLinear(8, 8, bias=False) self.v_proj = ColumnParallelLinear(8, 8, bias=False) + self.f_proj = ColumnParallelLinear(8, 8, bias=False) + self.g_proj = ColumnParallelLinear(8, 8, bias=False) + self.b_proj = ColumnParallelLinear(8, 1, bias=False) self.o_proj = RowParallelLinear(8, 8, bias=False) @@ -62,6 +65,17 @@ def install_lora_component(self, component: str, slot: nn.Module) -> None: self.lora_slots[component] = slot +class _ActiveLoraSlot(nn.Module): + def __init__(self) -> None: + super().__init__() + self.lora_A = nn.Parameter(torch.randn(1, 3)) + self.lora_B = nn.Parameter(torch.randn(4, 1)) + + @property + def enabled(self) -> bool: + return True + + class _SparseMLP(nn.Module): def __init__(self) -> None: super().__init__() @@ -100,6 +114,8 @@ def test_bailing_v3_full_profile_attaches_native_slots(monkeypatch) -> None: "q_proj", "k_proj", "v_proj", + "f_proj", + "g_proj", "o_proj", "q_a_proj", "q_b_proj", @@ -114,7 +130,13 @@ def test_bailing_v3_full_profile_attaches_native_slots(monkeypatch) -> None: registry = initialize_lora(model, LoraConfig(rank=4, alpha=4, target_modules=profile), seed=42) assert set(profile) <= set(BAILING_V3_TARGETS) - assert "layers.0.attention.q_proj" in registry.slots + kda_names = { + f"layers.0.attention.{component}" + for component in ("q_proj", "k_proj", "v_proj", "f_proj", "g_proj") + } + assert kda_names <= registry.slots.keys() + assert len({id(registry.slots[name]) for name in kda_names}) == len(kda_names) + assert "layers.0.attention.b_proj" not in registry.slots assert "layers.1.attention.q_a_proj" in registry.slots assert "layers.1.mlp.shared_experts.gate_proj" in registry.slots assert "layers.1.mlp.experts.{expert}.down_proj" in registry.slots @@ -132,3 +154,44 @@ def test_bailing_v3_requires_non_factorized_kda(monkeypatch) -> None: LoraConfig(target_modules=("q_proj",)), seed=42, ) + + +def test_bailing_v3_empty_route_keeps_expert_router_and_lora_gradients(monkeypatch) -> None: + from areno.models.bailing_v3 import model as bailing_model + + experts = bailing_model.BailingGroupedExperts.__new__(bailing_model.BailingGroupedExperts) + nn.Module.__init__(experts) + experts.linear_fc1 = nn.Linear(3, 4, bias=False) + experts.linear_fc2 = nn.Linear(2, 3, bias=False) + experts.lora_slots = nn.ModuleDict({"gate_proj": _ActiveLoraSlot()}) + experts.local_expert_start = 0 + experts.local_num_experts = 1 + + flat = torch.randn(2, 3, requires_grad=True) + router_weight = nn.Parameter(torch.randn(2, 1)) + topk_weight = router_weight.sigmoid() + topk_idx = torch.ones(2, 1, dtype=torch.long) + empty = flat.new_empty((0, 3)) + monkeypatch.setattr( + bailing_model, + "_areno_moe_topk_permute_no_compile", + lambda *args: ( + empty, + flat.new_empty((0,)), + torch.empty(0, dtype=torch.long), + torch.zeros(1, dtype=torch.long), + ), + ) + monkeypatch.setattr(bailing_model, "all_reduce", lambda value: value) + + experts(flat, topk_idx, topk_weight).sum().backward() + + parameters = ( + experts.linear_fc1.weight, + experts.linear_fc2.weight, + router_weight, + experts.lora_slots["gate_proj"].lora_A, + experts.lora_slots["gate_proj"].lora_B, + ) + assert all(parameter.grad is not None for parameter in parameters) + assert all(torch.count_nonzero(parameter.grad) == 0 for parameter in parameters) diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 60440327..c64fc347 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -218,6 +218,22 @@ def test_engine_config_allows_replicated_kv_lora_targets(self): lora=LoraConfig(), ) + def test_engine_config_rejects_router_bias_updates_only_for_lora(self): + """LoRA keeps router bias in the frozen base while fullweight may update it.""" + model = ModelConfig( + model_type="bailing_moe_v3", + num_attention_heads=4, + num_key_value_heads=4, + intermediate_size=16, + vocab_size=32, + moe_router_bias_update_rate=1e-3, + ) + + with self.assertRaisesRegex(ValueError, "moe_router_bias_update_rate=0"): + EngineConfig(model=model, tp_size=1, devices=[0], lora=LoraConfig()) + + EngineConfig(model=model, tp_size=1, devices=[0]) + def test_reference_view_requires_lora_at_engine_boundary(self): """The actor base can be reused only when the actor owns a native adapter.""" model = ModelConfig(num_attention_heads=4, num_key_value_heads=4, intermediate_size=16, vocab_size=32) diff --git a/tests/test_peft_adapter_cpu.py b/tests/test_peft_adapter_cpu.py new file mode 100644 index 00000000..eab37a3f --- /dev/null +++ b/tests/test_peft_adapter_cpu.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import json + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +import areno.engine.api as engine_api +from areno.adapters import LoraConfig +from areno.adapters.lora import AdapterRegistry, LoraSlot, RoutedExpertLoraSlot, _AdapterRuntimeState +from areno.adapters.peft import export_peft_adapter, load_peft_adapter +from areno.engine.config import ModelConfig, RuntimeConfig +from areno.engine.parallel.context import TPContext, set_tp_context + +_PREFIX = "base_model.model.model." + + +def _set_cpu_tp(*, rank: int = 0, world_size: int = 1) -> None: + set_tp_context(TPContext(rank=rank, world_size=world_size, device=torch.device("cpu"), group=None)) + + +def _dense_registry() -> AdapterRegistry: + _set_cpu_tp() + state = _AdapterRuntimeState() + config = LoraConfig(rank=2, alpha=4.0, target_modules=("q_proj",)) + slot = LoraSlot( + logical_name="layers.0.self_attn.q_proj", + base_weight=nn.Parameter(torch.zeros(1)), + global_in_features=4, + global_out_features=3, + local_in_features=4, + local_out_features=3, + row_parallel=False, + config=config, + seed=1, + runtime_state=state, + ) + return AdapterRegistry({slot.logical_name: slot}, config, state) + + +def _expert_registry() -> AdapterRegistry: + _set_cpu_tp(rank=1, world_size=2) + state = _AdapterRuntimeState() + config = LoraConfig(rank=2, alpha=4.0, target_modules=("down_proj",)) + slot = RoutedExpertLoraSlot( + logical_name="layers.0.mlp.experts.{expert}.down_proj", + base_weight=nn.Parameter(torch.zeros(1)), + local_num_experts=2, + local_expert_start=2, + in_features=3, + out_features=4, + config=config, + seed=1, + runtime_state=state, + ) + return AdapterRegistry({slot.logical_name: slot}, config, state) + + +def _key(logical_name: str, component: str) -> str: + return f"{_PREFIX}{logical_name}.lora_{component}.weight" + + +def _dense_state() -> dict[str, torch.Tensor]: + logical_name = "layers.0.self_attn.q_proj" + return { + _key(logical_name, "A"): torch.arange(8, dtype=torch.float32).reshape(2, 4), + _key(logical_name, "B"): torch.arange(6, dtype=torch.float32).reshape(3, 2), + } + + +def test_lora_config_rejects_semantic_modifiers_and_allows_neutral_values(tmp_path) -> None: + base_config = { + "peft_type": "LORA", + "r": 2, + "lora_alpha": 4, + "lora_dropout": 0, + "bias": "none", + "target_modules": ["q_proj"], + } + modifiers = { + "alora_invocation_tokens": [1, 2], + "layer_replication": [[0, 1]], + "trainable_token_indices": [0], + "target_parameters": ["layers.0.weight"], + "use_qalora": True, + } + for modifier, value in modifiers.items(): + config = {**base_config, modifier: value} + (tmp_path / "adapter_config.json").write_text(json.dumps(config), encoding="utf-8") + with pytest.raises(ValueError, match=modifier): + LoraConfig(adapter_path=str(tmp_path)) + + neutral_config = { + **base_config, + "alora_invocation_tokens": None, + "layer_replication": [], + "trainable_token_indices": None, + "target_parameters": [], + "use_qalora": False, + "qalora_group_size": 16, + } + (tmp_path / "adapter_config.json").write_text(json.dumps(neutral_config), encoding="utf-8") + + loaded = LoraConfig(adapter_path=str(tmp_path)) + + assert loaded.target_modules == ("q_proj",) + + +def test_peft_load_requires_exact_expert_key_set(tmp_path) -> None: + registry = _expert_registry() + state = {} + for expert_id in range(4): + logical_name = f"layers.0.mlp.experts.{expert_id}.down_proj" + state[_key(logical_name, "A")] = torch.full((2, 3), float(expert_id + 1)) + state[_key(logical_name, "B")] = torch.full((4, 2), float(expert_id + 2)) + save_file(state, tmp_path / "adapter_model.safetensors") + load_peft_adapter(registry, tmp_path) + + slot = registry.slots["layers.0.mlp.experts.{expert}.down_proj"] + torch.testing.assert_close(slot.lora_A[0], state[_key("layers.0.mlp.experts.2.down_proj", "A")]) + torch.testing.assert_close(slot.lora_B[1], state[_key("layers.0.mlp.experts.3.down_proj", "B")]) + + del state[_key("layers.0.mlp.experts.3.down_proj", "B")] + state[_key("layers.0.mlp.experts.4.down_proj", "A")] = torch.zeros(2, 3) + (tmp_path / "adapter_model.safetensors").unlink() + save_file(state, tmp_path / "adapter_model.safetensors") + + with pytest.raises(ValueError, match="tensor keys"): + load_peft_adapter(registry, tmp_path) + + +def test_peft_load_rejects_broadcastable_shape_before_copy(tmp_path) -> None: + registry = _dense_registry() + slot = registry.slots["layers.0.self_attn.q_proj"] + original_A = slot.lora_A.detach().clone() + original_B = slot.lora_B.detach().clone() + state = _dense_state() + state[_key("layers.0.self_attn.q_proj", "A")] = torch.full((2, 1), 9.0) + state[_key("layers.0.self_attn.q_proj", "B")].fill_(7.0) + save_file(state, tmp_path / "adapter_model.safetensors") + + with pytest.raises(ValueError, match="has shape"): + load_peft_adapter(registry, tmp_path) + + torch.testing.assert_close(slot.lora_A, original_A) + torch.testing.assert_close(slot.lora_B, original_B) + + +def test_export_preserves_model_reference_before_resolution(tmp_path, monkeypatch) -> None: + resolved_path = "/cache/models--example--base/snapshots/revision" + monkeypatch.setattr(engine_api, "resolve_model_path", lambda _model: resolved_path) + monkeypatch.setattr(engine_api, "config_from_hf", lambda _path: ModelConfig()) + + def fake_init(self, config, **_kwargs): + self.config = config + + monkeypatch.setattr(engine_api.ArenoEngine, "__init__", fake_init) + engine = engine_api.ArenoEngine.from_pretrained( + "example/base", + devices=[0], + start=False, + loss_fn=lambda _pack, logprobs: logprobs.sum(), + runtime_config=RuntimeConfig(attn_backend="native", compile_model=False), + ) + + assert engine.config.model_path == resolved_path + assert engine.config.base_model_name_or_path == "example/base" + + registry = _dense_registry() + exported = export_peft_adapter( + registry, + tmp_path, + base_model_name_or_path=engine.config.base_model_name_or_path, + ) + adapter_config = json.loads((tmp_path / "adapter_config.json").read_text(encoding="utf-8")) + + assert exported == str(tmp_path) + assert adapter_config["base_model_name_or_path"] == "example/base" From a56409fe432e88938cc2f25144ac7436658d7d47 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Fri, 21 Aug 2026 10:40:50 +0800 Subject: [PATCH 17/29] fix(lora): correct replicated norm and base metadata --- areno/api/backend/cuda/backend.py | 3 +++ areno/api/config.py | 5 +++++ areno/api/trainer_config.py | 5 ++++- areno/cli/model_refs.py | 2 ++ areno/cli/serve.py | 8 ++++++++ areno/engine/api.py | 3 ++- areno/engine/runtime/train_step.py | 20 ++++++++++++++++---- tests/test_cli_model_refs_cpu.py | 12 ++++++++++++ tests/test_config_data_cpu.py | 23 +++++++++++++++++++++++ tests/test_peft_adapter_cpu.py | 3 ++- tests/test_serve_cli_cpu.py | 2 ++ 11 files changed, 79 insertions(+), 7 deletions(-) diff --git a/areno/api/backend/cuda/backend.py b/areno/api/backend/cuda/backend.py index 80a4bfca..328f7937 100644 --- a/areno/api/backend/cuda/backend.py +++ b/areno/api/backend/cuda/backend.py @@ -170,6 +170,7 @@ def initialize(self, ctx: Context): policy_sync_bucket_mb=cfg.policy_sync_bucket_mb, lora_config=cfg.lora, reference_mode=cfg.reference_mode, + base_model_name_or_path=cfg.base_model_name_or_path, ) return self._policy_sync_bucket_bytes = cfg.policy_sync_bucket_mb * 1024 * 1024 @@ -225,6 +226,7 @@ def initialize(self, ctx: Context): role="train", lora_config=cfg.lora, reference_mode=cfg.reference_mode, + base_model_name_or_path=cfg.base_model_name_or_path, cluster_kwargs={"world_spec": world_spec, "partition": train_partition}, **common, ) @@ -239,6 +241,7 @@ def initialize(self, ctx: Context): loss_fn=None, role="rollout", lora_config=cfg.lora, + base_model_name_or_path=cfg.base_model_name_or_path, policy_sync_bucket_mb=cfg.policy_sync_bucket_mb, start=False, cluster_kwargs={"world_spec": world_spec, "partition": rollout_partition}, diff --git a/areno/api/config.py b/areno/api/config.py index 3f37485b..cfd5b27b 100644 --- a/areno/api/config.py +++ b/areno/api/config.py @@ -21,9 +21,14 @@ class CudaConfig: The `optimizer` and `runtime` dicts are passed verbatim to the engine's `OptimizerConfig`/`RuntimeConfig` so any new tuning knob can be added without changing this file. + + `base_model_name_or_path` keeps the caller-facing model reference for + portable PEFT metadata when `model_path` has already resolved to a local + cache path. """ model_path: str | None = None + base_model_name_or_path: str | None = field(default=None, kw_only=True) tp_size: int = 1 sequence_parallel: bool | None = None dp_size: int | None = None diff --git a/areno/api/trainer_config.py b/areno/api/trainer_config.py index 940b227d..ad056774 100644 --- a/areno/api/trainer_config.py +++ b/areno/api/trainer_config.py @@ -11,7 +11,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Literal from areno.adapters.config import LoraConfig @@ -30,6 +30,7 @@ class TrainerConfig: ckpt: str dataset_path: str backend: str | None = None + base_model_name_or_path: str | None = field(default=None, kw_only=True) model_hub: str = "modelscope" dataset_loader_fn: str | None = None save_path: str | None = None @@ -203,6 +204,7 @@ def cuda_config(self): from areno.api.config import CudaConfig return CudaConfig( + base_model_name_or_path=self.base_model_name_or_path, tp_size=self.tp_size, sequence_parallel=self.sequence_parallel, devices=self.train_devices, @@ -248,6 +250,7 @@ def cuda_config(self): from areno.api.config import CudaConfig return CudaConfig( + base_model_name_or_path=self.base_model_name_or_path, tp_size=self.tp_size, sequence_parallel=self.sequence_parallel, devices=self.train_devices, diff --git a/areno/cli/model_refs.py b/areno/cli/model_refs.py index 44192e1a..6feed1e1 100644 --- a/areno/cli/model_refs.py +++ b/areno/cli/model_refs.py @@ -46,6 +46,8 @@ def resolve_model_refs_for_config(config: ConfigT) -> ConfigT: cache: dict[str, str] = {} model_hub = str(getattr(config, "model_hub", "modelscope")) + if getattr(config, "base_model_name_or_path", None) is None: + config.base_model_name_or_path = config.ckpt config.ckpt = resolve_model_ref(config.ckpt, cache, model_hub=model_hub) algo = str(getattr(config, "algo", "")).lower() if algo == "dpo" and getattr(config, "ref_ckpt", None) is not None: diff --git a/areno/cli/serve.py b/areno/cli/serve.py index aa9aeb9d..28d13c41 100644 --- a/areno/cli/serve.py +++ b/areno/cli/serve.py @@ -206,6 +206,7 @@ def __init__( eager_decode: bool, attn_backend: str, lora: LoraConfig | None, + base_model_name_or_path: str | None, ) -> None: from areno.engine.config import RuntimeConfig @@ -221,6 +222,7 @@ def __init__( runtime_config=RuntimeConfig(eager_decode=bool(eager_decode), attn_backend=attn_backend), loss_fn=_serve_loss_fn, lora_config=lora, + base_model_name_or_path=base_model_name_or_path, ) self.max_model_len = int(self._engine.config.model.max_position_embeddings) @@ -310,6 +312,7 @@ def _create_serve_runtime( eager_decode: bool, attn_backend: str, lora: LoraConfig | None, + base_model_name_or_path: str | None, ) -> _CudaServeRuntime | _MlxServeRuntime: if backend_type == MLX: if lora is not None: @@ -329,6 +332,7 @@ def _create_serve_runtime( eager_decode=eager_decode, attn_backend=attn_backend, lora=lora, + base_model_name_or_path=base_model_name_or_path, ) @@ -344,6 +348,7 @@ def create_app( attn_backend: Literal["flash", "native"] = "flash", chat_template_enable_thinking: bool | None = None, lora: LoraConfig | None = None, + base_model_name_or_path: str | None = None, ) -> FastAPI: """Construct the FastAPI app: load tokenizer/engine, install routes and lifecycle hooks.""" if world_size < 1: @@ -374,6 +379,7 @@ def create_app( eager_decode=eager_decode, attn_backend=attn_backend, lora=lora, + base_model_name_or_path=base_model_name_or_path, ) if backend_type == MLX: tokenizer = engine.tokenizer @@ -1024,6 +1030,7 @@ def serve_command( """Click entry point: build the app and hand it to uvicorn.""" import uvicorn + base_model_name_or_path = model_path model_path = resolve_model_ref(model_path, model_hub=model_hub) lora = None if lora_rank is not None or lora_adapter_path is not None: @@ -1064,6 +1071,7 @@ def serve_command( attn_backend=attn_backend, chat_template_enable_thinking=False if disable_thinking else None, lora=lora, + base_model_name_or_path=base_model_name_or_path, ) uvicorn.run(app, host=host, port=port) diff --git a/areno/engine/api.py b/areno/engine/api.py index 71b52187..818bab11 100644 --- a/areno/engine/api.py +++ b/areno/engine/api.py @@ -213,6 +213,7 @@ def from_pretrained( policy_sync_bucket_mb: int = 64, lora_config: LoraConfig | None = None, reference_mode: str = "independent", + base_model_name_or_path: str | None = None, ) -> ArenoEngine: """Build an engine by reading model config from a checkpoint path. @@ -233,7 +234,7 @@ def from_pretrained( cfg = EngineConfig( model=model_config, model_path=model_path, - base_model_name_or_path=model, + base_model_name_or_path=(model if base_model_name_or_path is None else base_model_name_or_path), train_loss_fn=loss_fn, tp_size=tp_size, sequence_parallel=sequence_parallel, diff --git a/areno/engine/runtime/train_step.py b/areno/engine/runtime/train_step.py index 5aa956ce..bed2a27d 100644 --- a/areno/engine/runtime/train_step.py +++ b/areno/engine/runtime/train_step.py @@ -413,19 +413,31 @@ def _grads_for_norm(parameters): """Yield params whose grads should contribute to the global TP grad norm. TP-sharded params contribute on every rank (they are different slices). - Replicated params only contribute on rank 0 so the TP all-reduce does not - double-count them. + Replicated params contribute from one owner per unique output shard so the + TP all-reduce does not double-count replicas. """ ctx = get_tp_context() for param in parameters: if _param_grad(param) is None: continue - is_tp_parallel = bool(getattr(param, "tensor_model_parallel", False)) - if is_tp_parallel or ctx.rank == 0: + if _is_tp_norm_owner(param, ctx): yield param +def _is_tp_norm_owner(param: torch.nn.Parameter, ctx) -> bool: + """Select one owner for replicated output ranges and every true TP shard.""" + + output_range = getattr(param, "tp_replicated_output_range", None) + if output_range is not None: + start, end, global_size = output_range + local_size = end - start + unique_shards = global_size // local_size + ranks_per_shard = ctx.world_size // unique_shards + return ctx.rank == (start // local_size) * ranks_per_shard + return bool(getattr(param, "tensor_model_parallel", False)) or ctx.rank == 0 + + def _param_grad(param: torch.nn.Parameter) -> torch.Tensor | None: """Prefer the FP32 `main_grad` accumulator used by the optimizer split.""" diff --git a/tests/test_cli_model_refs_cpu.py b/tests/test_cli_model_refs_cpu.py index 2cde5e21..d4eb9e9d 100644 --- a/tests/test_cli_model_refs_cpu.py +++ b/tests/test_cli_model_refs_cpu.py @@ -7,6 +7,7 @@ from types import SimpleNamespace from unittest.mock import patch +from areno.api.trainer_config import TrainerConfig from areno.cli import model_refs @@ -87,6 +88,17 @@ def test_resolve_model_refs_for_config_honors_modelscope_for_all_roles(self): self.assertEqual(resolved.reward_ckpt, "/ms/org/reward") self.assertEqual(calls, ["org/actor", "org/reward"]) + def test_resolve_model_refs_preserves_actor_reference_for_adapter_metadata(self): + fake_modelscope = types.SimpleNamespace(snapshot_download=lambda repo: f"/ms/{repo}") + config = TrainerConfig(algo="sft", ckpt="org/actor", dataset_path="/data") + + with patch.dict(sys.modules, {"modelscope": fake_modelscope}): + resolved = model_refs.resolve_model_refs_for_config(config) + + self.assertEqual(resolved.ckpt, "/ms/org/actor") + self.assertEqual(resolved.base_model_name_or_path, "org/actor") + self.assertEqual(resolved.areno_config().base_model_name_or_path, "org/actor") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index c64fc347..1dd54672 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -33,6 +33,7 @@ ) from areno.engine.layers.attention_backend.infer import FlashAttnInferBackend, _native_prefill from areno.engine.layers.attention_backend.train import _native_train, _native_train_areno +from areno.engine.runtime import train_step as train_step_runtime from areno.engine.runtime.metadata import InferMeta, TrainMeta from areno.engine.training import _actor_train_model @@ -234,6 +235,28 @@ def test_engine_config_rejects_router_bias_updates_only_for_lora(self): EngineConfig(model=model, tp_size=1, devices=[0]) + def test_replicated_output_ranges_have_one_grad_norm_owner(self): + """Replicated B ranges contribute once while true TP shards contribute everywhere.""" + + def owners(output_range): + selected = [] + for rank in range(4): + parameter = torch.nn.Parameter(torch.ones(2, 2)) + parameter.grad = torch.ones_like(parameter) + parameter.tensor_model_parallel = True + if output_range is not None: + parameter.tp_replicated_output_range = output_range + ctx = SimpleNamespace(rank=rank, world_size=4) + with patch.object(train_step_runtime, "get_tp_context", return_value=ctx): + if list(train_step_runtime._grads_for_norm((parameter,))): + selected.append(rank) + return selected + + self.assertEqual(owners((0, 4, 4)), [0]) + self.assertEqual(owners((0, 2, 4)), [0]) + self.assertEqual(owners((2, 4, 4)), [2]) + self.assertEqual(owners(None), [0, 1, 2, 3]) + def test_reference_view_requires_lora_at_engine_boundary(self): """The actor base can be reused only when the actor owns a native adapter.""" model = ModelConfig(num_attention_heads=4, num_key_value_heads=4, intermediate_size=16, vocab_size=32) diff --git a/tests/test_peft_adapter_cpu.py b/tests/test_peft_adapter_cpu.py index eab37a3f..889e9235 100644 --- a/tests/test_peft_adapter_cpu.py +++ b/tests/test_peft_adapter_cpu.py @@ -158,11 +158,12 @@ def fake_init(self, config, **_kwargs): monkeypatch.setattr(engine_api.ArenoEngine, "__init__", fake_init) engine = engine_api.ArenoEngine.from_pretrained( - "example/base", + resolved_path, devices=[0], start=False, loss_fn=lambda _pack, logprobs: logprobs.sum(), runtime_config=RuntimeConfig(attn_backend="native", compile_model=False), + base_model_name_or_path="example/base", ) assert engine.config.model_path == resolved_path diff --git a/tests/test_serve_cli_cpu.py b/tests/test_serve_cli_cpu.py index 4aa1ebbc..9258f95d 100644 --- a/tests/test_serve_cli_cpu.py +++ b/tests/test_serve_cli_cpu.py @@ -35,10 +35,12 @@ def from_pretrained(cls, *args, **kwargs): decode_progress_interval_s=0.0, eager_decode=True, attn_backend="native", + base_model_name_or_path="org/base", ) assert captured["runtime_config"].eager_decode is True assert captured["runtime_config"].attn_backend == "native" + assert captured["base_model_name_or_path"] == "org/base" def test_create_app_can_disable_chat_template_thinking(monkeypatch): From 26e5cf2ffff126f8cb86c8d95a10a1c8ed3b046b Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Fri, 21 Aug 2026 11:01:36 +0800 Subject: [PATCH 18/29] test(serve): capture adapter base reference --- tests/test_serve_cli_cpu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_serve_cli_cpu.py b/tests/test_serve_cli_cpu.py index 9258f95d..9f18fdde 100644 --- a/tests/test_serve_cli_cpu.py +++ b/tests/test_serve_cli_cpu.py @@ -21,6 +21,7 @@ class FakeEngine: def from_pretrained(cls, *args, **kwargs): del args captured["runtime_config"] = kwargs["runtime_config"] + captured["base_model_name_or_path"] = kwargs["base_model_name_or_path"] return cls() monkeypatch.setattr(serve_mod, "load_tokenizer", lambda model_path: SimpleNamespace(eos_token_id=1)) From 2b239cbde91da66f4838f45cb9b61a26f68843cd Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Fri, 21 Aug 2026 11:06:52 +0800 Subject: [PATCH 19/29] fix(peft): expose stable base model reference --- areno/cli/serve.py | 9 ++++++++- areno/cli/train.py | 12 ++++++++++++ tests/test_config_data_cpu.py | 11 +++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/areno/cli/serve.py b/areno/cli/serve.py index 28d13c41..ee23c3d5 100644 --- a/areno/cli/serve.py +++ b/areno/cli/serve.py @@ -963,6 +963,11 @@ def _normalize_stop(stop: str | list[str] | None) -> list[str]: show_default=True, help="Remote hub for non-local model refs. Use 'modelscope' for ModelScope or 'hf' for Hugging Face.", ) +@click.option( + "--base-model-name-or-path", + default=None, + help="Stable base model reference associated with the PEFT adapter; defaults to the original model path.", +) @click.option("--tp-size", type=int, default=1, show_default=True, help="Tensor parallel size.") @click.option("--world-size", type=int, default=1, show_default=True, help="Total number of local worker ranks.") @click.option("--host", default="0.0.0.0", show_default=True, help="HTTP bind host.") @@ -1011,6 +1016,7 @@ def _normalize_stop(stop: str | list[str] | None) -> list[str]: def serve_command( model_path: str, model_hub: Literal["hf", "modelscope"], + base_model_name_or_path: str | None, tp_size: int, world_size: int, host: str, @@ -1030,7 +1036,8 @@ def serve_command( """Click entry point: build the app and hand it to uvicorn.""" import uvicorn - base_model_name_or_path = model_path + if base_model_name_or_path is None: + base_model_name_or_path = model_path model_path = resolve_model_ref(model_path, model_hub=model_hub) lora = None if lora_rank is not None or lora_adapter_path is not None: diff --git a/areno/cli/train.py b/areno/cli/train.py index 9ba0b3ef..71cff5a8 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -82,6 +82,7 @@ def flash_attention_unsupported_model_reason(model_config): ( "algo", "ckpt", + "base_model_name_or_path", "dataset_path", "model_hub", "dataset_loader_fn", @@ -232,6 +233,7 @@ def _trainer_config_from_options(**options) -> TrainerConfig: args.max_steps = getattr(args, "max_steps", None) args.score_micro_bs = getattr(args, "score_micro_bs", 8) args.model_hub = getattr(args, "model_hub", "modelscope") + args.base_model_name_or_path = getattr(args, "base_model_name_or_path", None) args.train_devices = getattr(args, "train_devices", None) args.sequence_parallel = getattr(args, "sequence_parallel", None) args.rollout_tp_size = getattr(args, "rollout_tp_size", None) @@ -828,6 +830,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: args.backend = getattr(args, "backend", None) args.score_micro_bs = getattr(args, "score_micro_bs", 8) args.model_hub = getattr(args, "model_hub", "modelscope") + args.base_model_name_or_path = getattr(args, "base_model_name_or_path", None) args.train_devices = getattr(args, "train_devices", None) args.sequence_parallel = getattr(args, "sequence_parallel", None) args.rollout_tp_size = getattr(args, "rollout_tp_size", None) @@ -853,6 +856,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: ckpt=args.ckpt, dataset_path=args.dataset_path, backend=args.backend, + base_model_name_or_path=args.base_model_name_or_path, model_hub=args.model_hub, dataset_loader_fn=args.dataset_loader_fn, save_path=args.save_path, @@ -912,6 +916,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: ckpt=args.ckpt, dataset_path=args.dataset_path, backend=args.backend, + base_model_name_or_path=args.base_model_name_or_path, model_hub=args.model_hub, dataset_loader_fn=args.dataset_loader_fn, save_path=args.save_path, @@ -969,6 +974,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: ckpt=args.ckpt, dataset_path=args.dataset_path, backend=args.backend, + base_model_name_or_path=args.base_model_name_or_path, model_hub=args.model_hub, dataset_loader_fn=args.dataset_loader_fn, reward_fn_path=args.reward_fn_path, @@ -1037,6 +1043,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: ckpt=args.ckpt, dataset_path=args.dataset_path, backend=args.backend, + base_model_name_or_path=args.base_model_name_or_path, model_hub=args.model_hub, dataset_loader_fn=args.dataset_loader_fn, reward_fn_path=args.reward_fn_path, @@ -1498,6 +1505,11 @@ def _dataset_builder_for_suffix(suffix: str) -> str: ) @click.option("--algo", type=str, default="gspo", show_default=True, help="Training algorithm registered in areno.api.") @click.option("--ckpt", default=None, help="Actor model/tokenizer checkpoint path or remote model repo ID.") +@click.option( + "--base-model-name-or-path", + default=None, + help="Stable base model reference written to PEFT adapter metadata; defaults to the original --ckpt value.", +) @click.option( "--dataset-path", default=None, help="Training dataset path, HF save_to_disk directory, or remote dataset ref." ) diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 1dd54672..98b3b115 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -943,6 +943,17 @@ def test_train_cli_preflight_skips_unused_reward_file_for_offline_algorithms(sel self.assertEqual(sft_cfg.algo, "sft") self.assertEqual(dpo_cfg.algo, "dpo") + def test_train_cli_accepts_stable_base_reference_for_adapter_metadata(self): + config = train_cli._trainer_config_from_options( + **_train_options( + ckpt="/pcache/local/base", + base_model_name_or_path="aistudio://project/base", + ) + ) + + self.assertEqual(config.base_model_name_or_path, "aistudio://project/base") + self.assertEqual(config.areno_config().base_model_name_or_path, "aistudio://project/base") + def test_train_cli_preflight_rejects_agent_file_without_callable_run_agent(self): """Agent hooks should fail before rollout/backend-heavy work.""" with tempfile.TemporaryDirectory() as tmp: From c99a40531af4393a875dfead4e6a138702c2c6d6 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Fri, 21 Aug 2026 11:08:37 +0800 Subject: [PATCH 20/29] style(tests): format Bailing LoRA coverage --- tests/test_bailing_v3_lora_cpu.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index 3997d22a..4165c582 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -130,10 +130,7 @@ def test_bailing_v3_full_profile_attaches_native_slots(monkeypatch) -> None: registry = initialize_lora(model, LoraConfig(rank=4, alpha=4, target_modules=profile), seed=42) assert set(profile) <= set(BAILING_V3_TARGETS) - kda_names = { - f"layers.0.attention.{component}" - for component in ("q_proj", "k_proj", "v_proj", "f_proj", "g_proj") - } + kda_names = {f"layers.0.attention.{component}" for component in ("q_proj", "k_proj", "v_proj", "f_proj", "g_proj")} assert kda_names <= registry.slots.keys() assert len({id(registry.slots[name]) for name in kda_names}) == len(kda_names) assert "layers.0.attention.b_proj" not in registry.slots From c20e8712224fa3b6564a6c6b301226bfb0edb92b Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Fri, 21 Aug 2026 20:42:41 +0800 Subject: [PATCH 21/29] perf(lora): optimize Bailing routed rollout --- areno/engine/config.py | 4 +- areno/models/bailing_v3/model.py | 35 +++++++++++------ tests/test_bailing_v3_lora_cpu.py | 65 ++++++++++++++++++++++++++++++- 3 files changed, 90 insertions(+), 14 deletions(-) diff --git a/areno/engine/config.py b/areno/engine/config.py index f6d6d5eb..f79d6de4 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -122,11 +122,11 @@ def resolve_compile_model(self, *, model: ModelConfig, devices: list[int]) -> No self.compile_model = False def resolve_eager_decode(self, *, model: ModelConfig, lora: LoraConfig | None) -> None: - """Use eager decode when routed-expert adapters need grouped execution.""" + """Use eager decode when routed-expert adapters lack a fused rollout path.""" if self.eager_decode or lora is None: return - if model.model_type in {"qwen3_moe", "bailing_moe_v3"} and { + if model.model_type == "qwen3_moe" and { "gate_proj", "up_proj", "down_proj", diff --git a/areno/models/bailing_v3/model.py b/areno/models/bailing_v3/model.py index 269fa06d..b7d523c4 100644 --- a/areno/models/bailing_v3/model.py +++ b/areno/models/bailing_v3/model.py @@ -289,7 +289,7 @@ def forward(self, hidden_states: torch.Tensor, num_padding_tokens: int = 0) -> t with sequence_parallel_region(False): topk_idx, topk_weight, _ = self.gate(hidden_states, num_padding_tokens) flat = expert_input.view(-1, hidden) - if self.training or self.experts.has_lora(): + if self.training or not self._infer_weights_ready: # Permute/unpermute path is autograd-friendly. out = self.experts(flat, topk_idx, topk_weight).view(bsz, seqlen, hidden) else: @@ -311,21 +311,15 @@ def prepare_infer_weights(self) -> None: down projection. Buffers are reused across calls if the shape/device already match to avoid reallocating on every weight refresh. """ - if self.experts.has_lora(): - self.clear_infer_weights() - return - gate_weights, up_weights, down_weights = self.experts.expert_weights() + gate_weights, up_weights, down_weights = self.experts.inference_weights() self._infer_gate_weight = self._updated_infer_weight( - self._infer_gate_weight, - torch.stack(gate_weights, dim=0).to(dtype=self.config.dtype).contiguous(), + self._infer_gate_weight, gate_weights.to(dtype=self.config.dtype).contiguous() ) self._infer_up_weight = self._updated_infer_weight( - self._infer_up_weight, - torch.stack(up_weights, dim=0).to(dtype=self.config.dtype).contiguous(), + self._infer_up_weight, up_weights.to(dtype=self.config.dtype).contiguous() ) self._infer_down_weight = self._updated_infer_weight( - self._infer_down_weight, - torch.stack(down_weights, dim=0).to(dtype=self.config.dtype).contiguous(), + self._infer_down_weight, down_weights.to(dtype=self.config.dtype).contiguous() ) # w1 = [gate || up] along the intermediate dim so SiLU(gate) * up can # be folded into a single fused kernel call. @@ -520,6 +514,25 @@ def expert_weights(self) -> tuple[list[torch.Tensor], list[torch.Tensor], list[t down_weights.append(_grouped_weight(self.linear_fc2, expert_id).detach()) return gate_weights, up_weights, down_weights + @torch.no_grad() + def inference_weights(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build derived fused-rollout weights without modifying the frozen base.""" + + gate_weights, up_weights, down_weights = self.expert_weights() + merged = { + "gate_proj": torch.stack(gate_weights, dim=0), + "up_proj": torch.stack(up_weights, dim=0), + "down_proj": torch.stack(down_weights, dim=0), + } + if self.has_active_lora(): + for component, weight in merged.items(): + if component not in self.lora_slots: + continue + slot = self.lora_slots[component] + delta = torch.bmm(slot.lora_B, slot.lora_A) + weight.add_(delta.mul_(slot.scale)) + return merged["gate_proj"], merged["up_proj"], merged["down_proj"] + @torch.no_grad() def full_expert_weights(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: """Gather all experts onto DP rank 0 for checkpoint saving (None elsewhere).""" diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index 4165c582..c007afe8 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -7,7 +7,8 @@ from torch import nn from areno.adapters.config import BAILING_V3_TARGETS, LoraConfig -from areno.adapters.lora import initialize_lora +from areno.adapters.lora import AdapterRegistry, RoutedExpertLoraSlot, _AdapterRuntimeState, initialize_lora +from areno.engine.config import ModelConfig, RuntimeConfig from areno.engine.layers import linear from areno.engine.layers.linear import ColumnParallelLinear, RowParallelLinear @@ -192,3 +193,65 @@ def test_bailing_v3_empty_route_keeps_expert_router_and_lora_gradients(monkeypat ) assert all(parameter.grad is not None for parameter in parameters) assert all(torch.count_nonzero(parameter.grad) == 0 for parameter in parameters) + + +def test_bailing_v3_expert_lora_merges_only_into_derived_infer_weights() -> None: + from areno.models.bailing_v3 import model as bailing_model + + experts = bailing_model.BailingGroupedExperts.__new__(bailing_model.BailingGroupedExperts) + nn.Module.__init__(experts) + experts.local_num_experts = 2 + experts.local_expert_start = 0 + experts.linear_fc1 = SimpleNamespace(weight=nn.Parameter(torch.arange(48, dtype=torch.float32).view(2, 8, 3))) + experts.linear_fc2 = SimpleNamespace(weight=nn.Parameter(torch.arange(24, dtype=torch.float32).view(2, 3, 4))) + runtime_state = _AdapterRuntimeState() + slots = nn.ModuleDict() + for component, in_features, out_features, base_weight in ( + ("gate_proj", 3, 4, experts.linear_fc1.weight), + ("up_proj", 3, 4, experts.linear_fc1.weight), + ("down_proj", 4, 3, experts.linear_fc2.weight), + ): + slot = RoutedExpertLoraSlot( + logical_name=f"layers.0.mlp.experts.{{expert}}.{component}", + base_weight=base_weight, + local_num_experts=2, + local_expert_start=0, + in_features=in_features, + out_features=out_features, + config=LoraConfig(rank=2, alpha=1, target_modules=(component,)), + seed=42, + runtime_state=runtime_state, + ) + slot.lora_A.data.fill_(0.5) + slot.lora_B.data.fill_(0.25) + slots[component] = slot + experts.lora_slots = slots + base_fc1 = experts.linear_fc1.weight.detach().clone() + base_fc2 = experts.linear_fc2.weight.detach().clone() + + gate, up, down = experts.inference_weights() + + expected_delta = 0.125 + torch.testing.assert_close(gate, base_fc1[:, :4] + expected_delta) + torch.testing.assert_close(up, base_fc1[:, 4:] + expected_delta) + torch.testing.assert_close(down, base_fc2 + expected_delta) + torch.testing.assert_close(experts.linear_fc1.weight, base_fc1) + torch.testing.assert_close(experts.linear_fc2.weight, base_fc2) + registry = AdapterRegistry(dict(slots.items()), LoraConfig(), runtime_state) + with registry.base_only(): + base_gate, base_up, base_down = experts.inference_weights() + torch.testing.assert_close(base_gate, base_fc1[:, :4]) + torch.testing.assert_close(base_up, base_fc1[:, 4:]) + torch.testing.assert_close(base_down, base_fc2) + + +def test_bailing_v3_routed_lora_keeps_cuda_graph_decode_enabled() -> None: + lora = LoraConfig(target_modules=("gate_proj", "up_proj", "down_proj")) + bailing_runtime = RuntimeConfig() + bailing_runtime.resolve_eager_decode(model=ModelConfig(model_type="bailing_moe_v3"), lora=lora) + assert not bailing_runtime.eager_decode + + qwen_runtime = RuntimeConfig() + with pytest.warns(RuntimeWarning, match="routed-expert LoRA"): + qwen_runtime.resolve_eager_decode(model=ModelConfig(model_type="qwen3_moe"), lora=lora) + assert qwen_runtime.eager_decode From cef13f43ff493a7f52532a12f5534f32587f667e Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Fri, 21 Aug 2026 20:46:02 +0800 Subject: [PATCH 22/29] perf(lora): offload merged Bailing expert train weights --- areno/models/bailing_v3/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/areno/models/bailing_v3/model.py b/areno/models/bailing_v3/model.py index b7d523c4..18e83bca 100644 --- a/areno/models/bailing_v3/model.py +++ b/areno/models/bailing_v3/model.py @@ -1575,7 +1575,7 @@ def clear_infer_weights(self) -> None: @torch.no_grad() def offload_train_weights(self) -> None: for layer in self.layers: - if isinstance(layer.mlp, BailingSparseMoeBlock) and not layer.mlp.experts.has_lora(): + if isinstance(layer.mlp, BailingSparseMoeBlock): layer.mlp.experts.offload_to_cpu() @torch.no_grad() From afb3657c8f41cad234a06ebbca4db3845b8165f6 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Sat, 22 Aug 2026 14:24:50 +0800 Subject: [PATCH 23/29] perf(lora): optimize Bailing KDA A projections --- areno/models/bailing_v3/model.py | 69 +++++++++++++++++++++++++++---- tests/test_bailing_v3_lora_cpu.py | 35 ++++++++++++++++ 2 files changed, 97 insertions(+), 7 deletions(-) diff --git a/areno/models/bailing_v3/model.py b/areno/models/bailing_v3/model.py index 18e83bca..dc9b988f 100644 --- a/areno/models/bailing_v3/model.py +++ b/areno/models/bailing_v3/model.py @@ -44,6 +44,7 @@ import torch.distributed as dist from fla.ops.lightning_attn import chunk_lightning_attn from torch import nn +from torch.nn import functional as F from areno.accel import ( areno_grouped_linear, @@ -1215,6 +1216,55 @@ def __init__(self, config: ModelConfig, layer_idx: int): self.eps = config.rms_norm_eps self.state_cache = torch.tensor([]) self.conv_cache = torch.tensor([]) + self.register_buffer("_infer_lora_A", torch.empty(0), persistent=False) + self._infer_lora_rank = 0 + + @torch.no_grad() + def prepare_lora_infer_weights(self) -> None: + """Pack the five KDA LoRA A projections for single-adapter inference.""" + + projections = (self.q_proj, self.k_proj, self.v_proj, self.f_proj, self.g_proj) + slots = tuple(projection.lora_slot for projection in projections) + if any(slot is None or not slot.enabled for slot in slots): + self._infer_lora_A = self._infer_lora_A.new_empty(0) + self._infer_lora_rank = 0 + return + value = torch.cat(tuple(slot.lora_A for slot in slots), dim=0).contiguous() + if ( + self._infer_lora_A.shape == value.shape + and self._infer_lora_A.device == value.device + and self._infer_lora_A.dtype == value.dtype + ): + self._infer_lora_A.copy_(value) + else: + self._infer_lora_A = value + self._infer_lora_rank = slots[0].rank + + @torch.no_grad() + def clear_lora_infer_weights(self) -> None: + self._infer_lora_A = self._infer_lora_A.new_empty(0) + self._infer_lora_rank = 0 + + def _project_qkvfg( + self, hidden_states: torch.Tensor, infer_meta: InferMeta | None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + projections = (self.q_proj, self.k_proj, self.v_proj, self.f_proj, self.g_proj) + slots = tuple(projection.lora_slot for projection in projections) + use_packed_lora = ( + infer_meta is not None + and self._infer_lora_A.numel() > 0 + and all(slot is not None and slot.enabled for slot in slots) + ) + if not use_packed_lora: + return tuple(projection(hidden_states) for projection in projections) + + packed_hidden = F.linear(hidden_states, self._infer_lora_A) + lora_inputs = packed_hidden.split(self._infer_lora_rank, dim=-1) + return tuple( + areno_linear(hidden_states, projection.weight, projection.bias) + + F.linear(lora_input, slot.lora_B) * slot.scale + for projection, slot, lora_input in zip(projections, slots, lora_inputs, strict=True) + ) @torch._dynamo.disable def forward( @@ -1226,15 +1276,16 @@ def forward( ) -> torch.Tensor: del position_ids hidden_states = hidden_states.to(dtype=self.q_proj.weight.dtype) - q = self._causal_conv(self.q_proj(hidden_states), self.q_conv1d_weight, 0, train_meta, infer_meta) - k = self._causal_conv(self.k_proj(hidden_states), self.k_conv1d_weight, 1, train_meta, infer_meta) - v = self._causal_conv(self.v_proj(hidden_states), self.v_conv1d_weight, 2, train_meta, infer_meta) + q, k, v, f, gate = self._project_qkvfg(hidden_states, infer_meta) batch, seqlen = q.shape[:2] + q = self._causal_conv(q, self.q_conv1d_weight, 0, train_meta, infer_meta) + k = self._causal_conv(k, self.k_conv1d_weight, 1, train_meta, infer_meta) + v = self._causal_conv(v, self.v_conv1d_weight, 2, train_meta, infer_meta) q = q.to(dtype=hidden_states.dtype) k = k.to(dtype=hidden_states.dtype) v = v.to(dtype=hidden_states.dtype) - f = self.f_proj(hidden_states).view(batch, seqlen, self.local_heads, self.head_dim) - gate = self.g_proj(hidden_states).view(batch, seqlen, self.local_heads, self.head_dim) + f = f.view(batch, seqlen, self.local_heads, self.head_dim) + gate = gate.view(batch, seqlen, self.local_heads, self.head_dim) beta = self.b_proj(hidden_states).view(batch, seqlen, self.local_heads) q = q.view(batch, seqlen, self.local_heads, self.head_dim) k = k.view(batch, seqlen, self.local_heads, self.head_dim) @@ -1560,15 +1611,19 @@ def set_kv_caches( @torch.no_grad() def prepare_infer_weights(self) -> None: - """Stack per-expert weights for fused-MoE inference on every MoE block.""" + """Prepare KDA LoRA and fused-MoE inference views.""" for layer in self.layers: + if isinstance(layer.attention, BailingKDAAttention): + layer.attention.prepare_lora_infer_weights() if isinstance(layer.mlp, BailingSparseMoeBlock): layer.mlp.prepare_infer_weights() @torch.no_grad() def clear_infer_weights(self) -> None: - """Drop fused-MoE inference tiles to reclaim memory before training.""" + """Drop KDA LoRA and fused-MoE inference views before training.""" for layer in self.layers: + if isinstance(layer.attention, BailingKDAAttention): + layer.attention.clear_lora_infer_weights() if isinstance(layer.mlp, BailingSparseMoeBlock): layer.mlp.clear_infer_weights() diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index c007afe8..2e978dc5 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -255,3 +255,38 @@ def test_bailing_v3_routed_lora_keeps_cuda_graph_decode_enabled() -> None: with pytest.warns(RuntimeWarning, match="routed-expert LoRA"): qwen_runtime.resolve_eager_decode(model=ModelConfig(model_type="qwen3_moe"), lora=lora) assert qwen_runtime.eager_decode + + +def test_bailing_v3_kda_packed_a_matches_canonical_slots(monkeypatch) -> None: + from areno.models.bailing_v3 import model as bailing_model + + monkeypatch.setattr(linear, "get_tp_context", _single_tp) + monkeypatch.setattr("areno.adapters.lora.get_tp_context", _single_tp) + model = _BailingModel() + registry = initialize_lora( + model, + LoraConfig(rank=4, alpha=4, target_modules=("q_proj", "k_proj", "v_proj", "f_proj", "g_proj")), + seed=42, + ) + attention = model.layers[0].attention + attention.register_buffer("_infer_lora_A", torch.empty(0), persistent=False) + attention._infer_lora_rank = 0 + for slot in registry.slots.values(): + slot.lora_B.data.normal_() + hidden_states = torch.randn(2, 3, 8) + components = ("q_proj", "k_proj", "v_proj", "f_proj", "g_proj") + expected = tuple( + getattr(attention, component)(hidden_states) for component in components + ) + + bailing_model.BailingKDAAttention.prepare_lora_infer_weights(attention) + actual = bailing_model.BailingKDAAttention._project_qkvfg(attention, hidden_states, SimpleNamespace()) + + assert attention._infer_lora_A.shape == (20, 8) + for packed, canonical in zip(actual, expected, strict=True): + torch.testing.assert_close(packed, canonical) + with registry.base_only(): + base = bailing_model.BailingKDAAttention._project_qkvfg(attention, hidden_states, SimpleNamespace()) + for component, output in zip(components, base, strict=True): + projection = getattr(attention, component) + torch.testing.assert_close(output, torch.nn.functional.linear(hidden_states, projection.weight)) From cc150591b037e878eedd3e728f464f5abbdfd32a Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Sat, 22 Aug 2026 14:41:17 +0800 Subject: [PATCH 24/29] test: stub CUDA linear in KDA packing oracle --- tests/test_bailing_v3_lora_cpu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index 2e978dc5..7db2040f 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -262,6 +262,7 @@ def test_bailing_v3_kda_packed_a_matches_canonical_slots(monkeypatch) -> None: monkeypatch.setattr(linear, "get_tp_context", _single_tp) monkeypatch.setattr("areno.adapters.lora.get_tp_context", _single_tp) + monkeypatch.setattr(bailing_model, "areno_linear", torch.nn.functional.linear) model = _BailingModel() registry = initialize_lora( model, From 3f7e33636a5558140c0057834c7c5378e3b2b163 Mon Sep 17 00:00:00 2001 From: Zhou FANG Date: Wed, 26 Aug 2026 19:25:49 +0800 Subject: [PATCH 25/29] fix(lora): align rebased tests with CUDA backend API --- areno/adapters/lora.py | 28 +++++++--------------------- areno/api/backend/cuda/roles.py | 16 ++++++++++------ areno/cli/serve.py | 2 +- areno/engine/worker.py | 2 +- areno/models/bailing_v3/model.py | 4 +--- tests/test_bailing_v3_lora_cpu.py | 4 +--- tests/test_cli_model_refs_cpu.py | 2 +- tests/test_config_data_cpu.py | 2 +- tests/test_qwen3_lora_e2e.py | 18 ++++++++---------- 9 files changed, 31 insertions(+), 47 deletions(-) diff --git a/areno/adapters/lora.py b/areno/adapters/lora.py index 8c91b5fc..cd4e9b5f 100644 --- a/areno/adapters/lora.py +++ b/areno/adapters/lora.py @@ -283,9 +283,7 @@ def _initialize_qwen3_lora( slots[logical_name] = slot if getattr(model_config, "enable_moe_block", False): - matched.update( - _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots) - ) + matched.update(_install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots)) else: gate_up = layer.mlp.gate_up_proj for component_index, component in enumerate(("gate_proj", "up_proj")): @@ -369,18 +367,14 @@ def _initialize_bailing_v3_lora( owner = getattr(attention, component, None) if component in requested and owner is not None: matched.add(component) - slot = _replicated_slot( - f"{attention_prefix}.{component}", owner, config, seed, runtime_state - ) + slot = _replicated_slot(f"{attention_prefix}.{component}", owner, config, seed, runtime_state) attention.install_lora_component(component, slot) slots[slot.logical_name] = slot for component in ("q_b_proj", "kv_b_proj"): owner = getattr(attention, component, None) if component in requested and owner is not None: matched.add(component) - _install_column_slot( - f"{attention_prefix}.{component}", owner, config, seed, runtime_state, slots - ) + _install_column_slot(f"{attention_prefix}.{component}", owner, config, seed, runtime_state, slots) if "dense" in requested: matched.add("dense") _install_row_slot( @@ -394,9 +388,7 @@ def _initialize_bailing_v3_lora( mlp_prefix = f"{prefix}.mlp" if hasattr(layer.mlp, "experts"): - matched.update( - _install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots) - ) + matched.update(_install_moe_slots(layer.mlp.experts, prefix, requested, config, seed, runtime_state, slots)) if layer.mlp.shared_experts is not None: matched.update( _install_dense_mlp_slots( @@ -411,9 +403,7 @@ def _initialize_bailing_v3_lora( ) else: matched.update( - _install_dense_mlp_slots( - layer.mlp, mlp_prefix, requested, config, seed, runtime_state, slots - ) + _install_dense_mlp_slots(layer.mlp, mlp_prefix, requested, config, seed, runtime_state, slots) ) return matched @@ -490,14 +480,10 @@ def _install_dense_mlp_slots( for component in ("gate_proj", "up_proj"): if component in requested: matched.add(component) - _install_column_slot( - f"{prefix}.{component}", getattr(mlp, component), config, seed, runtime_state, slots - ) + _install_column_slot(f"{prefix}.{component}", getattr(mlp, component), config, seed, runtime_state, slots) if "down_proj" in requested: matched.add("down_proj") - _install_row_slot( - f"{prefix}.down_proj", mlp.down_proj, config, seed, runtime_state, slots - ) + _install_row_slot(f"{prefix}.down_proj", mlp.down_proj, config, seed, runtime_state, slots) return matched diff --git a/areno/api/backend/cuda/roles.py b/areno/api/backend/cuda/roles.py index 8768eb8b..fa17252e 100644 --- a/areno/api/backend/cuda/roles.py +++ b/areno/api/backend/cuda/roles.py @@ -396,12 +396,16 @@ def score_logprobs(self, payload: ScorePayload) -> list[list[float]] | None: try: token_rows = payload.token_rows_by_dp[ctx.dp_rank] features = payload.features_by_dp[ctx.dp_rank] if payload.features_by_dp is not None else None - local = [] if not token_rows else self._score_logprob_rows( - model, - token_rows, - payload, - features=features, - sequence_parallel=sequence_parallel, + local = ( + [] + if not token_rows + else self._score_logprob_rows( + model, + token_rows, + payload, + features=features, + sequence_parallel=sequence_parallel, + ) ) return local if ctx.rank == 0 else None finally: diff --git a/areno/cli/serve.py b/areno/cli/serve.py index ee23c3d5..7dc3f548 100644 --- a/areno/cli/serve.py +++ b/areno/cli/serve.py @@ -20,8 +20,8 @@ from fastapi import FastAPI, HTTPException, Request from pydantic import BaseModel, Field -from areno.api import MLX, BackendType, MlxConfig, SamplingParams, Trainer, default_backend_type from areno.adapters import LoraConfig +from areno.api import MLX, BackendType, MlxConfig, SamplingParams, Trainer, default_backend_type from areno.api.multimodal import ( expand_image_tokens, image_token_counts_from_features, diff --git a/areno/engine/worker.py b/areno/engine/worker.py index 5e33c715..c08afe96 100644 --- a/areno/engine/worker.py +++ b/areno/engine/worker.py @@ -20,9 +20,9 @@ import torch import torch.distributed as dist -from areno.api.backend.cuda.roles import RoleManager, WorkerRole from areno.adapters import initialize_lora from areno.adapters.peft import export_peft_adapter, load_peft_adapter +from areno.api.backend.cuda.roles import RoleManager, WorkerRole from areno.engine.config import EngineConfig from areno.engine.data import RolloutOutput from areno.engine.data.sampling import _truncate_generated diff --git a/areno/models/bailing_v3/model.py b/areno/models/bailing_v3/model.py index dc9b988f..bbf26573 100644 --- a/areno/models/bailing_v3/model.py +++ b/areno/models/bailing_v3/model.py @@ -891,9 +891,7 @@ def _project( bsz, seqlen = q.shape[:2] q = q.view(bsz, seqlen, self.local_heads, self.head_dim) q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - kv_a = self._with_lora( - "kv_a_proj_with_mqa", mla_input, self.kv_a_proj_with_mqa(mla_input) - ) + kv_a = self._with_lora("kv_a_proj_with_mqa", mla_input, self.kv_a_proj_with_mqa(mla_input)) compressed_kv, k_rope = kv_a.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) if is_sequence_parallel_active(): k_rope = gather_from_sequence_parallel_region(k_rope) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index 7db2040f..a9682f3d 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -276,9 +276,7 @@ def test_bailing_v3_kda_packed_a_matches_canonical_slots(monkeypatch) -> None: slot.lora_B.data.normal_() hidden_states = torch.randn(2, 3, 8) components = ("q_proj", "k_proj", "v_proj", "f_proj", "g_proj") - expected = tuple( - getattr(attention, component)(hidden_states) for component in components - ) + expected = tuple(getattr(attention, component)(hidden_states) for component in components) bailing_model.BailingKDAAttention.prepare_lora_infer_weights(attention) actual = bailing_model.BailingKDAAttention._project_qkvfg(attention, hidden_states, SimpleNamespace()) diff --git a/tests/test_cli_model_refs_cpu.py b/tests/test_cli_model_refs_cpu.py index d4eb9e9d..0edd7757 100644 --- a/tests/test_cli_model_refs_cpu.py +++ b/tests/test_cli_model_refs_cpu.py @@ -97,7 +97,7 @@ def test_resolve_model_refs_preserves_actor_reference_for_adapter_metadata(self) self.assertEqual(resolved.ckpt, "/ms/org/actor") self.assertEqual(resolved.base_model_name_or_path, "org/actor") - self.assertEqual(resolved.areno_config().base_model_name_or_path, "org/actor") + self.assertEqual(resolved.cuda_config().base_model_name_or_path, "org/actor") if __name__ == "__main__": diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 98b3b115..b864da6c 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -952,7 +952,7 @@ def test_train_cli_accepts_stable_base_reference_for_adapter_metadata(self): ) self.assertEqual(config.base_model_name_or_path, "aistudio://project/base") - self.assertEqual(config.areno_config().base_model_name_or_path, "aistudio://project/base") + self.assertEqual(config.cuda_config().base_model_name_or_path, "aistudio://project/base") def test_train_cli_preflight_rejects_agent_file_without_callable_run_agent(self): """Agent hooks should fail before rollout/backend-heavy work.""" diff --git a/tests/test_qwen3_lora_e2e.py b/tests/test_qwen3_lora_e2e.py index 392833d7..4ede0922 100644 --- a/tests/test_qwen3_lora_e2e.py +++ b/tests/test_qwen3_lora_e2e.py @@ -12,7 +12,7 @@ from safetensors.torch import load_file from areno.adapters import LoraConfig -from areno.api import ArenoConfig, SamplingParams, Trainer +from areno.api import CudaConfig, SamplingParams, Trainer from areno.api.algorithms import get_algorithm from areno.api.roles import ModelRole from areno.api.trainer_config import DPOTrainerConfig, PolicyTrainerConfig, PPOTrainerConfig @@ -179,7 +179,7 @@ def reward_fn(record) -> float: ] reward_fn = None - backend_config = config.areno_config() + backend_config = config.cuda_config() backend_config.dp_size = 2 backend_config.runtime["compile_model"] = False observed = _ObservedReferenceTrainer( @@ -239,7 +239,7 @@ def test_qwen3_lora_tp2_dp2_rollout_train_peft(tmp_path: Path, model_env: str, m final_path = checkpoint_path / "step_000002" reexported_path = tmp_path / "adapter-reexported" lora = LoraConfig(rank=8, alpha=16.0) - backend_config = ArenoConfig( + backend_config = CudaConfig( tp_size=2, dp_size=2, devices=[0, 1, 2, 3], @@ -350,7 +350,7 @@ def reward_fn(record) -> float: imported = Trainer( 4, os.fspath(model_path), - custom_config=ArenoConfig( + custom_config=CudaConfig( tp_size=2, dp_size=2, devices=[0, 1, 2, 3], @@ -393,7 +393,7 @@ def test_qwen3_moe_lora_tp8_replicated_kv_roundtrip(tmp_path: Path) -> None: final_path = checkpoint_path / "step_000001" reexported_path = tmp_path / "adapter-reexported" lora = LoraConfig(rank=8, alpha=16.0) - backend_config = ArenoConfig( + backend_config = CudaConfig( tp_size=8, dp_size=1, devices=list(range(8)), @@ -471,7 +471,7 @@ def reward_fn(record) -> float: imported = Trainer( 8, os.fspath(model_path), - custom_config=ArenoConfig( + custom_config=CudaConfig( tp_size=8, dp_size=1, devices=list(range(8)), @@ -488,9 +488,7 @@ def reward_fn(record) -> float: reexported = load_file(reexported_path / "adapter_model.safetensors") assert reexported.keys() == final.keys() assert all(torch.equal(reexported[name], final[name]) for name in final) - torch.testing.assert_close( - torch.tensor(imported_logprobs), torch.tensor(trained_logprobs), rtol=0.0, atol=1.0e-5 - ) + torch.testing.assert_close(torch.tensor(imported_logprobs), torch.tensor(trained_logprobs), rtol=0.0, atol=1.0e-5) @pytest.mark.parametrize( @@ -544,7 +542,7 @@ def test_qwen3_lora_independent_rollout_two_step( metrics_log_dir=None, lora=lora, ) - backend_config = config.areno_config() + backend_config = config.cuda_config() backend_config.dp_size = 1 backend_config.runtime["compile_model"] = False inner = Trainer(2, os.fspath(model_path), custom_config=backend_config) From bbc59ef5a844cd30de6c291e388ceecd73e1681b Mon Sep 17 00:00:00 2001 From: xsuler Date: Thu, 27 Aug 2026 07:42:44 +0800 Subject: [PATCH 26/29] docs(lora): document native adapter workflows --- README.md | 37 ++++++++ areno/cli/train.py | 11 +++ docs/concepts/native-lora.rst | 131 +++++++++++++++++++++++++++++ docs/index.rst | 1 + tests/test_train_cli_config_cpu.py | 29 +++++++ 5 files changed, 209 insertions(+) create mode 100644 docs/concepts/native-lora.rst diff --git a/README.md b/README.md index c93bc947..e82238f9 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ AReno's mission is to make LLM RL **accessible** for a broad community of resear - 🪶 **Lightweight**: one self-contained train/serve stack that installs and loads only the native backend needed by the host—CUDA on Linux or MLX on Apple Silicon. - 🧰 **Agentic RL ready**: run an agent function against AReno's local OpenAI-compatible proxy, return explicit trajectories, and train from tokens, logprobs, rewards, and loss masks derived by the trainer. - 🎞️ **Multimodal**: use image, audio, and video content with compatible model processors through the same OpenAI-style message format in serving and agentic training. +- 🧩 **Native LoRA**: train TP-aware adapters for Qwen3, Qwen3-MoE, and Bailing-MoE V3, save standard PEFT artifacts, and reload them for training or serving. - 🧩 **Extensible**: easily register new algorithms, model adapters, reward functions, and hardware backends without changing the core. ## Installation @@ -306,6 +307,31 @@ reward function, OpenAI-compatible agent, and browser UI. For the full list of training options, run `areno train --help`. +#### Native LoRA + +Enable CUDA-native LoRA by passing a rank. AReno freezes the base model and +trains the selected projection adapters through the same rollout and training +engine, including agentic RL: + +```bash +areno train \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path gsm8k:main \ + --dataset-loader-fn examples/math/dataset_loader.py \ + --reward-fn-path examples/math/math_verify_reward.py \ + --algo gspo \ + --lora-rank 8 \ + --lora-alpha 16 \ + --save-path outputs/qwen3-lora \ + --save-interval 100 +``` + +Saved checkpoints contain standard PEFT `adapter_config.json` and +`adapter_model.safetensors` files. Resume training or serve an adapter by +supplying the frozen base checkpoint together with `--lora-adapter-path`. +See the [native LoRA guide](docs/concepts/native-lora.rst) for supported +models and targets, agentic training, save/reload, and serving examples. + ### Serving Serve a trained checkpoint as an OpenAI-compatible endpoint with continuous batching: @@ -318,6 +344,17 @@ areno serve \ --port 8000 ``` +To serve a saved native LoRA adapter without merging it into the base model: + +```bash +areno serve \ + --model-path Qwen/Qwen3-0.6B \ + --lora-adapter-path outputs/qwen3-lora/step_000100 \ + --tp-size 1 \ + --world-size 1 \ + --port 8000 +``` + The command selects CUDA or MLX from the host platform. MLX serving is single-process (`--tp-size 1 --world-size 1`) and uses the same long-lived continuous-batch request scheduler and HTTP API. diff --git a/areno/cli/train.py b/areno/cli/train.py index 71cff5a8..5991be5f 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -529,6 +529,17 @@ def _format_training_config_summary( ("mini_bs", str(config.mini_bs)), ("score_micro_bs", str(config.score_micro_bs)), ("gradient_accumulation_steps", _format_optional(config.gradient_accumulation_steps, default="auto")), + ("lora_rank", str(config.lora.rank) if config.lora is not None else "disabled"), + ("lora_alpha", str(config.lora.alpha) if config.lora is not None else "n/a"), + ("lora_dropout", str(config.lora.dropout) if config.lora is not None else "n/a"), + ( + "lora_target_modules", + ",".join(config.lora.target_modules) if config.lora is not None else "n/a", + ), + ( + "lora_adapter_path", + _format_optional(config.lora.adapter_path) if config.lora is not None else "n/a", + ), ( "optimizer", ( diff --git a/docs/concepts/native-lora.rst b/docs/concepts/native-lora.rst new file mode 100644 index 00000000..d95d128d --- /dev/null +++ b/docs/concepts/native-lora.rst @@ -0,0 +1,131 @@ +Native LoRA +=========== + +AReno can train and serve LoRA adapters directly in its CUDA engine. The base +model stays frozen while the LoRA A and B parameters participate in the same +tensor-parallel, data-parallel, sequence-parallel, rollout, and optimizer +paths as full-parameter training. No external PEFT runtime is required during +training or inference. + +Support +------- + +Native LoRA currently supports these CUDA model adapters: + +* Qwen3 +* Qwen3-MoE +* Bailing-MoE V3 checkpoints with ``no_kda_lora=true`` + +The default target modules are ``q_proj``, ``k_proj``, ``v_proj``, +``o_proj``, ``gate_proj``, ``up_proj``, and ``down_proj``. Select a subset +with ``--lora-target-modules``. Bailing-MoE V3 additionally supports its +native attention projection names, including ``q_a_proj``, ``q_b_proj``, +``kv_a_proj_with_mqa``, and ``kv_b_proj``. + +LoRA dropout must currently be zero. Standard PEFT LoRA adapters are accepted, +but options that change the adapter structure, such as DoRA, RS-LoRA, bias +training, rank patterns, alpha patterns, or ``modules_to_save``, are rejected +with a configuration error. Bailing-MoE V3 router-bias updates must also be +disabled so the base policy remains frozen. + +Train +----- + +Set ``--lora-rank`` to enable native LoRA. ``--lora-alpha`` defaults to 16 and +the default target list covers attention and MLP projections: + +.. code-block:: bash + + areno train \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path gsm8k:main \ + --dataset-loader-fn examples/math/dataset_loader.py \ + --reward-fn-path examples/math/math_verify_reward.py \ + --algo gspo \ + --world-size 1 \ + --tp-size 1 \ + --lora-rank 8 \ + --lora-alpha 16 \ + --lora-target-modules q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj \ + --save-path outputs/qwen3-lora \ + --save-interval 100 + +The resolved LoRA rank, alpha, dropout, target modules, and adapter path are +shown in the configuration summary printed before model loading. + +Agentic LoRA uses the normal agent hooks. For example, this trains the +Tic-Tac-Toe tool-calling policy: + +.. code-block:: bash + + python examples/agentic/tictactoe/dataset_generator.py \ + --output /tmp/areno-tictactoe.jsonl \ + --count 2048 \ + --seed 2026 + + areno train \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path /tmp/areno-tictactoe.jsonl \ + --dataset-loader-fn examples/agentic/tictactoe/dataset_loader.py \ + --reward-fn-path examples/agentic/tictactoe/reward.py \ + --agent-fn examples/agentic/tictactoe/run_agent.py \ + --algo gspo \ + --batch-size 1 \ + --n-samples 8 \ + --max-running-prompts 8 \ + --max-new-tokens 3071 \ + --lora-rank 8 \ + --lora-alpha 16 \ + --save-path outputs/tictactoe-lora \ + --save-interval 100 + +Save and reload +--------------- + +Each LoRA save directory contains PEFT-compatible +``adapter_config.json`` and ``adapter_model.safetensors`` files. The save is +adapter-only: continue to pass the original base checkpoint with ``--ckpt``. +To initialize a new training run from a saved adapter: + +.. code-block:: bash + + areno train \ + --ckpt Qwen/Qwen3-0.6B \ + --lora-adapter-path outputs/qwen3-lora/step_000100 \ + --dataset-path gsm8k:main \ + --dataset-loader-fn examples/math/dataset_loader.py \ + --reward-fn-path examples/math/math_verify_reward.py \ + --algo gspo \ + --save-path outputs/qwen3-lora-continued + +Adapter metadata is authoritative when ``--lora-adapter-path`` is present, so +its rank, alpha, dropout, and target modules replace the corresponding CLI +defaults. Adapter-only saves do not contain optimizer, scheduler, or RNG state; +loading one initializes the policy weights for a new run rather than exactly +resuming the old trainer state. + +Serve +----- + +Serve the frozen base and saved adapter together; merging is not required: + +.. code-block:: bash + + areno serve \ + --model-path Qwen/Qwen3-0.6B \ + --lora-adapter-path outputs/qwen3-lora/step_000100 \ + --world-size 1 \ + --tp-size 1 \ + --port 8000 + +The endpoint remains OpenAI compatible. ``/v1/models`` reports the base model, +and chat completion requests use the loaded adapter. + +Reference model reuse +--------------------- + +For algorithms that require a frozen reference policy, use +``--reference-mode reuse_actor_base`` when the reference is exactly the +actor's frozen base checkpoint. AReno temporarily disables the adapter to +evaluate the base policy, avoiding a second model copy. Keep the default +``independent`` mode when the reference checkpoint is different. diff --git a/docs/index.rst b/docs/index.rst index f6cb17eb..7cc230ed 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -28,6 +28,7 @@ AReno documentation Chat Templates Dataset Formats Multimodal Inputs + Native LoRA Reward Functions .. toctree:: diff --git a/tests/test_train_cli_config_cpu.py b/tests/test_train_cli_config_cpu.py index 4fc9e69c..f495f1ee 100644 --- a/tests/test_train_cli_config_cpu.py +++ b/tests/test_train_cli_config_cpu.py @@ -386,6 +386,35 @@ def test_training_config_summary_can_colorize_output(): assert "AReno training config" in summary +def test_training_config_summary_shows_lora_parameters(): + cfg = _trainer_config_from_options( + **_options( + lora_rank=8, + lora_alpha=16.0, + lora_target_modules="q_proj,v_proj", + lora_adapter_path=None, + ) + ) + + summary = _format_training_config_summary(cfg) + + assert re.search(r"(?m)^ lora_rank\s+8$", summary) + assert re.search(r"(?m)^ lora_alpha\s+16\.0$", summary) + assert re.search(r"(?m)^ lora_dropout\s+0\.0$", summary) + assert re.search(r"(?m)^ lora_target_modules\s+q_proj,v_proj$", summary) + assert re.search(r"(?m)^ lora_adapter_path\s+none$", summary) + + +def test_training_config_summary_marks_lora_disabled(): + cfg = _trainer_config_from_options(**_options(lora_rank=None, lora_adapter_path=None)) + + summary = _format_training_config_summary(cfg) + + assert re.search(r"(?m)^ lora_rank\s+disabled$", summary) + assert re.search(r"(?m)^ lora_alpha\s+n/a$", summary) + assert re.search(r"(?m)^ lora_adapter_path\s+n/a$", summary) + + def test_dashboard_run_config_serializes_lora(tmp_path): cfg = _trainer_config_from_options(**_options(lora_rank=8, lora_alpha=16.0, metrics_log_dir=str(tmp_path))) From 0bb83fea41903f1662da3b6b2128ddf6bff564cf Mon Sep 17 00:00:00 2001 From: xsuler Date: Thu, 27 Aug 2026 07:57:10 +0800 Subject: [PATCH 27/29] test(lora): isolate optional CPU dependencies --- tests/test_bailing_v3_lora_cpu.py | 41 ++++++++++++++++++++++-------- tests/test_train_cli_config_cpu.py | 6 ++++- 2 files changed, 36 insertions(+), 11 deletions(-) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index a9682f3d..e23f7ebf 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -1,6 +1,7 @@ from __future__ import annotations -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace import pytest import torch @@ -107,6 +108,29 @@ def _single_tp() -> SimpleNamespace: return SimpleNamespace(rank=0, world_size=1, group=None) +@pytest.fixture +def bailing_model_module(monkeypatch: pytest.MonkeyPatch): + """Import the model without requiring optional FLA kernels in CPU CI.""" + + fla = ModuleType("fla") + fla.__path__ = [] + fla_ops = ModuleType("fla.ops") + fla_ops.__path__ = [] + lightning_attn = ModuleType("fla.ops.lightning_attn") + lightning_attn.chunk_lightning_attn = lambda *args, **kwargs: None + kda = ModuleType("areno.accel.kda") + kda.areno_kda_chunk = lambda *args, **kwargs: None + kda.areno_kda_recurrent_update = lambda *args, **kwargs: None + monkeypatch.setitem(sys.modules, "fla", fla) + monkeypatch.setitem(sys.modules, "fla.ops", fla_ops) + monkeypatch.setitem(sys.modules, "fla.ops.lightning_attn", lightning_attn) + monkeypatch.setitem(sys.modules, "areno.accel.kda", kda) + + from areno.models.bailing_v3 import model as bailing_model + + return bailing_model + + def test_bailing_v3_full_profile_attaches_native_slots(monkeypatch) -> None: monkeypatch.setattr(linear, "get_tp_context", _single_tp) monkeypatch.setattr("areno.adapters.lora.get_tp_context", _single_tp) @@ -154,9 +178,8 @@ def test_bailing_v3_requires_non_factorized_kda(monkeypatch) -> None: ) -def test_bailing_v3_empty_route_keeps_expert_router_and_lora_gradients(monkeypatch) -> None: - from areno.models.bailing_v3 import model as bailing_model - +def test_bailing_v3_empty_route_keeps_expert_router_and_lora_gradients(monkeypatch, bailing_model_module) -> None: + bailing_model = bailing_model_module experts = bailing_model.BailingGroupedExperts.__new__(bailing_model.BailingGroupedExperts) nn.Module.__init__(experts) experts.linear_fc1 = nn.Linear(3, 4, bias=False) @@ -195,9 +218,8 @@ def test_bailing_v3_empty_route_keeps_expert_router_and_lora_gradients(monkeypat assert all(torch.count_nonzero(parameter.grad) == 0 for parameter in parameters) -def test_bailing_v3_expert_lora_merges_only_into_derived_infer_weights() -> None: - from areno.models.bailing_v3 import model as bailing_model - +def test_bailing_v3_expert_lora_merges_only_into_derived_infer_weights(bailing_model_module) -> None: + bailing_model = bailing_model_module experts = bailing_model.BailingGroupedExperts.__new__(bailing_model.BailingGroupedExperts) nn.Module.__init__(experts) experts.local_num_experts = 2 @@ -257,9 +279,8 @@ def test_bailing_v3_routed_lora_keeps_cuda_graph_decode_enabled() -> None: assert qwen_runtime.eager_decode -def test_bailing_v3_kda_packed_a_matches_canonical_slots(monkeypatch) -> None: - from areno.models.bailing_v3 import model as bailing_model - +def test_bailing_v3_kda_packed_a_matches_canonical_slots(monkeypatch, bailing_model_module) -> None: + bailing_model = bailing_model_module monkeypatch.setattr(linear, "get_tp_context", _single_tp) monkeypatch.setattr("areno.adapters.lora.get_tp_context", _single_tp) monkeypatch.setattr(bailing_model, "areno_linear", torch.nn.functional.linear) diff --git a/tests/test_train_cli_config_cpu.py b/tests/test_train_cli_config_cpu.py index f495f1ee..0d0542cd 100644 --- a/tests/test_train_cli_config_cpu.py +++ b/tests/test_train_cli_config_cpu.py @@ -928,7 +928,9 @@ def test_train_help_places_epochs_under_basic_not_checkpointing(): def test_train_help_remains_complete_and_groups_every_declared_option(): ctx = train_cli.click.Context(train_cli.train_command) declared = { - param.name for param in train_cli.train_command.get_params(ctx) if param.get_help_record(ctx) is not None + param.name + for param in train_cli.train_command.get_params(ctx) + if param.get_help_record(ctx) is not None and not param.name.startswith("_click_") } grouped = [name for _, names in TRAIN_OPTION_GROUPS for name in names] @@ -939,6 +941,8 @@ def test_train_help_remains_complete_and_groups_every_declared_option(): output = _help_output() for param in train_cli.train_command.get_params(ctx): + if param.name.startswith("_click_"): + continue record = param.get_help_record(ctx) if record is not None: assert record[0].split()[0].rstrip(",") in output, f"option dropped from help: {param.name}" From ec24a6801734ddbedd673703472acd1ce9d22c84 Mon Sep 17 00:00:00 2001 From: xsuler Date: Thu, 27 Aug 2026 08:02:04 +0800 Subject: [PATCH 28/29] test(lora): stub Triton CPU imports --- tests/test_bailing_v3_lora_cpu.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index e23f7ebf..d278c7ea 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -121,10 +121,23 @@ def bailing_model_module(monkeypatch: pytest.MonkeyPatch): kda = ModuleType("areno.accel.kda") kda.areno_kda_chunk = lambda *args, **kwargs: None kda.areno_kda_recurrent_update = lambda *args, **kwargs: None + accel_ops = ModuleType("areno.accel.ops") + + class _KernelConfig: + def __init__(self, *args, **kwargs) -> None: + pass + + accel_ops.FusedMoeConfig = _KernelConfig + accel_ops.SegLaMeta = _KernelConfig + accel_ops.areno_fused_experts = lambda *args, **kwargs: None + accel_ops.areno_silu_and_mul = lambda *args, **kwargs: None + accel_ops.log_once = lambda *args, **kwargs: None + accel_ops.seg_la_fwd = lambda *args, **kwargs: None monkeypatch.setitem(sys.modules, "fla", fla) monkeypatch.setitem(sys.modules, "fla.ops", fla_ops) monkeypatch.setitem(sys.modules, "fla.ops.lightning_attn", lightning_attn) monkeypatch.setitem(sys.modules, "areno.accel.kda", kda) + monkeypatch.setitem(sys.modules, "areno.accel.ops", accel_ops) from areno.models.bailing_v3 import model as bailing_model From 5cb320f5cc4d679b74b3fe3adec1281544f124f6 Mon Sep 17 00:00:00 2001 From: xsuler Date: Thu, 27 Aug 2026 08:06:06 +0800 Subject: [PATCH 29/29] test(lora): complete CPU kernel stubs --- tests/test_bailing_v3_lora_cpu.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_bailing_v3_lora_cpu.py b/tests/test_bailing_v3_lora_cpu.py index d278c7ea..08cd28f6 100644 --- a/tests/test_bailing_v3_lora_cpu.py +++ b/tests/test_bailing_v3_lora_cpu.py @@ -131,7 +131,9 @@ def __init__(self, *args, **kwargs) -> None: accel_ops.SegLaMeta = _KernelConfig accel_ops.areno_fused_experts = lambda *args, **kwargs: None accel_ops.areno_silu_and_mul = lambda *args, **kwargs: None + accel_ops.can_use_cuda_kernel = lambda *args, **kwargs: False accel_ops.log_once = lambda *args, **kwargs: None + accel_ops.rms_norm_gate_fwd = lambda *args, **kwargs: None accel_ops.seg_la_fwd = lambda *args, **kwargs: None monkeypatch.setitem(sys.modules, "fla", fla) monkeypatch.setitem(sys.modules, "fla.ops", fla_ops)