feat: multi lora async

[mathewjhan]

Summary

Allow disagg fully async training on multiple loras (all-linear, excluding expert) per step in Miles using megatron-bridge.

Two LoRAs trained together (DAPO + GSM8K)

Feature

• Train multiple LoRAs in a single training step
• Use as a long running service, with in_place pause generation and LoRA upsert updates
• Use as normal training job, stopping when no more LoRAs are left to train
• FastAPI control plane to register and deregister adapter runs

Dependent PRs

SGLang

• rid prefix abort: feat(sglang-miles): Support aborting requests by rid prefix - multi-lora needs sgl-project/sglang#30912
• upsert LoRA loading: [lora-related] feat(sglang-miles): Support upsert when loading adapters from tensors/distributed sgl-project/sglang#30913

Unified above: sgl-project/sglang#31253 (for testing - do not need merge)

Megatron-Bridge (radixark)

• multi lora support: feat: multi-lora support Megatron-Bridge#20

Dev setup

Tested image: radixark/miles:dev-202607090055

You need to install the 3 forks:

SGLang

This fork has two small changes to SGLang (has both PRs):

• Upsert semantics for loading LoRAs with the same lora name, to avoid having to unload; currently, unloading requires there to be zero requests for that LoRA to be loaded. When running pause generation before unload, we cannot complete more requests, when there are still running requests on unload, it is impossible to unload the adapter and miles will hang
• Abort by rid prefix; this is useful to target specific requests under a "namespace", such as when a LoRA run has been deregistered from miles and we want to abort all of the requests for that LoRA adapter. The alternative is to track all requests for that LoRA adapter individually and send in individual abort requests, since SGLang currently doesn't support batch abort

git clone https://github.com/mathewjhan/sglang
cd sglang
git checkout sglang-miles-mathew-dev-2
uv pip install --system --break-system-packages --no-deps -e python

Megatron-Bridge

This fork adds a multi-lora support to Megatron-Bridge along with some helper methods/user api

git clone https://github.com/mathewjhan/Megatron-Bridge
cd Megatron-Bridge
git checkout radixark/bridge/multilora-2
uv pip install --system --break-system-packages --no-deps -e .

Miles

git clone https://github.com/mathewjhan/miles
cd miles
git checkout feat/multi-lora-async
uv pip install --system --break-system-packages --no-deps -e .

Running

see: examples/multi_lora

For normal training (not as a service):

1. run examples/multi_lora/provision.sh
2. configure W&B credentials and settings in examples/multi_lora/run_job.sh
3. run examples/multi_lora/run_job.sh |& tee run.log

For multi-lora training as a long running service:

1. run examples/multi_lora/provision.sh
2. configure W&B credentials and settings in examples/multi_lora/run_service.sh
3. run examples/multi_lora/run_service.sh |& tee run.log in one shell
4. wait for the service to be ready ("No adapters; sleeping for 5.0s...")
5. optionally run the smoke test in another shell

python3 examples/multi_lora/service_smoke.py \
      --api-url http://127.0.0.1:8068 \
      --data /root/gsm8k/train.parquet \
      --input-key messages \
      --label-key label \
      --rm-type math \
      --steps 5

Minor features added to existing code (backwards compatible)

• Add AdapterRef and RewardSpec to Sample type so individual samples can access their own reward functions and adapter names during rollout

Missing features for future

• Colocated + update_weight_from_tensor currently unsupported, only update_weight_from_distributed
• Per adapter batch size currently unsupported, all adapters stepped at once (will be added later when decoupling the optimizer)
• Currently doesn't checkpoint optimizer + scheduler state yet, but can be added later as future PR
• Dataset checkpoint loading per adapter
• MultiLoRA not applied to experts as of now due to more complex bookkeeping required (need to keep track of the [adapter index, routed experts] together for grouped gemm)
• Doesn't support pure LoRA load from *.bin/*.safetensors yet, can be added later as a future PR, only resume from a megatron checkpoint or train from scratch
• Support add_samples to recycle stale/aborted samples

[gemini-code-assist]

Code Review

This pull request introduces a fully-async multi-LoRA training framework, enabling concurrent training of multiple LoRA adapters against a shared base model. It adds a background rollout worker, an async round-robin data source, a Ray-based controller with an HTTP control-plane API, and integrates slot-keyed adapter management into the Megatron-LM and SGLang backends. The review feedback identifies a critical distributed deadlock risk in actor.py caused by non-deterministic set iteration during collective operations, which can be fixed by sorting. Additionally, the reviewer recommends replacing several assert statements with ValueError for input validation and closing the httpx.Client in the smoke test to prevent resource leaks.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[yushengsu-thu]

Ran a deep verification pass over this branch together with its two sglang dependency PRs (sgl-project#30912/#30913); every finding below was adversarially re-verified against the code. Two issues to flag on the miles side — one is a hard crash/correctness bug in the logprob-recompute path, the other is a lifecycle race that silently corrupts rollout data across adapter slot reuse. (Side note: the deeper sglang-side findings are filed on the two dependency PRs directly.)

[yushengsu-thu]

@mathewjhan examples/multi_lora/tests/test_multi_lora_async_rollout.py::test_process_group_recycles_aborted fails on this branch (pre-existing at e371b2e, unrelated to the M1/M2 fixes): the test expects aborted groups to be re-queued, but the code path documents that re-queuing is not wired up. Which direction do you want — wire up re-queuing for aborted groups, or adjust the test to pin the current drop-and-count behavior? I'm happy to pick it up once the direction is set.

[yushengsu-thu]

Thanks for confirming — marked test_process_group_recycles_aborted as xfail(strict=True) in 50f6ff1 so the suite is green on this PR while still pinning the intended end-state behavior (wiring up re-queuing will XPASS and force the marker off). Two things worth bundling into that future PR: (1) the re-queuing itself (per-adapter source needs a write path), and (2) aborted/stale-dropped rows currently still count toward num_row, so adapters quietly train on fewer samples than configured — probably the more impactful half. Also worth skipping re-queue for RETIRING adapters (their slot is being freed; nothing can serve the regenerated group).

[yushengsu-thu]

Merged radixark/main into this branch to resolve the conflict (the branch was 36 commits behind). Single conflict in miles/rollout/sglang_rollout.py — an import-line collision between the multi-LoRA imports and main's new call_agent_abort_hook; resolved as the union. compileall clean; PR is now MERGEABLE.

Updated with Claude Code.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.