[multi-lora] /v1 post-training API with per-adapter token metering - #1822
Closed
yushengsu-thu wants to merge 12 commits into
Closed
[multi-lora] /v1 post-training API with per-adapter token metering#1822yushengsu-thu wants to merge 12 commits into
yushengsu-thu wants to merge 12 commits into
Conversation
Training-side support for putting multi-LoRA adapters on MoE experts, on top of MultiLoRAGroupedExpertLinear in megatron-bridge. Force moe_permute_fusion off for multi-LoRA runs that target expert leaves. Most bridge MoE providers default it on (Qwen3-MoE among them), and with it on the dispatcher records TE's row_id_map instead of a token gather index, which expert adapters cannot replay. Disabling it only costs the fused permute kernel, so this turns off an optimization rather than refusing to build. Fix two things that would have corrupted or lost expert adapters: * slice_lora_to_rank addressed the rank axis from the front, so a packed grouped-expert export ([num_experts, rank, in] / [num_experts, out, rank]) was sliced on its expert or output axis. Address it from the end instead. * Megatron-native adapter shards were named by (tp, pp) only and written by DP rank 0. Expert-parallel ranks sharing a (tp, pp) coordinate hold different local experts, so their shards collided and a resume loaded one EP rank's experts onto every rank. Shards are now keyed by (tp, pp, ep) — the suffix is omitted at ep_size 1 so existing checkpoints stay loadable — and both the writer election and the completeness check come from one cached gloo all-gather of the realized coordinates. Neither can be derived locally: no single group rank elects one writer per coordinate, and the realized coordinates are not the cross product of the group sizes once expert TP is smaller than tensor TP, which expert multi-LoRA requires. Validate the configurations the adapter cannot serve. Model-dependent checks live in _validate_multi_lora_moe_support, where the provider values are concrete: expert TP must be 1, experts must be grouped (SequentialMLP linears are skipped, so the experts would train nothing), no fp8/fp4 experts, no capacity padding, and both expert projections must be targeted since sglang applies the expert delta after gate_up and after down together. Launch-time checks cover pipeline size 1 and --qkv-format thd; the latter closes a pre-existing hole, as per-slot token spans assume samples pack contiguously in the sequence-major flattening, which bshd's [b, s] batch does not. Expert TP is deliberately NOT checked against the CLI: its default is None and Megatron only resolves it to tensor_model_parallel_size in its own validate_args, which runs after miles' — so comparing the raw value rejected runs that simply omitted the flag. Also stop forwarding exclude_modules to MultiLoRA: ModuleMatcher asserts it is empty whenever target_modules is set, and the exclusion has already been subtracted from target_modules during argument validation.
Comment-only change addressing the 19 review comments: every flagged docstring/comment block keeps its load-bearing constraint and drops the derivation. No code change; the four touched fast-test files pass 48/48. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…sentence per review
…me, leaf-name mapping, squashed comments
yushengsu-thu
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23 and
yueming-yuan
as code owners
July 27, 2026 08:44
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
yushengsu-thu
force-pushed
the
radixark/multi-lora-api
branch
from
July 27, 2026 09:00
2d18968 to
700dd80
Compare
yushengsu-thu
changed the title
Every adapter registration carries a monotonic TokenUsage meter split by
class (prefill / cached-prefill / sample / scoring-prefill / train) and
side (rollout vs training). Rollout tokens are counted at engine-response
receipt (aborted/stale/filter/retire drops included, itemized in detail
buckets); train tokens bank through the existing exactly-once
record_batch_adapters -> mark_batch_trained commit. Cross-process
reporting uses incarnation-keyed cumulative snapshots (idempotent under
at-least-once delivery). A timestamped JSONL journal under
{save}/multi_lora_controller/ is the durable ledger: replayed on restart,
frozen per registration at slot free, retained across job deletion. An
external billing backend consumes the ledger; no rates live in miles.
Declarative resource surface over the existing controller verbs: postTrainingJobs (create/get/list/cancel/delete/batchGetState), datasets, evaluators, models (checkpoint scan + :download), /v1/info and /v1/usage. Typed error envelope with requestId, idempotent create, reserved names, 429 on slot capacity, shadow EXTERNAL jobs for ops-plane registrations, stopReason derivation (user/ops/max-steps). Slot indices never leak. The legacy /adapter_runs ops plane is unchanged, and the gateway never reaches into registry internals, so a Tinker-style primitives API can mount later as a sibling driver.
ui/console.html: zero-dependency console over /v1 — launch/stop server, one-click job creation with prefilled defaults (auto-provisions the demo dataset/evaluator), per-LoRA cancel, live slot usage (a/n), token-usage ledger and model/checkpoint views. Served at GET /ui by the gateway. ui/launcher.py: standalone host so the UI exists before the server does — serves the console, proxies /v1, starts/stops the miles server. ui/run.sh: one-command entry (tunnel + launcher + server + browser); --fresh wipes the usage ledger and checkpoints for a clean demo slate.
examples/multi_lora/v1_api_demo.py drives the public surface end to end: dataset/evaluator registration, N job submissions, step/usage polling, mid-run cancellation, final ledger and model/checkpoint readout. Same flow verified live on an 8xH200 devbox (4 jobs x 5 steps, 2 cancelled).
yushengsu-thu
force-pushed
the
radixark/multi-lora-api
branch
from
July 27, 2026 09:06
700dd80 to
ffc5229
Compare
yushengsu-thu
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Declarative post-training API over the multi-LoRA controller: submit LoRA jobs against a shared base model, meter every adapter's token usage (rollout/training, counting only — no pricing), and operate everything from a single-file web console. Tinker-style execution core untouched; legacy
/adapter_runsops plane unchanged.APIs
POST /v1/postTrainingJobs— create a LoRA post-training jobGET /v1/postTrainingJobs[/{id}]— list/inspect jobs, live progress and usagePOST /v1/postTrainingJobs/{id}:cancel— stop one adapter's trainingDELETE /v1/postTrainingJobs/{id}— remove terminal job, ledger keptGET /v1/postTrainingJobs:batchGetState— cheap bulk state pollPOST|GET|DELETE /v1/datasets— register cluster datasetsPOST|GET|DELETE /v1/evaluators— register reward configsGET /v1/models[/{id}]— adapter states and checkpointsGET /v1/models/{id}:download— checkpoint file pathsGET /v1/usage— per-registration token ledger (billing substrate)GET /v1/info— capacity and limitsGET /ui— web console: launch server, create/cancel jobsTests: 74 control-plane tests (
tests/fast/ray/multi_lora/,--noconftest); e2e-verified on an 8×H200 devbox (4 jobs × 5 steps, 2 cancelled mid-run).To-do