feat: multi lora async#1638
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
left a comment
yushengsu-thu
left a comment
There was a problem hiding this comment.
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.)
_build_prefill_scoring_payload keyed LoRA solely on is_lora_enabled(args)
(always true in multi-LoRA) and sent the single-adapter name miles_lora,
which is never registered on multi-LoRA engines -> every step crashed with
'adapters are not loaded' when --recompute-logprobs-via-prefill is on; the
batched path also applied payloads[0]'s lora_path to a mixed-adapter batch.
- resolve lora_path per sample: adapter samples score under their own
__miles_slot_{N}, single-adapter LoRA keeps miles_lora, base keeps none
- batched scoring groups by (logprob_start_len, lora_path) and rejects
mixed-adapter batch payloads
- centralize the engine-side slot adapter name as slot_lora_name(), shared
by rollout, prefill scoring, and the weight-push mixin
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
The retire-time prefix abort fires exactly once (RETIRING->CLEANUP); requests
can survive it (a multi-turn group POSTs its next turn after the round, or a
request sits in the engine's tokenizer window) and, once the slot is reused
and the next adapter's weights are upserted into the same __miles_slot_{N},
keep decoding under the wrong weights with no error. Batch-time collection
filters already drop such groups from training data; these changes shrink the
window and stop wasting decode on them:
- MultiLoRABackend.free_slot: re-run the prefix abort right before the slot
is released for reuse (second round, after a full step of settling)
- sglang_rollout.generate: refuse to POST for an adapter that is no longer
sampleable (deregistered/cleaned up) and abort the sample locally instead
- the tokenizer-window escape itself is fixed engine-side (sgl-project#30912)
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
|
@mathewjhan |
Re-queuing aborted groups (add_samples) is intended but deferred to a future PR (confirmed by mathewjhan; left unimplemented for lack of testing time). strict=True so wiring it up forces marker removal. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
|
Thanks for confirming — marked |
Each sglang tokenizer worker process holds its own LoRA registry with no cross-worker sync (sgl-project/sglang#31084), so per-step adapter upserts resolve on whichever worker the router picks and fail non-deterministically. sglang rejects the upsert at runtime anyway; failing at launch avoids burning GPU time until the first weight push. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
fdefa55 to
f115610
Compare
1fbe335 to
75d7538
Compare
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
… bshd slice Two context-parallel regressions on shared (non-multi-LoRA) paths: - thd + allgather_cp: a chunk->batch rename also hit the torch method calls, leaving tokens.batch(...) / loss_masks.batch(...) which do not exist on torch.Tensor -- every --allgather-cp run crashed with AttributeError on the first micro-batch. Restore Tensor.chunk(). - bshd: tokens were sliced with slice_with_cp at the top of the branch and again in the non-allgather arm. chunk_size derives from max_seqlen, not the input length, so the second pass silently replaces the back half of each CP rank's tokens with pad zeros (shape unchanged, no error) and hands the allgather arm pre-sliced input. Drop the top slice, matching the structure of the same branch on main.
…e-slice bshd layout get_batch's CP paths had no coverage, which let both regressions in the previous commit ship undetected. Pin the normal-CP (zigzag) contract at cp_size=2 with CUDA stubbed so the tests run on CPU workers: - per-rank adapter_token_counts equal the post-slice per-sample lengths summed per slot, stream padding attributed to the last slot, and sum to the rank-local token count; - counts are identical across CP ranks (zigzag gives each rank the same padded share of every sample); - unsorted adapter_slots micro-batches are rejected; - bshd tokens match a single zigzag slice (regression test for the double slice_with_cp).
--global-batch-size defaults to None (Megatron derives it later), but the adapter-cap fallback guarded only hasattr, which is always true, so validation died with 'TypeError: int * None'. Leave the cap unset when the global batch size is not given; the existing is-not-None branch below already treats None as no cap. Fixes the shipped test_accepts_default_single_tokenizer_worker, which fails at HEAD.
…'s buffered state get_groups only discards a retired adapter's buffered groups and partial step stats when a later generate call observes the name missing from the snapshot. When the last adapter retires, the driver stops generating and idles, so that sweep never runs; re-registering the same name then ships the old tenant's groups (old prompts, rewards, slot_version) into the new tenant's first batch and merges the old rewards into its step statistics. The staleness filter does not help: same-slot reuse leaves a version delta of 1-2, within the run scripts' --max-weight-staleness 3. Give every registration a unique identity and fence rollout state on it: - AdapterRecord/AdapterRun carry a registration_id (uuid per register()); - process_group stamps it on samples next to the slot_version stamp; - get_groups drops a name's buffer and partial stats when its snapshot registration_id changes, and GroupBuffer.drop_foreign() drops stragglers of the old registration that land after that sweep (an in-flight generation of the retired tenant can complete at any time). Covered by two new batch-collection tests (re-registration reset, foreign straggler filtering) plus a registration stamp assertion in the process_group test.
…producer Two stacked defects let one bad registration stall every adapter behind a misleading empty-batch timeout: - a nonexistent data path was accepted at registration and only blew up when the producer's get_samples built the RolloutDataSource — killing the shared producer thread for ALL adapters; the driver then looped on EmptyBatchTimeoutError while get_or_create resurrected the worker into the same crash; - every multi-LoRA sample carries a RewardSpec, and _resolve_reward_config treated any spec as final, so an adapter without its own rm_type never fell back to the sample metadata or the process-wide --rm-type / --custom-rm-path: every generated group failed reward computation and was silently dropped, starving the batch collector the same way. Fixes, in three layers: - resolve_adapter_config rejects a registration whose data path does not exist or whose reward config resolves to nothing anywhere (adapter fields or process args), so the API returns a clear error up front; - _resolve_reward_config resolves per field: spec values win, unset fields fall through to sample metadata then args (the non-multi-LoRA spec-less chain is unchanged); - the producer records an unexpected death in worker.failure and collect_batch raises it immediately with the original cause instead of waiting out the empty-batch timeout. Covered by registration-validation tests, a reward-resolution precedence test file, and a dead-producer surfacing test; existing controller test fixtures now use a real data file and an explicit rm_type.
…per-param state Slot cleanup zeroed exp_avg/exp_avg_sq and the per-param state['step'], but TE/apex FusedAdam keeps the Adam clock per param GROUP: after a tenant trained N steps, group['step'] stayed N through retirement, so the slot's next tenant started with the previous adapter's bias-correction clock (its early updates scaled as if N steps had already happened). The per-param reset only ever covered torch's AdamW fallback, whose clock lives in state['step']. Reset the group clocks of the retired slot's param groups (matched by the miles_multi_lora_slot tag the per-slot optimizer builder stamps on them); co-tenant slots keep their clocks since every slot owns its own Adam child. Covered by slot-cleanup tests for both clock layouts, stubbing the bridge module so they run against any bridge build.
The model allocates every slot at --lora-rank, so export_adapter_weights yields max-rank-padded tensors, while adapter_config.json declares the adapter's real r — PEFT refuses such a checkpoint with a shape mismatch (the example's gsm8k adapter at rank 16 under --lora-rank 32 hits this on every save). The weight-sync push path already trims with slice_lora_to_rank; apply the same trim on the checkpoint export. The Megatron shard keeps the padded layout for slot resume, unchanged. Adds unit tests for slice_lora_to_rank (per-dim trims for lora_A/lora_B, hard failure on trained padding, pass-throughs), which had no coverage.
save_multi_lora_checkpoints elected writers with intra_dp.rank == 0, but
LoRA params are replicated across DP AND CP: with context_parallel_size >
1 every CP rank passed the gate, so cp_size ranks wrote the same
adapter_megatron_tp{t}_pp{p}.pt concurrently and raced each other's
shutil.rmtree/os.replace promotion of the same checkpoint dir. Gate on
the combined intra_dp_cp group instead: exactly one writer per (tp, pp)
shard and one global promoter.
The bridge LoRA model setup built DistributedDataParallelConfig with only use_distributed_optimizer set, leaving grad_reduce_in_fp32 at its False default — silently ignoring --accumulate-allreduce-grads-in-fp32 (which the multi-LoRA run scripts pass). This matters doubly for multi-LoRA: per-adapter accumulation RETAINS gradients in this DDP buffer across train batches, so a bf16 buffer compounds rounding error with every accumulated batch and every idempotent re-reduce. Wire the flag through.
…sions Version bumps drive the rollout-side staleness filter. A new/recovered engine triggers a full resync of every loaded adapter, and the bump set was the push set, so adapters whose weights had not changed aged one version per recovery: their buffered and in-flight groups drifted toward --max-weight-staleness and were dropped despite being perfectly fresh. Keep the full push, bump only the adapters that actually stepped (or were newly loaded) since the last push.
Argument validation requires --num-rollout (when num_epoch is unset) and the README says the run stops there, but the driver loop never read it — the only exits were adapter exhaustion or service-mode idling. Cap the loop at num_rollout as documented.
collect_batch already read multi_lora_max_empty_wait_s (falling back to the 30s default), but the argument was never registered, so the timeout was not actually configurable from the command line.
The per-slot optimizer machinery is Adam-only: _adam_init_state_fn seeds exp_avg/exp_avg_sq, and slot retirement cleanup resets exactly those moments plus the Adam step clocks. Muon already had a dedicated rejection, but any other non-Adam optimizer (e.g. sgd) built fine and silently leaked its state across slot tenants at retirement. Guard at argument validation and again in build_multi_lora_optimizer.
The v2 train group has no reconcile_adapters, so the multi-LoRA driver's first loop iteration dies with AttributeError after engines and models are already up. Fail at argument validation with a clear message instead. (Full v2 support also needs train-outcome propagation and reconcile fan-out to every cell; until then the combination is unsupported.)
The LR schedule budget is a single global sample count that every adapter's steps drain together: under a decaying --lr-decay-style, an adapter registered later in a service-mode run starts at an already-decayed LR and can sit at --min-lr for its whole life. A per-adapter schedule position is a larger design change; until then, surface the sharp edge at launch (the shipped examples use constant LR and are unaffected).
|
Note on Muon support (currently rejected at launch for multi-LoRA): it would be a small change (~50 LOC). Megatron's |
2 participants
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
in_placepause generation and LoRA upsert updatesDependent PRs
SGLang
ridprefix abort: feat(sglang-miles): Support aborting requests by rid prefix - multi-lora needs sgl-project/sglang#30912Unified above: sgl-project/sglang#31253 (for testing - do not need merge)
Megatron-Bridge (radixark)
Dev setup
Tested image:
radixark/miles:dev-202607090055You need to install the 3 forks:
SGLang
This fork has two small changes to SGLang (has both PRs):
Megatron-Bridge
This fork adds a multi-lora support to Megatron-Bridge along with some helper methods/user api
Miles
Running
see:
examples/multi_loraFor normal training (not as a service):
examples/multi_lora/provision.shexamples/multi_lora/run_job.shexamples/multi_lora/run_job.sh |& tee run.logFor multi-lora training as a long running service:
examples/multi_lora/provision.shexamples/multi_lora/run_service.shexamples/multi_lora/run_service.sh |& tee run.login one shellMinor features added to existing code (backwards compatible)
AdapterRefandRewardSpectoSampletype so individual samples can access their own reward functions and adapter names during rolloutMissing features for future