Skip to content

[multi-lora] 2/n Tinker frontend: tinker==0.24.1 JSON/REST adapter - #2346

Open
yushengsu-thu wants to merge 67 commits into
tinker-compatible-backendfrom
tinker-frontend
Open

[multi-lora] 2/n Tinker frontend: tinker==0.24.1 JSON/REST adapter#2346
yushengsu-thu wants to merge 67 commits into
tinker-compatible-backendfrom
tinker-frontend

Conversation

@yushengsu-thu

@yushengsu-thu yushengsu-thu commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #2273; review this PR against tinker-compatible-backend.
#2365 adds the full-stack example and client-owned RL/SFT walkthrough.

Summary

Adds miles/ray/tinker_frontend, a protocol adapter for the exact
tinker==0.24.1 JSON core-loop surface. An unmodified official SDK client can
create LoRA training clients, submit explicit training operations,
checkpoint/resume, publish sampler weights, and sample through the existing
SGLang router.

The client still owns data construction, sampling policy, rewards,
advantages, and batch scheduling. This PR translates protocol requests;
#2273 owns operation ordering and Multi-LoRA execution.

Reviewer

official tinker==0.24.1 SDK
             │
             ▼
TinkerFrontendHTTPServer
  ├── /api/v1 routes, auth, and HTTP error mapping
  ├── wire.py: exact 0.24.1 JSON request/response models
  └── TinkerFrontend
      ├── state: sessions, models, futures, checkpoints, samplers
      ├── translation: Datum / Adam / results ↔ backend payloads
      │
      ├── TRAINING
      │     └── MultiLoraOperationController (#2273)
      │         └── MultiLoraOperationBackend
      │             └── ordered ledger
      │                 ├── forward[/backward]
      │                 │   └── MultiLoraOperationBatchFn → Megatron
      │                 └── optim/save/load/publish
      │                     └── Multi-LoRA control handlers
      │
      └── SAMPLING
            └── SamplingAdmission (weighted by num_samples)
                └── bounded SamplingTransport → SGLang router /generate

Sampling never enters MultiLoraOperationBatchFn or the legacy
dataset/reward rollout path.

Protocol and operation identity

  • Health/readiness, capabilities, sessions, heartbeats, telemetry, model
    creation/info/unload, and asynchronous /retrieve_future polling.
  • forward, forward_backward, per-call optim_step, immutable full-state
    save/load with optimizer, and latest-weight sampler publication.
  • Base-model and ephemeral-LoRA /asample.
  • Text-only shifted targets for cross_entropy, importance_sampling, and
    ppo; cross-entropy reports chunk-additive loss:sum,
    unmasked_tokens:sum, and loss_weight:sum.
  • /api/v1/client/config keeps the 0.24.1 client on its JSON/httpx path.
    Other SDK versions are rejected at bootstrap; this does not implement the
    0.25+ protobuf protocol.

One SDK training client maps to one backend registration. Every training verb
forwards its SDK seq_id as that registration's operation ordinal.
Out-of-order chunks are gap-buffered by #2273. A frontend-rejected request is
recorded as terminal FAILED(user) so its already-spent ordinal cannot stall
later work.

Request IDs and fingerprints make identical retries replay-safe. Conflicting
content maps to 422, retryable pressure maps to 429 + Retry-After, and
expired futures map to 410. Terminal response bytes are retained before the
backend record is acknowledged; bounded fingerprint tombstones prevent an
expired identity from silently executing again.

Sampling safety and lifecycle

  • Global admission is counted in sub-generations: one request consumes
    num_samples permits across all SDK clients.
  • The router transport uses the same hard concurrency bound and matching
    connection pool, with no pool-acquisition timeout behind the admission gate.
  • prompt + max_tokens is checked against an explicitly configured engine
    limit or one lazily discovered through the router and a healthy worker.
  • Sampling parameters are rebuilt from an allowlist; client-supplied backend
    registration/routing identity never reaches SGLang.
  • A multi-sample failure cancels and awaits sibling generations. Shutdown
    stops admission, drains sampling tasks, and then closes the transport.
  • Session/future TTL maintenance cancels orphaned samples, resolves abandoned
    operation futures, and retires undelivered results into tombstones without
    freeing their execution identity.
  • Ephemeral-LoRA sampling is version-checked before and after generation, so a
    stale or mid-flight republished sampler fails instead of knowingly returning
    cross-version output.

For non-loopback binds, an API key is mandatory. When configured, it protects
the SDK surface except health probes. Operator routes such as
/adapter_runs* and /info remain loopback-only even with the SDK key.

Supporting changes

  • Adds backend facade projections needed by the frontend without exposing
    registry/ledger internals or the router URL.
  • Adds terminal rejected-ordinal recording for failures after the SDK has
    already consumed a sequence number.
  • Adds frontend launch, admission, context-limit, and retention flags.
  • Pins tinker==0.24.1 in CPU CI so the real SDK contract tests run, and keeps
    pytest package arguments contiguous under pytest 9.1.

Validation

At current stacked head e87e4ef2d:

  • H200 focused gate covering the frontend, operation backend, trainer/model
    initialization, arguments, predicates, and CI runner: 452 passed,
    3 skipped
    .
  • Frontend CLI surface smoke passed: all seven current flags are present and
    the two removed legacy flags are absent.
  • Pre-commit passed over all 40 effective PR files.
  • Current GitHub pre-commit, CPU, 2-GPU, and all but one selected 4/8-GPU
    shards pass. The remaining 4-GPU H200 shard timed out after 1,800 seconds in
    test_inkling_small_4layer_lora_ci.py, which is outside this PR's effective
    diff; no Tinker frontend test failed.

Earlier official-SDK GPU acceptance covered Qwen pure SFT/RL and GPT-OSS
multi-adapter SFT/checkpoint and RL on 2×H200. The current head subsequently
adds sampling-lifecycle hardening; that code is covered by the focused H200
gate above, but the complete SDK GPU acceptance matrix was not rerun at this
exact head.

Scope and limits

  • Requires [multi-lora] 1/n operation backend: explicit training with Tinker compatibility #2273. Multi-LoRA is the only wired training target; the unwired
    FullParameterExecutor seam is not exposed here and has no launch or
    GPU/E2E support claim.
  • Text-only dense 1-D targets; no multimodal input, sparse CSR, nested top-K,
    or additional loss families.
  • Full-state restore requires optimizer state. Named persistent sampler
    checkpoints and checkpoint TTL expiry are not implemented.
  • Sampling serves latest weights only and admission is frontend-global, not
    per tenant. Request/output byte quotas and token-ID vocabulary upper-bound
    validation remain out of scope.
  • Session, future, and tinker:// checkpoint catalogs are in-memory; process
    restart reconciliation for previously claimed operations is unchanged.

0821 liveness fixes merged forward from #2273: capped generate-failure tolerance in the driver, claimed-operation TTL backstop, FAILED child-runtime cooldown self-heal.

  • 0821: SamplingClient.compute_logprobs() / sample(include_prompt_logprobs=True) now served natively (sglang logprob_start_len=0; topk_prompt_logprobs still rejected typed); README matrix updated; real-SDK contract tests + live GPU parity vs TrainingClient.forward (mean |Δ|=0.032, max 0.138).

@yushengsu-thu yushengsu-thu changed the title tinker-frontend: SDK-compatible REST frontend [multi-lora] tinker-frontend: SDK-compatible REST frontend Aug 10, 2026
@yushengsu-thu
yushengsu-thu force-pushed the tinker-compatible-backend branch from a3d8038 to f442377 Compare August 10, 2026 20:42
@yushengsu-thu
yushengsu-thu force-pushed the tinker-frontend branch 2 times, most recently from e5b4196 to cfa35d4 Compare August 10, 2026 23:13
@yushengsu-thu
yushengsu-thu force-pushed the tinker-compatible-backend branch from f442377 to f969743 Compare August 10, 2026 23:13
@yushengsu-thu yushengsu-thu changed the title [multi-lora] tinker-frontend: SDK-compatible REST frontend [multi-lora] tinker-frontend: SDK-compatible REST frontend (tinker==0.24.1 JSON subset) Aug 10, 2026
@yushengsu-thu
yushengsu-thu marked this pull request as ready for review August 11, 2026 00:26
…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 (#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 (#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.
@yushengsu-thu
yushengsu-thu force-pushed the tinker-compatible-backend branch from f969743 to 0486b5a Compare August 11, 2026 03:23
yushengsu-thu and others added 12 commits August 17, 2026 18:48
… (codex-0817-sft-fix)

# Conflicts:
#	tests/fast/ray/tinker_backend/test_backend.py
#	tests/fast/ray/tinker_backend/test_operations.py
…or; permanent SFT/pre-HTTP contract suite (codex-0817-sft-fix §3.2/§7)

The golden-acceptance mini-loop trained on all-ones-weight plain-LM
datums, so unmasked_tokens:sum == loss_weight:sum held by construction
and the §7 denominator bug was invisible in the GPU smoke that exists to
catch exactly this class. The corpus is now teacher-forced prompt-masked
SFT (codex-0817-sft-fix §2: prompt weight 0, completion weight 1), the
per-token display divides loss:sum by loss_weight:sum (guarded — weights
are arbitrary finite floats), and the loop asserts the two metrics stay
DISTINCT and exact against the locally computed weight/position counts,
making the denominator a live GPU regression.

tests/fast/ray/tinker_backend/frontend/test_sdk_sft_contract.py makes
the report's §3.2 scratch adversarial suite a permanent regression over
the real tinker==0.24.1 wheel and the live HTTP stack:

- SFT training contract: three-fb accumulation stepping exactly once,
  prompt-masked and fractional CE weights (loss_weight:sum vs
  unmasked_tokens:sum separation on the wire), dirty-save rejection
  keeping its gradients, rejected Adam params not dropping the window,
  no-grad forward remaining checkpointable, non-destructive
  .result(timeout);
- pre-HTTP failure modes: the NaN-serialization hole now runs the FULL
  gap-timeout chain — typed RequestFailedError naming missing ordinal 2,
  step clock proving the sealed identity never executed, the SAME
  TrainingClient resubmitting successfully (turn counter advanced), and
  the ledger holding exactly one SealedGap; the immediate-cancel probe
  stays as SDK characterization (server queue empty — nothing
  server-side can terminalize it, per the 覆核-confirmed §5 verdict).

Fixtures are reused from test_sdk_contract by module reference (a
fixture import would F811 against the test parameters).
…ner's initialize hook on the canonical module instance

Post-main-merge, CI partition packing newly placed
test_model_initialize.py directly before the tinker trainer tests and
test_master_reload_skips_restored_slots started failing with
'ParallelState not initialized' — a deterministic cross-file pollution
chain, reproduced and bisected on the gate box:

- test_model_initialize's module fixture tore down via
  sys.modules.clear() + snapshot restore. That evicts EVERY module first
  imported during its window — including torch internals whose module
  bodies hold one-shot registrations (a later fresh re-import trips
  torch's mega-cache 'artifact already registered' assert) — and leaves
  stale submodule attributes on retained parent packages. The teardown
  now drops only the namespaces its stubs poisoned (miles/megatron/
  sglang) and restores the originals over the stubs; real third-party
  modules imported during the window stay put.
- test_trainer's monkeypatch targeted the initialize function by STRING
  path; pytest resolves that by walking package attributes from the top,
  so with a stale parent-package attribute it patched the evicted module
  instance while load_adapters' function-level import fetched the fresh
  one and called the real function. Both call sites now import the
  module and patch the canonical sys.modules instance directly.

Verified on the gate box: the exact failing partition pairing
(test_model_initialize.py + test_trainer.py) now passes 20/20; each file
still passes standalone.
…nager changes (RolloutFnHandoff, direct-await invocation)

Propagates the backend reverts b95bc46 + 51c9e14 to the frontend.
Conflict in validate_tinker_args resolved by dropping the three
dispatch-bypass asserts introduced by the reverted handoff commit while
keeping the frontend-side flag guards and the --tinker-frontend HTTP
server default (fe4). Round-3 cleanup, Tau sampling fixes,
preflight/metrics/reaper, and the SFT round all stay.
Signed-off-by: Ethan (Yusheng) Su <yushengsu.thu@gmail.com>
Signed-off-by: Ethan (Yusheng) Su <yushengsu.thu@gmail.com>
@yushengsu-thu yushengsu-thu changed the title [multi-lora] 2/n tinker-frontend: SDK-compatible REST frontend (tinker==0.24.1 JSON subset) [multi-lora] 2/n Tinker frontend: tinker==0.24.1 JSON/REST adapter Aug 18, 2026
yushengsu-thu and others added 17 commits August 18, 2026 14:56
# Conflicts:
#	docs/examples/index.md
#	docs/examples/multi-lora-operations.md
#	examples/README.md
#	examples/multi_lora_operations/README.md
#	miles/utils/tinker_backend.py
#	tests/e2e/multi_lora_operations/tinker_sdk_mini_loop.py
#	tests/e2e/multi_lora_operations/tinker_sdk_poison_window.py
#	tests/e2e/multi_lora_operations/tinker_sdk_rl_quality.py
The 0.24.1 SDK's compute_logprobs() is a 1-sample, 1-token asample with
prompt_logprobs=true, which the frontend answered with a typed v1
rejection. sglang scores prompts natively — logprob_start_len=0 returns
input_token_logprobs on the same generate — so translate the wire flag to
that router call and map the per-token scores (position 0 has no context
and stays null) into SampleResponse.prompt_logprobs. The request costs one
admission sub-generation and the spent-seq fence is unchanged; an engine
response missing or mis-sizing the scores resolves as a typed server
terminal. topk_prompt_logprobs remains rejected.
# Conflicts:
#	tests/ci/requirements-ci-cpu.txt
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant