[Feature] Multi-Lora support to allow asynchronous RL finetuning for multiple tasks.

[HwVanICI]

Checklist

• This feature will maintain backward compatibility with the current APIs in
  areal/api/. If not, please raise a refactor issue first.

Motivation

Currently, AReaL mainly focuses on single-task RL finetuning. However, if there are multiple RL finetuning tasks to be run together, the memory, computation, and hardware requirements are significantly increased. This is because the only viable solution right now is to run each finetuning task separately. Nevertheless, in many scenarios, these tasks can be executed more efficiently through the use of Parameter-Efficient Fine-Tuning (PEFT) techniques, such as LoRA.

Therefore, there is a pressing need to support efficient PEFT finetuning techniques to enable a cloud-based system for finetuning multiple RL tasks in a multi-tenant, multi-user Finetuning as a Service (FaaS) setting.

In this RFC, we propose RL-based multi-LoRA finetuning that allows for the parallelized finetuning of multiple LoRA adapters for different RL tasks. This will enable resources to be shared more efficiently, thereby avoiding the inefficiencies associated with running them separately.

Proposed change

From a high level, we propose to add the following functionality to support multi-Lora RL finetuning support

• Multi-lora rollout support in vLLM inference engine.
• Multi-lora training support in the FSDP finetuning engine.
• Optimizations and performance enhancements.

To this regard, we propose a multi-milestone RFC to enable full support for multi-LoRA RL finetuning in AReaL:

Milestone 1: Basic LoRA Functionalities for Ascend-vLLM

• Support single-LoRA weight updates via load-from-disk
• Support single-LoRA weight updates via broadcast

Milestone 2: Advanced Multi-LoRA Features

• Enable multi-LoRA weight updates via load-from-disk or broadcast
• Support training multiple LoRAs either sequentially or concurrently

Milestone 3: Inference-side Optimizations

• Implement task-specific rollout interruption
• Improve rollout management and scheduling for load balancing and reduced rollout staleness

Milestone 4: Training-side Optimizations

• Optimize resource utilization via bubble reduction and balanced task-level allocation
• Implement job scheduling for training
• Optimize LoRA and optimizer state loading/offloading

Each milestone will correspond to concrete PR deliverables, with details as follows:

Milestone	Description
M1 — Basic LoRA Functionalities	Single-LoRA weights update via load-from-disk/broadcast
M2 — Advanced Multi-LoRA	Multi-LoRA updates & concurrent/sequential training support
M3 — Inference Optimizations	Task-specific rollout interruption + scheduling improvements
M4 — Training Optimizations	Bubble reduction, allocation, job scheduling, state offloading

Additional Information

See the following prior art and references:

• Motivation for LoRA PEFT: https://arxiv.org/abs/2106.09685
• Multi-LoRA kernel research: https://arxiv.org/abs/2505.14620

[garrett4wade]

Hi @HwVanICI, thanks for the proposal! I agree that the multi-LoRA feature is important for an RFT service, and we really want to integrate this feature within AReaL.

I have the following questions for clarification:

1. How would a job join and depart from the service? Do they have to be launched together, or can new training jobs join asynchronously? FYI, Tinker uses a clock-cycle-based solution.
2. What is the user interface? How will users know whether their job belongs to an existing multi-LoRA job, or whether they should start a new job?

From my perspective, I would suggest the following design:

1. The inference engine can remain as-is, supporting only "single-LoRA" inference. However, we would attach experiment and trial names to the LoRA's name. Different LoRA tasks would be bound to the same set of remote inference servers. Since inference servers (e.g., vLLM and SGLang) typically support serving different LoRA parameters, we can launch the server in advance and connect it with multiple LoRA fine-tuning tasks. This design makes implementing multi-LoRA serving inference clients unnecessary.

2. For the train engine, we should follow a similar design pattern as inference: launching training engines in advance as a service and passing the service address to the user's training job to let them create and update the LoRA parameters. In other words, we should design and implement a RemoteTrainEngine with corresponding HTTP servers. An example would look like:

# TrainEngine server
@app.post("/train_batch")
async def train_batch(request):
    lora_name = request.json()['lora_name']
    batch = request.json()['batch']
    global lora_server
    # The server will determine when to execute the real forward/backward pass
    # only after enough batches have been accumulated
    stats = await lora_server.train_batch_async(lora_name=lora_name, batch=batch)
    return stats

# Client/User code
engine = RemoteTrainEngine(config)
engine.initialize(addr=addr, lora_name=lora_name)  # The server's address
stats = engine.train_batch(batch)  # Send request and wait

Let's get feedback from @nuzant and @rchardx on the API design. I'll follow up by next Monday with our decision.

[HwVanICI]

Thanks a lot for the thoughtful questions and design suggestions! Below is our clarification and proposed alignment.

How jobs join or depart the system

Our intent is to allow asynchronous joining and departing of LoRA training jobs.

• A training job can join an existing multi-LoRA training pool at any time
• A job can finish and leave without disturbing others
• The scheduler assigns new LoRA tasks to a shared pool of workers when capacity allows
  This model is consistent with Tinker’s clock-cycle-based design, where a worker pool repeatedly performs forward/backward/optimizer operations and multiple LoRA tasks time-share the same pool.

User interface

Users do not need to explicitly choose between:

• joining an existing multi-LoRA pool, or
• starting a new pool
  Instead, the system will manage this automatically:
• Users provide only lora_name, experiment_name, and training config
• The scheduler routes each training job to an appropriate pool
• The pool’s RemoteTrainEngine exposes a stable address for users
  This keeps the interface simple and makes multi-tenancy transparent to end users.

Design Adaptation

We agree with your direction and propose the following alignment:

Inference
We will keep inference as “single-LoRA per rollout request,” as you suggested. Inference servers (vLLM/SGLang) will preload multiple LoRAs, and the client specifies the LoRA name. This keeps inference simple and avoids unnecessary multi-LoRA logic on the rollout side.

Training
Your suggestion of introducing a RemoteTrainEngine and a long-running shared worker pool is very much aligned with our vision for Milestones 3 and 4.

We propose:

• M1–M2: implement basic multi-LoRA weight update / sequential or concurrent training
• M3–M4: extend the training backend into a shared Training Pool, similar to Tinker’s clock-cycle design
  • Forward/backward requests from multiple LoRAs are queued
  • Worker pool executes them using a shared clock cycle
  • Optimizer step scheduling is coordinated inside the pool

[garrett4wade]

@HwVanICI The plan looks great! I think we are ready to move on.

[rchardx]

Based on my analysis of the AReaL codebase and the discussion in this issue, here are some suggestions:

---

I. Design Recommendations

1. Prioritizing RemoteTrainEngine Design

The RemoteTrainEngine design proposed by @garrett4wade is crucial. I suggest laying the groundwork for this abstraction in M1/M2 rather than waiting until M3/M4:

Rationale:

• The existing RemoteInfEngine has proven the feasibility of HTTP-based remote engines
• Introducing this abstraction early provides a clear upgrade path from M1/M2's disk-based and broadcast LoRA updates
• Reference the RemoteInfBackendProtocol pattern in areal/core/remote_inf_engine.py

2. LoRA Identification and Routing Mechanism

I suggest extending LoRA-related fields in WeightUpdateMeta:

@dataclass
class WeightUpdateMeta:
    type: Literal["disk", "nccl"]
    path: str | None = None
    # Existing fields...
    use_lora: bool = False

    # New fields
    lora_name: str | None = None  # Unique identifier for multi-lora scenarios
    experiment_name: str | None = None  # Associated experiment name
    trial_name: str | None = None  # Associated trial name

This enables dynamic lora_name support in SGLangBackend.build_disk_weight_update_requests, instead of the hardcoded "lora_1".

3. Thoughts on Clock-Cycle Scheduling

Tinker's clock-cycle design is elegant, but implementing it in AReaL requires consideration of:

• Integration with existing StalenessManager: The current StalenessManager manages staleness for a single task and needs to be extended for multi-task version awareness
• Propose adding MultiTaskStalenessManager:

  class MultiTaskStalenessManager:
      def __init__(self):
          self.task_versions: dict[str, int] = {}  # lora_name -> version
          self.task_stats: dict[str, RolloutStat] = {}
  
      def get_capacity_for_task(self, lora_name: str) -> int:
          """Calculate available capacity for a specific LoRA task."""
          ...

---

II. Early Execution Steps

Milestone 1 (Basic LoRA) Recommendations

1. Improve test coverage for examples/lora/gsm8k_grpo_lora.py

   • Current code shows that in LoRA scenarios only rank 0 submits rollouts (# NOTE: special design for lora, only rank 0 submits rollout)
   • Suggest adding end-to-end tests under areal/tests/

2. Unify LoRA configuration entry point

   • Currently LoRA configuration is scattered in FSDPEngine.__init__ using PEFT's LoraConfig
   • Suggest adding a unified LoRATrainConfig dataclass in areal/api/cli_args.py:

     @dataclass
     class LoRATrainConfig:
         r: int = 8
         lora_alpha: int = 16
         target_modules: list[str] = field(default_factory=lambda: ["q_proj", "v_proj"])
         lora_dropout: float = 0.0
         # ... other PEFT parameters

3. Enhance SGLang LoRA load/unload logging

   • Add more detailed logging in SGLangBackend.build_disk_weight_update_requests for easier debugging

Milestone 2 (Multi-LoRA) Recommendations

1. Design a LoRA Registry component

   # areal/core/lora_registry.py
   class LoRARegistry:
       """Central registry for managing multiple LoRA adapters."""
   
       def __init__(self):
           self.adapters: dict[str, LoRAAdapterMeta] = {}
           self.lock = threading.Lock()
   
       def register(self, name: str, config: LoRATrainConfig, owner: str) -> None:
           ...
   
       def get_active_adapters(self) -> list[str]:
           ...
   
       def get_adapter_for_request(self, request_id: str) -> str | None:
           """Route a request to the appropriate LoRA adapter."""
           ...

2. Extend ModelRequest to support LoRA routing

   @dataclass
   class ModelRequest:
       # Existing fields...
       lora_name: str | None = None  # Specify which LoRA to use

Milestone 3/4 Architecture Recommendations

1. Training Pool Service Architecture

   Suggest following the existing areal/launcher/ pattern and adding:

   areal/
     service/
       __init__.py
       training_pool.py      # Training pool HTTP server
       training_client.py    # Client for RemoteTrainEngine
       pool_scheduler.py     # Clock-cycle based scheduler
   

2. Integration strategy with vLLM/SGLang

   • Inference side: Reuse existing RemoteSGLangEngine, differentiate by lora_name parameter
   • Training side: Create new RemoteTrainEngine service, independent of inference

---

III. Risks and Challenges

1. Memory management: Multiple LoRAs loaded simultaneously on GPU may cause OOM. Suggest implementing LoRA offloading strategy in M4

2. Version consistency: In async multi-tenant scenarios, ensuring each LoRA's rollout uses the correct weight version is a key challenge

3. Test coverage: Multi-LoRA scenario testing requires simulating multiple concurrent tasks. Suggest establishing a dedicated multi-lora test directory under areal/tests/

---

IV. Documentation Recommendations

Suggest adding multi_lora.md under docs/customization/, including:

• Multi-LoRA architecture diagram
• Configuration examples
• Migration guide from single-task LoRA

---

cc @HwVanICI @garrett4wade @nuzant