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

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.

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",

areno/accel/csrc/extension.cpp

@@ -15,7 +15,10 @@ std::vector<torch::Tensor> 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<torch::Tensor> areno_grouped_linear_backward_cuda(
     torch::Tensor grad_output,
     torch::Tensor input,
     torch::Tensor weight,
-    std::vector<int64_t> tokens_per_expert);
+    std::vector<int64_t> tokens_per_expert,
+    bool need_grad_input,
+    bool need_grad_weight);
 std::vector<torch::Tensor> 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<torch::Tensor> areno_depthwise_causal_conv1d_silu_forward_cuda(torch::Tensor input, torch::Tensor weight);
 std::vector<torch::Tensor> areno_depthwise_causal_conv1d_silu_decode_cuda(
     torch::Tensor current,

areno/accel/csrc/linear.cu

@@ -193,57 +193,64 @@ std::vector<torch::Tensor> 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<torch::Tensor> areno_grouped_linear_backward_cuda(
     torch::Tensor grad_output,
     torch::Tensor input,
     torch::Tensor weight,
-    std::vector<int64_t> tokens_per_expert) {
+    std::vector<int64_t> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> areno_grouped_linear_backward_counts_cuda(
       counts[static_cast<size_t>(i)] = static_cast<int64_t>(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);
 }