[multi-lora] 3/n Tinker full stack: backend, frontend, and RL/SFT docs - #2365
Open
yushengsu-thu wants to merge 214 commits into
Open
[multi-lora] 3/n Tinker full stack: backend, frontend, and RL/SFT docs#2365yushengsu-thu wants to merge 214 commits into
yushengsu-thu wants to merge 214 commits into
Conversation
…ate serialization
Sample gains optional per-token float channels (loss_weights, advantages)
for client-supplied training data: response-aligned like loss_mask, merged
across turns like the OPD lists (zeros over injected observation spans),
carried on the wire as float32 typed_ragged, CP-sliced like
rollout_log_probs. The binary int32 loss_masks stay untouched.
miles/backends/megatron_utils/tinker_backend/checkpoint.py holds the slot
training-state serialization for the tinker-compatible backend: bf16
adapter weights + positional per-child optimizer state (fp32 masters, Adam
moments, step counters), per-rank atomic shards, rank-0 manifest committed
after a barrier, optional manifest ttl_seconds, and named immutable states
at states/{tag}. Loading fences on format, world topology, and LoRA
rank/alpha shape — never on the display name, so a new registration may
restore another run's state (create-from-checkpoint).
Provenance: radixark#2242 data-channel and checkpoint commits, renamed to the
tinker namespace, minus swap-in/out (they belong to the residency layer).
yushengsu-thu
force-pushed
the
agent/tinker-full-stack-example
branch
from
August 11, 2026 00:19
41145cf to
b96fdee
Compare
yushengsu-thu
marked this pull request as ready for review
August 11, 2026 00:26
yushengsu-thu
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23 and
yueming-yuan
as code owners
August 11, 2026 00:26
Fixed residency: a registration binds the lowest free slot for its whole life or queues behind a full pool (bootstrap drains the queue at retirement); there is no eviction, no bind-at-selection, and therefore no reservation transactions — tenancy changes only on the driver-sequenced register/deregister path. Pins mark slots whose state is immovable (dirty-grads: accumulated gradients no checkpoint carries). The run lifecycle is PENDING -> READY -> RETIRING -> CLEANUP -> COMPLETED, where READY comes from the trainer finishing the slot load — never from a weight publish: serving is a separate axis (serving_version stays 0 until save_weights_for_sampler) and record_weight_update no longer promotes. commit_tinker_step advances the per-run step clock, releases the dirty pin, and honors the optional client-set num_step bound; set_step repositions the baseline for state resume. AdapterRunConfig is the client-driven minimum: rank (server ceiling --lora-rank), optional save/num_step/metadata; alpha is server-resolved and never client-settable. Provenance: radixark#2137 slot pool/registry reworked for fixed residency and readiness/serving decoupling; radixark#2242 tinker lifecycle methods.
…tries, strict execution order One registration is strictly serialized: an operation is claimable only when every earlier ordinal has ARRIVED and reached a terminal state, which carries the client's per-model ordering end to end and keeps an optim_step from ever overtaking its forward_backward batches. Arrival may be out of order — the tinker SDK deliberately posts the first chunk of a large forward_backward last — so operations buffer by ordinal (consecutive from 1 per registration) and a gap below the head blocks all claims until it fills. NOTE: this reorder buffer moves to the tinker frontend when one lands. Retries are fingerprinted (sha256 over kind + canonical payload): re-enqueueing a known operation_id with identical content returns the original operation; different content is a conflict error, never silently swallowed. Cancel applies to QUEUED only and the cancelled ordinal still counts for contiguity; retirement fences open operations; terminal results are retained until acked (enqueue backpressure — mapped to HTTP 429 — is the capacity knob, never result eviction). Provenance: radixark#2242 operation ledger + the arrival/fingerprint upgrades from the design review.
…roller and HTTP surface
…t sums, all-rank NaN veto
…llection, forward-only batches
…oncile, publish staging, SDK-format metrics
…election, BatchPlan conversion, DP zero-weight padding
…rbs, arg defaults, example service
…, resume phases Phase A drives one adapter through the full operation lifecycle against a live service at DP=2 (register -> forward_backward x3 + odd-count fbs -> optim_step -> save_weights_for_sampler + router sampling -> save_state -> load_state -> post-restore fb/optim -> deregister), asserting result shapes (DP zero-weight padding never leaks rows), finite loss:sum/grad_norm, the publish barrier's serving identity, weight movement across optim_step, and post-deregister operation fencing. Phase B: forward operations return logprobs identical to a forward_backward of the same payload, take no dirty pin (save_state right after passes the unstepped-gradients gate), and an optim_step with nothing accumulated is an empty step (grad_norm 0.0, clock advances). Phase C: LayerWise DP sharding is real (disjoint per-rank ownership); a cross-slot restore is allowed exactly when the per-rank ownership signatures match (they coincide on this deployment) and bitwise-correct; a state whose shards carry a foreign signature (rank-swapped) is refused unanimously as a clean user error with the trainer staying healthy; a foreign-signature sidecar falls back to a fresh init at re-registration. Phase D: deregister writes the final sidecar; re-registering the same name auto-resumes it — step clock restored, probe logprobs identical, weights and optimizer fp32 masters/moments bitwise-preserved (no re-quantization), and training continues.
…ver the operation API Four adapters run independent client-driven RL loops against a live service (disjoint GSM8K shards, ranks 8/16/16/32, lr 1e-5/2e-5/4e-5/1e-5), 50 optimizer steps each: sample through the router with the adapter's serving name and rollout logprobs, score with the math grader, GRPO advantages (per-prompt mean baseline, std-normalized, sample-mean token scaling), forward_backward with loss_fn=importance_sampling, optim_step with grad_clip_norm 1.0, save_weights_for_sampler — the publish barrier keeps every loop on-policy. Per-step CSVs record reward, loss:sum, grad_norm, train-vs-rollout logprob abs-diff, and serving version; the final summary carries first/last-10 reward means, least-squares slopes, step clocks, and serving versions. H200 evidence (2 train DP=2 + 2 rollout GPUs, Qwen3-4B, thinking mode @ 512 new tokens): reward first-10 -> last-10 over 50 steps — rl_a 0.094 -> 0.481, rl_b 0.100 -> 0.603, rl_c 0.247 -> 0.803, rl_d 0.056 -> 0.275 (4/4 growing); step clocks exactly 50, serving versions 51/51/51/51 advancing independently, zero operation failures; ~149 optimizer steps/h per adapter (~595/h aggregate).
…on still fills its arrival slot A tinker HTTP frontend rejects some submissions AFTER the client has spent the (model, seq) ordinal: the SDK never resends a rejected sequence number, so refusing to record it would leave a permanent arrival gap and every later operation of the registration would buffer forever. record_rejected() inserts the operation born terminal FAILED(user) — same identity rules as enqueue (idempotent identical retry, conflict otherwise), hole-filler treatment for backpressure — and reject_operation() exposes it on the backend + controller.
The protocol layer of the SDK frontend, verified against the tinker==0.24.1 wheel source and captured live traffic (JSON path: proto_write_fwdbwd is the wheel's own default False, so the frontend serves pure JSON and vendors no protobuf). wire.py mirrors the request shapes the SDK actually POSTs and the client-config flags that pin it to this protocol; translation.py bridges the official next-token Datum to the backend's trailing-response-span sample (tokens + [target[-1]], response_length = N), requires true next-token alignment wherever a position carries loss, and types every v1 boundary rejection (sparse/top-K/multimodal/CISPO/DRO/seed) as UserInputError.
One SDK training client == one backend registration; every training verb forwards ordinal = seq_id verbatim (the 0.24.1 per-model counter is exactly the ledger's per-registration ordinal contract — the D5 note), and a submission the frontend rejects still consumes its ordinal as terminal FAILED(user). Request ids are deterministic in the SDK's own coordinates so retries replay and true conflicts 422 (never 409, which the SDK retries). retrieve_future long-polls, stores terminal bodies for replay BEFORE acking the backend record, and translates results per kind; save_state mints tinker:// paths into an in-memory catalog; ephemeral sampler publishes bind (name, registration_id, serving_version) and go stale loudly on republish; asample proxies to the sglang router under the registration-scoped serving name with the versioned KV extra_key. Tests drive the service against a real TinkerBackend executed by a fake driver that speaks only the documented trainer verbs.
TinkerFrontendHTTPServer extends the controller's registration server with the /api/v1 routes tinker==0.24.1 speaks, selected via --tinker-frontend (or --multi-lora-http-server-path). Backpressure maps to 429 + Retry-After (retryable to the SDK), conflicts to 422, expired/unknown futures to 410 (the SDK raises a retryable "promise expired" toward the caller — it does NOT re-run training requests, so delivered results answer 410 from fingerprint tombstones). With --tinker-api-key/$MILES_TINKER_API_KEY set, every route except the health probes requires X-API-Key (constant-time compare); the operator plane (/adapter_runs*, /info) additionally accepts loopback peers only, whatever the bind — the SDK key is a client credential, never an operator one. A non-loopback bind without a key refuses to start; --tinker-frontend without --tinker-backend (or a key without the frontend) fails loud at validation.
…e HTTP tinker.ServiceClient(base_url, api_key) against a real localhost uvicorn: capabilities/auth, create -> fb -> optim -> forward chain with metrics and future pipelining, >1024-datum forward_backward (the SDK splits chunks and posts the first one LAST — gap-buffered reorder + combiner reassembly), CE/IS/PPO, typed user failures that consume their seq, 429 backpressure retried to success, save_state -> create_training_client_from_state_with_ optimizer resume chain, weights-only resume as a typed rejection, immutable states, ephemeral publish -> sample with serving identity on the router payload, stale-after-republish fail-loud, base-model sessions, and low-level models.unload. Skipped where the tinker wheel is absent (hosted CPU CI).
The README gains the /api/v1 frontend section (launch flags, SDK pin, the seq->ordinal mapping, and the frontend-level v1 rejections); the ledger's frontend note now records the decided design: the frontend forwards the SDK's per-model seq_id verbatim, so the backend gap buffer IS the reorder point, and rejected submissions consume their ordinal terminally.
The serving fork accepts a per-request sampling_seed, so a client seed can be honored instead of rejected: each fanned-out sample i gets seed + i — deterministic per request, still diverse across num_samples.
The golden-acceptance clients that drive the live H200 deployments through the UNMODIFIED official tinker==0.24.1 SDK (base_url + api_key only): - tinker_sdk_mini_loop.py — cookbook-style supervised loop: capabilities -> create_lora_training_client(rank=16) -> 10x fb(cross_entropy)+optim(1e-4) with decreasing loss:sum -> publish + sample (coherent continuation) -> save_state -> load_state_with_optimizer -> fb/optim resumes -> a >1024-datum fb (SDK chunks, posts the first chunk last; the ledger reorders) -> a channel-mismatch datum surfacing as a typed RequestFailedError that consumes its seq AND poisons its gradient window (radixark#2258 §5): the window's optim_step fails as a discard, and the next round steps normally. - tinker_sdk_rl_quality.py — the SDK port of tinker_rl_quality.py: four concurrent GRPO loops on disjoint GSM8K shards (ranks 8/16/16/32, lrs 1e-5/2e-5/4e-5/1e-5), thinking mode, 50 optimizer steps each; per step sample -> grade -> grouped advantages -> fb(importance_sampling) -> optim(grad_clip_norm=1.0) -> save_weights_and_get_sampling_client as the on-policy publish barrier. Step clock / serving version evidence comes from the operator /adapter_runs routes on the same uvicorn. - tinker_sdk_poison_window.py — the poison-window collective semantics on a live DP=2 deployment (radixark#2258 §5): a good fb EXECUTES into the window, a failed chunk poisons it, and the window's optim_step must discard on every rank — probe forward logprobs re-read EXACTLY (bit-for-bit) after the discard, the recovery step's grad_norm equals a clean-window reference for the same batch (residue would double it), step/serving clocks hold, a concurrently-training neighbor adapter never perturbs, and a 1030-datum fb whose LATE chunk fails after the 1024-datum chunk landed discards just the same.
…le tests Multi-LoRA requires --tinker-backend at launch and the tinker rollout fn always attaches the batch execution lease, so the else-branch that trusted per-sample stamped slots (and the two tests exercising it: stamped-slot fallback, non-tinker heterogeneous reward normalization) can never run in production. Convert now fails loudly when an adapter-stamped batch arrives without a lease, a new guard test pins that, and the one legacy-channel test that stamped adapters incidentally now uses plain samples as production legacy batches do.
test_rollout_data_conversion.py::test_unaligned_input_is_trimmed_to_multiple already pins the same trim-to-multiple behavior on the same production branch; keeping a second copy in the padding suite adds maintenance cost without coverage.
The abort_all_requests=False behavior was introduced on main (radixark#2589), not by this stack, so its regression test belongs in a standalone test-only PR against main rather than riding the tinker backend; the file returns to its main-tree content.
_thread_main records a dead loop in run.error and the harness previously still exited 0 after printing the summary, so a wrapper (or a human checking $?) would read an aborted run as a pass; the summary and CSVs still land first, then the process fails if any loop aborted.
All 35 sites flagged by review 5003387723 on radixark#2273: each multi-line comment, docstring, assert/error message, or argparse help string is now a single line keeping the load-bearing invariant; no behavior change (test-matched substrings preserved).
…into tinker-frontend
…r-full-stack-example
Unacked terminal results of a retired registration lived in the ledger forever (probe: 50 dead registrations retained 14.1 MB, monotonic). fence() now strips request payloads (fingerprints keep retry identity, results stay pollable), and evicting a COMPLETED record from the registry ring fires drop_tenant, purging the tenant's queue and by_id entries. The eviction slice also clamps at zero so under-cap rings no longer evict early. A dropped operation polls as None; the frontend already maps missing operations to typed tombstones.
…uted sync LoRA sync sends only adapter tensors and never refills base weights; opening the session anyway makes begin/end_weight_update restore and re-pack the quantized base buffers with nothing loaded in between, corrupting the frozen base (reproduced on Kimi-K2.5 W4A16, TP8). Also re-adds the update_weight_version abort_all_requests=False wire pin so main radixark#2589's no-abort behavior cannot silently regress. Absorbed from closed PRs radixark#2715 and radixark#2713.
Move the import-integrity checks out of the fast suite and into tests/ci/verify_source_resolution.py, which every CPU and GPU CI job runs before pytest: statically resolve every miles-internal import site (including function-local ones) across miles/ and examples/, walk the optional namespaces in full when present, and import the update_weight lazy-import targets. Failures raise RuntimeError with the offending file:line and import target.
…bridge The launch-time probe (_bridge_recompute_patch_recognizes_multi_lora and its source-inspection helper) rejected full recompute and expert-target MoE recompute on a Megatron-Bridge without radixark#27. The deployment now tracks the bridge branch, which carries radixark#27, so the pre-radixark#27 shape can no longer reach launch; the guard and its test are retired. The CI-level LayerWise dependency canary remains the guard against a stale image. README and docs mirror drop the bridge-version requirement wording; supported recompute combos stay documented.
Absorbs the recently merged split PRs and maintainer changes. One conflict: miles/backends/megatron_utils/multi_lora_optimizer.py was modified on main (Megatron-LM bump renamed enable_gloo_process_groups to use_gloo_process_groups) but the stack deletes the legacy adapter-sample-level path entirely; resolved as deleted, and the same rename is applied to the stack's replacement (api_backends/multi_lora/optimizer.py and its test) so it matches the bumped Megatron argument name.
5 tasks
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.
Summary
This is the documentation/integration leaf of the Multi-LoRA Tinker stack.
Because this PR formally targets
main, GitHub's Files changed tabincludes the backend and frontend inherited from the first two PRs. The
logical review order is:
Relative to the current #2346 head, this PR is documentation-only:
examples/multi_lora_operations/README.mdis the source of truth.docs/examples/multi-lora-operations.mdis its generated docs mirror.No runtime or test code is unique to this layer.
Reviewer
Miles executes ordered training operations and serves explicitly published
adapter weights. The application still owns its data, sampling loop, reward
function, scoring, advantages, and update schedule.
What this layer adds
The guide shows:
--tinker-frontend;sample → score → build Datum → forward_backward → optim_step → publish;Quick start:
python examples/multi_lora_operations/run_multi_lora_operations.py prepare python examples/multi_lora_operations/run_multi_lora_operations.py serve \ --extra-args "--tinker-frontend"The full runnable RL version is
tests/e2e/tinker_frontend/tinker_sdk_rl_quality.py.Validation
At current head
83b55d13c:e87e4ef2dis an ancestor.generated docs mirror: two files, +188/-2 total.
git diff --check: passed.current GitHub run.
documentation-only leaf; runtime GPU/E2E evidence remains in [multi-lora] 1/n operation backend: explicit training with Tinker compatibility #2273 and
[multi-lora] 2/n Tinker frontend: tinker==0.24.1 JSON/REST adapter #2346.
Scope and limits
tinker==0.24.1JSON core-loop subset, not the 0.25+ protobufprotocol.
cross_entropy,importance_sampling, andppo.not implemented.
Datumconstruction.FullParameterExecutoris an implemented but unwired optimizer-control seamfrom [multi-lora] 1/n operation backend: explicit training with Tinker compatibility #2273. Full-parameter launch, training flow, checkpointing,
publication, and GPU/E2E support are outside this stack.
See #2273 and #2346 for the detailed backend and frontend contracts.
0821 liveness fixes merged forward from #2273/#2346: capped generate-failure tolerance in the driver, claimed-operation TTL backstop, FAILED child-runtime cooldown self-heal.