[WIP][multi-lora] rewrite: tinker-compatible multi-LoRA backend by yushengsu-thu · Pull Request #2775 · radixark/miles

yushengsu-thu · 2026-08-27T04:16:40Z

Clean and refactor now (#2365): split into

Empty commit so a draft PR against main can open before the first implementation commit lands.

@bridge

Multi-LoRA v2 needs the MultiLoRA fixes on Megatron-Bridge @bridge (the pinned 7f0fb345 revision predates them; moving ref, so rebuild with --no-cache).

AdapterIdentity (AdapterRef's name/slot plus the never-reused registration_id) and the committed ServingRef give registry, rollout, and trainer one exact-identity shape without a new datatype module; the two optimizer step policy literals live with their implementer so tinker's request-driven AdamW-on-sum and E2E's scheduled mean share one dispatch vocabulary.

yushengsu-thu · 2026-08-30T06:53:04Z



+@dataclass(frozen=True)
+class AdapterIdentity:


maybe change name to LoraIdentity

In all multi-lora related files (including e2e or tinker) class, function, var with adapter prefix or postfix need to be changed to lora

install_adapter_slot/release_adapter_slot carry only the rank-local mechanism (optional state_dict load-then-init, seed forked off the global RNG stream, clear plus optimizer/grad/scheduler cleanup) so the tinker frontend can drive slots without inheriting checkpoint auto-resume, controller RPCs, or the final-save policy, which stay in the _register/_deregister e2e wrappers unchanged.

A batch may carry request_loss_fn (+config) to pick the loss per call; without it dispatch is unchanged. Rejected under recompute_loss_function because checkpoint reruns the loss fn in backward.

forward_with_request_loss reuses the forward_only pipeline (new default-off extra_batch_keys/include_batch_in_f params) and a collector that returns detached per-sample logprobs in input order plus the request-scoped loss.

forward_backward_with_request_loss runs train() with optimizer=None so grads accumulate under keep-grads; request loss keys ride the batch via the new extra_loss_batch param and per-step losses come back via loss_sink.

step_adapter_slots now takes per-slot grad scale and clip (E2E passes 1/batch + args.clip_grad, tinker passes 1.0 + AdamParams.grad_clip_norm) plus an opt-in finite vote that zeroes and skips a non-finite slot; apply_adam_params writes per-call AdamW hyperparameters, and slots without a scheduler are skipped.

Each rank saves its owned share of one slot's fp32 masters, Adam moments, and step clocks under slot-neutral param keys (per-global-rank files, since LayerWise scatters whole params across DP); full load restores them and syncs master to model instead of reload_model_params. save_multi_lora_checkpoints gains an opt-in include_optimizer_state.

export_adapter_slot gathers one slot's HF-named adapter tensors sliced to the real rank as an immutable CPU snapshot (no push, no ServingRef commit); the checkpoint HF export now goes through the same function.

One preflight/execute pair per backend verb on MegatronTrainRayActor, each validating the exact slot binding first, delegating to the existing primitives; the pure validators live in utils/multi_lora.

The GPU gate showed init_adapter_slot only sets rank/alpha masks while weights come from the previous clear's unseeded re-init, so a seeded create was not deterministic; the seeded path now clears and re-inits under a scope that also reseeds the megatron CUDA-RNG tracker.

install_adapter_slot now passes seed straight to the bridge's init_adapter_slot (radixark/Megatron-Bridge#34), which owns LoRA weight init and the RNG-tracker reseeding; the local _seeded_rng_scope dance is gone. A seed is ignored when loading a state_dict.

One fast-gpu test drives every rank-local handler on a real tiny MultiLoRALinear + per-slot optimizer: create (seeded), forward/FB delegation, optim_step, save_state, load_state with and without optimizer, export_adapter, and release.

yushengsu-thu · 2026-08-30T23:01:05Z

+    for slot, scale in step_scales.items():
        children = _slot_children(optimizer, slot)
-        # Copy accumulated main_grads into the owned masters' grads, then scale the sum to the adapter-batch mean.
+        # Copy accumulated main_grads into the owned masters' grads, then apply the caller's scale to the sum.
        for child in children:
            child.prepare_grads()
-            for main_param in child.get_parameters():
-                if main_param.grad is not None:
-                    main_param.grad.mul_(1.0 / batch_size)
+            if scale != 1.0:
+                for main_param in child.get_parameters():
+                    if main_param.grad is not None:
+                        main_param.grad.mul_(scale)


note: check if need any optimization in computing on later

yushengsu-thu · 2026-08-30T23:01:30Z

(Above) Backend, Trainer:

Refactor the backend and trainer to add interactive (Tinker) support to the existing non-interactive (E2E) multi-LoRA backend.
The interactive (tinker) operation/SDK require the following functions: (tinker operations SDK → miles multi-lora functions):

ServiceClient.create_lora_training_client() → multi_lora_create_model()
TrainingClient.forward() → multi_lora_forward()
TrainingClient.forward_backward() → multi_lora_forward_backward()
TrainingClient.optim_step() → multi_lora_optim_step()
TrainingClient.save_state() → multi_lora_save_state()
TrainingClient.load_state() → multi_lora_load_state(with_optimizer=False)
TrainingClient.load_state_with_optimizer() → multi_lora_load_state(with_optimizer=True)
TrainingClient.save_weights_for_sampler() → multi_lora_export_adapter() — only the export candidate is currently implemented; SGLang publication and commit are not yet implemented.

Serving-only consumers need /generate without Sample, reward, or GenerateState coupling; generate() keeps admission and sample mutation and now delegates payload assembly, POST, and logprob parsing.

Per-adapter sources become RolloutDataSourceWithBuffer so recycled groups actually return, reconcile rebuilds a source when its name is re-registered so a new tenant never reads the old tenant's data, and get_samples takes an eligibility filter for admission-gated producers.

get_groups now pops an adapter only as one complete batch under its exact ServingRef and blocks it until publication advances the version; the producer caps buffered+inflight groups at the open batch quota. The old partial-pop unit tests are removed in favor of a GPU CI rollout-to-handoff E2E (test_multi_lora_rollout.py) driving real SGLang multi-LoRA generation through admission to the handoff.

# Conflicts: # miles/backends/megatron_utils/model.py

yushengsu-thu · 2026-08-31T07:31:38Z

(Above) rollout (e2e):

abstract async def call_sglang_generate_endpoint(...) and def parse_output_token_logprobs(...) in miles/rollout/sglang_rollout.py for tinker re-using
Changed codes in async_rollout.py and data_source.py: each LoRA must complete its full rollout before Trainer admission.
Preserved fully asynchronous multi-LoRA rollout and multi-slot batch coalescing in async_rollout.py.

Re-add the tinker==0.26.1 Dockerfile pin and add ensure_tinker_runtime(), failing fast on Python < 3.11 or a missing or mismatched SDK without importing tinker.

Add --multi-lora-frontend with an e2e/tinker validation split, a serving-only RolloutManager branch that skips dataset, rollout-fn, reward, and eval wiring, an epoch-math guard, and a launcher branch that boots the SGLang fleet and hands the router URL to the controller. Per-request sampling stays off the RolloutManager per the direct-to-SGLang design.

Serve the unmodified tinker SDK sampling plane from the controller process: sessions, long-poll futures, dedupe, two-level admission, and proto sample terminals in miles/tinker, dispatched straight to SGLang through a one-shot post seam on the neutral endpoint. Generalize MultiLoRABackend aborts to exact-rid/prefix and mount the routes on the existing control-plane app.

Add tinker/datum_batch.py turning ForwardBackwardInput wire dicts into the existing rollout_data keys with a plain loss spec and (request_id, row) ownership sideband, and tinker/losses.py validating loss names, configs, and required inputs; cispo and dro reject as unsupported.

yushengsu-thu · 2026-08-31T23:18:44Z

(Above) Rollout (tinker): - most effort on miles/tinker

Pin Tinker 0.26.1 and validate the runtime only when the Tinker frontend starts.
Keep RolloutManager responsible only for SGLang fleet/router lifecycle; create no E2E or Trainer resources.
Implement sampling HTTP, Futures, admission limits, and abort handling; call SGLang directly.
Decode Datum/loss inputs while preserving row order, adapter identity, and request ownership.
Validate the unmodified SDK through Trainer-ready input without executing training.

yushengsu-thu commented Aug 29, 2026

View reviewed changes

Comment thread docker/Dockerfile Outdated

yushengsu-thu force-pushed the new2-multi-lora-tinker branch 2 times, most recently from 5812cd7 to 16d544e Compare August 29, 2026 22:09

chore: anchor multi-lora v2 development branch

3d104c8

Empty commit so a draft PR against main can open before the first implementation commit lands.

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from e1abe87 to 3d104c8 Compare August 30, 2026 06:33

docker: prepare Megatron-Bridge upgrade

1499baa

Multi-LoRA v2 needs the MultiLoRA fixes on Megatron-Bridge @bridge (the pinned 7f0fb345 revision predates them; moving ref, so rebuild with --no-cache).

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from ab443ed to 0961dad Compare August 30, 2026 06:41

yushengsu-thu commented Aug 30, 2026

View reviewed changes

Comment thread miles/backends/training_utils/loss_hub/losses.py Outdated

multi-lora: request-scoped loss seam in loss dispatch

76f2699

A batch may carry request_loss_fn (+config) to pick the loss per call; without it dispatch is unchanged. Rejected under recompute_loss_function because checkpoint reruns the loss fn in backward.

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from 912823b to a56c6b8 Compare August 30, 2026 07:22

yushengsu-thu added 7 commits August 30, 2026 00:36

multi-lora: exact-slot export candidate

29e21bc

export_adapter_slot gathers one slot's HF-named adapter tensors sliced to the real rank as an immutable CPU snapshot (no push, no ServingRef commit); the checkpoint HF export now goes through the same function.

multi-lora: rank-local operation handlers on the megatron actor

5b6b8d5

One preflight/execute pair per backend verb on MegatronTrainRayActor, each validating the exact slot binding first, delegating to the existing primitives; the pure validators live in utils/multi_lora.

yushengsu-thu mentioned this pull request Aug 30, 2026

[peft] Optional seeded re-init in model-level init_adapter_slot radixark/Megatron-Bridge#34

Merged

yushengsu-thu commented Aug 30, 2026

View reviewed changes

Comment thread miles/backends/megatron_utils/multi_lora_utils.py Outdated

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from a121c17 to f5a0666 Compare August 30, 2026 18:15

yushengsu-thu commented Aug 30, 2026

View reviewed changes

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from b05765e to 201f09d Compare August 31, 2026 00:30

yushengsu-thu commented Aug 31, 2026

View reviewed changes

Comment thread miles/rollout/sglang_rollout.py Outdated

multi-lora: extract call_sglang_generate_endpoint from rollout generate

fc71184

Serving-only consumers need /generate without Sample, reward, or GenerateState coupling; generate() keeps admission and sample mutation and now delegates payload assembly, POST, and logprob parsing.

yushengsu-thu force-pushed the new2-multi-lora-tinker branch 6 times, most recently from 115ed08 to 43f9e0f Compare August 31, 2026 04:49

yushengsu-thu force-pushed the new2-multi-lora-tinker branch 2 times, most recently from 806c765 to 061763a Compare August 31, 2026 06:31

yushengsu-thu added 2 commits August 31, 2026 00:17

Merge remote-tracking branch 'upstream/main' into new2-multi-lora-tinker

c9f58d3

# Conflicts: # miles/backends/megatron_utils/model.py

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from 061763a to b8dd437 Compare August 31, 2026 07:17

yushengsu-thu force-pushed the new2-multi-lora-tinker branch 3 times, most recently from bb0aca3 to c9f58d3 Compare August 31, 2026 17:44

yushengsu-thu changed the title ~~[multi-lora] v2: tinker-compatible multi-LoRA backend~~ [multi-lora] rewrite: tinker-compatible multi-LoRA backend Aug 31, 2026

multi-lora: tinker dependency pin and runtime gate

4e0efcc

Re-add the tinker==0.26.1 Dockerfile pin and add ensure_tinker_runtime(), failing fast on Python < 3.11 or a missing or mismatched SDK without importing tinker.

yushengsu-thu force-pushed the new2-multi-lora-tinker branch from 47bbcce to 4e0efcc Compare August 31, 2026 19:40

yushengsu-thu added 3 commits August 31, 2026 14:54

yushengsu-thu changed the title ~~[multi-lora] rewrite: tinker-compatible multi-LoRA backend~~ [WIP][multi-lora] rewrite: tinker-compatible multi-LoRA backend Sep 1, 2026

Navigation Menu

Uh oh!

[WIP][multi-lora] rewrite: tinker-compatible multi-LoRA backend - #2775

[WIP][multi-lora] rewrite: tinker-compatible multi-LoRA backend#2775
yushengsu-thu wants to merge 22 commits into
mainfrom
new2-multi-lora-tinker

yushengsu-thu commented Aug 27, 2026 •

edited

Loading

Uh oh!

Uh oh!

yushengsu-thu Aug 30, 2026

Uh oh!

yushengsu-thu Aug 30, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

yushengsu-thu Aug 30, 2026

Uh oh!

yushengsu-thu commented Aug 30, 2026 •

edited

Loading

Uh oh!

Uh oh!

yushengsu-thu commented Aug 31, 2026 •

edited

Loading

Uh oh!

yushengsu-thu commented Aug 31, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant



		@dataclass(frozen=True)
		class AdapterIdentity:

Uh oh!

Conversation

yushengsu-thu commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

yushengsu-thu Aug 30, 2026

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Aug 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yushengsu-thu Aug 30, 2026

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu commented Aug 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

yushengsu-thu commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yushengsu-thu commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

yushengsu-thu commented Aug 27, 2026 •

edited

Loading

yushengsu-thu Aug 30, 2026 •

edited

Loading

yushengsu-thu commented Aug 30, 2026 •

edited

Loading

yushengsu-thu commented Aug 31, 2026 •

edited

Loading

yushengsu-thu commented Aug 31, 2026 •

edited

Loading