diff --git a/.github/workflows/policies/comment-command-access.json b/.github/workflows/policies/comment-command-access.json
index ac00985a58..28f2dcf0d0 100644
--- a/.github/workflows/policies/comment-command-access.json
+++ b/.github/workflows/policies/comment-command-access.json
@@ -36,6 +36,7 @@
"run-ci-image",
"run-ci-long",
"run-ci-lora",
+ "run-ci-multi-lora",
"run-ci-megatron",
"run-ci-miles-plugin",
"run-ci-model-scripts",
diff --git a/docs/advanced/lora.md b/docs/advanced/lora.md
index 28ed547364..38484d31c0 100644
--- a/docs/advanced/lora.md
+++ b/docs/advanced/lora.md
@@ -18,9 +18,9 @@ Rollout requests then select the live named adapter with `lora_path`.
Multi-LoRA currently supports disaggregated rollout only and rejects
`--colocate` at launch. Each SGLang engine keeps the base checkpoint resident;
-miles selectively exports and NCCL-broadcasts newly loaded or optimizer-stepped
-adapters into their corresponding SGLang slots. New or restarted engines receive
-every loaded adapter, while unchanged adapters are not resent.
+adapter weights reach the engines through explicit pushes: each push registers
+an immutable version under its own engine-side name, so in-flight sampling
+against an older version is never disturbed.
Model support is therefore a three-way contract rather than a hard-coded
allowlist:
@@ -68,7 +68,8 @@ full-scale experiment evidence. It is not an exhaustive model whitelist.
| Implementation | Model family | Architecture exercised | Evidence | Notes |
|---|---|---|---|---|
| Bridge | Qwen2.5 0.5B / 3B | Dense | [0.5B CUDA and ROCm E2E](https://github.com/radixark/miles/blob/main/tests/e2e/lora/test_lora_qwen2.5_0.5B.py), [3B disaggregated recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-qwen2.5-3B-megatron-lora-disaggregated.sh) | The simplest starting point; `all-linear` works. |
-| Bridge | Qwen3 4B | Dense | [Single-LoRA recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-qwen3-4B-megatron-lora.sh), [AMD launcher](https://github.com/radixark/miles/blob/main/scripts/amd/run_qwen3_4b_lora.py), [multi-LoRA recipe](https://github.com/radixark/miles/tree/main/examples/multi_lora) | The AMD launcher uses Triton attention and LoRA backends. |
+| Bridge | Qwen3 4B | Dense | [Single-LoRA recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-qwen3-4B-megatron-lora.sh), [AMD launcher](https://github.com/radixark/miles/blob/main/scripts/amd/run_qwen3_4b_lora.py) | The AMD launcher uses Triton attention and LoRA backends. |
+| Bridge | Qwen3-30B-A3B | MoE | [Multi-LoRA gateway example](https://github.com/radixark/miles/tree/main/examples/multi_lora) | Disaggregated Tinker training and versioned sampling. |
| Bridge | GPT-OSS 20B | MoE | [Recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh), [MoE LoRA E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_gpt_oss_20b_moe_lora_ci.py) | Uses the SGLang `triton` LoRA backend. |
| Bridge | Kimi K2.5 | Multimodal MoE + MLA | [16-node recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-kimi-k25-megatron-lora.sh) | Demonstrates shared-outer expert LoRA and an INT4 rollout / fake-QAT setup. |
| Bridge | GLM-5 / 5.1 / 5.2 744B-A40B | MoE + MLA + DSA | [GLM-5.1 launcher](https://github.com/radixark/miles/blob/main/scripts/run_glm5_1_744b_a40b_lora.py), [GLM-5.2 launcher](https://github.com/radixark/miles/blob/main/scripts/run_glm5_2_744b_a40b_lora.py) | CI covers reduced 6-layer / 5-layer checkpoints; historical full-744B results are described below. |
@@ -289,121 +290,29 @@ stability evidence rather than a released benchmark.
## Multi-LoRA training
-### Current dataset-driven backend
-
-The implementation on `main` trains multiple adapters against one shared base
-model through the [fully async example](https://github.com/radixark/miles/tree/main/examples/multi_lora).
-Each registered adapter supplies its own dataset, reward, rollout batch shape,
-rank/alpha, and checkpoint directory, with most fields inheriting process-wide
-defaults. LR/WD hyperparameters come from the global CLI; each fixed slot has
-its own Adam state and independently clocked scheduler. The trainer coalesces
-ready prompt-group slices or partial adapter batches and selectively upserts
-only changed adapters into SGLang.
-
-Set the slot capacity with `--multi-lora-n-adapters N`. A bounded run registers
-repeatable `--multi-lora-adapter NAME PATH` entries at startup; service mode can
-start with empty slots and register adapters through the controller HTTP API.
-This path currently forces Megatron-Bridge LoRA and requires disaggregated NCCL
-broadcast, PP1, THD, Adam, and no train offload. Shared-outer expert adapters are
-unsupported, and MoE expert adapters cannot use FP8/FP4 experts.
-
-Native multi-LoRA is not implied by the native single-adapter work: both current
-`main` and the Tinker-oriented branch below still build multi-LoRA through
-Megatron-Bridge. Native multi-LoRA is tracked separately in
-[issue #2141](https://github.com/radixark/miles/issues/2141).
-
-### Future Tinker-compatible operation backend
-
-[PR #2273](https://github.com/radixark/miles/pull/2273) is the active
-Tinker-oriented backend proposal. It changes ownership of the training loop:
-instead of the server owning a dataset, reward function, and one-step schedule,
-clients submit explicit operations against a registered adapter. Its primary
-intended consumer is a Tinker-compatible training service rather than a generic
-server-owned dataset scheduler.
-
-```text
-Tinker-style client
- | register + ordered operations
- v
-controller / operation ledger
- | bind one fixed LoRA slot
- v
-Megatron-Bridge multi-LoRA trainer
- | forward, backward, optimizer, checkpoint
- | save_weights_for_sampler publish barrier
- v
-SGLang router + registration-scoped adapter identity
-```
-
-The operation surface separates compute, optimization, and publication:
-
-| Operation | Contract |
-|---|---|
-| `forward` / `forward_backward` | Return per-datum log probabilities; backward calls accumulate client-scaled gradient sums. |
-| `optim_step` | Apply client-supplied Adam parameters and clipping to one slot, with an all-rank non-finite veto. |
-| `save_weights_for_sampler` | Publish the latest adapter and complete only after the new serving version is live. |
-| `save_state` / `load_state` | Save or restore immutable per-adapter weights and optimizer state behind shape, world-size, and ownership fences. |
-
-The design uses fixed residency rather than transparent LRU eviction. Operations
-are strictly serialized per registration, while idempotent retries, gap-buffered
-arrival, acknowledgements, and backpressure make execution retry-safe and
-order-safe. A registration-scoped serving identity prevents an old request from
-using a slot after that slot has been reassigned. Authenticated remote access is
-the responsibility of the future frontend, not the Ray operation API in #2273.
-
-The v1 scope in the PR is deliberately narrow: text-only synchronous training,
-one shared base model, shifted 1-D targets, `cross_entropy`, importance-sampling,
-and PPO losses, per-call Adam, and latest-only sampler weights. Multimodal,
-top-K/SDFT targets, CISPO/DRO, asynchronous or pinned-snapshot off-policy
-training, and cross-world-size restore are outside v1.
-
-
-This backend is implemented in an open PR, not released on `main`; the PR
-reports H200 validation. PR #2273 provides the operation backend, but its v1
-training operations are still exposed through the controller's Ray API. The
-stacked [PR #2346](https://github.com/radixark/miles/pull/2346) adds a REST
-frontend compatible with the official `tinker==0.24.1` client; its GPU frontend
-E2E is still pending. If #2273 lands as proposed, it replaces the current
-dataset-driven driver.
-
-
-## Compatibility and limitations
-
-- **Training backend:** Megatron only; FSDP has no LoRA training path.
-- **Implementation path on `main`:** Bridge is the general path; native/raw LoRA
- is model-specific to Inkling. General native coverage is pending PR #1792.
-- **Remote transport:** NCCL broadcast only, with PP1. P2P/RDMA and disk-delta
- reject LoRA.
-- **PPO:** shared actor/critic PPO with Bridge LoRA is untested; the critic never
- gets adapters.
-- **Resume:** miles adapter shards are resumable with the matching parallel
- topology. Direct HF PEFT import into the Bridge model is not yet implemented;
- native Inkling has a custom importer.
-- **Memory optimizations:** `--rematerialize-param-from-master-weight` and
- streamed optimizer state on NVMe reject LoRA. Ordinary actor disk offload is a
- different feature and is used by the draft agentic recipe.
-- **Multi-LoRA:** the current and proposed Tinker operation paths both require
- Bridge; native multi-LoRA remains roadmap work. Evaluation, colocate, PP,
- train offload, shared-outer experts, and FP8/FP4 MoE expert adapters are not
- supported by the current multi-adapter path.
-- **Agentic sessions:** the current session integration selects the fixed
- `miles_lora` adapter, not a multi-LoRA slot; `--lora-train-only` is also not a
- supported combination for this path.
-
-## Internals
-
-- `miles/backends/megatron_utils/bridge_lora_helpers.py` builds and wraps the
- general Bridge LoRA model.
-- `miles/backends/megatron_utils/lora_utils.py` resolves module names, creates
- standard/canonical adapters, and implements adapter checkpoint helpers.
-- `miles_plugins/models/inkling/lora.py` implements the native/raw LoRA path
- available on current `main`.
-- `miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py`
- handles colocated adapter export and IPC loading.
-- `miles/backends/megatron_utils/update_weight/update_weight_from_distributed/`
- gathers and broadcasts adapters to remote SGLang engines.
-- `miles/rollout/session/core.py` attaches the single adapter to agentic session
- requests.
-- `miles/ray/multi_lora/`, `miles/rollout/multi_lora/`, and
- `miles/backends/megatron_utils/multi_lora_*.py` implement the multi-adapter
- controller, routing, scheduling, optimization, and checkpoint path.
+Multi-LoRA training is served through the Tinker protocol: `serve_tinker.py`
+turns a miles deployment into a training service where each client drives its
+own adapter slot with explicit forward_backward / optim_step operations (the
+server owns no datasets or schedules). The gateway lives in `miles/tinker/`,
+the slot mechanics in `miles/backends/megatron_utils/lora/`. The dataset-driven
+multi-LoRA v1 backend (fully-async driver, adapter controller, per-adapter data
+sources) has been removed.
+
+Clients use the official `tinker` SDK with the gateway URL and a distinct API key
+per tenant. The key identifies ownership of models, futures, and checkpoints.
+Each adapter accumulates gradients until its client submits an optimizer step;
+saving weights for sampling publishes an immutable adapter version.
+
+`save_state` saves adapter parameters and optimizer state, including FP32 masters.
+It waits for preceding commands but neither steps nor saves pending gradients;
+call it after `optim_step` to save the effect of the accumulated training work.
+Sampler saves commit an immutable adapter directory on shared storage, without
+contacting inference engines. Engines load the same frozen base at startup and
+load adapter snapshots from disk on demand. A failed engine load fails sampling
+without invalidating the saved snapshot or model training.
+
+Futures, model leases, and sequence deduplication are in memory and are lost on
+gateway restart. Saved checkpoints retain their ownership and adapter metadata
+and can be used to create a new training or sampling client.
+
+See the [gateway example](/examples/multi-lora) for the launcher and smoke client.
diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md
index 15a49bca7e..e675f892f4 100644
--- a/docs/developer/architecture.md
+++ b/docs/developer/architecture.md
@@ -68,7 +68,7 @@ is loaded only when a run names its import path in a flag (`--spec`, or one of t
per-architecture weight bridges, `megatron_bridge/` the `megatron.bridge` shims, and
`optimizers/` optimizer plugins.
-`train.py`, `train_async.py` and `train_multi_lora_async.py` are the entry points. They are
+`train.py` and `train_async.py` are the entry points. They are
thin; most logic lives in the modules above.
## A request's life
diff --git a/docs/developer/contributor-guide.md b/docs/developer/contributor-guide.md
index e0857ba5d2..b428ed5d96 100644
--- a/docs/developer/contributor-guide.md
+++ b/docs/developer/contributor-guide.md
@@ -11,7 +11,6 @@ welcome: bug reports, doc fixes, new model recipes, full features.
miles/
├── train.py # synchronous entry point
├── train_async.py # fully-async entry point
-├── train_multi_lora_async.py # multi-LoRA async entry point
├── miles/ # the package
│ ├── backends/
│ │ ├── megatron_utils/ # Megatron actor, weight sync, checkpointing, fp32 markers
diff --git a/docs/examples/index.md b/docs/examples/index.md
index 016fc9d518..0b3f097444 100644
--- a/docs/examples/index.md
+++ b/docs/examples/index.md
@@ -13,7 +13,7 @@ End-to-end training workflows — the place to start.
- **[geo3k_vlm](/examples/geo3k-vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset.
- **[multi_turn](/examples/geo3k-vlm/multi-turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment.
- **[lora](https://github.com/radixark/miles/tree/main/examples/lora)**: LoRA fine-tuning with the Megatron backend.
-- **[multi_lora](/examples/multi-lora)**: Fully-async multi-adapter LoRA training with a slot-keyed adapter page table.
+- **[multi_lora](/examples/multi-lora)**: Serve concurrent LoRA fine-tuning clients on one shared base model through the Tinker protocol.
- **[on_policy_distillation](/examples/on-policy-distillation)**: Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop.
- **[qwen3_5_35b_selfdistill](/examples/on-policy-distillation/qwen3-5-35b-selfdistill)**: Two-phase self-distillation of Qwen3.5-35B-A3B on one 8xH200 node, with an in-process Megatron teacher.
- **[ppo](/examples/ppo)**: Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs.
diff --git a/docs/examples/multi-lora.md b/docs/examples/multi-lora.md
index 91778cb21c..44adaa0715 100644
--- a/docs/examples/multi-lora.md
+++ b/docs/examples/multi-lora.md
@@ -1,149 +1,98 @@
---
-title: "Multi-LoRA Training Example (fully-async)"
-description: "Fully-async multi-adapter LoRA training with a slot-keyed adapter page table."
+title: "Multi-LoRA Tinker Gateway"
+description: "Serve concurrent LoRA fine-tuning clients on one shared base model through the Tinker protocol."
# Generated from examples/multi_lora/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file.
---
-Train multiple LoRA adapters concurrently against a shared base model, using a
-fully-async rollout (continuous producer) + a slot-keyed LoRA page table on the
-SGLang engines (in-place upsert, no unload, no drain).
+> **Read the docs:** [Multi-LoRA training](https://miles.radixark.com/docs/advanced/lora#multi-lora-training).
-This example trains two adapters on Qwen3-4B:
-
-- **gsm8k** — grade-school math, `rm_type: math`
-- **dapo_math** — competition math (DAPO-Math-17k), `rm_type: deepscaler`
+- `serve_qwen3_30b_a3b_tinker.py`: prepare Qwen3-30B-A3B and launch the gateway.
+- `run_multi_tenant_example.py`: check marker memorization for one client or adapter isolation across concurrent tenants.
## Layout
-```
-run_multi_lora.py # launcher: prepare / train / full-train / serve
-service_smoke.py # register/deregister smoke test against the API
-adapters/
- gsm8k.yaml
- dapo_math.yaml
-```
-
-The implementation lives in the library: the driver is `train_multi_lora_async.py`
-at the repo root (next to `train.py`/`train_async.py`), the rollout fn and data
-source are `miles/rollout/multi_lora/`, and the controller is
-`miles/ray/multi_lora/` (registry + backend + HTTP API, plus the named Ray
-actor pinned to the head node).
-
-## Design (decoupled per-adapter optimizers)
-
-- **Controller** (Ray actor + control-plane HTTP API) is the source of truth:
- `POST/GET/DELETE /adapter_runs` plus `GET /adapter_runs/state`. The data source
- reads it; the trainer reads it. Generation traffic goes straight to the router;
- on deregister the controller aborts the adapter's in-flight requests
- engine-side by rid prefix (`rid = {adapter}::{uuid}`, set in `generate`).
-- **Per-adapter gradient accumulation.** Each adapter has its own batch shape:
- `rollout_batch_size` prompt groups per optimizer step, each group holding
- `n_samples_per_prompt` responses (`adapter_global_batch_size =
- rollout_batch_size x n_samples_per_prompt` samples per step). Completed
- prompt groups flow into training continuously in multiples of the
- adapter's `min_groups_per_dp_split` (the smallest group count whose samples
- split evenly across data-parallel ranks), gradients
- accumulate in the DDP buffers across train batches, and an adapter's
- optimizer steps exactly when its adapter batch fills — independent of every other
- adapter. The controller tracks adapter batch progress (`accumulated_groups`) and commits
- it only after a successful train call.
-- **Per-slot optimizers.** One Adam per adapter slot under Megatron's
- `LayerWiseDistributedOptimizer` (whole-parameter ZeRO-1): per-slot state,
- step counts, and gradient clipping; optimizer state sharded across DP ranks;
- plain DDP all-reduce (no distributed optimizer) makes cross-batch gradient
- retention idempotent.
-- **Batch collection.** The collection loop (same shape as fully_async's)
- pops groups from the per-adapter buffers round-robin, one
- `min_groups_per_dp_split` at a time, capped at each adapter's remaining
- batch, until the batch reaches `--global-batch-size` samples or a non-empty
- batch makes no progress for `--multi-lora-max-coalesce-wait-s` (the target
- can be permanently unreachable, so it trains on whatever is ready) — a
- single adapter with a small batch trains alone without waiting for
- anyone. Samples enter the gradient buffers with weight 1; at step time the
- slot's accumulated gradient is scaled by `1/adapter_global_batch_size`
- (a constant known in advance), so an adapter's update is identical to what
- it would get training alone.
-- **Selective weight sync.** Only adapters whose optimizer stepped are pushed
- to the engines (upsert into the slot-keyed page table); only their slot
- versions bump, keeping staleness filtering per-adapter accurate.
-- Adapters deregister on committed optimizer-step count (`num_step`) in the
- controller's train-commit path (`mark_batch_trained`), so stop checks happen
- exactly when steps advance. `num_step` is relative to the adapter's
- start/resume step. When an adapter doesn't set `num_step`, it is derived
- from `num_epoch` (default 1) as `num_epoch x len(dataset) //
- rollout_batch_size` once the data source loads the dataset (post-filter
- length). The trainer's
- `reconcile_adapters` (before each generate) retires it at the next sync
- point and cleans up (save ckpt + clear Megatron slot + zero its optimizer
- state and retained gradients). The adapter's untrained tail — buffered
- groups and any partially accumulated gradients — is discarded.
-- **Batch ⊆ loaded property:** `reconcile_adapters` runs before `generate`, so the
- batch is fetched with loaded = active; active only shrinks during generate, so every
- adapter in the batch is live on the trainer.
-
-## Provision (once)
+One 8-GPU node, disaggregated (multi-LoRA forbids `--colocate`):
-```bash
-python examples/multi_lora/run_multi_lora.py prepare
-```
+- 4 training GPUs: TP2 for the dense layers, EP4 for the 128 routed experts.
+- 4 sampling GPUs: two SGLang engines of 2 GPUs each, serving adapter versions by name.
+- 4 adapter slots (`--multi-lora-n-adapters`), rank up to 32, covering attention
+ (`linear_qkv`, `linear_proj`), the per-expert MoE projections (`linear_fc1`, `linear_fc2`),
+ and the output layer (`output_layer`) so the cookbook's default `train_unembed=True` is servable.
-Downloads `Qwen/Qwen3-4B` (to `/root/models`), `zhuzilin/dapo-math-17k`, and
-`zhuzilin/gsm8k` (to `/root/datasets`).
+The gateway currently resolves Tinker training groups for `qwen3` and `qwen3_moe`.
+`--tinker-train-attn`, `--tinker-train-mlp`, and `--tinker-train-unembed` default to enabled;
+use `--no-tinker-train-attn`, `--no-tinker-train-mlp`, or `--no-tinker-train-unembed` to disable a group.
+Every client's corresponding SDK flags must match the server layout. Tinker startup rejects
+`--target-modules` and `--exclude-modules`; native Miles training still accepts them.
## Run
+The gateway implements the `tinker==0.26.2` wire schema (newer SDKs renamed protobuf fields); install that exact version on the serving node and the client:
+
```bash
-python examples/multi_lora/run_multi_lora.py train # or: full-train (prepare + train)
+pip install "tinker==0.26.2"
```
-Registers the two adapters from CLI flags and trains until each hits its `num_step`,
-then exits.
-
-## Service mode
+Start the gateway:
```bash
-python examples/multi_lora/run_multi_lora.py serve
+python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py prepare # once per node
+python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py serve # Tinker API on :10613
```
-Starts with no adapters and idles; register/deregister at runtime through the
-control-plane API (port 8068):
+Checkpoints default to `/checkpoints/`; use `--save-dir` to choose another root.
+
+Install `tinker` on the client, then run the marker checks:
```bash
-python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \
- --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math
-```
+# one client: train, save for sampler, sample back the marker
+python examples/multi_lora/run_multi_tenant_example.py --base-model /root/models/Qwen3-30B-A3B --mode single
-## Multi-LoRA CLI flags
-
-| Flag | Purpose |
-| --- | --- |
-| `--multi-lora-n-adapters N` | Max concurrent adapter slots. `0` disables (default); `> 0` enables. |
-| `--multi-lora-adapter NAME PATH` | Register an adapter at startup. Repeatable. `PATH` → an `adapter.yaml`. |
-
-Per-adapter `rank` in `adapter.yaml` must be `<= --lora-rank`.
-
-## adapter.yaml
-
-```yaml
-rank: 16
-alpha: 16
-rollout_batch_size: 32 # prompt groups per optimizer step (defaults to --rollout-batch-size)
-n_samples_per_prompt: 4 # group shape (defaults to --n-samples-per-prompt)
-data: /root/datasets/gsm8k/train.parquet
-input_key: messages
-label_key: label
-rm_type: math
-num_step: 400 # stop adapter after N optimizer steps
- # (default: derived from num_epoch, itself default 1)
-# optional: save, num_epoch, custom_rm_path, ...
+# four tenants training concurrently on the same prompt with different markers;
+# passing means the adapters stayed isolated end to end
+python examples/multi_lora/run_multi_tenant_example.py --base-model /root/models/Qwen3-30B-A3B --mode multi --clients 4
```
-The derived `adapter_global_batch_size = rollout_batch_size x
-n_samples_per_prompt` is the adapter's samples-per-optimizer-step (the
-per-adapter analog of `--global-batch-size`).
-
-Batch-shape constraints (validated at registration, not at runtime):
-`n_samples_per_prompt` must be a divisor or multiple of the trainer's
-data-parallel size; `rollout_batch_size` must be a multiple of the adapter's
-`min_groups_per_dp_split`;
-`adapter_global_batch_size` is capped by
-`--multi-lora-max-adapter-global-batch-size` (default 4x `--global-batch-size`).
+## Supported inputs
+
+Training accepts text with 1-D loss inputs. 2-D soft targets, including SDFT,
+are not supported. Sampling requires a `/sampler_weights/` path returned by
+`save_weights_for_sampler()`; `/weights/` training checkpoints cannot be sampled directly.
+
+## Failure handling
+
+A terminal failure of `forward_backward`, `optim_step`, or `load_state` ends
+training for that model, including commands already queued behind it.
+This includes content validation failures with a valid model and sequence.
+Create a new model and restore a saved checkpoint to continue; completed
+futures and published checkpoints keep their results.
+
+Known request-local failures of `forward` or sampling leave model training
+available. Checkpoint load/save execution failures, including filesystem errors,
+invalidate the shared trainer cell and stop the server.
+Saving sampler weights commits an immutable directory;
+it does not call the inference engines. Sampling loads that snapshot from disk
+on demand, including after cache eviction. An engine load failure fails the
+sampling request; it leaves the snapshot and training state intact. Unknown
+trainer execution failures invalidate the shared trainer cell and stop the server.
+
+This gateway provides failure isolation, not automatic training recovery.
+Checkpoints persist; futures, deduplication, and unsaved accumulation do not
+survive a server restart.
+
+## Sampler snapshots
+
+Training and inference must use the same base checkpoint. Tinker engines load
+that frozen base at startup and serve without trainer weight updates; dummy
+loading and `update_weights: true` are rejected. Ordinary full-model and
+single-LoRA training continue to use the existing weight updater.
+
+`--tinker-checkpoint-root` must be on storage shared by the trainers, gateway,
+and every inference engine. A sampler save exports the current adapter weights,
+then publishes its tensors, adapter config, and `META.json` together through an
+atomic symlink replacement. Existing sampler versions cannot be overwritten. Saving between
+`forward_backward` and `optim_step` neither applies nor discards pending gradients.
+
+Training checkpoint names also point to immutable version directories. Overwriting
+atomically switches the link; older versions remain on disk for active readers.
+Legacy directory checkpoints can still be loaded; save under a new name instead
+of overwriting them.
diff --git a/examples/README.md b/examples/README.md
index c8ef2d7cdb..4577f4d997 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -9,7 +9,7 @@ End-to-end training workflows — the place to start.
- **[geo3k_vlm](./geo3k_vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset.
- **[multi_turn](./geo3k_vlm/multi_turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment.
- **[lora](./lora)**: LoRA fine-tuning with the Megatron backend.
-- **[multi_lora](./multi_lora)**: Fully-async multi-adapter LoRA training with a slot-keyed adapter page table.
+- **[multi_lora](./multi_lora)**: Serve concurrent LoRA fine-tuning clients on one shared base model through the Tinker protocol.
- **[on_policy_distillation](./on_policy_distillation)**: Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop.
- **[qwen3_5_35b_selfdistill](./on_policy_distillation/qwen3_5_35b_selfdistill)**: Two-phase self-distillation of Qwen3.5-35B-A3B on one 8xH200 node, with an in-process Megatron teacher.
- **[ppo](./ppo)**: Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs.
diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md
index 8423ddae1d..9aa3d2f1e2 100644
--- a/examples/multi_lora/README.md
+++ b/examples/multi_lora/README.md
@@ -1,146 +1,95 @@
-# Multi-LoRA Training Example (fully-async)
+# Multi-LoRA Tinker Gateway
-Train multiple LoRA adapters concurrently against a shared base model, using a
-fully-async rollout (continuous producer) + a slot-keyed LoRA page table on the
-SGLang engines (in-place upsert, no unload, no drain).
+> **Read the docs:** [Multi-LoRA training](https://miles.radixark.com/docs/advanced/lora#multi-lora-training).
-This example trains two adapters on Qwen3-4B:
-
-- **gsm8k** — grade-school math, `rm_type: math`
-- **dapo_math** — competition math (DAPO-Math-17k), `rm_type: deepscaler`
+- `serve_qwen3_30b_a3b_tinker.py`: prepare Qwen3-30B-A3B and launch the gateway.
+- `run_multi_tenant_example.py`: check marker memorization for one client or adapter isolation across concurrent tenants.
## Layout
-```
-run_multi_lora.py # launcher: prepare / train / full-train / serve
-service_smoke.py # register/deregister smoke test against the API
-adapters/
- gsm8k.yaml
- dapo_math.yaml
-```
-
-The implementation lives in the library: the driver is `train_multi_lora_async.py`
-at the repo root (next to `train.py`/`train_async.py`), the rollout fn and data
-source are `miles/rollout/multi_lora/`, and the controller is
-`miles/ray/multi_lora/` (registry + backend + HTTP API, plus the named Ray
-actor pinned to the head node).
-
-## Design (decoupled per-adapter optimizers)
-
-- **Controller** (Ray actor + control-plane HTTP API) is the source of truth:
- `POST/GET/DELETE /adapter_runs` plus `GET /adapter_runs/state`. The data source
- reads it; the trainer reads it. Generation traffic goes straight to the router;
- on deregister the controller aborts the adapter's in-flight requests
- engine-side by rid prefix (`rid = {adapter}::{uuid}`, set in `generate`).
-- **Per-adapter gradient accumulation.** Each adapter has its own batch shape:
- `rollout_batch_size` prompt groups per optimizer step, each group holding
- `n_samples_per_prompt` responses (`adapter_global_batch_size =
- rollout_batch_size x n_samples_per_prompt` samples per step). Completed
- prompt groups flow into training continuously in multiples of the
- adapter's `min_groups_per_dp_split` (the smallest group count whose samples
- split evenly across data-parallel ranks), gradients
- accumulate in the DDP buffers across train batches, and an adapter's
- optimizer steps exactly when its adapter batch fills — independent of every other
- adapter. The controller tracks adapter batch progress (`accumulated_groups`) and commits
- it only after a successful train call.
-- **Per-slot optimizers.** One Adam per adapter slot under Megatron's
- `LayerWiseDistributedOptimizer` (whole-parameter ZeRO-1): per-slot state,
- step counts, and gradient clipping; optimizer state sharded across DP ranks;
- plain DDP all-reduce (no distributed optimizer) makes cross-batch gradient
- retention idempotent.
-- **Batch collection.** The collection loop (same shape as fully_async's)
- pops groups from the per-adapter buffers round-robin, one
- `min_groups_per_dp_split` at a time, capped at each adapter's remaining
- batch, until the batch reaches `--global-batch-size` samples or a non-empty
- batch makes no progress for `--multi-lora-max-coalesce-wait-s` (the target
- can be permanently unreachable, so it trains on whatever is ready) — a
- single adapter with a small batch trains alone without waiting for
- anyone. Samples enter the gradient buffers with weight 1; at step time the
- slot's accumulated gradient is scaled by `1/adapter_global_batch_size`
- (a constant known in advance), so an adapter's update is identical to what
- it would get training alone.
-- **Selective weight sync.** Only adapters whose optimizer stepped are pushed
- to the engines (upsert into the slot-keyed page table); only their slot
- versions bump, keeping staleness filtering per-adapter accurate.
-- Adapters deregister on committed optimizer-step count (`num_step`) in the
- controller's train-commit path (`mark_batch_trained`), so stop checks happen
- exactly when steps advance. `num_step` is relative to the adapter's
- start/resume step. When an adapter doesn't set `num_step`, it is derived
- from `num_epoch` (default 1) as `num_epoch x len(dataset) //
- rollout_batch_size` once the data source loads the dataset (post-filter
- length). The trainer's
- `reconcile_adapters` (before each generate) retires it at the next sync
- point and cleans up (save ckpt + clear Megatron slot + zero its optimizer
- state and retained gradients). The adapter's untrained tail — buffered
- groups and any partially accumulated gradients — is discarded.
-- **Batch ⊆ loaded property:** `reconcile_adapters` runs before `generate`, so the
- batch is fetched with loaded = active; active only shrinks during generate, so every
- adapter in the batch is live on the trainer.
-
-## Provision (once)
+One 8-GPU node, disaggregated (multi-LoRA forbids `--colocate`):
-```bash
-python examples/multi_lora/run_multi_lora.py prepare
-```
+- 4 training GPUs: TP2 for the dense layers, EP4 for the 128 routed experts.
+- 4 sampling GPUs: two SGLang engines of 2 GPUs each, serving adapter versions by name.
+- 4 adapter slots (`--multi-lora-n-adapters`), rank up to 32, covering attention
+ (`linear_qkv`, `linear_proj`), the per-expert MoE projections (`linear_fc1`, `linear_fc2`),
+ and the output layer (`output_layer`) so the cookbook's default `train_unembed=True` is servable.
-Downloads `Qwen/Qwen3-4B` (to `/root/models`), `zhuzilin/dapo-math-17k`, and
-`zhuzilin/gsm8k` (to `/root/datasets`).
+The gateway currently resolves Tinker training groups for `qwen3` and `qwen3_moe`.
+`--tinker-train-attn`, `--tinker-train-mlp`, and `--tinker-train-unembed` default to enabled;
+use `--no-tinker-train-attn`, `--no-tinker-train-mlp`, or `--no-tinker-train-unembed` to disable a group.
+Every client's corresponding SDK flags must match the server layout. Tinker startup rejects
+`--target-modules` and `--exclude-modules`; native Miles training still accepts them.
## Run
+The gateway implements the `tinker==0.26.2` wire schema (newer SDKs renamed protobuf fields); install that exact version on the serving node and the client:
+
```bash
-python examples/multi_lora/run_multi_lora.py train # or: full-train (prepare + train)
+pip install "tinker==0.26.2"
```
-Registers the two adapters from CLI flags and trains until each hits its `num_step`,
-then exits.
-
-## Service mode
+Start the gateway:
```bash
-python examples/multi_lora/run_multi_lora.py serve
+python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py prepare # once per node
+python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py serve # Tinker API on :10613
```
-Starts with no adapters and idles; register/deregister at runtime through the
-control-plane API (port 8068):
+Checkpoints default to `/checkpoints/`; use `--save-dir` to choose another root.
+
+Install `tinker` on the client, then run the marker checks:
```bash
-python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \
- --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math
-```
+# one client: train, save for sampler, sample back the marker
+python examples/multi_lora/run_multi_tenant_example.py --base-model /root/models/Qwen3-30B-A3B --mode single
-## Multi-LoRA CLI flags
-
-| Flag | Purpose |
-| --- | --- |
-| `--multi-lora-n-adapters N` | Max concurrent adapter slots. `0` disables (default); `> 0` enables. |
-| `--multi-lora-adapter NAME PATH` | Register an adapter at startup. Repeatable. `PATH` → an `adapter.yaml`. |
-
-Per-adapter `rank` in `adapter.yaml` must be `<= --lora-rank`.
-
-## adapter.yaml
-
-```yaml
-rank: 16
-alpha: 16
-rollout_batch_size: 32 # prompt groups per optimizer step (defaults to --rollout-batch-size)
-n_samples_per_prompt: 4 # group shape (defaults to --n-samples-per-prompt)
-data: /root/datasets/gsm8k/train.parquet
-input_key: messages
-label_key: label
-rm_type: math
-num_step: 400 # stop adapter after N optimizer steps
- # (default: derived from num_epoch, itself default 1)
-# optional: save, num_epoch, custom_rm_path, ...
+# four tenants training concurrently on the same prompt with different markers;
+# passing means the adapters stayed isolated end to end
+python examples/multi_lora/run_multi_tenant_example.py --base-model /root/models/Qwen3-30B-A3B --mode multi --clients 4
```
-The derived `adapter_global_batch_size = rollout_batch_size x
-n_samples_per_prompt` is the adapter's samples-per-optimizer-step (the
-per-adapter analog of `--global-batch-size`).
-
-Batch-shape constraints (validated at registration, not at runtime):
-`n_samples_per_prompt` must be a divisor or multiple of the trainer's
-data-parallel size; `rollout_batch_size` must be a multiple of the adapter's
-`min_groups_per_dp_split`;
-`adapter_global_batch_size` is capped by
-`--multi-lora-max-adapter-global-batch-size` (default 4x `--global-batch-size`).
+## Supported inputs
+
+Training accepts text with 1-D loss inputs. 2-D soft targets, including SDFT,
+are not supported. Sampling requires a `/sampler_weights/` path returned by
+`save_weights_for_sampler()`; `/weights/` training checkpoints cannot be sampled directly.
+
+## Failure handling
+
+A terminal failure of `forward_backward`, `optim_step`, or `load_state` ends
+training for that model, including commands already queued behind it.
+This includes content validation failures with a valid model and sequence.
+Create a new model and restore a saved checkpoint to continue; completed
+futures and published checkpoints keep their results.
+
+Known request-local failures of `forward` or sampling leave model training
+available. Checkpoint load/save execution failures, including filesystem errors,
+invalidate the shared trainer cell and stop the server.
+Saving sampler weights commits an immutable directory;
+it does not call the inference engines. Sampling loads that snapshot from disk
+on demand, including after cache eviction. An engine load failure fails the
+sampling request; it leaves the snapshot and training state intact. Unknown
+trainer execution failures invalidate the shared trainer cell and stop the server.
+
+This gateway provides failure isolation, not automatic training recovery.
+Checkpoints persist; futures, deduplication, and unsaved accumulation do not
+survive a server restart.
+
+## Sampler snapshots
+
+Training and inference must use the same base checkpoint. Tinker engines load
+that frozen base at startup and serve without trainer weight updates; dummy
+loading and `update_weights: true` are rejected. Ordinary full-model and
+single-LoRA training continue to use the existing weight updater.
+
+`--tinker-checkpoint-root` must be on storage shared by the trainers, gateway,
+and every inference engine. A sampler save exports the current adapter weights,
+then publishes its tensors, adapter config, and `META.json` together through an
+atomic symlink replacement. Existing sampler versions cannot be overwritten. Saving between
+`forward_backward` and `optim_step` neither applies nor discards pending gradients.
+
+Training checkpoint names also point to immutable version directories. Overwriting
+atomically switches the link; older versions remain on disk for active readers.
+Legacy directory checkpoints can still be loaded; save under a new name instead
+of overwriting them.
diff --git a/examples/multi_lora/adapters/dapo_math.yaml b/examples/multi_lora/adapters/dapo_math.yaml
deleted file mode 100644
index 3a1a3ff8a8..0000000000
--- a/examples/multi_lora/adapters/dapo_math.yaml
+++ /dev/null
@@ -1,9 +0,0 @@
-rank: 32
-alpha: 32
-rollout_batch_size: 8 # prompt groups per optimizer step
-n_samples_per_prompt: 8 # -> 64 samples per step
-data: /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
-input_key: prompt
-label_key: label
-rm_type: deepscaler
-num_step: 500
diff --git a/examples/multi_lora/adapters/gsm8k.yaml b/examples/multi_lora/adapters/gsm8k.yaml
deleted file mode 100644
index 2290611464..0000000000
--- a/examples/multi_lora/adapters/gsm8k.yaml
+++ /dev/null
@@ -1,9 +0,0 @@
-rank: 16
-alpha: 16
-rollout_batch_size: 32 # prompt groups per optimizer step
-n_samples_per_prompt: 4 # -> 128 samples per step
-data: /root/datasets/gsm8k/train.parquet
-input_key: messages
-label_key: label
-rm_type: math
-num_step: 400
diff --git a/examples/multi_lora/run_client_recipes.py b/examples/multi_lora/run_client_recipes.py
new file mode 100644
index 0000000000..e23be03f09
--- /dev/null
+++ b/examples/multi_lora/run_client_recipes.py
@@ -0,0 +1,73 @@
+"""Run the official tinker-cookbook recipes against a running gateway.
+
+The cookbook's sl_loop (SFT) and rl_loop (GRPO) are the executable definition
+of the Tinker wire contract; passing them is the gateway's acceptance bar.
+
+Requires, next to the pinned SDK (tinker==0.26.2):
+
+ pip install git+https://github.com/thinking-machines-lab/tinker-cookbook@1f962eda3a2c
+
+``--base-model`` must be both the name this gateway serves (--tinker-base-model)
+and a HuggingFace name the cookbook can resolve a tokenizer and renderer for.
+"""
+
+import argparse
+import os
+import tempfile
+
+
+def run_sft(base_url: str, base_model: str, steps: int) -> None:
+ from tinker_cookbook.recipes import sl_loop
+
+ sl_loop.main(
+ sl_loop.Config(
+ base_url=base_url,
+ model_name=base_model,
+ log_path=tempfile.mkdtemp(prefix="cookbook-sl-"),
+ batch_size=4,
+ max_length=1024,
+ lora_rank=8,
+ save_every=0,
+ ttl_seconds=None,
+ max_steps=steps,
+ )
+ )
+
+
+def run_rl(base_url: str, base_model: str, steps: int) -> None:
+ from tinker_cookbook.recipes import rl_loop
+
+ rl_loop.main(
+ rl_loop.Config(
+ base_url=base_url,
+ model_name=base_model,
+ log_path=tempfile.mkdtemp(prefix="cookbook-rl-"),
+ batch_size=2,
+ group_size=4,
+ lora_rank=8,
+ save_every=0,
+ ttl_seconds=None,
+ max_tokens=512,
+ max_steps=steps,
+ )
+ )
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--base-url", required=True)
+ parser.add_argument("--base-model", required=True)
+ parser.add_argument("--mode", choices=["sft", "rl", "both"], default="both")
+ parser.add_argument("--steps", type=int, default=2)
+ args = parser.parse_args()
+
+ os.environ.setdefault("TINKER_API_KEY", "tml-cookbook-acceptance")
+ if args.mode in ("sft", "both"):
+ run_sft(args.base_url, args.base_model, args.steps)
+ if args.mode in ("rl", "both"):
+ run_rl(args.base_url, args.base_model, args.steps)
+ print(f"cookbook acceptance passed: mode={args.mode} steps={args.steps}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/multi_lora/run_multi_lora.py b/examples/multi_lora/run_multi_lora.py
deleted file mode 100644
index 417d14af65..0000000000
--- a/examples/multi_lora/run_multi_lora.py
+++ /dev/null
@@ -1,204 +0,0 @@
-"""Multi-LoRA fully-async GRPO example (Qwen3-4B, disaggregated 4 train + 4 rollout GPUs).
-
-Trains multiple LoRA adapters concurrently on a shared base model. Two example
-adapters ship in ``adapters/``: gsm8k (rm_type=math) and dapo_math
-(rm_type=deepscaler); each carries its own rank/alpha, batch shape, dataset,
-reward, and ``num_step`` stop condition. The driver is
-``train_multi_lora_async.py`` at the repo root; fully-async training forbids
-``--colocate`` (generation needs continuous GPU).
-
-Usage:
- python examples/multi_lora/run_multi_lora.py prepare # download Qwen3-4B + both datasets (once per node)
- python examples/multi_lora/run_multi_lora.py train # bounded run: registers the two adapters, exits when each hits num_step
- python examples/multi_lora/run_multi_lora.py full-train # prepare + train
- python examples/multi_lora/run_multi_lora.py serve # service mode: no adapters preloaded, idles for registrations (API on :8068)
-
-Service mode pairs with the smoke client:
- python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \\
- --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math
-"""
-
-from dataclasses import dataclass
-
-import typer
-
-import miles.utils.external_utils.command_utils as U
-
-app = typer.Typer()
-
-_ADAPTER_DIR = f"{U.repo_base_dir}/examples/multi_lora/adapters"
-
-
-@dataclass
-class ScriptArgs(U.ExecuteTrainConfig):
- run_id: str = U.create_run_id()
-
- hf_checkpoint: str | None = None
- model_dir: str = "/root/models"
- data_dir: str = "/root/datasets"
- save_dir: str = "/tmp/multi_lora"
- megatron_path: str = "/root/Megatron-LM"
-
- # Disaggregated split (fully-async forbids colocate).
- num_gpus_per_node: int = 8
- actor_num_gpus: int = 4
- rollout_num_gpus: int = 4
- tp: int = 2
-
- # LoRA slot pool. Per-adapter rank/alpha come from adapter.yaml, capped by lora_rank.
- lora_rank: int = 32
- lora_alpha: int = 32
- lora_dropout: float = 0.0
- target_modules: str = "all-linear"
- n_adapters: int = 4
- # Comma-separated adapter names; each resolves to adapters/{name}.yaml (train mode only).
- adapters: str = "dapo_math,gsm8k"
-
- # Global rollout defaults; the per-adapter batch shapes live in the yamls.
- num_rollout: int = 50
- rollout_batch_size: int = 32
- n_samples_per_prompt: int = 8
- rollout_max_response_len: int = 4096
- global_batch_size: int = 256
- max_weight_staleness: int = 3
-
- # Service mode.
- api_port: int = 8068
-
- save_interval: int = 5
- enable_wandb: bool = False
- extra_args: str = ""
-
- def __post_init__(self):
- if self.hf_checkpoint is None:
- self.hf_checkpoint = f"{self.model_dir}/Qwen3-4B"
-
-
-def _prepare_download(args: ScriptArgs):
- U.exec_command_cpu(f"mkdir -p {args.data_dir} {args.model_dir}")
- U.exec_command_cpu(f"hf download Qwen/Qwen3-4B --local-dir {args.model_dir}/Qwen3-4B")
- U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir)
- U.hf_download_dataset("zhuzilin/gsm8k", data_dir=args.data_dir)
-
-
-def _train(args: ScriptArgs, service: bool):
- mode = "service" if service else "bounded"
- print(
- f"[run] multi-LoRA ({mode}): {args.actor_num_gpus} train + {args.rollout_num_gpus} rollout GPUs, tp={args.tp}"
- )
-
- ckpt_args = f"--hf-checkpoint {args.hf_checkpoint} --megatron-to-hf-mode bridge "
-
- lora_args = (
- f"--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} "
- f'--lora-dropout {args.lora_dropout} --target-modules "{args.target_modules}" '
- )
-
- multi_lora_args = f"--multi-lora-n-adapters {args.n_adapters} --multi-lora-idle-poll-s 5 "
- if service:
- # No adapters preloaded; the control-plane API accepts registrations at runtime.
- multi_lora_args += f"--multi-lora-api-port {args.api_port} "
- else:
- for name in args.adapters.split(","):
- multi_lora_args += f'--multi-lora-adapter "{name}" "{_ADAPTER_DIR}/{name}.yaml" '
- multi_lora_args += "--multi-lora-disable-service-mode "
-
- # in_place pause + upsert weight push is what lets adapters refresh without
- # unloading (an unload would deadlock behind paused in-flight requests).
- sync_args = f"--pause-generation-mode in_place --max-weight-staleness {args.max_weight_staleness} --use-tis "
-
- rollout_args = (
- "--apply-chat-template --rollout-shuffle "
- f"--num-rollout {args.num_rollout} "
- f"--rollout-batch-size {args.rollout_batch_size} "
- f"--n-samples-per-prompt {args.n_samples_per_prompt} "
- f"--rollout-max-response-len {args.rollout_max_response_len} "
- "--rollout-temperature 1 "
- f"--global-batch-size {args.global_batch_size} "
- )
-
- grpo_args = (
- "--advantage-estimator grpo --kl-loss-coef 0.00 --kl-coef 0.00 "
- "--entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 "
- )
-
- optimizer_args = (
- "--optimizer adam --lr 1e-5 --lr-decay-style constant --weight-decay 0.1 "
- "--adam-beta1 0.9 --adam-beta2 0.98 "
- )
-
- perf_args = (
- f"--tensor-model-parallel-size {args.tp} --sequence-parallel "
- "--pipeline-model-parallel-size 1 --context-parallel-size 1 "
- "--expert-model-parallel-size 1 --expert-tensor-parallel-size 1 "
- "--use-dynamic-batch-size --max-tokens-per-gpu 9216 "
- )
-
- sglang_args = "--rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.8 "
-
- topology_args = (
- f"--actor-num-nodes 1 --actor-num-gpus-per-node {args.actor_num_gpus} "
- f"--rollout-num-gpus {args.rollout_num_gpus} "
- )
-
- save_args = f"--save {args.save_dir} --save-interval {args.save_interval} "
-
- misc_args = (
- "--attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 "
- "--attention-softmax-in-fp32 --attention-backend flash "
- )
-
- wandb_args = U.get_default_wandb_args(__file__, run_id=args.run_id) if args.enable_wandb else ""
-
- train_args = (
- f"{ckpt_args} {lora_args} {multi_lora_args} {sync_args} {rollout_args} {grpo_args} "
- f"{optimizer_args} {perf_args} {sglang_args} {topology_args} {save_args} {misc_args} "
- f"{wandb_args} {args.extra_args} "
- )
-
- U.execute_train(
- train_args=train_args,
- config=args,
- num_gpus_per_node=args.num_gpus_per_node,
- megatron_model_type="qwen3-4B",
- train_script="train_multi_lora_async.py",
- megatron_path=args.megatron_path,
- )
-
-
-@app.command()
-@U.dataclass_cli
-def prepare(args: ScriptArgs):
- """Download Qwen3-4B and both task datasets. Run once per node before training."""
- _prepare_download(args)
-
-
-@app.command()
-@U.dataclass_cli
-def train(args: ScriptArgs):
- """Bounded run: register the adapters from adapters/, train until each hits num_step, exit."""
- _train(args, service=False)
-
-
-@app.command()
-@U.dataclass_cli
-def full_train(args: ScriptArgs):
- """Download model + datasets, then run the bounded training."""
- _prepare_download(args)
- _train(args, service=False)
-
-
-@app.command()
-@U.dataclass_cli
-def serve(args: ScriptArgs):
- """Service mode: no adapters preloaded; register/deregister via the HTTP API while it idles."""
- _train(args, service=True)
-
-
-@app.callback()
-def _callback() -> None:
- pass
-
-
-if __name__ == "__main__":
- app()
diff --git a/examples/multi_lora/run_multi_tenant_example.py b/examples/multi_lora/run_multi_tenant_example.py
new file mode 100644
index 0000000000..796a0623cf
--- /dev/null
+++ b/examples/multi_lora/run_multi_tenant_example.py
@@ -0,0 +1,99 @@
+"""Check marker memorization and concurrent adapter isolation through the Tinker SDK."""
+
+import argparse
+import asyncio
+
+import tinker
+from tinker import types
+
+PROMPT = "The secret word is"
+
+MARKERS = [
+ "cobalt-lantern-73",
+ "velvet-otter-19",
+ "amber-glacier-58",
+ "crimson-abacus-31",
+ "silver-nebula-86",
+ "indigo-walrus-44",
+ "maroon-zeppelin-27",
+ "golden-isotope-65",
+]
+
+
+def build_datums(tokenizer, marker: str) -> list[types.Datum]:
+ texts = [
+ f"{PROMPT} {marker}.",
+ f"Remember this: the secret word is {marker}.",
+ f"Q: What is the secret word? A: The secret word is {marker}.",
+ f"Note for later. The secret word is {marker}.",
+ ]
+ datums = []
+ for text in texts:
+ tokens = tokenizer.encode(text)
+ datums.append(
+ types.Datum(
+ model_input=types.ModelInput.from_ints(tokens[:-1]),
+ loss_fn_inputs={"target_tokens": tokens[1:], "weights": [1.0] * (len(tokens) - 1)},
+ )
+ )
+ return datums
+
+
+async def run_client(index: int, args) -> dict:
+ marker = MARKERS[index]
+ tag = f"client-{index}"
+ # tml- prefix required by the SDK; the gateway reads the key as the tenant id
+ service = tinker.ServiceClient(base_url=args.base_url, api_key=f"tml-smoke-{tag}")
+ training = await service.create_lora_training_client_async(base_model=args.base_model, rank=args.lora_rank)
+ tokenizer = training.get_tokenizer()
+
+ datums = build_datums(tokenizer, marker)
+ target_tokens = sum(datum.loss_fn_inputs["target_tokens"].shape[0] for datum in datums)
+
+ losses = []
+ for _ in range(args.steps):
+ fb_future = await training.forward_backward_async(datums, loss_fn="cross_entropy")
+ optim_future = await training.optim_step_async(types.AdamParams(learning_rate=args.lr))
+ fb = await fb_future
+ await optim_future
+ losses.append(fb.metrics["loss:sum"] / target_tokens)
+
+ save_future = await training.save_weights_for_sampler_async(name="final")
+ sampler_path = (await save_future).path
+ sampler = await service.create_sampling_client_async(model_path=sampler_path)
+
+ response = await sampler.sample_async(
+ prompt=types.ModelInput.from_ints(tokenizer.encode(PROMPT)),
+ num_samples=1,
+ sampling_params=types.SamplingParams(max_tokens=args.max_tokens, temperature=0.0),
+ )
+ completion = tokenizer.decode(response.sequences[0].tokens)
+
+ ok = losses[-1] < losses[0] and marker in completion
+ print(
+ f"[{tag}] {'PASS' if ok else 'FAIL'} loss {losses[0]:.3f} -> {losses[-1]:.3f}, "
+ f"marker {marker!r}, completion {completion!r}"
+ )
+ return {"tag": tag, "ok": ok}
+
+
+async def main(args):
+ n_clients = 1 if args.mode == "single" else args.clients
+ assert n_clients <= len(MARKERS), f"at most {len(MARKERS)} clients"
+ results = await asyncio.gather(*(run_client(i, args) for i in range(n_clients)))
+ failed = [r["tag"] for r in results if not r["ok"]]
+ assert not failed, f"failed clients: {failed}"
+ print(f"all {n_clients} client(s) passed")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--base-url", default="http://127.0.0.1:10613")
+ parser.add_argument("--base-model", required=True)
+ parser.add_argument("--mode", choices=["single", "multi"], default="single")
+ parser.add_argument("--clients", type=int, default=4)
+ parser.add_argument("--lora-rank", type=int, default=32)
+ parser.add_argument("--steps", type=int, default=12)
+ parser.add_argument("--lr", type=float, default=1e-3)
+ parser.add_argument("--max-tokens", type=int, default=24)
+ asyncio.run(main(parser.parse_args()))
diff --git a/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py b/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py
new file mode 100644
index 0000000000..1b7352dc44
--- /dev/null
+++ b/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py
@@ -0,0 +1,127 @@
+"""Prepare Qwen3-30B-A3B and serve the Tinker gateway on separate training and sampling GPUs."""
+
+from dataclasses import dataclass, field
+
+import typer
+
+import miles.utils.external_utils.command_utils as U
+
+app = typer.Typer()
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+ run_id: str = field(default_factory=U.create_run_id)
+
+ hf_checkpoint: str | None = None
+ model_type: str = "qwen3-30B-A3B"
+ model_dir: str = "/root/models"
+ save_dir: str | None = None
+ megatron_path: str = "/root/Megatron-LM"
+
+ num_gpus_per_node: int = 8
+ actor_num_gpus: int = 4
+ rollout_num_gpus: int = 4
+ tp: int = 2
+ ep: int = 4
+
+ # LoRA slot pool; per-client rank comes from the SDK, capped by lora_rank.
+ lora_rank: int = 32
+ lora_alpha: int = 64
+ n_adapters: int = 4
+ tinker_train_attn: bool = True
+ tinker_train_mlp: bool = True
+ tinker_train_unembed: bool = True
+
+ tinker_port: int = 10613
+ rollout_num_gpus_per_engine: int = 2
+ sglang_mem_fraction_static: float = 0.7
+
+ extra_args: str = ""
+
+ def __post_init__(self):
+ if self.save_dir is None:
+ self.save_dir = f"{self.output_dir}/checkpoints"
+ if self.hf_checkpoint is None:
+ self.hf_checkpoint = f"{self.model_dir}/Qwen3-30B-A3B"
+
+
+@app.command()
+@U.dataclass_cli
+def prepare(args: ScriptArgs):
+ """Download the Qwen3-30B-A3B checkpoint. Run once per node before serving."""
+ U.exec_command_cpu(f"mkdir -p {args.model_dir}")
+ U.exec_command_cpu(f"hf download Qwen/Qwen3-30B-A3B --local-dir {args.model_dir}/Qwen3-30B-A3B")
+
+
+@app.command()
+@U.dataclass_cli
+def serve(args: ScriptArgs):
+ """Serve the Tinker gateway (idles until clients connect)."""
+ print(
+ f"[run] tinker gateway: {args.actor_num_gpus} train + {args.rollout_num_gpus} rollout GPUs, "
+ f"{args.n_adapters} adapter slots, port {args.tinker_port}"
+ )
+
+ ckpt_args = f"--hf-checkpoint {args.hf_checkpoint} --megatron-to-hf-mode bridge "
+
+ lora_args = (
+ f"--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} --lora-dropout 0.0 "
+ "--no-gradient-accumulation-fusion "
+ f"--multi-lora-n-adapters {args.n_adapters} "
+ )
+
+ tinker_args = f"--tinker-server-port {args.tinker_port} " f"--tinker-checkpoint-root {args.save_dir}/{args.run_id}"
+
+ for group in ("attn", "mlp", "unembed"):
+ enabled = getattr(args, f"tinker_train_{group}")
+ tinker_args += f" --{'' if enabled else 'no-'}tinker-train-{group}"
+
+ # initial config only; AdamParams come per optim_step request
+ optimizer_args = "--optimizer adam --lr 1e-4 "
+
+ perf_args = (
+ f"--tensor-model-parallel-size {args.tp} --sequence-parallel "
+ "--pipeline-model-parallel-size 1 --context-parallel-size 1 "
+ f"--expert-model-parallel-size {args.ep} --expert-tensor-parallel-size 1 "
+ "--use-dynamic-batch-size --max-tokens-per-gpu 8192 "
+ )
+
+ sglang_args = (
+ f"--rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine} "
+ f"--sglang-mem-fraction-static {args.sglang_mem_fraction_static} "
+ "--sglang-lora-backend triton "
+ )
+
+ topology_args = (
+ f"--actor-num-nodes 1 --actor-num-gpus-per-node {args.actor_num_gpus} "
+ f"--rollout-num-gpus {args.rollout_num_gpus} "
+ )
+
+ misc_args = (
+ "--attention-dropout 0.0 --hidden-dropout 0.0 "
+ "--accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 --attention-backend flash "
+ )
+
+ train_args = (
+ f"{ckpt_args} {lora_args} {tinker_args} {optimizer_args} "
+ f"{perf_args} {sglang_args} {topology_args} {misc_args} {args.extra_args} "
+ )
+
+ U.execute_train(
+ train_args=train_args,
+ config=args,
+ num_gpus_per_node=args.num_gpus_per_node,
+ megatron_model_type=args.model_type,
+ train_script="serve_tinker.py",
+ megatron_path=args.megatron_path,
+ )
+
+
+@app.callback()
+def _callback() -> None:
+ pass
+
+
+if __name__ == "__main__":
+ app()
diff --git a/examples/multi_lora/service_smoke.py b/examples/multi_lora/service_smoke.py
deleted file mode 100644
index fb6685e34f..0000000000
--- a/examples/multi_lora/service_smoke.py
+++ /dev/null
@@ -1,170 +0,0 @@
-"""Smoke test for multi-LoRA service mode: register/deregister against a running
-trainer, using step counts as the race-free progress signal.
-
-Usage: python examples/multi_lora/service_smoke.py --api-url http://HOST:8068 \\
- --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math
-"""
-
-import argparse
-import sys
-import time
-
-import httpx
-
-POLL_INTERVAL_S = 5.0
-
-
-class SmokeFailure(Exception):
- pass
-
-
-class ServiceClient:
- def __init__(self, api_url: str, timeout_s: float):
- self.api_url = api_url.rstrip("/")
- self.timeout_s = timeout_s
- self.http = httpx.Client(timeout=30.0)
-
- def adapters(self, states: set[str] | None = None) -> dict:
- response = self.http.get(f"{self.api_url}/adapter_runs")
- response.raise_for_status()
- wanted_states = states if states is not None else {"ACTIVE"}
- return {
- status["name"]: {
- "slot": status["slot"],
- "version": status["version"],
- "step": status["step"],
- "state": status["state"],
- }
- for status in response.json()["adapters"]
- if status["state"] in wanted_states
- }
-
- def active_adapters(self) -> dict:
- return self.adapters(states={"ACTIVE"})
-
- def register(self, name: str, config: dict) -> httpx.Response:
- return self.http.post(f"{self.api_url}/adapter_runs", json={"name": name, "config": config})
-
- def deregister(self, name: str) -> None:
- response = self.http.delete(f"{self.api_url}/adapter_runs/{name}")
- response.raise_for_status()
-
- def wait_for(self, description: str, predicate) -> dict:
- deadline = time.time() + self.timeout_s
- while time.time() < deadline:
- try:
- adapters = self.active_adapters()
- except httpx.HTTPError as e:
- print(f" ... api not reachable yet ({e})")
- adapters = None
- if adapters is not None:
- if predicate(adapters):
- print(f" ok: {description} (active={adapters})")
- return adapters
- print(f" waiting for {description} (active={adapters})")
- time.sleep(POLL_INTERVAL_S)
- raise SmokeFailure(f"timed out after {self.timeout_s}s waiting for: {description}")
-
- def wait_for_step(self, name: str, min_step: int) -> None:
- # Step-triggered deregistration can move an adapter to RETIRING quickly;
- # count both ACTIVE and RETIRING for progress waits.
- self.wait_for(
- f"'{name}' to reach step {min_step}",
- lambda _active: (
- (adapters := self.adapters(states={"ACTIVE", "RETIRING"}))
- and name in adapters
- and adapters[name]["step"] >= min_step
- ),
- )
-
- def register_when_allowed(self, name: str, config: dict) -> None:
- """Registration is rejected while a same-named adapter is cleaning up;
- retry until the name frees."""
- deadline = time.time() + self.timeout_s
- while time.time() < deadline:
- response = self.register(name, config)
- if response.status_code == 200:
- print(f" ok: registered '{name}'")
- return
- print(f" register '{name}' rejected ({response.status_code}): {response.text[:200]}")
- time.sleep(POLL_INTERVAL_S)
- raise SmokeFailure(f"timed out registering '{name}'")
-
-
-def main() -> int:
- parser = argparse.ArgumentParser(description=__doc__)
- parser.add_argument("--api-url", required=True, help="controller API listener, e.g. http://host:8068")
- parser.add_argument("--data", required=True, help="prompt dataset path for the test adapters")
- parser.add_argument("--input-key", default="text")
- parser.add_argument("--label-key", default="label")
- parser.add_argument("--rm-type", default="math")
- parser.add_argument("--rank", type=int, default=16)
- parser.add_argument("--alpha", type=int, default=16)
- parser.add_argument("--save", default=None, help="per-adapter save dir root override (default: trainer --save)")
- parser.add_argument("--steps", type=int, default=2, help="training steps to wait for per phase")
- parser.add_argument(
- "--num-step-smoke",
- type=int,
- default=1,
- help="num_step used by the auto-deregister smoke adapter",
- )
- parser.add_argument("--timeout", type=float, default=1800.0, help="per-phase timeout in seconds")
- args = parser.parse_args()
-
- def config(name: str) -> dict:
- cfg = {
- "rank": args.rank,
- "alpha": args.alpha,
- "data": args.data,
- "input_key": args.input_key,
- "label_key": args.label_key,
- "rm_type": args.rm_type,
- }
- if args.save:
- cfg["save"] = f"{args.save}/{name}"
- return cfg
-
- client = ServiceClient(args.api_url, args.timeout)
- try:
- print("phase 1: api reachable, no active adapters expected")
- client.wait_for("api reachable", lambda adapters: True)
-
- print("phase 2: register smoke_auto with num_step; expect auto-deregister after committed steps")
- auto_cfg = config("smoke_auto")
- auto_cfg["num_step"] = args.num_step_smoke
- client.register_when_allowed("smoke_auto", auto_cfg)
- client.wait_for_step("smoke_auto", args.num_step_smoke)
- client.wait_for("'smoke_auto' auto-deregistered", lambda adapters: "smoke_auto" not in adapters)
-
- print("phase 3: register smoke_a; expect promotion + training progress")
- client.register_when_allowed("smoke_a", config("smoke_a"))
- client.wait_for_step("smoke_a", args.steps)
-
- print("phase 4: register smoke_b mid-run; both must train")
- client.register_when_allowed("smoke_b", config("smoke_b"))
- client.wait_for_step("smoke_b", args.steps)
-
- print("phase 5: deregister smoke_a mid-run; smoke_b must keep training")
- step_b = client.active_adapters()["smoke_b"]["step"]
- client.deregister("smoke_a")
- client.wait_for("'smoke_a' gone from active set", lambda adapters: "smoke_a" not in adapters)
- client.wait_for_step("smoke_b", step_b + 1)
-
- print("phase 6: re-register the name smoke_a (waits out cleanup, reuses slot)")
- client.register_when_allowed("smoke_a", config("smoke_a"))
- client.wait_for_step("smoke_a", 1)
-
- print("phase 7: deregister everything; service should drain to idle")
- client.deregister("smoke_a")
- client.deregister("smoke_b")
- client.wait_for("no active adapters", lambda adapters: not adapters)
-
- print("SMOKE TEST PASSED")
- return 0
- except SmokeFailure as failure:
- print(f"SMOKE TEST FAILED: {failure}", file=sys.stderr)
- return 1
-
-
-if __name__ == "__main__":
- sys.exit(main())
diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py
index 0bdc30cda7..ee9cd15765 100644
--- a/miles/backends/megatron_utils/actor.py
+++ b/miles/backends/megatron_utils/actor.py
@@ -13,7 +13,10 @@
from torch_memory_saver import torch_memory_saver
from miles.backends.megatron_utils.ft.types import TrainStepOutput
+from miles.backends.megatron_utils.lora.utils import build_lora_sync_config, is_lora_enabled, lora_rollout_enabled
from miles.backends.megatron_utils.rematerialize_utils import build_main_cast_context
+from miles.backends.megatron_utils.update_weight.hf_weight_iterator import get_hf_weight_iterator
+from miles.backends.training_utils.weight_update.updater import WeightUpdater
from miles.dashboard import hooks as dashboard_hooks
from miles.ray.specs.train import compute_trainer_pool_id
from miles.ray.train_actor import TrainRayActor
@@ -26,7 +29,6 @@
from miles.utils.ft_utils.indep_dp import IndepDPInfo
from miles.utils.hf_config import load_hf_config
from miles.utils.memory_utils import clear_memory, print_memory
-from miles.utils.multi_lora import is_multi_lora_enabled
from miles.utils.processing_utils import load_tokenizer
from miles.utils.ray_utils import Box
from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups
@@ -56,7 +58,6 @@
from .ft.in_memory_checkpoint import InMemoryCheckpointManager
from .ft.indep_dp import reconfigure_indep_dp_group
from .initialize import init, is_first_replica_megatron_main_rank
-from .lora_utils import is_lora_enabled, lora_rollout_enabled
from .model import TrainStepOutcome, forward_only, initialize_model_and_optimizer, save, train
from .named_weights import named_params_and_buffers
from .parallel import verify_megatron_parallel_state
@@ -247,35 +248,7 @@ def init(
if self.args.vocab_size is None:
self.args.vocab_size = self.tokenizer.vocab_size
- from miles.backends.training_utils.weight_update.updater import WeightUpdater
-
- from .lora_utils import build_lora_sync_config
- from .update_weight.hf_weight_iterator import get_hf_weight_iterator
-
- is_lora = lora_rollout_enabled(args)
- uses_colocate_protocol = self.args.colocate
- if is_lora and not uses_colocate_protocol:
- assert args.megatron_to_hf_mode == "bridge", (
- "LoRA weight sync over distributed engines requires "
- f"--megatron-to-hf-mode bridge (got {args.megatron_to_hf_mode!r})."
- )
- self.weight_updater = WeightUpdater(
- self.args,
- self.model,
- weights_getter=self._get_actor_weights,
- model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
- quantization_config=getattr(self.hf_config, "quantization_config", None),
- iterator_factory=get_hf_weight_iterator,
- parallel_state=get_parallel_state(),
- is_lora=is_lora,
- lora_sync_config=build_lora_sync_config(self.args) if is_lora else None,
- )
-
- # Adapters currently loaded into Megatron slots on this rank.
- self.loaded_adapters: dict[str, object] = {}
- # Adapters with stale engine-side weights (newly loaded or just trained);
- # consumed by the next update_weights. Identical on every rank.
- self._multi_lora_pending_push: set[str] = set()
+ self._init_training_state()
# empty cache after initialization
clear_memory()
@@ -294,6 +267,29 @@ def init(
return start_rollout_id
+ def _init_training_state(self) -> None:
+ args = self.args
+ is_lora = lora_rollout_enabled(args)
+ uses_colocate_protocol = self.args.colocate
+ if is_lora and not uses_colocate_protocol:
+ assert args.megatron_to_hf_mode == "bridge", (
+ "LoRA weight sync over distributed engines requires "
+ f"--megatron-to-hf-mode bridge (got {args.megatron_to_hf_mode!r})."
+ )
+ model_name = type(self.hf_config).__name__.lower() if args.model_name is None else args.model_name
+ quantization_config = getattr(self.hf_config, "quantization_config", None)
+ self.weight_updater = WeightUpdater(
+ args,
+ self.model,
+ weights_getter=self._get_actor_weights,
+ model_name=model_name,
+ quantization_config=quantization_config,
+ iterator_factory=get_hf_weight_iterator,
+ parallel_state=get_parallel_state(),
+ is_lora=is_lora,
+ lora_sync_config=build_lora_sync_config(args) if is_lora else None,
+ )
+
def _clear_quantized_weight_workspaces(self) -> None:
if not (
self.args.clear_quantized_weight_workspaces_on_offload
@@ -667,68 +663,11 @@ def train_actor(
logger.info(f"Updating ref model at rollout_id {rollout_id}")
self.weights_backuper.backup("ref")
- if train_step_outcome == TrainStepOutcome.NORMAL and is_multi_lora_enabled(self.args):
- from miles.backends.megatron_utils.multi_lora_utils import commit_trained_batch
-
- commit_trained_batch(rollout_data, rollout_id, self._multi_lora_pending_push)
-
log_perf_data(rollout_id, self.args, extra_metrics=self.weight_updater.pop_metrics())
self._heartbeat.bump()
return TrainStepOutput(outcome=train_step_outcome)
- @with_logs
- @timer
- def reconcile_adapters(self) -> None:
- """Load adapters the controller wants served; retire deregistered ones, dropping their untrained tail."""
- if not is_multi_lora_enabled(self.args):
- return
- from miles.backends.megatron_utils.multi_lora_utils import cleanup_adapters as _cleanup_adapters
- from miles.backends.megatron_utils.multi_lora_utils import load_adapters as _load_adapters
- from miles.ray.multi_lora.controller import get_multi_lora_controller
-
- broadcast_buffer = [None]
- if is_first_replica_megatron_main_rank():
- controller = get_multi_lora_controller()
- ray.get(controller.retire_adapters.remote())
- broadcast_buffer[0] = ray.get(controller.snapshot.remote())
- if dist.is_initialized():
- dist.broadcast_object_list(broadcast_buffer, src=0, group=get_gloo_group())
- snapshot = broadcast_buffer[0]
- should_be_loaded = {**snapshot["active"], **snapshot["pending"], **snapshot["retiring"]}
- cleanup_names = set(snapshot["cleanup"])
-
- loaded_names = set(self.loaded_adapters)
- # Sorted so per-adapter collectives (checkpoint export) run in the same
- # order on every rank; set iteration order is process-specific.
- adapters_to_load = sorted(
- (adapter for name, adapter in should_be_loaded.items() if name not in loaded_names),
- key=lambda adapter: adapter.name,
- )
- adapters_to_clean_up = sorted(
- (self.loaded_adapters[n] for n in loaded_names if n in cleanup_names or n not in should_be_loaded),
- key=lambda adapter: adapter.name,
- )
- if adapters_to_load:
- _load_adapters(self.args, self.model, self.optimizer, adapters_to_load)
- for adapter in adapters_to_load:
- self.loaded_adapters[adapter.name] = adapter
- self._multi_lora_pending_push.add(adapter.name)
- if self._enable_weight_backup:
- self.weights_backuper.backup("actor")
- if adapters_to_clean_up:
- _cleanup_adapters(self.args, self.model, self.optimizer, adapters_to_clean_up)
- for adapter in adapters_to_clean_up:
- self.loaded_adapters.pop(adapter.name, None)
- self._multi_lora_pending_push.discard(adapter.name)
- if self._enable_weight_backup:
- self.weights_backuper.backup("actor")
-
- # Deregistered before ever being loaded: nothing to save or clear.
- if is_first_replica_megatron_main_rank():
- for name in cleanup_names - loaded_names:
- ray.get(get_multi_lora_controller().free_slot.remote(name))
-
@timer
def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
self._heartbeat.bump()
@@ -740,13 +679,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
maybe_finalize_async_save(blocking=True)
- if is_multi_lora_enabled(self.args):
- from miles.backends.megatron_utils.multi_lora_utils import save_due_adapter_checkpoints
-
- if not save_due_adapter_checkpoints(self.args, self.model):
- return
- else:
- save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler)
+ save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler)
if force_sync and self.args.async_save:
maybe_finalize_async_save(blocking=True)
@@ -843,28 +776,11 @@ def update_weights(self, info: "UpdatableEngines") -> int | None:
destroy_process_groups()
return None
- version_update_names: list[str] = []
- if is_multi_lora_enabled(self.args):
- from miles.backends.megatron_utils.multi_lora_utils import select_adapters_to_push
-
- self.weight_updater.multi_lora_adapters, version_update_names = select_adapters_to_push(
- # TODO: may improve condition (currently needs_reconnect) after yusheng's refactor
- self.loaded_adapters,
- self._multi_lora_pending_push,
- needs_reconnect,
- )
-
with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
print_memory("before update_weights")
self.weight_updater.update_weights()
print_memory("after update_weights")
- if is_multi_lora_enabled(self.args):
- from miles.backends.megatron_utils.multi_lora_utils import commit_weight_push
-
- self._multi_lora_pending_push.clear()
- commit_weight_push(version_update_names, self._is_first_replica_megatron_main_rank)
-
if self.args.ci_test and len(rollout_engines) > 0 and not is_lora_enabled(self.args):
engine = random.choice(rollout_engines)
engine_version = async_utils.run(engine.get_weight_version())
diff --git a/miles/backends/megatron_utils/checkpoint.py b/miles/backends/megatron_utils/checkpoint.py
index 1a34f6e807..756520501a 100644
--- a/miles/backends/megatron_utils/checkpoint.py
+++ b/miles/backends/megatron_utils/checkpoint.py
@@ -15,7 +15,7 @@
from miles.utils import megatron_bridge_utils
from miles_plugins.models.deepseek_v4.arguments import assert_checkpoint_is_current, is_dsv4_model
-from .lora_utils import is_lora_enabled, is_lora_model, load_lora_adapter, save_lora_checkpoint
+from .lora.utils import is_lora_enabled, is_lora_model, load_lora_adapter, save_lora_checkpoint
from .model_provider import LinearForLastLayer
try:
diff --git a/miles/backends/megatron_utils/hf_export.py b/miles/backends/megatron_utils/hf_export.py
index 90e4383c0f..1f921e170a 100644
--- a/miles/backends/megatron_utils/hf_export.py
+++ b/miles/backends/megatron_utils/hf_export.py
@@ -18,7 +18,7 @@
import torch
from megatron.core.distributed import DistributedDataParallel as DDP
-from miles.backends.megatron_utils.lora_utils import is_lora_model, save_lora_checkpoint
+from miles.backends.megatron_utils.lora.utils import is_lora_model, save_lora_checkpoint
from miles.backends.megatron_utils.named_weights import named_params_and_buffers
from miles.backends.megatron_utils.update_weight.hf_weight_iterator_direct import HfWeightIteratorDirect
from miles.backends.training_utils.parallel import get_parallel_state
diff --git a/miles/ray/multi_lora/__init__.py b/miles/backends/megatron_utils/lora/__init__.py
similarity index 100%
rename from miles/ray/multi_lora/__init__.py
rename to miles/backends/megatron_utils/lora/__init__.py
diff --git a/miles/backends/megatron_utils/lora/actor.py b/miles/backends/megatron_utils/lora/actor.py
new file mode 100644
index 0000000000..efa74957dc
--- /dev/null
+++ b/miles/backends/megatron_utils/lora/actor.py
@@ -0,0 +1,75 @@
+from contextlib import ExitStack
+
+from miles.backends.megatron_utils.actor import MegatronTrainRayActor
+from miles.backends.megatron_utils.lora import checkpoint as lora_checkpoint
+from miles.backends.megatron_utils.lora import model as lora_model
+from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer
+from miles.backends.megatron_utils.lora.utils import build_lora_sync_config
+from miles.backends.megatron_utils.update_weight.hf_weight_iterator import get_hf_weight_iterator
+from miles.backends.training_utils.data import get_rollout_data
+from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement
+from miles.backends.training_utils.weight_update.snapshot_publisher import WeightPublisher
+from miles.utils.multi_lora import AdapterSpec
+from miles.utils.ray_utils import Box
+from miles.utils.tracking_utils.structured_log import with_logs
+
+
+class MultiLoRATrainRayActor(MegatronTrainRayActor):
+ def _init_training_state(self) -> None:
+ args = self.args
+ self.slot_optimizers: dict[int, SlotOptimizer] = {}
+ iterator = get_hf_weight_iterator(
+ args,
+ self.model,
+ required_placement=WeightUpdatePlacement(gather_pp=True),
+ model_name=type(self.hf_config).__name__.lower() if args.model_name is None else args.model_name,
+ quantization_config=getattr(self.hf_config, "quantization_config", None),
+ )
+ self.weight_publisher = WeightPublisher(iterator, build_lora_sync_config(args))
+
+ @with_logs
+ def forward_backward(self, batch_id: int, rollout_data_ref: Box) -> dict:
+ self._heartbeat.bump()
+ with ExitStack() as stack:
+ rollout_data, store_get_result = get_rollout_data(self.args, rollout_data_ref)
+ stack.enter_context(store_get_result)
+ return lora_model.run_forward_backward(self.args, batch_id, self.model, rollout_data)
+
+ @with_logs
+ def optim_step(self, adam_params_by_slot: dict[int, dict]) -> dict[int, dict]:
+ self._heartbeat.bump()
+ return lora_model.optim_step(self.slot_optimizers, adam_params_by_slot)
+
+ @with_logs
+ def forward_only(self, batch_id: int, rollout_data_ref: Box) -> dict:
+ """Same loss pass as forward_backward, without the backward: the Tinker
+ forward() contract returns the requested loss per datum."""
+ self._heartbeat.bump()
+ with ExitStack() as stack:
+ rollout_data, store_get_result = get_rollout_data(self.args, rollout_data_ref)
+ stack.enter_context(store_get_result)
+ return lora_model.run_forward_backward(self.args, batch_id, self.model, rollout_data, forward_only=True)
+
+ @with_logs
+ def load_slot(
+ self, slot: int, rank: int, alpha: float, ckpt_path: str | None = None, load_optimizer: bool = True
+ ) -> None:
+ self.slot_optimizers[slot] = lora_model.load_slot(self.args, self.model, slot, rank, alpha)
+ if ckpt_path is not None:
+ lora_checkpoint.load_slot(self.model, self.slot_optimizers[slot], ckpt_path, load_optimizer)
+
+ @with_logs
+ def save_slot(self, slot: int, path: str, metadata: dict | None = None) -> None:
+ lora_checkpoint.save_slot(self.model, self.slot_optimizers[slot], path, metadata=metadata)
+
+ @with_logs
+ def export_slot(self, slot: int, rank: int, alpha: float, path: str, metadata: dict | None = None) -> None:
+ """Write the slot's adapter as an engine-loadable dir."""
+ self._heartbeat.bump()
+ self.weight_publisher.publish_adapter(AdapterSpec(slot=slot, rank=rank, alpha=alpha), path, metadata=metadata)
+
+ @with_logs
+ def unload_slot(self, slot: int) -> dict | None:
+ slot_optimizer = self.slot_optimizers.pop(slot)
+ lora_model.unload_slot(self.model, slot_optimizer)
+ return None
diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/lora/bridge.py
similarity index 96%
rename from miles/backends/megatron_utils/bridge_lora_helpers.py
rename to miles/backends/megatron_utils/lora/bridge.py
index d40335d066..b3f2779aaa 100644
--- a/miles/backends/megatron_utils/bridge_lora_helpers.py
+++ b/miles/backends/megatron_utils/lora/bridge.py
@@ -15,7 +15,7 @@
from miles.utils.hf_config import load_hf_config
from miles.utils.multi_lora import is_multi_lora_enabled, targets_expert_leaves
-from .lora_utils import convert_target_modules_to_hf, patch_param_grad_buffer_for_colocate_mode_lora
+from .utils import convert_target_modules_to_hf, patch_param_grad_buffer_for_colocate_mode_lora
logger = logging.getLogger(__name__)
@@ -37,7 +37,7 @@ def _make_value_model_hook(hidden_size: int):
"""Create a pre-wrap hook that replaces the output layer with a value head."""
from megatron.core import parallel_state
- from .model_provider import LinearForLastLayer
+ from ..model_provider import LinearForLastLayer
def hook(model):
model_post_process = []
@@ -168,11 +168,11 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list:
if is_multi_lora_enabled(args):
_validate_multi_lora_moe_support(args, provider)
- from miles.backends.megatron_utils.multi_lora_utils import create_multi_lora_instance
+ from miles.backends.megatron_utils.lora.slots import create_multi_lora_instance
lora = create_multi_lora_instance(args)
else:
- from .lora_utils import create_lora_instance
+ from .utils import create_lora_instance
lora = create_lora_instance(args)
diff --git a/miles/backends/megatron_utils/lora/checkpoint.py b/miles/backends/megatron_utils/lora/checkpoint.py
new file mode 100644
index 0000000000..a4ed50db5f
--- /dev/null
+++ b/miles/backends/megatron_utils/lora/checkpoint.py
@@ -0,0 +1,61 @@
+"""Per-slot checkpoints in megatron's torch_dist format.
+
+Every tensor carries its global shard coordinates and a slot-agnostic key,
+so a checkpoint reloads under any tp/pp/ep/world-size layout and into any
+free slot."""
+
+from collections.abc import Sequence
+from pathlib import Path
+
+from megatron.core import dist_checkpointing
+from megatron.core.dist_checkpointing.dict_utils import nested_values
+from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.core.utils import unwrap_model
+
+from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer
+from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+
+_WEIGHTS_KEY = "adapter_weights"
+_OPTIM_KEY = "adapter_optimizer"
+
+
+def _slot_weights_sharded_state_dict(model: Sequence[DDP], slot: int) -> dict:
+ sharded: dict = {}
+ for chunk in unwrap_model(model):
+ sharded |= chunk.sharded_state_dict()
+ marker = f".adapters.{slot}."
+ slot_sharded = {key: value for key, value in sharded.items() if marker in key}
+ assert slot_sharded, f"slot {slot} exposed no adapter tensors"
+ return slot_sharded
+
+
+def _canonicalize_slot_keys(tree: dict, slot: int) -> dict:
+ """Rewrite storage keys so a checkpoint saved from slot i loads into slot j."""
+ marker, canonical = f".adapters.{slot}.", ".adapter."
+ for sharded in nested_values(tree):
+ if hasattr(sharded, "key"):
+ sharded.key = sharded.key.replace(marker, canonical)
+ return tree
+
+
+def save_slot(model: Sequence[DDP], slot_optimizer: SlotOptimizer, path: str, metadata: dict | None = None) -> None:
+ weights = _slot_weights_sharded_state_dict(model, slot_optimizer.slot)
+ sharded = {_WEIGHTS_KEY: weights, _OPTIM_KEY: slot_optimizer.sharded_state(weights, is_loading=False)}
+ _canonicalize_slot_keys(sharded, slot_optimizer.slot)
+ write_checkpoint_dir(path, lambda tmp_dir: dist_checkpointing.save(sharded, str(tmp_dir)), metadata=metadata)
+
+
+def load_slot(model: Sequence[DDP], slot_optimizer: SlotOptimizer, path: str, load_optimizer: bool) -> None:
+ checkpoint_dir = str(Path(path))
+ weights = _slot_weights_sharded_state_dict(model, slot_optimizer.slot)
+ shells = {_WEIGHTS_KEY: weights}
+ if load_optimizer:
+ shells[_OPTIM_KEY] = slot_optimizer.sharded_state(weights, is_loading=True)
+ _canonicalize_slot_keys(shells, slot_optimizer.slot)
+
+ loaded = dist_checkpointing.load(shells, checkpoint_dir)
+ if load_optimizer:
+ slot_optimizer.load_sharded_state(loaded[_OPTIM_KEY])
+ else:
+ # weights-only load keeps the fresh Adam state the slot init just created
+ slot_optimizer.reload_masters()
diff --git a/miles/backends/megatron_utils/lora/model.py b/miles/backends/megatron_utils/lora/model.py
new file mode 100644
index 0000000000..4935f2db62
--- /dev/null
+++ b/miles/backends/megatron_utils/lora/model.py
@@ -0,0 +1,81 @@
+"""Multi-LoRA slot operations extending the shared Megatron model execution."""
+
+from argparse import Namespace
+from collections.abc import Sequence
+
+from megatron.core.distributed import DistributedDataParallel as DDP
+
+from miles.backends.megatron_utils.lora.optimizer import (
+ SlotOptimizer,
+ reset_grad_metadata_keep_grads,
+ step_slot_optimizers,
+)
+from miles.backends.megatron_utils.model import run_forward_backward_pass, setup_train_iteration_config
+from miles.backends.training_utils.data import get_data_iterator
+from miles.backends.training_utils.log_utils import aggregate_train_losses
+from miles.backends.training_utils.parallel import get_parallel_state
+from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase
+from miles.utils.types import RolloutBatch
+
+
+def run_forward_backward(
+ args: Namespace,
+ batch_id: int,
+ model: Sequence[DDP],
+ rollout_data: RolloutBatch,
+ *,
+ forward_only: bool = False,
+) -> dict:
+ data_iterator, num_microbatches = get_data_iterator(args, model, rollout_data)
+ assert len(num_microbatches) == 1, "a work unit is a single forward/backward pass"
+
+ for iterator in data_iterator:
+ iterator.reset()
+ for model_chunk in model:
+ model_chunk.train()
+ # disable_optimizer: bf16 loss scaling is the identity, so the pass needs no optimizer
+ setup_train_iteration_config(args, model, None, disable_optimizer=True)
+ if not forward_only:
+ reset_grad_metadata_keep_grads(model)
+
+ dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=batch_id)
+ losses_reduced = run_forward_backward_pass(
+ args,
+ dumper_phase_util,
+ data_iterator,
+ model,
+ num_microbatches[0],
+ num_rollouts=None,
+ forward_only=forward_only,
+ )
+ per_datum_outputs = [output for microbatch in losses_reduced for output in microbatch["per_datum"]]
+ dumper_phase_util.finalize(model)
+
+ if get_parallel_state().is_pp_last_stage:
+ return {"metrics": aggregate_train_losses(losses_reduced, None), "per_datum": per_datum_outputs}
+ return {"metrics": {}, "per_datum": per_datum_outputs}
+
+
+def optim_step(
+ slot_optimizers: dict[int, SlotOptimizer],
+ adam_params_by_slot: dict[int, dict],
+) -> dict[int, dict]:
+ stepped = {slot: slot_optimizers[slot] for slot in adam_params_by_slot}
+ return step_slot_optimizers(stepped, adam_params_by_slot)
+
+
+def load_slot(args: Namespace, model: Sequence[DDP], slot: int, rank: int, alpha: float) -> SlotOptimizer:
+ from megatron.bridge.peft.multi_lora_layers import init_adapter_slot
+
+ init_adapter_slot(model, slot, rank=rank, alpha=alpha)
+ slot_optimizer = SlotOptimizer(args, model, slot)
+ slot_optimizer.zero_grads()
+ return slot_optimizer
+
+
+def unload_slot(model: Sequence[DDP], slot_optimizer: SlotOptimizer) -> None:
+ """The caller drops the instance afterwards; its state dies with it."""
+ from megatron.bridge.peft.multi_lora_layers import clear_adapter_slot
+
+ clear_adapter_slot(model, slot_optimizer.slot)
+ slot_optimizer.zero_grads()
diff --git a/miles/backends/megatron_utils/lora/optimizer.py b/miles/backends/megatron_utils/lora/optimizer.py
new file mode 100644
index 0000000000..514c772002
--- /dev/null
+++ b/miles/backends/megatron_utils/lora/optimizer.py
@@ -0,0 +1,217 @@
+"""Per-tenant slot optimizers for multi-LoRA.
+
+Each live slot owns an independent LayerWiseDistributedOptimizer over exactly
+its adapter parameters, built at load and destroyed with the slot, so a fresh
+tenant never inherits optimizer state. Requires plain DDP all-reduce
+(use_distributed_optimizer OFF) so cross-batch gradient retention stays
+idempotent."""
+
+import logging
+import math
+from argparse import Namespace
+from contextlib import contextmanager
+from dataclasses import fields
+
+import torch
+from megatron.core.optimizer import get_megatron_optimizer
+from megatron.core.optimizer.clip_grads import clip_grad_by_total_norm_fp32, get_grad_norm_fp32
+from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer
+from megatron.core.optimizer.optimizer_config import OptimizerConfig
+from megatron.core.process_groups_config import ProcessGroupCollection
+
+
+logger = logging.getLogger(__name__)
+
+
+def validate_multi_lora_optimizer_args(args: Namespace) -> None:
+ """Reject launch options the per-slot optimizers cannot honor."""
+ assert not args.use_distributed_optimizer, (
+ "multi-LoRA per-slot optimizers require use_distributed_optimizer=False: "
+ "gradient retention relies on all-reduce idempotency, and LayerWise "
+ "sharding replaces byte-level ZeRO"
+ )
+ assert args.bf16 and not args.fp16, "multi-LoRA per-slot optimizers require bf16 (no dynamic loss scaler)"
+ assert (
+ args.optimizer or ""
+ ).lower() == "adam", (
+ f"multi-LoRA per-slot optimizers only implement Adam semantics; got optimizer={args.optimizer!r}"
+ )
+ for flag in ("optimizer_cpu_offload", "stream_optimizer_state_to_disk", "rematerialize_param_from_master_weight"):
+ assert not getattr(args, flag, False), f"--{flag.replace('_', '-')} is not supported with multi-LoRA slots"
+
+
+def adapter_slot_parameters(model, slot: int) -> list[torch.nn.Parameter]:
+ """All parameters belonging to one adapter slot, across model chunks."""
+ from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear
+
+ parameters = []
+ seen = set()
+ model_chunks = model if isinstance(model, (list, tuple)) else [model]
+ for model_chunk in model_chunks:
+ for module in model_chunk.modules():
+ if not isinstance(module, MultiLoRALinear):
+ continue
+ for param in module.adapters[slot].parameters():
+ if id(param) not in seen:
+ parameters.append(param)
+ seen.add(id(param))
+ return parameters
+
+
+def _adam_init_state_fn(opt, config=None):
+ for group in opt.param_groups:
+ for p in group["params"]:
+ if len(opt.state[p]) == 0:
+ opt.state[p]["exp_avg"] = torch.zeros_like(p.data)
+ opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data)
+
+
+@contextmanager
+def _only_slot_trainable(model_chunks, slot_params: list[torch.nn.Parameter]):
+ """Temporarily freeze every trainable param outside ``slot_params`` so the
+ stock param-group builder sees exactly one slot (the Muon construction
+ pattern from megatron's ``get_megatron_muon_optimizer``)."""
+ slot_ids = {id(p) for p in slot_params}
+ frozen = []
+ for model_chunk in model_chunks:
+ for param in model_chunk.parameters():
+ if param.requires_grad and id(param) not in slot_ids:
+ param.requires_grad = False
+ frozen.append(param)
+ try:
+ yield
+ finally:
+ for param in frozen:
+ param.requires_grad = True
+
+
+def _build_slot_base_optimizers(config, model, slot_params, *, use_gloo_process_groups: bool) -> list:
+ # create FP32 master wrappers only after LayerWise shards the base optimizers
+ config.bf16 = False
+ with _only_slot_trainable(model, slot_params):
+ chained = get_megatron_optimizer(config, list(model), use_gloo_process_groups=use_gloo_process_groups)
+ optimizers = [
+ child.optimizer
+ for child in chained.chained_optimizers
+ if getattr(child, "optimizer", None) is not None and child.get_parameters()
+ ]
+ config.bf16 = True
+ return optimizers
+
+
+class SlotOptimizer:
+ """One tenant's optimizer over one adapter slot.
+
+ Wraps a per-slot LayerWiseDistributedOptimizer, so every method touches
+ only this slot's parameters and state: reloads cannot round other tenants'
+ FP32 masters, and the parameter all-gather moves only this slot."""
+
+ def __init__(self, args: Namespace, model, slot: int) -> None:
+ self.slot = slot
+ self._model = model
+ slot_params = adapter_slot_parameters(model, slot)
+ assert slot_params, f"adapter slot {slot} has no parameters; is this a multi-LoRA model?"
+
+ config = OptimizerConfig(
+ **{f.name: getattr(args, f.name) for f in fields(OptimizerConfig) if hasattr(args, f.name)}
+ )
+ config.timers = None
+ base_optimizers = _build_slot_base_optimizers(
+ config, model, slot_params, use_gloo_process_groups=args.use_gloo_process_groups
+ )
+ assert base_optimizers, f"adapter slot {slot} produced no optimizer children"
+ self._inner = LayerWiseDistributedOptimizer(
+ base_optimizers,
+ config,
+ ProcessGroupCollection.use_mpu_process_groups(),
+ init_state_fn_list=[_adam_init_state_fn] * len(base_optimizers),
+ model_chunks=list(model),
+ )
+ # params are scattered whole across DP ranks; per-child norm/clip reductions must span the world
+ for child in self._inner.chained_optimizers:
+ child.grad_stats_parallel_group = None
+
+ def apply_adam_params(self, adam_params: dict) -> None:
+ for child in self._inner.chained_optimizers:
+ for group in child.param_groups:
+ group["lr"] = adam_params["learning_rate"]
+ group["betas"] = (adam_params["beta1"], adam_params["beta2"])
+ group["eps"] = adam_params["eps"]
+ group["weight_decay"] = adam_params["weight_decay"]
+
+ def prepare_grads(self) -> None:
+ for child in self._inner.chained_optimizers:
+ child.prepare_grads()
+
+ def clip_and_step(self, clip_grad: float) -> dict:
+ """-> {"grad_norm": x} stepped, {"skipped_nonfinite": 1.0} dropped.
+ Collective: the slot's grad norm is all-reduced over the world."""
+ grads_for_norm = []
+ slot_params = []
+ for child in self._inner.chained_optimizers:
+ grads_for_norm += child.get_grads_for_grad_norm()
+ slot_params += child.get_parameters()
+ slot_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None)
+ # BF16 has no grad scaler; the all-reduced norm detects inf/nan on every rank
+ if not math.isfinite(slot_norm):
+ return {"skipped_nonfinite": 1.0}
+ if clip_grad > 0.0 and slot_params:
+ clip_grad_by_total_norm_fp32(slot_params, clip_grad, slot_norm, False)
+ for child in self._inner.chained_optimizers:
+ child.step_with_ready_grads()
+ return {"grad_norm": float(slot_norm)}
+
+ def allgather_params(self) -> None:
+ self._inner.allgather_params()
+
+ def reload_masters(self) -> None:
+ """Refresh this slot's FP32 masters from its model parameters."""
+ self._inner.reload_model_params()
+
+ def sharded_state(self, model_sharded_state_dict: dict, *, is_loading: bool) -> dict:
+ """torch_dist form: every state tensor carries its param's global coordinates."""
+ return self._inner.sharded_state_dict(model_sharded_state_dict, is_loading=is_loading)
+
+ def load_sharded_state(self, loaded: dict) -> None:
+ self._inner.load_state_dict(loaded)
+
+ def zero_grads(self) -> None:
+ """Zero the slot's gradients everywhere they live: the DDP ``main_grad``
+ buffer views and any lingering ``grad``/``main_param.grad`` references."""
+ for param in adapter_slot_parameters(self._model, self.slot):
+ if (main_grad := getattr(param, "main_grad", None)) is not None:
+ main_grad.zero_()
+ param.grad = None
+ if (main_param := getattr(param, "main_param", None)) is not None:
+ main_param.grad = None
+
+
+def reset_grad_metadata_keep_grads(model_chunks) -> None:
+ """Reset DDP bookkeeping while retaining each slot's gradient accumulation window."""
+ for model_chunk in model_chunks:
+ if getattr(model_chunk.config, "cuda_graph_impl", "none") != "transformer_engine":
+ for param in model_chunk.params_with_grad:
+ param.grad_added_to_main_grad = False
+ for bucket_group in model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups:
+ bucket_group.reset()
+
+
+def step_slot_optimizers(
+ slot_optimizers: dict[int, SlotOptimizer],
+ adam_params_by_slot: dict[int, dict],
+) -> dict[int, dict]:
+ """Step slots in the same order on every rank; execution failures invalidate the cell."""
+ slots = sorted(adam_params_by_slot)
+ for slot in slots:
+ slot_optimizers[slot].apply_adam_params(adam_params_by_slot[slot])
+ slot_optimizers[slot].prepare_grads()
+
+ outcomes = {
+ slot: slot_optimizers[slot].clip_and_step(adam_params_by_slot[slot]["grad_clip_norm"]) for slot in slots
+ }
+ for slot in slots:
+ slot_optimizers[slot].zero_grads()
+ for slot in slots:
+ if "grad_norm" in outcomes[slot]:
+ slot_optimizers[slot].allgather_params()
+ return outcomes
diff --git a/miles/backends/megatron_utils/lora/slots.py b/miles/backends/megatron_utils/lora/slots.py
new file mode 100644
index 0000000000..4f51795b4e
--- /dev/null
+++ b/miles/backends/megatron_utils/lora/slots.py
@@ -0,0 +1,65 @@
+"""Slot mechanics for multi-LoRA: model instrumentation, shard addressing, and
+per-slot optimizer state hygiene. Kept deliberately free of any control-plane
+or data-plane logic."""
+
+import logging
+from argparse import Namespace
+
+import torch
+
+
+logger = logging.getLogger(__name__)
+
+
+def create_multi_lora_instance(args: Namespace):
+ """Create a MultiLoRA instance from training args."""
+ from megatron.bridge.peft.multi_lora import MultiLoRA
+
+ from miles.backends.megatron_utils.lora.utils import convert_target_modules_to_megatron
+
+ lora_type_name = getattr(args, "lora_type", "lora").lower()
+ if lora_type_name == "canonical_lora":
+ from megatron.bridge.peft.canonical_lora import CanonicalLoRA
+
+ lora_cls = CanonicalLoRA
+ else:
+ from megatron.bridge.peft.lora import LoRA
+
+ lora_cls = LoRA
+
+ # exclude_modules was already folded into target_modules during arg validation.
+ return MultiLoRA(
+ target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls),
+ n_adapters=args.multi_lora_n_adapters,
+ dim=args.lora_rank,
+ alpha=args.lora_alpha,
+ dropout=getattr(args, "lora_dropout", 0.0),
+ lora_A_init_method=getattr(args, "lora_A_init_method", "xavier"),
+ lora_B_init_method=getattr(args, "lora_B_init_method", "zero"),
+ )
+
+
+def slice_lora_to_rank(hf_name: str, tensor: torch.Tensor, adapter_rank: int) -> torch.Tensor:
+ """Trim a max-rank-padded LoRA tensor to ``adapter_rank`` on the rank axis, addressed
+ from the end so packed grouped-expert exports are not sliced on the expert axis."""
+ if "lora_A" in hf_name:
+ rank_dim = tensor.ndim - 2
+ if adapter_rank < tensor.shape[rank_dim]:
+ remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank)
+ assert remainder.abs().max() == 0, (
+ f"lora_A padded dims are non-zero: {hf_name}, "
+ f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
+ )
+ return tensor.narrow(rank_dim, 0, adapter_rank)
+ return tensor
+ if "lora_B" in hf_name:
+ rank_dim = tensor.ndim - 1
+ if adapter_rank < tensor.shape[rank_dim]:
+ remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank)
+ assert remainder.abs().max() == 0, (
+ f"lora_B padded dims are non-zero: {hf_name}, "
+ f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
+ )
+ return tensor.narrow(rank_dim, 0, adapter_rank)
+ return tensor
+ return tensor
diff --git a/miles/backends/megatron_utils/lora_utils.py b/miles/backends/megatron_utils/lora/utils.py
similarity index 99%
rename from miles/backends/megatron_utils/lora_utils.py
rename to miles/backends/megatron_utils/lora/utils.py
index 1008f06319..68a090244c 100644
--- a/miles/backends/megatron_utils/lora_utils.py
+++ b/miles/backends/megatron_utils/lora/utils.py
@@ -28,6 +28,7 @@
"gate_proj": "linear_fc1",
"up_proj": "linear_fc1",
"down_proj": "linear_fc2",
+ "lm_head": "output_layer",
# GDN (Qwen3.5/Qwen3-Next): both slices live in the single fused megatron in_proj
"in_proj_qkvz": "in_proj",
"in_proj_ba": "in_proj",
@@ -44,6 +45,7 @@
"gate_proj": "linear_fc1_gate",
"up_proj": "linear_fc1_up",
"down_proj": "linear_fc2",
+ "lm_head": "output_layer",
"in_proj_qkvz": "in_proj",
"in_proj_ba": "in_proj",
}
@@ -66,6 +68,7 @@
"linear_proj": ["o_proj"],
"linear_fc1": ["gate_proj", "up_proj"],
"linear_fc2": ["down_proj"],
+ "output_layer": ["lm_head"],
# CanonicalLoRA (split layers)
"linear_q": ["q_proj"],
"linear_k": ["k_proj"],
@@ -84,6 +87,7 @@
"gate_proj",
"up_proj",
"down_proj",
+ "lm_head",
"in_proj_qkvz",
"in_proj_ba",
}
@@ -236,8 +240,8 @@ def convert_target_modules_to_megatron(
) -> list[str]:
"""Convert HuggingFace LoRA target module names to Megatron format.
- HF: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
- Megatron (LoRA): linear_qkv, linear_proj, linear_fc1, linear_fc2
+ HF: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj, lm_head
+ Megatron (LoRA): linear_qkv, linear_proj, linear_fc1, linear_fc2, output_layer
Megatron (CanonicalLoRA): linear_q, linear_k, linear_v, linear_proj,
linear_fc1_up, linear_fc1_gate, linear_fc2
@@ -284,7 +288,7 @@ def convert_target_modules_to_hf(megatron_modules: list[str]) -> list[str]:
Supports both standard LoRA and CanonicalLoRA module names.
- Megatron standard: linear_qkv, linear_proj, linear_fc1, linear_fc2
+ Megatron standard: linear_qkv, linear_proj, linear_fc1, linear_fc2, output_layer
Megatron canonical: linear_q, linear_k, linear_v, linear_proj,
linear_fc1_up, linear_fc1_gate, linear_fc2
HF: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py
index d05643598a..aff3307f55 100644
--- a/miles/backends/megatron_utils/model.py
+++ b/miles/backends/megatron_utils/model.py
@@ -51,7 +51,7 @@
save_model_hashes,
)
from .initialize import is_first_replica_megatron_main_rank
-from .lora_utils import is_lora_enabled, is_lora_model
+from .lora.utils import is_lora_enabled, is_lora_model
from .model_provider import get_model_provider_func
from .parallel import get_packed_seq_params
@@ -63,10 +63,10 @@ def _has_loadable_ckpt(load_dir: str | None) -> bool:
return bool(load_dir) and Path(load_dir).is_dir() and any(Path(load_dir).iterdir())
-from .bridge_lora_helpers import _ensure_model_list, _setup_lora_model_via_bridge # noqa: F401
+from .lora.bridge import _ensure_model_list, _setup_lora_model_via_bridge # noqa: F401
-def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler:
+def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler | None:
"""Create and configure the optimizer learning-rate/weight-decay scheduler.
This configures iteration-based schedules derived from the global batch size
@@ -198,9 +198,11 @@ def setup_model_and_optimizer(
layer_wise_distributed_optimizer="dist" in config.optimizer.lower(),
)
elif is_multi_lora_enabled(args):
- from miles.backends.megatron_utils.multi_lora_optimizer import build_multi_lora_optimizer
+ from miles.backends.megatron_utils.lora.optimizer import validate_multi_lora_optimizer_args
- optimizer = build_multi_lora_optimizer(args, config, model)
+ # per-tenant SlotOptimizers are built at load_slot; there is no pool optimizer
+ validate_multi_lora_optimizer_args(args)
+ return model, None, None
else:
optimizer = get_megatron_optimizer(
config=config,
@@ -300,21 +302,8 @@ def forward_only(
def forward_step(
data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False
) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
- """Forward step used by Megatron's pipeline engine.
+ """Return the model output and its batch-bound loss callback."""
- Args:
- data_iterator (DataIterator): Input data iterator.
- model (GPTModel): The GPT model chunk to execute.
-
- Returns:
- tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
- Output tensor(s) and a callable that computes and packages results
- to be collected by the engine.
- """
-
- assert not return_schedule_plan, "forward_only step should never return schedule plan"
-
- # Get the batch.
batch = get_batch(
data_iterator,
[
@@ -413,68 +402,15 @@ def _zero_grads(model: Sequence[DDP], optimizer: MegatronOptimizer | None, disab
optimizer.zero_grad()
-def train_one_step(
- args: Namespace,
- rollout_id: int,
- step_id: int,
- data_iterator: Sequence[DataIterator],
- model: Sequence[DDP],
- optimizer: MegatronOptimizer | None,
- opt_param_scheduler: OptimizerParamScheduler | None,
- num_microbatches: int,
- num_rollouts: int,
- witness_info: WitnessInfo | None,
- attempt: int,
- ft_test_action_executor: FTTestActionActorExecutor | None = None,
-) -> tuple[dict[str, float], float, TrainStepOutcome]:
- """Execute a single pipeline-parallel training step.
-
- Runs forward/backward over ``num_microbatches``, applies optimizer step and
- one scheduler step when gradients are valid.
-
- Multi-LoRA: gradients are retained across train calls (per-adapter
- gradient accumulation); only the slots in the batch's ``step_slots`` step,
- and only their gradients are zeroed.
-
- Args:
- args: Runtime arguments.
- rollout_id: Rollout identifier.
- step_id: Step index within the current rollout.
- data_iterator: Iterable(s) yielding training batches.
- model: Sequence of DDP-wrapped model chunks.
- optimizer: Optimizer instance.
- opt_param_scheduler: LR/WD scheduler.
- num_microbatches: Number of microbatches to process.
- num_rollouts: This step's rollout count (loss normalizer + LR increment).
-
- Returns:
- Tuple of (reduced loss dict, gradient norm, step outcome).
- """
- args = get_args()
- parallel_state = get_parallel_state()
- dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=rollout_id)
- disable_optimizer = args.debug_disable_optimizer or optimizer is None
- multi_lora = is_multi_lora_enabled(args)
-
- if multi_lora:
- from miles.backends.megatron_utils.multi_lora_optimizer import reset_grad_metadata_keep_grads
-
- # Retain accumulated per-adapter gradients; reset only the per-iteration
- # DDP bookkeeping. Slot grads are zeroed selectively at step time.
- reset_grad_metadata_keep_grads(model)
- else:
- _zero_grads(model, optimizer, disable_optimizer)
-
- if args.custom_megatron_before_train_step_hook_path:
- from miles.utils.function_registry import load_function
-
- custom_before_train_step_hook = load_function(args.custom_megatron_before_train_step_hook_path)
- custom_before_train_step_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler)
+def run_forward_backward_pass(
+ args, dumper_phase_util, data_iterator, model, num_microbatches, num_rollouts, forward_only=False
+):
+ """One pipeline forward/backward pass over the microbatches; no optimizer interaction."""
@dumper_phase_util.wrap_forward_step
def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[
torch.Tensor,
- Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]],
+ Callable[[torch.Tensor], tuple[torch.Tensor, int, dict]],
]:
"""Forward step used by Megatron's pipeline engine during training.
@@ -483,7 +419,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
model (GPTModel): The GPT model chunk to execute.
Returns:
- tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]]]:
+ tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, int, dict]]]:
Output tensor(s) and the loss function, which returns
(loss, num_elems, {"keys": list[str], "values": torch.Tensor}).
"""
@@ -508,6 +444,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
"witness_ids",
"opd_reverse_kl",
"rollout_mask_sums",
+ "loss_weights",
+ "target_tokens",
+ "sample_indices",
],
args.data_pad_size_multiplier,
args.qkv_format,
@@ -568,15 +507,48 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
# Forward pass.
forward_backward_func = get_forward_backward_func()
- losses_reduced = forward_backward_func(
- forward_step_func=forward_step,
- data_iterator=data_iterator,
- model=model,
- num_microbatches=num_microbatches,
- seq_length=args.seq_length,
- micro_batch_size=args.micro_batch_size,
- decoder_seq_length=args.decoder_seq_length,
- forward_only=False,
+ with torch.no_grad() if forward_only else nullcontext():
+ return forward_backward_func(
+ forward_step_func=forward_step,
+ data_iterator=data_iterator,
+ model=model,
+ num_microbatches=num_microbatches,
+ seq_length=args.seq_length,
+ micro_batch_size=args.micro_batch_size,
+ decoder_seq_length=args.decoder_seq_length,
+ forward_only=forward_only,
+ )
+
+
+def train_one_step(
+ args: Namespace,
+ rollout_id: int,
+ step_id: int,
+ data_iterator: Sequence[DataIterator],
+ model: Sequence[DDP],
+ optimizer: MegatronOptimizer | None,
+ opt_param_scheduler: OptimizerParamScheduler | None,
+ num_microbatches: int,
+ num_rollouts: int,
+ witness_info: WitnessInfo | None,
+ attempt: int,
+ ft_test_action_executor: FTTestActionActorExecutor | None = None,
+) -> tuple[dict[str, float], float, TrainStepOutcome]:
+ """Run pipeline forward/backward, then step the optimizer and scheduler when gradients are valid."""
+ args = get_args()
+ parallel_state = get_parallel_state()
+ dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=rollout_id)
+ disable_optimizer = args.debug_disable_optimizer or optimizer is None
+ _zero_grads(model, optimizer, disable_optimizer)
+
+ if args.custom_megatron_before_train_step_hook_path:
+ from miles.utils.function_registry import load_function
+
+ custom_before_train_step_hook = load_function(args.custom_megatron_before_train_step_hook_path)
+ custom_before_train_step_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler)
+
+ losses_reduced = run_forward_backward_pass(
+ args, dumper_phase_util, data_iterator, model, num_microbatches, num_rollouts
)
outcome = TrainStepOutcome.NORMAL
@@ -598,7 +570,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
outcome = TrainStepOutcome.DISCARDED_SHOULD_RETRY
valid_step = False
- if (not disable_optimizer) and (not multi_lora) and (not getattr(args, "check_for_nan_in_loss_and_grad", True)):
+ if (not disable_optimizer) and (not getattr(args, "check_for_nan_in_loss_and_grad", True)):
found_inf_flag = optimizer.prepare_grads()
if found_inf_flag:
valid_step = False
@@ -623,24 +595,11 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
dumper_phase_util.finalize(model)
if not disable_optimizer and valid_step:
- if multi_lora:
- from miles.backends.megatron_utils.multi_lora_utils import step_stepped_adapter_slots
-
- grad_norm = step_stepped_adapter_slots(
- args, model, optimizer, data_iterator[0].rollout_data, rollout_id, step_id
- )
- else:
- # Update parameters.
- update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
+ update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
+ assert update_successful
+ opt_param_scheduler.step(increment=num_rollouts)
- # Update learning rate.
- assert update_successful
- opt_param_scheduler.step(increment=num_rollouts)
-
- # release grad (multi-LoRA retains accumulated grads; stepped slots were
- # zeroed selectively inside step_adapter_slots)
- if not multi_lora:
- _zero_grads(model, optimizer, disable_optimizer)
+ _zero_grads(model, optimizer, disable_optimizer)
log_structured(
logger.info,
@@ -676,12 +635,30 @@ def finalize_model_grads_with_empty_cache(*args, **kwargs):
free, total = torch.cuda.mem_get_info(device)
if free / total < 0.1:
clear_memory()
- from .lora_utils import reduce_marked_lora_grads
+ from .lora.utils import reduce_marked_lora_grads
reduce_marked_lora_grads(args[0])
return finalize_model_grads(*args, **kwargs)
+def setup_train_iteration_config(args, model, optimizer, disable_optimizer):
+ """Set the DDP gradient and parameter sync hooks for this training iteration."""
+ config = get_model_config(model[0])
+ config.grad_scale_func = None if disable_optimizer else optimizer.scale_loss
+ config.timers = None
+ if isinstance(model[0], DDP) and args.overlap_grad_reduce:
+ no_sync_funcs = [model_chunk.no_sync for model_chunk in model]
+ config.no_sync_func = no_sync_funcs[0] if len(model) == 1 else no_sync_funcs
+ if args.align_grad_reduce:
+ grad_sync_funcs = [model_chunk.start_grad_sync for model_chunk in model]
+ config.grad_sync_func = grad_sync_funcs[0] if len(model) == 1 else grad_sync_funcs
+ if args.overlap_param_gather and args.align_param_gather:
+ param_sync_funcs = [model_chunk.start_param_sync for model_chunk in model]
+ config.param_sync_func = param_sync_funcs[0] if len(model) == 1 else param_sync_funcs
+ config.finalize_model_grads_func = finalize_model_grads_with_empty_cache
+ return config
+
+
def train(
rollout_id: int,
model: Sequence[DDP],
@@ -724,27 +701,7 @@ def train(
for model_module in model:
model_module.train()
- # Setup some training config params.
- config = get_model_config(model[0])
- config.grad_scale_func = None if disable_optimizer else optimizer.scale_loss
- config.timers = None
- if isinstance(model[0], DDP) and args.overlap_grad_reduce:
- assert config.no_sync_func is None, (
- "When overlap_grad_reduce is True, config.no_sync_func must be None; "
- "a custom no_sync_func is not supported when overlapping grad-reduce"
- )
- config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
- if len(model) == 1:
- config.no_sync_func = config.no_sync_func[0]
- if args.align_grad_reduce:
- config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]
- if len(model) == 1:
- config.grad_sync_func = config.grad_sync_func[0]
- if args.overlap_param_gather and args.align_param_gather:
- config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]
- if len(model) == 1:
- config.param_sync_func = config.param_sync_func[0]
- config.finalize_model_grads_func = finalize_model_grads_with_empty_cache
+ config = setup_train_iteration_config(args, model, optimizer, disable_optimizer)
pre_hook_enabled = False
diff --git a/miles/backends/megatron_utils/multi_lora_optimizer.py b/miles/backends/megatron_utils/multi_lora_optimizer.py
deleted file mode 100644
index 8cd96b2a29..0000000000
--- a/miles/backends/megatron_utils/multi_lora_optimizer.py
+++ /dev/null
@@ -1,195 +0,0 @@
-"""Per-slot decoupled Adam optimizers for multi-LoRA, chained under Megatron's LayerWiseDistributedOptimizer;
-requires plain DDP all-reduce (use_distributed_optimizer OFF) so cross-batch gradient retention stays idempotent."""
-
-import logging
-from argparse import Namespace
-from collections.abc import Sequence
-from contextlib import contextmanager
-
-import torch
-from megatron.core.optimizer import get_megatron_optimizer
-from megatron.core.optimizer.clip_grads import clip_grad_by_total_norm_fp32, get_grad_norm_fp32
-from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer
-from megatron.core.optimizer.optimizer import MegatronOptimizer
-from megatron.core.optimizer.optimizer_config import OptimizerConfig
-from megatron.core.process_groups_config import ProcessGroupCollection
-
-logger = logging.getLogger(__name__)
-
-
-def adapter_slot_parameters(model, slot: int) -> list[torch.nn.Parameter]:
- """All parameters belonging to one adapter slot, across model chunks."""
- from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear
-
- parameters = []
- seen = set()
- model_chunks = model if isinstance(model, (list, tuple)) else [model]
- for model_chunk in model_chunks:
- for module in model_chunk.modules():
- if not isinstance(module, MultiLoRALinear):
- continue
- for param in module.adapters[slot].parameters():
- if id(param) not in seen:
- parameters.append(param)
- seen.add(id(param))
- return parameters
-
-
-def _adam_init_state_fn(opt, config=None):
- for group in opt.param_groups:
- for p in group["params"]:
- if len(opt.state[p]) == 0:
- opt.state[p]["exp_avg"] = torch.zeros_like(p.data)
- opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data)
-
-
-@contextmanager
-def _only_slot_trainable(model_chunks, slot_params: list[torch.nn.Parameter]):
- """Temporarily freeze every trainable param outside ``slot_params`` so the
- stock param-group builder sees exactly one slot (the Muon construction
- pattern from megatron's ``get_megatron_muon_optimizer``)."""
- slot_ids = {id(p) for p in slot_params}
- frozen = []
- for model_chunk in model_chunks:
- for param in model_chunk.parameters():
- if param.requires_grad and id(param) not in slot_ids:
- param.requires_grad = False
- frozen.append(param)
- try:
- yield
- finally:
- for param in frozen:
- param.requires_grad = True
-
-
-def build_multi_lora_optimizer(
- args: Namespace,
- config: OptimizerConfig,
- model_chunks: Sequence,
-) -> MegatronOptimizer:
- """Build one Float16-wrapped Adam per adapter slot under a LayerWiseDistributedOptimizer (ChainedOptimizer);
- each child's param groups are tagged with ``miles_multi_lora_slot`` and narrowed to this rank's shard."""
- assert not config.use_distributed_optimizer, (
- "multi-LoRA per-slot optimizers require use_distributed_optimizer=False: "
- "gradient retention relies on all-reduce idempotency, and LayerWise "
- "sharding replaces byte-level ZeRO"
- )
- assert not config.fp16, "multi-LoRA per-slot optimizers require bf16 (no dynamic loss scaler)"
- assert (config.optimizer or "").lower() == "adam", (
- "multi-LoRA per-slot optimizers only implement Adam semantics (state init, "
- f"slot retirement cleanup, step clocks); got optimizer={config.optimizer!r}"
- )
-
- pg_collection = ProcessGroupCollection.use_mpu_process_groups()
-
- # Defer bf16 master-weight creation into LayerWise (post-sharding) so fp32 masters exist only for owned params.
- reset_bf16 = config.bf16
- config.bf16 = False
-
- base_optimizers: list = []
- init_fns: list = []
- slot_child_indices: dict[int, list[int]] = {}
- try:
- for slot in range(args.multi_lora_n_adapters):
- slot_params = adapter_slot_parameters(model_chunks, slot)
- assert slot_params, f"adapter slot {slot} has no parameters; is this a multi-LoRA model?"
- with _only_slot_trainable(model_chunks, slot_params):
- chained = get_megatron_optimizer(
- config,
- list(model_chunks),
- use_gloo_process_groups=args.use_gloo_process_groups,
- )
- children = [
- child
- for child in chained.chained_optimizers
- if getattr(child, "optimizer", None) is not None and child.get_parameters()
- ]
- assert children, f"adapter slot {slot} produced no optimizer children"
- slot_child_indices[slot] = list(range(len(base_optimizers), len(base_optimizers) + len(children)))
- for child in children:
- for group in child.param_groups:
- group["miles_multi_lora_slot"] = slot
- base_optimizers.append(child)
- init_fns.append(_adam_init_state_fn)
- finally:
- config.bf16 = reset_bf16
-
- optimizer = LayerWiseDistributedOptimizer(base_optimizers, config, pg_collection, init_state_fn_list=init_fns)
-
- # Params are scattered whole across DP ranks, so per-child norm/clip reductions must span the world.
- for child in optimizer.chained_optimizers:
- child.grad_stats_parallel_group = None
-
- optimizer.miles_slot_child_indices = slot_child_indices
- logger.info(
- f"Built multi-LoRA LayerWise optimizer: {args.multi_lora_n_adapters} slots, "
- f"{len(optimizer.chained_optimizers)} chained children"
- )
- return optimizer
-
-
-def _slot_children(optimizer, slot: int):
- return [optimizer.chained_optimizers[i] for i in optimizer.miles_slot_child_indices[slot]]
-
-
-def reset_grad_metadata_keep_grads(model_chunks) -> None:
- """Reset DDP per-iteration grad bookkeeping WITHOUT zeroing grad buffers, so per-adapter accumulation
- survives across train batches (replaces ``DistributedDataParallel.zero_grad_buffer``)."""
- for model_chunk in model_chunks:
- if getattr(model_chunk.config, "cuda_graph_impl", "none") != "transformer_engine":
- for param in model_chunk.params_with_grad:
- param.grad_added_to_main_grad = False
- for bucket_group in model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups:
- bucket_group.reset()
-
-
-def zero_adapter_slot_grads(model, slot: int) -> None:
- """Zero one slot's gradients everywhere they live: the DDP ``main_grad`` buffer views
- and any lingering ``grad``/``main_param.grad`` references."""
- for param in adapter_slot_parameters(model, slot):
- if (main_grad := getattr(param, "main_grad", None)) is not None:
- main_grad.zero_()
- param.grad = None
- if (main_param := getattr(param, "main_param", None)) is not None:
- main_param.grad = None
-
-
-def step_adapter_slots(
- optimizer,
- model,
- step_batch_sizes: dict[int, int],
- clip_grad: float,
-) -> dict[int, float]:
- """Step exactly the slots in ``step_batch_sizes`` (slot -> batch size), retaining all other slots' gradients;
- scales each slot's accumulated grad sum by 1/batch_size and returns the grad norm per stepped slot."""
- grad_norms: dict[int, float] = {}
-
- for slot, batch_size in step_batch_sizes.items():
- children = _slot_children(optimizer, slot)
- # Copy accumulated main_grads into the owned masters' grads, then scale the sum to the adapter-batch mean.
- 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)
-
- # Per-slot grad norm over the slot's children, reduced across the whole world (whole-param DP scatter).
- grads_for_norm = []
- slot_params = []
- for child in children:
- grads_for_norm += child.get_main_grads_for_grad_norm()
- slot_params += child.get_parameters()
- slot_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None)
- if clip_grad > 0.0 and slot_params:
- clip_grad_by_total_norm_fp32(slot_params, clip_grad, slot_norm, False)
- grad_norms[slot] = float(slot_norm)
-
- for child in children:
- child.step_with_ready_grads()
-
- zero_adapter_slot_grads(model, slot)
-
- if step_batch_sizes:
- optimizer.allgather_params()
-
- return grad_norms
diff --git a/miles/backends/megatron_utils/multi_lora_scheduler.py b/miles/backends/megatron_utils/multi_lora_scheduler.py
deleted file mode 100644
index 4f148eb599..0000000000
--- a/miles/backends/megatron_utils/multi_lora_scheduler.py
+++ /dev/null
@@ -1,88 +0,0 @@
-"""Per-adapter LR/WD schedules for multi-LoRA: one ``OptimizerParamScheduler`` per adapter slot, positioned by
-the adapter's own trained samples. Adapters without a known ``num_step`` warm up, then hold ``--lr`` constant."""
-
-import logging
-from argparse import Namespace
-
-from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
-
-logger = logging.getLogger(__name__)
-
-
-class _SlotParamGroups:
- """Minimal optimizer facade: the scheduler only reads ``.param_groups``."""
-
- def __init__(self, param_groups: list[dict]):
- self.param_groups = param_groups
-
-
-def build_slot_scheduler(args: Namespace, optimizer, adapter, resume_step: int) -> OptimizerParamScheduler:
- """Build the slot's scheduler and position it at the adapter's committed
- samples. Rebuilt on every adapter load, so slot reuse starts fresh."""
- from miles.backends.megatron_utils.multi_lora_optimizer import _slot_children
-
- groups = [group for child in _slot_children(optimizer, adapter.slot) for group in child.param_groups]
- samples_per_step = adapter.config.adapter_global_batch_size
- num_step = adapter.config.num_step
-
- decay_steps = num_step * samples_per_step if num_step is not None else None
- if args.lr_warmup_fraction is not None and decay_steps is not None:
- lr_warmup_steps = args.lr_warmup_fraction * decay_steps
- else:
- lr_warmup_steps = args.lr_warmup_iters * samples_per_step
- if decay_steps is None:
- # No horizon: warm up, then hold constant. The decay steps only need
- # to satisfy the scheduler's warmup < decay invariant.
- lr_decay_style = "constant"
- decay_steps = int(lr_warmup_steps) + 1
- else:
- lr_decay_style = args.lr_decay_style
-
- scheduler = OptimizerParamScheduler(
- _SlotParamGroups(groups),
- init_lr=args.lr_warmup_init,
- max_lr=args.lr,
- min_lr=args.min_lr,
- lr_warmup_steps=lr_warmup_steps,
- lr_decay_steps=decay_steps,
- lr_decay_style=lr_decay_style,
- start_wd=args.start_weight_decay,
- end_wd=args.end_weight_decay,
- wd_incr_steps=decay_steps,
- wd_incr_style=args.weight_decay_incr_style,
- use_checkpoint_opt_param_scheduler=False,
- override_opt_param_scheduler=False,
- wsd_decay_steps=(
- args.lr_wsd_decay_iters * samples_per_step
- if lr_decay_style == "WSD" and args.lr_wsd_decay_iters is not None
- else None
- ),
- lr_wsd_decay_style=args.lr_wsd_decay_style,
- )
- if resume_step:
- scheduler.step(increment=resume_step * samples_per_step)
- return scheduler
-
-
-def install_slot_scheduler(args: Namespace, optimizer, adapter, resume_step: int) -> None:
- """Attach the adapter's scheduler to the optimizer, keyed by slot."""
- if not hasattr(optimizer, "miles_slot_schedulers"):
- optimizer.miles_slot_schedulers = {}
- optimizer.miles_slot_schedulers[adapter.slot] = build_slot_scheduler(args, optimizer, adapter, resume_step)
-
-
-def drop_slot_scheduler(optimizer, slot: int) -> None:
- """Detach a retired slot's scheduler (the next tenant installs its own)."""
- getattr(optimizer, "miles_slot_schedulers", {}).pop(slot, None)
-
-
-def step_slot_schedulers(optimizer, step_batch_sizes: dict[int, int]) -> dict[int, float]:
- """Advance exactly the stepped slots' schedules by their batch samples.
- Returns slot -> new learning rate, for logging."""
- lr_by_slot: dict[int, float] = {}
- for slot, batch_size in step_batch_sizes.items():
- scheduler = optimizer.miles_slot_schedulers[slot]
- scheduler.step(increment=batch_size)
- if scheduler.optimizer.param_groups: # empty on ranks owning none of the slot's params
- lr_by_slot[slot] = scheduler.optimizer.param_groups[0]["lr"]
- return lr_by_slot
diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py
deleted file mode 100644
index 5c4a5f2bfa..0000000000
--- a/miles/backends/megatron_utils/multi_lora_utils.py
+++ /dev/null
@@ -1,490 +0,0 @@
-import json
-import logging
-import os
-from argparse import Namespace
-from collections.abc import Mapping
-from pathlib import Path
-
-import ray
-import torch
-import torch.distributed as dist
-
-from miles.backends.training_utils.parallel import get_parallel_state
-from miles.ray.multi_lora.controller import get_multi_lora_controller
-from miles.utils.adapter_config import AdapterRun
-from miles.utils.distributed_utils import get_gloo_group
-
-logger = logging.getLogger(__name__)
-
-# Cached by adapter_shard_topology(); the topology is fixed for the run.
-_shard_topology: tuple[bool, tuple[tuple[int, int, int], ...]] | None = None
-
-
-def create_multi_lora_instance(args: Namespace):
- """Create a MultiLoRA instance from training args."""
- from megatron.bridge.peft.multi_lora import MultiLoRA
-
- from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_megatron
-
- lora_type_name = getattr(args, "lora_type", "lora").lower()
- if lora_type_name == "canonical_lora":
- from megatron.bridge.peft.canonical_lora import CanonicalLoRA
-
- lora_cls = CanonicalLoRA
- else:
- from megatron.bridge.peft.lora import LoRA
-
- lora_cls = LoRA
-
- # exclude_modules was already folded into target_modules during arg validation.
- return MultiLoRA(
- target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls),
- n_adapters=args.multi_lora_n_adapters,
- dim=args.lora_rank,
- alpha=args.lora_alpha,
- dropout=getattr(args, "lora_dropout", 0.0),
- lora_A_init_method=getattr(args, "lora_A_init_method", "xavier"),
- lora_B_init_method=getattr(args, "lora_B_init_method", "zero"),
- )
-
-
-def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int) -> str:
- """Adapter shard name for one (tp, pp, ep) coordinate; EP ranks hold different local
- experts. The ep suffix is omitted at ep_size == 1 so legacy checkpoints stay loadable."""
- name = f"adapter_megatron_tp{tp_rank}_pp{pp_rank}"
- if ep_size > 1:
- name += f"_ep{ep_rank}"
- return name + ".pt"
-
-
-def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
- """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` via one cached gloo all-gather."""
- global _shard_topology
- if _shard_topology is not None:
- return _shard_topology
- parallel_state = get_parallel_state()
- coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank)
- if not dist.is_initialized():
- _shard_topology = (True, (coords,))
- return _shard_topology
-
- current_rank = dist.get_rank()
- group = get_gloo_group()
- gathered: list[object] = [None] * dist.get_world_size(group=group)
- dist.all_gather_object(gathered, (coords, current_rank), group=group)
- is_writer = current_rank == min(rank for entry_coords, rank in gathered if entry_coords == coords)
- _shard_topology = (is_writer, tuple(sorted({entry_coords for entry_coords, _ in gathered})))
- return _shard_topology
-
-
-def all_megatron_checkpoints_exist(step_dir: Path, shard_names) -> bool:
- return all((step_dir / name).exists() for name in shard_names)
-
-
-def find_latest_checkpoint(ckpt_dir: Path) -> tuple[Path | None, int]:
- _, coords = adapter_shard_topology()
- if not ckpt_dir.exists():
- return None, 0
-
- parallel_state = get_parallel_state()
- ep_size = parallel_state.ep.size
- my_coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank)
-
- expected = {megatron_shard_name(*coord, ep_size) for coord in coords}
- my_shard = megatron_shard_name(*my_coords, ep_size)
- # Legacy pre-expert-adapter layout: no ep suffix; safe for all EP ranks to read (EP-replicated).
- legacy = {megatron_shard_name(tp, pp, 0, 1) for tp, pp, _ in coords}
- my_legacy = megatron_shard_name(my_coords[0], my_coords[1], 0, 1)
-
- def get_step(d):
- return int(d.name.split("_")[1])
-
- step_dirs = sorted(
- [d for d in ckpt_dir.iterdir() if d.is_dir() and d.name.startswith("step_")],
- key=get_step,
- reverse=True,
- )
- for step_dir in step_dirs:
- step = get_step(step_dir)
- if all_megatron_checkpoints_exist(step_dir, expected):
- return step_dir / my_shard, step
- if ep_size > 1 and all_megatron_checkpoints_exist(step_dir, legacy):
- logger.info(f"[multilora] resuming from pre-expert-parallel shard layout in {step_dir}")
- return step_dir / my_legacy, step
-
- return None, 0
-
-
-def zero_optimizer_state_for_adapter(optimizer, model, idx: int) -> None:
- from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear, _iter_multi_lora_modules
-
- target_main_params = set()
- for module in _iter_multi_lora_modules(model):
- if not isinstance(module, MultiLoRALinear):
- continue
- adapter = module.adapters[idx]
- for param in adapter.parameters():
- main = getattr(param, "main_param", None)
- target_main_params.add(id(main if main is not None else param))
-
- chained = getattr(optimizer, "chained_optimizers", [optimizer])
- for chained_optimizer in chained:
- inner = getattr(chained_optimizer, "optimizer", chained_optimizer)
- if inner is None:
- continue
- # TE/apex FusedAdam tracks the Adam step per param GROUP, not per param;
- # reset the retired slot's groups so the next tenant restarts bias correction.
- for group in inner.param_groups:
- if group.get("miles_multi_lora_slot") == idx and "step" in group:
- if isinstance(group["step"], torch.Tensor):
- group["step"].zero_()
- else:
- group["step"] = 0
- for param, state in inner.state.items():
- if id(param) not in target_main_params:
- continue
- if "exp_avg" in state:
- state["exp_avg"].zero_()
- if "exp_avg_sq" in state:
- state["exp_avg_sq"].zero_()
- # Bias correction restarts for the slot's next tenant.
- if "step" in state:
- if isinstance(state["step"], torch.Tensor):
- state["step"].zero_()
- else:
- state["step"] = 0
-
-
-def slice_lora_to_rank(hf_name: str, tensor: torch.Tensor, adapter_rank: int) -> torch.Tensor:
- """Trim a max-rank-padded LoRA tensor to ``adapter_rank`` on the rank axis, addressed
- from the end so packed grouped-expert exports are not sliced on the expert axis."""
- if "lora_A" in hf_name:
- rank_dim = tensor.ndim - 2
- if adapter_rank < tensor.shape[rank_dim]:
- remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank)
- assert remainder.abs().max() == 0, (
- f"lora_A padded dims are non-zero: {hf_name}, "
- f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
- )
- return tensor.narrow(rank_dim, 0, adapter_rank)
- return tensor
- if "lora_B" in hf_name:
- rank_dim = tensor.ndim - 1
- if adapter_rank < tensor.shape[rank_dim]:
- remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank)
- assert remainder.abs().max() == 0, (
- f"lora_B padded dims are non-zero: {hf_name}, "
- f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
- )
- return tensor.narrow(rank_dim, 0, adapter_rank)
- return tensor
- return tensor
-
-
-def save_multi_lora_checkpoints(
- args,
- model,
- adapter_steps: Mapping[str, int],
- adapters: Mapping[str, AdapterRun],
-):
- """Save per-adapter checkpoints in two formats per adapter.
-
- Layout (per adapter)::
-
- {adapter.save}/checkpoints/step_{iteration}/
- ├── adapter_megatron_tp{tp}_pp{pp}[_ep{ep}].pt ← per-rank shard, fast resume
- ├── adapter_model.safetensors ← gathered HF, inference / external
- └── adapter_config.json ← HF PEFT metadata (r, alpha, ...)
- """
- from megatron.bridge import AutoBridge
- from megatron.bridge.peft.multi_lora_layers import expose_adapter_slot
- from safetensors.torch import save_file as save_safetensors
-
- from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_hf
- from miles.utils import megatron_bridge_utils
-
- parallel_state = get_parallel_state()
- tp_rank = parallel_state.tp.rank
- pp_rank = parallel_state.pp.rank
- ep_rank = parallel_state.ep.rank
- ep_size = parallel_state.ep.size
- # Exactly one writer per (tp, pp, ep) shard; see adapter_shard_topology.
- is_shard_writer, _ = adapter_shard_topology()
- is_global_writer = is_shard_writer and tp_rank == 0 and pp_rank == 0 and ep_rank == 0
-
- target_modules_hf = (
- convert_target_modules_to_hf(list(args.target_modules))
- if args.target_modules
- else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
- )
-
- bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
-
- for adapter_name, adapter in adapters.items():
- config = adapter.config
- log_prefix = f"[multilora] ({adapter_name})"
- iteration = adapter_steps[adapter_name]
-
- if config.save is None:
- logger.info(f"{log_prefix} skipping checkpoint (no save dir configured)")
- continue
-
- final_dir = config.save / "checkpoints" / f"step_{iteration}"
- tmp_dir = config.save / "checkpoints" / f"_tmp_step_{iteration}"
- if is_shard_writer:
- tmp_dir.mkdir(parents=True, exist_ok=True)
- if dist.is_initialized():
- dist.barrier()
-
- with expose_adapter_slot(model, adapter.slot):
- # Megatron checkpoints
- if is_shard_writer:
- shard: dict[str, torch.Tensor] = {
- name: param.data.cpu()
- for batch in model
- for name, param in batch.named_parameters()
- if ".adapter." in name
- }
- native_path = tmp_dir / megatron_shard_name(tp_rank, pp_rank, ep_rank, ep_size)
- torch.save(shard, native_path)
- logger.info(f"{log_prefix} saved Megatron shard " f"({len(shard)} tensors) to {native_path}")
-
- hf_state: dict[str, torch.Tensor] = {}
- with megatron_bridge_utils.patch_megatron_model(model):
- for hf_name, weight, _megatron_name in bridge.export_adapter_weights(
- model,
- cpu=True,
- show_progress=False,
- ):
- # Slice from the shared --lora-rank down to this adapter's real rank to
- # match adapter_config's r; clone() since safetensors rejects aliased views.
- hf_state[hf_name] = slice_lora_to_rank(hf_name, weight, config.rank).clone()
-
- if is_global_writer:
- save_safetensors(
- hf_state,
- str(tmp_dir / "adapter_model.safetensors"),
- metadata={"format": "pt"},
- )
- adapter_config_json = {
- "peft_type": "LORA",
- "r": config.rank,
- "lora_alpha": config.alpha,
- "target_modules": target_modules_hf,
- "lora_dropout": getattr(args, "lora_dropout", 0.0),
- "bias": "none",
- "task_type": "CAUSAL_LM",
- }
- with open(tmp_dir / "adapter_config.json", "w") as f:
- json.dump(adapter_config_json, f, indent=2)
- os.sync()
- logger.info(f"{log_prefix} saved HF PEFT to {tmp_dir} " f"({len(hf_state)} tensors)")
-
- if dist.is_initialized():
- dist.barrier()
-
- # Write to a temp dir and move into place so readers never see a
- # partially written checkpoint.
- if is_global_writer:
- if final_dir.exists():
- import shutil
-
- shutil.rmtree(final_dir)
- os.replace(tmp_dir, final_dir)
- logger.info(f"{log_prefix} promoted checkpoint to {final_dir}")
- if dist.is_initialized():
- dist.barrier()
-
-
-def _register_adapter(adapter: AdapterRun, model) -> int:
- """Install one adapter on this rank's local model shard. Returns the step
- of the checkpoint it resumed from (0 for a fresh adapter)."""
- from megatron.bridge.peft.multi_lora_layers import init_adapter_slot, load_adapter
-
- name = adapter.name
- config = adapter.config
- slot = adapter.slot
- log_prefix = f"[multilora] ({name})"
-
- step = 0
- if config.save is not None:
- ckpt_root = config.save / "checkpoints"
- ckpt, step = find_latest_checkpoint(ckpt_root)
- else:
- ckpt = None
-
- if ckpt is None:
- logger.info(f"{log_prefix} no checkpoint, starting from random init")
- step = 0
- else:
- state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
- loaded = load_adapter(model, slot, state_dict)
- assert loaded > 0, (
- f"{log_prefix} loaded 0 tensors from {ckpt} "
- f"(state_dict has {len(state_dict)} entries) — name mismatch?"
- )
- logger.info(f"{log_prefix} loaded from {ckpt} ({loaded} tensors)")
-
- init_adapter_slot(model, slot, rank=config.rank, alpha=config.alpha)
- logger.info(f"{log_prefix} installed at slot {slot}")
- return step
-
-
-def _deregister_adapter(adapter: AdapterRun, args, model, optimizer) -> None:
- """Model-side cleanup for one adapter."""
- from megatron.bridge.peft.multi_lora_layers import clear_adapter_slot
-
- name = adapter.name
- slot = adapter.slot
- log_prefix = f"[multilora] ({name})"
-
- if args.save_interval is not None:
- # The controller still holds the step count until free_slot runs.
- step = ray.get(get_multi_lora_controller().adapter_step.remote(name))
- save_multi_lora_checkpoints(args, model, {name: step}, {name: adapter})
- logger.info(f"{log_prefix} saved final checkpoint at step {step}")
- else:
- logger.info(f"{log_prefix} save_interval unset; skipping final checkpoint")
-
- clear_adapter_slot(model, slot)
- logger.info(f"{log_prefix} cleared adapter slot {slot}")
-
- # Prevent future slot tenants from inheriting optimizer momentum or the
- # previous tenant's partially accumulated gradients.
- from miles.backends.megatron_utils.multi_lora_optimizer import zero_adapter_slot_grads
-
- from miles.backends.megatron_utils.multi_lora_scheduler import drop_slot_scheduler
-
- zero_optimizer_state_for_adapter(optimizer, model, slot)
- zero_adapter_slot_grads(model, slot)
- drop_slot_scheduler(optimizer, slot)
- optimizer.reload_model_params()
- logger.info(f"{log_prefix} cleared optimizer state and retained grads for slot {slot}")
-
-
-def load_adapters(args, model, optimizer, adapters) -> int:
- """Load adapters into Megatron slots; resumes step counts from checkpoints."""
- from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
- from miles.utils.distributed_utils import get_gloo_group
-
- if dist.is_initialized():
- dist.barrier(group=get_gloo_group())
- if not adapters:
- return 0
- from miles.backends.megatron_utils.multi_lora_scheduler import install_slot_scheduler
-
- resume_steps: dict[str, int] = {}
- for adapter in adapters:
- resume_steps[adapter.name] = _register_adapter(adapter, model)
- # Per-adapter LR/WD schedule, positioned at the resumed step count.
- install_slot_scheduler(args, optimizer, adapter, resume_steps[adapter.name])
- if dist.is_initialized():
- dist.barrier(group=get_gloo_group())
- optimizer.reload_model_params()
- if is_first_replica_megatron_main_rank():
- for name, step in resume_steps.items():
- if step > 0:
- ray.get(get_multi_lora_controller().set_adapter_step.remote(name, step))
- return len(adapters)
-
-
-def cleanup_adapters(args, model, optimizer, adapters) -> int:
- """Save final ckpt + clear Megatron slot, then free_slot on the controller."""
- from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
- from miles.utils.distributed_utils import get_gloo_group
-
- if dist.is_initialized():
- dist.barrier(group=get_gloo_group())
- if not adapters:
- return 0
- for adapter in adapters:
- _deregister_adapter(adapter, args, model, optimizer)
- if dist.is_initialized():
- dist.barrier(group=get_gloo_group())
- if is_first_replica_megatron_main_rank():
- for adapter in adapters:
- ray.get(get_multi_lora_controller().free_slot.remote(adapter.name))
- return len(adapters)
-
-
-def step_stepped_adapter_slots(args, model, optimizer, rollout_data, rollout_id: int, step_id: int) -> float:
- """Optimizer-step the slots whose adapter batch completes with this train batch and advance
- their per-adapter LR/WD schedules. Returns the max grad norm across stepped slots (0.0 if none)."""
- from miles.backends.megatron_utils.multi_lora_optimizer import step_adapter_slots
- from miles.backends.megatron_utils.multi_lora_scheduler import step_slot_schedulers
- from miles.utils.tracking_utils.structured_log import log_structured
-
- # slot -> adapter_global_batch_size for adapter batches completing now.
- step_batch_sizes = dict(rollout_data.get("step_adapter_batch_sizes", {}))
- grad_norms_by_slot = step_adapter_slots(
- optimizer,
- model,
- step_batch_sizes,
- clip_grad=args.clip_grad,
- )
-
- if lr_by_slot := step_slot_schedulers(optimizer, step_batch_sizes):
- log_structured(
- logger.info,
- tag="train",
- op="adapter_lr",
- rollout=rollout_id,
- step=step_id,
- **{f"slot_{slot}": lr for slot, lr in lr_by_slot.items()},
- )
- return max(grad_norms_by_slot.values(), default=0.0)
-
-
-def commit_trained_batch(rollout_data, rollout_id: int, pending_push: set) -> None:
- """A train call landed: schedule the stepped adapters' engine push and
- commit the batch on the controller (main rank only). The stepped set ships
- with the train data, identical on all ranks."""
- from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
-
- pending_push.update(rollout_data.get("step_adapter_names", []))
- if is_first_replica_megatron_main_rank():
- ray.get(get_multi_lora_controller().mark_batch_trained.remote(rollout_id))
-
-
-def save_due_adapter_checkpoints(args, model) -> bool:
- """Save per-adapter checkpoints for adapters at a save-interval multiple
- without a checkpoint on disk. Rank 0 picks and broadcasts, so the
- collective export lines up. Returns False when nothing is due."""
- from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
- from miles.utils.distributed_utils import get_gloo_group
-
- due_buffer = [None]
- if is_first_replica_megatron_main_rank() and args.save_interval is not None:
- snapshot = ray.get(get_multi_lora_controller().snapshot.remote())
- adapters = {**snapshot["active"], **snapshot["retiring"]}
- due_buffer[0] = {
- name: adapter
- for name, adapter in adapters.items()
- if adapter.step > 0
- and adapter.step % args.save_interval == 0
- and adapter.config.save is not None
- and not (Path(adapter.config.save) / "checkpoints" / f"step_{adapter.step}").exists()
- }
- if dist.is_initialized():
- dist.broadcast_object_list(due_buffer, src=0, group=get_gloo_group())
- due_adapters = due_buffer[0]
- if not due_adapters:
- return False
- adapter_steps = {name: adapter.step for name, adapter in due_adapters.items()}
- save_multi_lora_checkpoints(args, model, adapter_steps, due_adapters)
- return True
-
-
-def select_adapters_to_push(loaded_adapters: dict, pending_push: set, has_new_engines: bool) -> tuple[dict, list]:
- """Pick the stale adapters to push (all loaded adapters when engines are new). Returns
- (adapters to push keyed by name, names to version-bump — only those whose weights changed)."""
- pending = pending_push & set(loaded_adapters)
- push_names = set(loaded_adapters) if has_new_engines else pending
- return {name: loaded_adapters[name] for name in sorted(push_names)}, sorted(pending)
-
-
-def commit_weight_push(version_update_names: list, is_main_rank: bool) -> None:
- """A weight push landed: bump the pushed adapters' slot versions on the
- controller (promotes PENDING adapters to ACTIVE)."""
- if version_update_names and is_main_rank:
- ray.get(get_multi_lora_controller().record_weight_update.remote(version_update_names))
diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py
index ca3c245704..3e95f52a0d 100644
--- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py
+++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py
@@ -35,7 +35,7 @@ def __init__(self, *args, **kwargs):
def _hf_atomic_update_groups(self):
return get_hf_atomic_update_groups(self.model_name, q_lora_rank=self.args.q_lora_rank)
- def _iter_hf_adapter_units(self, lora_name, adapter, *, materialize):
+ def _iter_hf_adapter_units(self, adapter, *, materialize):
"""Both megatron exporters are PP-local after gathering TP/EP; the PP
gather runs only where the resolved placement asks for it."""
named_tensors = self._export_pp_local_lora(adapter)
@@ -55,7 +55,7 @@ def _iter_hf_adapter_units(self, lora_name, adapter, *, materialize):
raise RuntimeError("LoRA weight sync failed: the adapter export contains no lora_A/lora_B names.")
while named_tensors:
hf_name, tensor = named_tensors.pop(0)
- yield [(f"{lora_name}:{hf_name}", tensor)]
+ yield [(hf_name, tensor)]
@abstractmethod
def _export_pp_local_lora(self, adapter) -> list[tuple[str, torch.Tensor]]:
diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
index a48db2d68a..32d4daebc3 100644
--- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
+++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
@@ -75,11 +75,11 @@ def _export_pp_local_lora(self, adapter):
from megatron.bridge.peft.multi_lora_layers import expose_adapter_slot
- from ..multi_lora_utils import slice_lora_to_rank
+ from ..lora.slots import slice_lora_to_rank
with expose_adapter_slot(self.model, adapter.slot):
named_tensors = self._export_current_adapter()
- return [(h, slice_lora_to_rank(h, w, adapter.config.rank)) for h, w in named_tensors]
+ return [(h, slice_lora_to_rank(h, w, adapter.rank)) for h, w in named_tensors]
def _export_current_adapter(self) -> list:
with megatron_bridge_utils.patch_megatron_model(self.model):
diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
index 046fd353e2..053e5b3bfe 100644
--- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
+++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
@@ -201,7 +201,7 @@ def _get_megatron_local_param_infos(
"""
pp_size = get_parallel_state().pp.size
- from ..lora_utils import _is_adapter_param_name
+ from ..lora.utils import _is_adapter_param_name
param_infos: dict[str, ParamInfo] = {}
rank = dist.get_rank()
diff --git a/miles/backends/sglang_utils/sglang_config.py b/miles/backends/sglang_utils/sglang_config.py
index 9b60cc2302..7649feb2cf 100644
--- a/miles/backends/sglang_utils/sglang_config.py
+++ b/miles/backends/sglang_utils/sglang_config.py
@@ -9,6 +9,7 @@
import yaml
from miles.backends.sglang_utils.arguments import collect_eval_sglang_overrides
+from miles.utils.multi_lora import is_multi_lora_enabled
from miles.utils.pydantic_utils import FrozenStrictBaseModel
logger = logging.getLogger(__name__)
@@ -226,7 +227,11 @@ def resolve(cls, raw: _RawModelConfig, args, offset_cursor: "_OffsetCursor") ->
effective_model_path = default_model_path
update_weights = raw.update_weights
- if update_weights is None:
+ if is_multi_lora_enabled(args):
+ assert update_weights is not True, "Tinker loads a frozen base; update_weights must be false"
+ assert effective_model_path == args.hf_checkpoint, "Tinker engines must load the trainer's base checkpoint"
+ update_weights = False
+ elif update_weights is None:
if effective_model_path != args.hf_checkpoint:
logger.warning(
f"Model '{raw.name}' uses model_path='{effective_model_path}' which differs "
diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py
index 4ee8094f1b..c01549bc48 100644
--- a/miles/backends/sglang_utils/sglang_engine.py
+++ b/miles/backends/sglang_utils/sglang_engine.py
@@ -7,7 +7,7 @@
from sglang.srt.server_args import ServerArgs
-from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_hf, sglang_lora_target_all_sentinel
+from miles.backends.megatron_utils.lora.utils import convert_target_modules_to_hf, sglang_lora_target_all_sentinel
from miles.backends.sglang_utils.server_args_utils import server_args_to_argv
from miles.utils.lora import LORA_ADAPTER_NAME, lora_base_cpu_backup_enabled, lora_rollout_enabled
from miles.utils.multi_lora import is_multi_lora_enabled
@@ -190,4 +190,11 @@ def _compute_server_args(
for key in unused_keys:
kwargs.pop(key)
+ if is_multi_lora_enabled(args):
+ assert kwargs.get("load_format") != "dummy", "Tinker engines must load the frozen base from disk"
+ if kwargs.get("max_loaded_loras") is None:
+ # use --sglang-max-loaded-loras to override
+ # TODO: dynamic allocation
+ kwargs["max_loaded_loras"] = 2 * kwargs["max_loras_per_batch"]
+
return kwargs
diff --git a/miles/backends/training_utils/checkpoint_io.py b/miles/backends/training_utils/checkpoint_io.py
new file mode 100644
index 0000000000..d82bc4c775
--- /dev/null
+++ b/miles/backends/training_utils/checkpoint_io.py
@@ -0,0 +1,70 @@
+"""Checkpoint directories: written collectively, complete at their final path."""
+
+# TODO: isolate checkpoint IO failures; they currently terminate the trainer cell.
+
+import json
+import os
+import shutil
+from collections.abc import Callable
+from pathlib import Path
+from uuid import uuid4
+
+import torch.distributed as dist
+
+from miles.utils.distributed_utils import get_gloo_group
+
+
+def write_checkpoint_dir(
+ path: str | Path,
+ write_shards: Callable[[Path], None],
+ metadata: dict | None = None,
+ *,
+ overwrite: bool = True,
+) -> None:
+ """Write collectively, then atomically point ``path`` at the completed version.
+
+ All ranks must call. Readers may still hold an older version, so retain it.
+ """
+ final_dir = Path(path)
+ tmp_dir = final_dir.parent / f"_tmp_{final_dir.name}"
+
+ def make_tmp_dir():
+ if _rank() == 0:
+ if not overwrite and final_dir.exists():
+ raise FileExistsError(f"checkpoint {final_dir} already exists")
+ if final_dir.exists() and not final_dir.is_symlink():
+ raise NotImplementedError(
+ f"cannot overwrite a legacy checkpoint directory {final_dir}; save under a new name"
+ )
+ # a crashed attempt may leave shards or an unpublished version link
+ if tmp_dir.is_symlink():
+ tmp_dir.unlink()
+ elif tmp_dir.exists():
+ shutil.rmtree(tmp_dir)
+ tmp_dir.mkdir(parents=True)
+
+ def publish_dir():
+ if _rank() != 0:
+ return
+ if metadata is not None:
+ (tmp_dir / "META.json").write_text(json.dumps(metadata, indent=2))
+ version_dir = final_dir.parent / f"_version_{final_dir.name}_{uuid4().hex}"
+ os.replace(tmp_dir, version_dir)
+ tmp_dir.symlink_to(version_dir.name, target_is_directory=True)
+ os.replace(tmp_dir, final_dir)
+
+ make_tmp_dir()
+ _barrier()
+ write_shards(tmp_dir)
+ _barrier()
+ publish_dir()
+ _barrier()
+
+
+def _rank() -> int:
+ return dist.get_rank() if dist.is_initialized() else 0
+
+
+def _barrier() -> None:
+ if dist.is_initialized():
+ dist.barrier(group=get_gloo_group())
diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py
index 0716074c76..45d0c049ef 100644
--- a/miles/backends/training_utils/data.py
+++ b/miles/backends/training_utils/data.py
@@ -210,6 +210,10 @@ def get_batch(
if "dynamic_global_batch_size" in data_iterator.rollout_data:
batch["dynamic_global_batch_size"] = data_iterator.rollout_data["dynamic_global_batch_size"]
+ for unit_key in ("loss_fn", "loss_fn_config"):
+ if unit_key in data_iterator.rollout_data:
+ batch[unit_key] = data_iterator.rollout_data[unit_key]
+
# No-op safety net if batches reach get_batch without rollout-level preprocessing.
expand_multimodal_rollout_data_in_place(batch, qkv_format=qkv_format)
diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py
index 5683c54b01..52a56acb7b 100644
--- a/miles/backends/training_utils/log_utils.py
+++ b/miles/backends/training_utils/log_utils.py
@@ -210,9 +210,6 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc
"num_rollouts",
"n_adapters",
"adapter_slots",
- "step_slots",
- "step_adapter_names",
- "step_adapter_batch_sizes",
"prompt_group_sizes",
]:
continue
diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py
index 92e89d9940..e68a253afc 100644
--- a/miles/backends/training_utils/loss.py
+++ b/miles/backends/training_utils/loss.py
@@ -132,11 +132,12 @@ def loss_function(
logits: torch.Tensor,
apply_megatron_loss_scaling: bool = False,
num_rollouts: int | None = None,
-) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]:
+) -> tuple[torch.Tensor, int | torch.Tensor, dict]:
"""Dispatch to the configured loss and rescale for Megatron integration.
Selects one of "policy_loss", "value_loss", "sft_loss", or a custom loss
- function based on `args.loss_type`, computes the loss and metrics, then
+ function based on `args.loss_type`, or the Tinker loss specified by
+ `batch["loss_fn"]`. Computes the loss and metrics, then
rescales the loss by micro-batch and parallelism factors to integrate with
Megatron's gradient accumulation.
@@ -157,6 +158,7 @@ def loss_function(
`args.calculate_per_token_loss` is True, else `1` (int).
- `logging_dict` has keys "keys" (list of str metric names) and
"values" (1D tensor: [count, metric1, metric2, ...]).
+ Tinker losses may also return "per_datum" outputs.
"""
parallel_state = get_parallel_state()
num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]])
@@ -172,7 +174,7 @@ def loss_function(
denominators=batch.get("rollout_mask_sums", None),
)
- func = get_loss_function(args)
+ func = get_loss_function(args, batch.get("loss_fn"))
if args.recompute_loss_function:
loss, log = checkpoint(
@@ -196,11 +198,8 @@ def loss_function(
else:
assert args.use_dynamic_global_batch_size == ("dynamic_global_batch_size" in batch)
global_batch_size = batch.get("dynamic_global_batch_size", args.global_batch_size)
- # Multi-LoRA: samples enter the gradient buffers with weight 1; per-adapter
- # normalization (1/adapter_global_batch_size, a constant known in advance)
- # is applied to the accumulated slot gradient at optimizer-step time.
- if is_multi_lora_enabled(args):
- global_batch_size = 1
+ # Tinker losses already carry the client's normalization
+ loss_normalizer = 1 if is_multi_lora_enabled(args) else global_batch_size
if not args.calculate_per_token_loss:
if apply_megatron_loss_scaling:
loss_parallel_size = (
@@ -208,17 +207,19 @@ def loss_function(
if args.true_on_policy_mode and parallel_state.is_ulysses_cp
else parallel_state.intra_dp_cp.size
)
- loss = loss * num_microbatches / global_batch_size * loss_parallel_size
+ loss = loss * num_microbatches / loss_normalizer * loss_parallel_size
else:
- loss = loss / global_batch_size * parallel_state.intra_dp.size
+ loss = loss / loss_normalizer * parallel_state.intra_dp.size
else:
if apply_megatron_loss_scaling:
loss = loss * parallel_state.cp.size
+ per_datum = log.pop("per_datum", None) if batch.get("loss_fn") is not None else None
return (
loss,
torch.tensor(num_tokens if args.calculate_per_token_loss else 1, device=logits.device),
{
+ **({"per_datum": per_datum} if per_datum is not None else {}),
"keys": list(log.keys()),
"values": torch.tensor(
[num_samples if not args.calculate_per_token_loss else num_tokens] + list(log.values()),
diff --git a/miles/backends/training_utils/loss_hub/losses.py b/miles/backends/training_utils/loss_hub/losses.py
index 2d5a207738..d2d40b0693 100644
--- a/miles/backends/training_utils/loss_hub/losses.py
+++ b/miles/backends/training_utils/loss_hub/losses.py
@@ -505,7 +505,11 @@ def sft_loss_function(
)
-def get_loss_function(args: Namespace) -> LossFunction:
+def get_loss_function(args: Namespace, loss_fn: str | None = None) -> LossFunction:
+ if loss_fn is not None:
+ from miles.backends.training_utils.loss_hub.tinker_losses import TINKER_LOSS_FUNCTIONS
+
+ return TINKER_LOSS_FUNCTIONS[loss_fn]
match args.loss_type:
case "policy_loss":
return policy_loss_function
diff --git a/miles/backends/training_utils/loss_hub/tinker_losses.py b/miles/backends/training_utils/loss_hub/tinker_losses.py
new file mode 100644
index 0000000000..27f2fade1c
--- /dev/null
+++ b/miles/backends/training_utils/loss_hub/tinker_losses.py
@@ -0,0 +1,163 @@
+"""Sum Tinker losses over datums without server-side normalization.
+
+Client loss inputs carry normalization; the trainer accumulates raw sums.
+The SDK represents custom-loss gradients as `weights = -dL/dlogprob` with `cross_entropy`.
+"""
+
+from argparse import Namespace
+from collections.abc import Callable
+
+import torch
+
+from miles.backends.training_utils.loss_hub.logit_processors import get_log_probs_and_entropy
+from miles.utils.types import RolloutBatch
+
+PPO_DEFAULTS = {"clip_low_threshold": 0.8, "clip_high_threshold": 1.2}
+CISPO_DEFAULTS = {"clip_low_threshold": 0.0, "clip_high_threshold": 4.0}
+DRO_DEFAULTS = {"beta": 0.05}
+
+
+def _target_logprobs(args: Namespace, batch: RolloutBatch, logits: torch.Tensor) -> list[torch.Tensor]:
+ # Tinker targets are explicit labels: splice them over the response region of the gather sequence
+ label_tokens = [
+ torch.cat([tokens[: len(tokens) - len(targets)], _as_tensor_like(targets, tokens)])
+ for tokens, targets in zip(batch["unconcat_tokens"], batch["target_tokens"], strict=True)
+ ]
+ outputs = get_log_probs_and_entropy(
+ logits,
+ args=args,
+ unconcat_tokens=label_tokens,
+ total_lengths=batch["total_lengths"],
+ response_lengths=batch["response_lengths"],
+ with_entropy=False,
+ max_seq_lens=batch.get("max_seq_lens", None),
+ )
+ return outputs["log_probs"]
+
+
+def _as_tensor_like(values, reference: torch.Tensor) -> torch.Tensor:
+ return torch.as_tensor(values, dtype=reference.dtype, device=reference.device)
+
+
+def _response_masks(batch: RolloutBatch, log_probs: list[torch.Tensor]) -> list[torch.Tensor]:
+ """Per-datum loss masks; a DP-padding datum is all zeros and must not reach the objective."""
+ return [_as_tensor_like(mask, log_prob) for mask, log_prob in zip(batch["loss_masks"], log_probs, strict=True)]
+
+
+def _sum_loss_and_outputs(
+ batch: RolloutBatch,
+ logits: torch.Tensor,
+ log_probs: list[torch.Tensor],
+ per_datum_losses: list[torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+ if per_datum_losses:
+ loss = torch.stack(per_datum_losses).sum()
+ else:
+ # a microbatch with no supervised tokens still needs the graph alive; fp32 sum avoids fp16 inf -> nan
+ loss = logits.sum(dtype=torch.float32) * 0
+ per_datum = [
+ {"sample_index": index, "logprobs": log_prob.detach().cpu(), "loss": sample_loss.detach().cpu()}
+ for index, log_prob, sample_loss in zip(batch["sample_indices"], log_probs, per_datum_losses, strict=True)
+ ]
+ return loss, {"loss": loss.detach(), "per_datum": per_datum}
+
+
+def cross_entropy_loss_function(
+ args: Namespace,
+ batch: RolloutBatch,
+ logits: torch.Tensor,
+ sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+ log_probs = _target_logprobs(args, batch, logits)
+ per_datum_losses = [
+ -(_as_tensor_like(weights, log_prob) * log_prob * mask).sum()
+ for log_prob, weights, mask in zip(
+ log_probs, batch["loss_weights"], _response_masks(batch, log_probs), strict=True
+ )
+ ]
+ return _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses)
+
+
+def importance_sampling_loss_function(
+ args: Namespace,
+ batch: RolloutBatch,
+ logits: torch.Tensor,
+ sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+ log_probs = _target_logprobs(args, batch, logits)
+ per_datum_losses = []
+ for log_prob, sampling_log_prob, advantage, mask in zip(
+ log_probs, batch["rollout_log_probs"], batch["advantages"], _response_masks(batch, log_probs), strict=True
+ ):
+ ratio = torch.exp(log_prob - _as_tensor_like(sampling_log_prob, log_prob))
+ per_datum_losses.append(-(ratio * _as_tensor_like(advantage, log_prob) * mask).sum())
+ return _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses)
+
+
+def ppo_loss_function(
+ args: Namespace,
+ batch: RolloutBatch,
+ logits: torch.Tensor,
+ sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+ config = batch.get("loss_fn_config") or {}
+ clip_low = config.get("clip_low_threshold", PPO_DEFAULTS["clip_low_threshold"])
+ clip_high = config.get("clip_high_threshold", PPO_DEFAULTS["clip_high_threshold"])
+ log_probs = _target_logprobs(args, batch, logits)
+ per_datum_losses = []
+ for log_prob, sampling_log_prob, advantage, mask in zip(
+ log_probs, batch["rollout_log_probs"], batch["advantages"], _response_masks(batch, log_probs), strict=True
+ ):
+ ratio = torch.exp(log_prob - _as_tensor_like(sampling_log_prob, log_prob))
+ advantages = _as_tensor_like(advantage, log_prob)
+ objective = torch.minimum(ratio * advantages, torch.clamp(ratio, clip_low, clip_high) * advantages)
+ per_datum_losses.append(-(objective * mask).sum())
+ return _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses)
+
+
+def cispo_loss_function(
+ args: Namespace,
+ batch: RolloutBatch,
+ logits: torch.Tensor,
+ sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+ config = batch.get("loss_fn_config") or {}
+ clip_low = config.get("clip_low_threshold", CISPO_DEFAULTS["clip_low_threshold"])
+ clip_high = config.get("clip_high_threshold", CISPO_DEFAULTS["clip_high_threshold"])
+ log_probs = _target_logprobs(args, batch, logits)
+ per_datum_losses = []
+ for log_prob, sampling_log_prob, advantage, mask in zip(
+ log_probs, batch["rollout_log_probs"], batch["advantages"], _response_masks(batch, log_probs), strict=True
+ ):
+ ratio = torch.exp(log_prob - _as_tensor_like(sampling_log_prob, log_prob))
+ coefficient = torch.clamp(ratio, clip_low, clip_high).detach()
+ per_datum_losses.append(-(coefficient * log_prob * _as_tensor_like(advantage, log_prob) * mask).sum())
+ return _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses)
+
+
+def dro_loss_function(
+ args: Namespace,
+ batch: RolloutBatch,
+ logits: torch.Tensor,
+ sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+ config = batch.get("loss_fn_config") or {}
+ beta = config.get("beta", DRO_DEFAULTS["beta"])
+ log_probs = _target_logprobs(args, batch, logits)
+ per_datum_losses = []
+ for log_prob, sampling_log_prob, advantage, mask in zip(
+ log_probs, batch["rollout_log_probs"], batch["advantages"], _response_masks(batch, log_probs), strict=True
+ ):
+ divergence = log_prob - _as_tensor_like(sampling_log_prob, log_prob)
+ objective = log_prob * _as_tensor_like(advantage, log_prob) - 0.5 * beta * divergence**2
+ per_datum_losses.append(-(objective * mask).sum())
+ return _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses)
+
+
+TINKER_LOSS_FUNCTIONS = {
+ "cross_entropy": cross_entropy_loss_function,
+ "importance_sampling": importance_sampling_loss_function,
+ "ppo": ppo_loss_function,
+ "cispo": cispo_loss_function,
+ "dro": dro_loss_function,
+}
diff --git a/miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py b/miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py
index c9eb03bc3f..5b52b37f52 100644
--- a/miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py
+++ b/miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py
@@ -30,6 +30,10 @@ class WeightUpdatePlacement:
gather_tp: bool = True
gather_ep: bool = True
+ @property
+ def is_full_gather(self) -> bool:
+ return self.gather_pp and self.gather_tp and self.gather_ep
+
def resolve_placement(required: WeightUpdatePlacement, forced: WeightUpdatePlacement | None) -> WeightUpdatePlacement:
"""Join of the protocol's required placement and the iterator's forced one:
@@ -87,11 +91,14 @@ def iter_hf_weights(
whose tensors join the stream under ``{lora_name}:{hf_key}`` names.
``materialize=False`` joins every collective but yields nothing.
"""
+
+ def prefixed_units(lora_name, adapter):
+ for unit in self._iter_hf_adapter_units(adapter, materialize=materialize):
+ yield [(f"{lora_name}:{name}", tensor) for name, tensor in unit]
+
hf_param_units = self._iter_hf_param_units(weights, materialize=materialize) if include_base else iter(())
for lora_name, adapter in adapters:
- hf_param_units = itertools.chain(
- hf_param_units, self._iter_hf_adapter_units(lora_name, adapter, materialize=materialize)
- )
+ hf_param_units = itertools.chain(hf_param_units, prefixed_units(lora_name, adapter))
atomic_update_groups = self._hf_atomic_update_groups() if include_base and materialize else []
hf_param_units = assemble_atomic_update_groups(hf_param_units, atomic_update_groups)
yield from pack_units_by_size(hf_param_units, self.args.update_weight_buffer_size)
@@ -112,11 +119,18 @@ def _hf_atomic_update_groups(self) -> list[AtomicUpdateGroup]:
"""Backend hook: HF-namespace atomic groups for this model. Default none."""
return []
+ def materialize_adapter(self, adapter, *, materialize: bool = True) -> dict[str, torch.Tensor]:
+ """One adapter as ``{hf_key: tensor}``, skipping the transport bucketing.
+ Collective: ``materialize=False`` joins the gathers, returns {}."""
+ return {
+ name: tensor
+ for unit in self._iter_hf_adapter_units(adapter, materialize=materialize)
+ for name, tensor in unit
+ }
+
@abstractmethod
- def _iter_hf_adapter_units(
- self, lora_name: str, adapter, *, materialize: bool
- ) -> Iterator[list[tuple[str, torch.Tensor]]]:
+ def _iter_hf_adapter_units(self, adapter, *, materialize: bool) -> Iterator[list[tuple[str, torch.Tensor]]]:
"""Backend hook: this rank's slice of the adapter per ``self.placement``,
- one unit per parameter, names ``{lora_name}:{hf_key}``, rank-trimmed.
- Collectives must run lockstep on every rank; ``materialize=False`` joins
- them but yields nothing."""
+ one unit per parameter, bare ``hf_key`` names, rank-trimmed. Collectives
+ must run lockstep on every rank; ``materialize=False`` joins them but
+ yields nothing."""
diff --git a/miles/backends/training_utils/weight_update/snapshot_publisher.py b/miles/backends/training_utils/weight_update/snapshot_publisher.py
new file mode 100644
index 0000000000..0669990e1c
--- /dev/null
+++ b/miles/backends/training_utils/weight_update/snapshot_publisher.py
@@ -0,0 +1,35 @@
+import json
+from pathlib import Path
+
+import safetensors.torch
+import torch
+import torch.distributed as dist
+
+from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+from miles.backends.training_utils.weight_update.hf_weight_iterator import HfWeightIteratorBase
+from miles.utils.multi_lora import AdapterSpec
+
+
+class WeightPublisher:
+ def __init__(self, iterator: HfWeightIteratorBase, adapter_config: dict) -> None:
+ assert iterator.placement.is_full_gather, "publishing requires the full adapter on rank 0"
+ self._iterator = iterator
+ self._adapter_config = adapter_config
+
+ @torch.no_grad()
+ def publish_adapter(self, adapter: AdapterSpec, path: str, metadata: dict | None = None) -> None:
+ is_writer = dist.get_rank() == 0
+
+ def write_shards(tmp_dir: Path):
+ tensors = {
+ name: tensor.detach().contiguous().cpu()
+ for name, tensor in self._iterator.materialize_adapter(adapter, materialize=is_writer).items()
+ }
+ data = safetensors.torch.save(tensors) if is_writer else None
+
+ if is_writer:
+ config = self._adapter_config | {"r": adapter.rank, "lora_alpha": adapter.alpha}
+ (tmp_dir / "adapter_config.json").write_text(json.dumps(config))
+ (tmp_dir / "adapter_model.safetensors").write_bytes(data)
+
+ write_checkpoint_dir(path, write_shards, metadata=metadata, overwrite=False)
diff --git a/miles/backends/training_utils/weight_update/updater.py b/miles/backends/training_utils/weight_update/updater.py
index ed819becdb..096fb246fb 100644
--- a/miles/backends/training_utils/weight_update/updater.py
+++ b/miles/backends/training_utils/weight_update/updater.py
@@ -29,7 +29,6 @@
from miles.backends.training_utils.weight_update.utils import record_lora_checksums
from miles.utils.distributed_utils import get_gloo_group
from miles.utils.lora import LORA_ADAPTER_NAME
-from miles.utils.multi_lora import is_multi_lora_enabled, slot_lora_name
from miles.utils.timer import timer
logger = logging.getLogger(__name__)
@@ -70,8 +69,6 @@ def __init__(
assert lora_sync_config is not None
self._lora_sync_config = lora_sync_config
self._registered_adapters: set[str] = set()
- # Set by the actor before each update_weights call (loaded map at reconcile).
- self.multi_lora_adapters = None
def connect_rollout_engines(
self,
@@ -152,10 +149,6 @@ def _get_updated_adapters(self) -> list[tuple[str, object]]:
identical on every rank so the iterator's collectives align."""
if not self.is_lora:
return []
- if is_multi_lora_enabled(self.args):
- adapters = self.multi_lora_adapters
- assert adapters is not None, "actor must set multi_lora_adapters before update_weights"
- return [(slot_lora_name(adapters[name].slot), adapters[name]) for name in sorted(adapters)]
return [(LORA_ADAPTER_NAME, None)]
def _register_new_lora_adapters(self, rollout_engines, adapters: list[tuple[str, object]]) -> None:
@@ -166,6 +159,6 @@ def _register_new_lora_adapters(self, rollout_engines, adapters: list[tuple[str,
continue
config = self._lora_sync_config
if adapter is not None:
- config = config | {"r": adapter.config.rank, "lora_alpha": adapter.config.alpha}
+ config = config | {"r": adapter.rank, "lora_alpha": adapter.alpha}
register_lora_adapter(rollout_engines, lora_name=lora_name, lora_config=config)
self._registered_adapters.add(lora_name)
diff --git a/miles/backends/training_utils/weight_update/utils.py b/miles/backends/training_utils/weight_update/utils.py
index 569beb2a42..6f72120e1c 100644
--- a/miles/backends/training_utils/weight_update/utils.py
+++ b/miles/backends/training_utils/weight_update/utils.py
@@ -20,7 +20,7 @@ def get_data_replica_rank_and_size(parallel_state: ParallelState, placement: Wei
def record_lora_checksums(bucket, checksums) -> None:
- """Accumulate the sha256 manifest the engines verify at end_weight_update."""
+ """Accumulate the sha256 checksums the engines verify at end_weight_update."""
for name, tensor in bucket:
if ":" not in name:
continue
diff --git a/miles/ray/multi_lora/backend.py b/miles/ray/multi_lora/backend.py
deleted file mode 100644
index bee97443c1..0000000000
--- a/miles/ray/multi_lora/backend.py
+++ /dev/null
@@ -1,189 +0,0 @@
-"""Multi-LoRA backend: the registry plus engine-facing aborts, shared by the
-controller Ray actor and the HTTP server. Subclass via
-``--multi-lora-backend-path``."""
-
-import asyncio
-import logging
-from dataclasses import replace
-from pathlib import Path
-from typing import Any
-
-import httpx
-
-from miles.ray.multi_lora.registry import AdapterRegistry, AdapterState
-from miles.utils.adapter_config import AdapterRunConfig
-from miles.utils.http_utils import router_worker_base_urls
-from miles.utils.multi_lora import RID_SEPARATOR, min_groups_per_dp_split
-
-logger = logging.getLogger(__name__)
-
-
-class MultiLoRABackend:
- """Registry + engine-facing aborts, shared by the Ray actor and HTTP server.
- Subclass via --multi-lora-backend-path."""
-
- def __init__(self, args: Any, router_url: str) -> None:
- self.args = args
- self.registry = AdapterRegistry(args.multi_lora_n_adapters)
- self.router_url = router_url.rstrip("/")
- self.client: httpx.AsyncClient | None = None
-
- async def init(self) -> None:
- self.client = httpx.AsyncClient(timeout=httpx.Timeout(30.0))
-
- async def close(self) -> None:
- if self.client is not None:
- await self.client.aclose()
- self.client = None
-
- async def validate_adapter(self, name: str, config: Any) -> None:
- """Override to reject adapter registrations (raise ValueError)."""
-
- def resolve_adapter_config(self, name: str, config: Any) -> Any:
- """Resolve optional adapter-local values against process-wide defaults
- and validate the batch shape against the trainer's DP layout.
-
- All batch-shape constraints are enforced here, at registration, so a
- bad config fails immediately instead of crashing an arbitrary later
- train batch.
- """
- if config is None or not isinstance(config, AdapterRunConfig):
- return config
-
- rank = config.rank if config.rank is not None else getattr(self.args, "lora_rank", 1)
- alpha = config.alpha if config.alpha is not None else getattr(self.args, "lora_alpha", rank)
- rollout_batch_size = (
- config.rollout_batch_size
- if config.rollout_batch_size is not None
- else getattr(self.args, "rollout_batch_size", None)
- )
- n_samples_per_prompt = (
- config.n_samples_per_prompt
- if config.n_samples_per_prompt is not None
- else getattr(self.args, "n_samples_per_prompt", 1)
- )
-
- if type(rank) is not int or rank <= 0:
- raise ValueError(f"Adapter '{name}' rank must be a positive integer")
- if rank > getattr(self.args, "lora_rank", rank):
- raise ValueError(f"Adapter '{name}' rank {rank} exceeds the allocated maximum rank {self.args.lora_rank}")
- if alpha is None or alpha <= 0:
- raise ValueError(f"Adapter '{name}' must have a positive alpha")
- if type(rollout_batch_size) is not int or rollout_batch_size <= 0:
- raise ValueError(f"Adapter '{name}' rollout_batch_size must be a positive integer (prompt groups)")
- if type(n_samples_per_prompt) is not int or n_samples_per_prompt <= 0:
- raise ValueError(f"Adapter '{name}' n_samples_per_prompt must be a positive integer")
- if config.num_step is not None and (type(config.num_step) is not int or config.num_step <= 0):
- raise ValueError(f"Adapter '{name}' num_step must be a positive integer")
- if config.num_epoch is not None and (type(config.num_epoch) is not int or config.num_epoch <= 0):
- raise ValueError(f"Adapter '{name}' num_epoch must be a positive integer")
- if config.num_step is not None and config.num_epoch is not None:
- logger.warning(f"Adapter '{name}' sets both num_step and num_epoch; num_step takes precedence")
-
- # A bad data path or unresolvable reward config does not fail at this
- # API otherwise: the data path kills the shared rollout producer thread
- # and an empty reward config burns every generated sample, either way
- # stalling ALL adapters behind a misleading empty-batch timeout.
- if not Path(config.data).expanduser().exists():
- raise ValueError(
- f"Adapter '{name}' data path '{config.data}' does not exist "
- "(checked from the controller process, which runs on the head node with the rollout data source)"
- )
- if (
- config.custom_rm_path is None
- and not (config.rm_type or "").strip()
- and getattr(self.args, "custom_rm_path", None) is None
- and not (getattr(self.args, "rm_type", None) or "").strip()
- ):
- raise ValueError(
- f"Adapter '{name}' has no reward config: set rm_type or custom_rm_path in the adapter "
- "config, or launch with --rm-type / --custom-rm-path"
- )
-
- adapter_global_batch_size = rollout_batch_size * n_samples_per_prompt
- if (max_batch := getattr(self.args, "multi_lora_max_adapter_global_batch_size", None)) is not None:
- if adapter_global_batch_size > max_batch:
- raise ValueError(
- f"Adapter '{name}' consumes {adapter_global_batch_size} samples per step "
- f"(rollout_batch_size {rollout_batch_size} x n_samples_per_prompt {n_samples_per_prompt}), "
- f"exceeding --multi-lora-max-adapter-global-batch-size {max_batch}"
- )
- if (dp_size := getattr(self.args, "multi_lora_dp_size", None)) is not None:
- try:
- group_multiple = min_groups_per_dp_split(n_samples_per_prompt, dp_size)
- except ValueError as e:
- raise ValueError(f"Adapter '{name}': {e}") from None
- if rollout_batch_size % group_multiple != 0:
- raise ValueError(
- f"Adapter '{name}' rollout_batch_size {rollout_batch_size} must be a multiple of "
- f"its min_groups_per_dp_split ({group_multiple} at dp_size={dp_size}), so the "
- f"adapter batch can complete from evenly-splitting takes"
- )
-
- save = Path(config.save) if config.save is not None else None
- if save is None:
- if getattr(self.args, "save", None) is None:
- raise ValueError(f"Adapter '{name}' has no save dir: set 'save' in the adapter config or pass --save")
- save = Path(self.args.save) / "adapters" / name
-
- return replace(
- config,
- rank=rank,
- alpha=alpha,
- rollout_batch_size=rollout_batch_size,
- n_samples_per_prompt=n_samples_per_prompt,
- save=save,
- )
-
- async def register(self, name: str, config: Any) -> dict:
- config = self.resolve_adapter_config(name, config)
- await self.validate_adapter(name, config)
- result = self.registry.register(name, config)
- resolved = getattr(config, "save", None)
- if resolved is not None:
- logger.info(f"Adapter '{name}' registered (slot {result['slot']}), checkpoints -> {resolved}")
- return result
-
- async def deregister(self, name: str) -> None:
- self.registry.deregister(name)
-
- async def retire_adapters(self) -> list[str]:
- names = self.registry.retire_adapters()
- for name in names:
- await self.abort_adapter_requests(name)
- return names
-
- async def free_slot(self, name: str) -> int:
- """Free the adapter's slot after one final abort round: requests can survive the
- ``retire_adapters`` abort (e.g. multi-turn groups), and must not leak to the slot's next tenant."""
- record = self.registry.records.get(name)
- if record is not None and record.state is AdapterState.CLEANUP:
- await self.abort_adapter_requests(name)
- return self.registry.free_slot(name)
-
- async def worker_urls(self) -> list[str]:
- assert self.client is not None
- for endpoint, extract in (
- ("/list_workers", lambda body: body["urls"]),
- ("/workers", lambda body: [worker["url"] for worker in body["workers"]]),
- ):
- try:
- resp = await self.client.get(f"{self.router_url}{endpoint}")
- if resp.status_code == 200:
- return router_worker_base_urls(extract(resp.json()))
- except Exception:
- continue
- return []
-
- async def abort_adapter_requests(self, adapter_name: str) -> None:
- prefix = f"{adapter_name}{RID_SEPARATOR}"
- urls = await self.worker_urls()
- if not urls:
- logger.warning(f"Abort for adapter '{adapter_name}': no workers discovered at {self.router_url}")
- return
- results = await asyncio.gather(
- *(self.client.post(f"{url}/abort_request", json={"rid": prefix, "prefix": True}) for url in urls),
- return_exceptions=True,
- )
- if failures := sum(isinstance(r, Exception) for r in results):
- logger.warning(f"Abort for adapter '{adapter_name}': {failures}/{len(results)} posts failed")
diff --git a/miles/ray/multi_lora/controller.py b/miles/ray/multi_lora/controller.py
deleted file mode 100644
index e44ef8b3b2..0000000000
--- a/miles/ray/multi_lora/controller.py
+++ /dev/null
@@ -1,123 +0,0 @@
-"""Named Ray actor wrapping the multi-LoRA backend + HTTP server."""
-
-import time
-from functools import cache
-from typing import Any
-
-import ray
-
-from miles.ray.multi_lora.backend import MultiLoRABackend
-from miles.ray.multi_lora.http_server import MultiLoRAHTTPServer
-from miles.utils.adapter_config import AdapterRun
-from miles.utils.function_registry import load_function
-from miles.utils.misc import SingletonMeta, get_current_node_ip
-from miles.utils.ray_utils import compute_ray_pin_head_options
-
-CONTROLLER_NAME = "miles_multi_lora_controller"
-CONTROLLER_NAMESPACE = "miles"
-
-
-@cache
-def get_multi_lora_controller():
- return ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE)
-
-
-class AdaptersCache(metaclass=SingletonMeta):
- """TTL-cached controller snapshot; get/get_all expose the sampleable
- projection (active + retiring)."""
-
- def __init__(self, ttl_s: float = 1.0) -> None:
- self.ttl_s = ttl_s
- self.snapshot: dict = {"pending": {}, "active": {}, "retiring": {}, "cleanup": []}
- self.last_refresh: float | None = None
-
- async def get_snapshot(self) -> dict:
- now = time.monotonic()
- if self.last_refresh is None or now - self.last_refresh >= self.ttl_s:
- try:
- self.snapshot = await get_multi_lora_controller().snapshot.remote()
- self.last_refresh = now
- except Exception:
- pass
- return self.snapshot
-
- async def get_all(self) -> dict[str, "AdapterRun"]:
- snapshot = await self.get_snapshot()
- return {**snapshot["active"], **snapshot["retiring"]}
-
- async def get(self, adapter_name: str) -> "AdapterRun | None":
- return (await self.get_all()).get(adapter_name)
-
-
-def _load_subclass(path: str | None, base_cls):
- if not path:
- return base_cls
- cls = load_function(path)
- assert issubclass(cls, base_cls), f"{path} must point to a {base_cls.__name__} subclass, got {cls}"
- return cls
-
-
-@ray.remote(num_cpus=0)
-class MultiLoRAController:
- def __init__(self, *, args, router_url: str, host: str = "0.0.0.0") -> None:
- backend_cls = _load_subclass(getattr(args, "multi_lora_backend_path", None), MultiLoRABackend)
- server_cls = _load_subclass(getattr(args, "multi_lora_http_server_path", None), MultiLoRAHTTPServer)
- self.backend = backend_cls(args, router_url)
- self.server = server_cls(self.backend, host, api_port=getattr(args, "multi_lora_api_port", 0))
-
- async def start(self) -> int:
- await self.backend.init()
- await self.server.start()
- return self.server.actual_api_port
-
- async def stop(self) -> None:
- await self.server.stop()
- await self.backend.close()
-
- async def register_adapter(self, name: str, config: Any) -> dict:
- return await self.backend.register(name, config)
-
- async def deregister_adapter(self, name: str) -> None:
- await self.backend.deregister(name)
-
- async def retire_adapters(self) -> list[str]:
- return await self.backend.retire_adapters()
-
- async def free_slot(self, name: str) -> int:
- return await self.backend.free_slot(name)
-
- def record_weight_update(self, names: list[str]) -> None:
- self.backend.registry.record_weight_update(names)
-
- def record_batch_adapters(self, rollout_id: int, groups: dict[str, int], step_names: list[str]) -> None:
- self.backend.registry.record_batch_adapters(rollout_id, groups, step_names)
-
- def mark_batch_trained(self, rollout_id: int) -> list[str]:
- return self.backend.registry.mark_batch_trained(rollout_id)
-
- def resolve_num_step(self, name: str, dataset_rows: int) -> None:
- self.backend.registry.resolve_num_step(name, dataset_rows)
-
- def set_adapter_step(self, name: str, step: int) -> None:
- self.backend.registry.set_step(name, step)
-
- def adapter_step(self, name: str) -> int:
- return self.backend.registry.step_count(name)
-
- def snapshot(self) -> dict:
- return self.backend.registry.snapshot()
-
- def http_host(self) -> str:
- return get_current_node_ip()
-
- def api_port(self) -> int:
- return self.server.actual_api_port
-
-
-def create_multilora_controller(args, router_url: str, host: str = "0.0.0.0"):
- # Pinned to the head node so the API sits at a port-forwardable address.
- return MultiLoRAController.options(
- name=CONTROLLER_NAME,
- namespace=CONTROLLER_NAMESPACE,
- **compute_ray_pin_head_options(),
- ).remote(args=args, router_url=router_url, host=host)
diff --git a/miles/ray/multi_lora/http_server.py b/miles/ray/multi_lora/http_server.py
deleted file mode 100644
index b209142e1e..0000000000
--- a/miles/ray/multi_lora/http_server.py
+++ /dev/null
@@ -1,129 +0,0 @@
-"""Multi-LoRA control-plane HTTP API over a MultiLoRABackend.
-
-Subclass via ``--multi-lora-http-server-path`` (override add_routes /
-create_app)."""
-
-import asyncio
-from dataclasses import asdict
-from pathlib import Path
-
-import uvicorn
-from fastapi import FastAPI, HTTPException, Query, Request
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-
-from miles.ray.multi_lora.registry import AdapterState
-from miles.utils.adapter_config import AdapterRunConfig, parse_adapter_run_yaml
-
-
-class RegisterAdapterRequest(BaseModel):
- """Exactly one of ``config`` (inline) or ``yaml_path`` must be set."""
-
- name: str
- config: AdapterRunConfig | None = None
- yaml_path: str | None = None
-
-
-_NAMES_QUERY = Query(default_factory=list)
-
-
-class MultiLoRAHTTPServer:
- """Control-plane API over a MultiLoRABackend. Subclass via
- --multi-lora-http-server-path (add_routes / create_app)."""
-
- def __init__(self, backend, host="127.0.0.1", api_port=0):
- self.backend = backend
- self.host = host
- self.api_port = api_port
- self.api_server: uvicorn.Server | None = None
- self.api_task: asyncio.Task | None = None
-
- @property
- def actual_api_port(self) -> int:
- if self.api_server is not None and self.api_server.started:
- return self.api_server.servers[0].sockets[0].getsockname()[1]
- return self.api_port
-
- def create_app(self) -> FastAPI:
- app = FastAPI(title="Miles Multi-LoRA Controller")
-
- @app.exception_handler(ValueError)
- async def value_error_handler(request: Request, exc: ValueError):
- return JSONResponse({"detail": str(exc)}, status_code=400)
-
- @app.exception_handler(RuntimeError)
- async def runtime_error_handler(request: Request, exc: RuntimeError):
- status = 409 if "No free adapter slots" in str(exc) else 500
- return JSONResponse({"detail": str(exc)}, status_code=status)
-
- return app
-
- def add_routes(self, app: FastAPI) -> None:
- app.get("/health")(self.health)
- app.get("/adapter_runs")(self.list_adapters)
- app.get("/adapter_runs/state")(self.adapter_states) # before /adapter_runs/{name}
- app.get("/adapter_runs/{name}")(self.get_adapter)
- app.post("/adapter_runs")(self.register_adapter)
- app.delete("/adapter_runs/{name}")(self.deregister_adapter)
-
- async def start(self) -> None:
- app = self.create_app()
- self.add_routes(app)
- config = uvicorn.Config(app, host=self.host, port=self.api_port, log_level="warning", access_log=False)
- self.api_server = uvicorn.Server(config)
- self.api_task = asyncio.create_task(self.api_server.serve())
- while not self.api_server.started:
- if self.api_task.done():
- self.api_task.result()
- raise RuntimeError("uvicorn exited before startup completed")
- await asyncio.sleep(0.01)
-
- async def stop(self) -> None:
- if self.api_server is not None:
- self.api_server.should_exit = True
- await self.api_task
- self.api_server = self.api_task = None
-
- async def health(self) -> dict:
- return {"status": "healthy"}
-
- def adapter_statuses(self) -> list[dict]:
- registry = self.backend.registry
- statuses = []
- for record in registry.records.values():
- flat = asdict(registry.view(record))
- flat |= flat.pop("config")
- flat["save"] = str(flat["save"])
- flat["state"] = record.state
- if record.state is AdapterState.COMPLETED:
- flat["version"] = None
- statuses.append(flat)
- return statuses
-
- async def list_adapters(self) -> dict:
- return {"adapters": self.adapter_statuses()}
-
- async def adapter_states(self, names: list[str] = _NAMES_QUERY) -> dict:
- return {"states": {name: self.backend.registry.adapter_state(name) for name in names}}
-
- async def get_adapter(self, name: str) -> dict:
- for status in self.adapter_statuses():
- if status["name"] == name:
- return status
- raise HTTPException(status_code=404, detail=f"Adapter '{name}' not registered")
-
- async def register_adapter(self, request: RegisterAdapterRequest) -> dict:
- if (request.config is None) == (request.yaml_path is None):
- raise HTTPException(status_code=400, detail="Exactly one of 'config' or 'yaml_path' must be set")
- if request.yaml_path is not None:
- config = parse_adapter_run_yaml(Path(request.yaml_path))
- else:
- config = request.config
- return await self.backend.register(request.name, config)
-
- async def deregister_adapter(self, name: str) -> dict:
- state = self.backend.registry.adapter_state(name)
- if state is None:
- raise HTTPException(status_code=404, detail=f"Adapter '{name}' not registered")
- await self.backend.deregister(name)
- return {"status": "ok", "name": name}
diff --git a/miles/ray/multi_lora/registry.py b/miles/ray/multi_lora/registry.py
deleted file mode 100644
index 4c8723c29d..0000000000
--- a/miles/ray/multi_lora/registry.py
+++ /dev/null
@@ -1,252 +0,0 @@
-"""Multi-LoRA adapter registry: the controller-owned lifecycle state machine.
-
-One record per adapter name, walking PENDING -> ACTIVE -> RETIRING -> CLEANUP
--> COMPLETED. Slots are reused across registrations but ``slot_versions``
-never reset, so a (slot, version) pair never recurs.
-"""
-
-import logging
-import re
-import uuid
-from dataclasses import dataclass, field, replace
-from enum import Enum
-from pathlib import Path
-from typing import Any
-
-from miles.utils.adapter_config import AdapterRun, AdapterRunConfig
-
-logger = logging.getLogger(__name__)
-
-VALID_ADAPTER_NAME = re.compile(r"^[A-Za-z0-9._-]+$")
-
-
-class AdapterState(str, Enum):
- PENDING = "PENDING"
- ACTIVE = "ACTIVE"
- RETIRING = "RETIRING"
- CLEANUP = "CLEANUP"
- COMPLETED = "COMPLETED"
-
-
-# States that hold a slot.
-LIVE_STATES = (
- AdapterState.PENDING,
- AdapterState.ACTIVE,
- AdapterState.RETIRING,
- AdapterState.CLEANUP,
-)
-
-
-@dataclass
-class AdapterRecord:
- name: str
- slot: int
- config: Any
- step: int = 0
- # Baseline step for relative num_step stopping (supports checkpoint resume).
- start_step: int = 0
- # Committed prompt groups accumulated toward the current optimizer step.
- # Only advanced by mark_batch_trained (after a successful train call).
- accumulated_groups: int = 0
- state: AdapterState = AdapterState.PENDING
- # Unique per registration: a re-registered name is a new tenant, and
- # rollout-side state stamped by the previous tenant must not carry over.
- registration_id: str = field(default_factory=lambda: uuid.uuid4().hex)
-
-
-MAX_BATCH_RECORDS = 16
-MAX_COMPLETED_RECORDS = 1024
-
-
-class AdapterRegistry:
- """One record per name; ``slot_versions`` never reset, so (slot, version)
- never recurs across slot reuse."""
-
- def __init__(self, max_adapters: int) -> None:
- self.max_adapters = max_adapters
- self.free_slots: set[int] = set(range(max_adapters))
- self.slot_versions: list[int] = [0] * max_adapters
- self.records: dict[str, AdapterRecord] = {}
- self.batch_records: dict[int, dict] = {}
-
- def in_state(self, *states: AdapterState) -> dict[str, AdapterRecord]:
- return {name: r for name, r in self.records.items() if r.state in states}
-
- def find(self, name: str) -> AdapterRecord | None:
- record = self.records.get(name)
- return record if record is not None and record.state in LIVE_STATES else None
-
- def is_active(self, name: str) -> bool:
- record = self.records.get(name)
- return record is not None and record.state in (AdapterState.ACTIVE, AdapterState.RETIRING)
-
- def register(self, name: str, config: Any) -> dict:
- if not VALID_ADAPTER_NAME.match(name) or name in (".", ".."):
- raise ValueError(f"Adapter name '{name}' is invalid: use only letters, digits, '.', '_' and '-'")
- if (existing := self.records.get(name)) is not None:
- if existing.state in (AdapterState.PENDING, AdapterState.ACTIVE):
- raise ValueError(f"Adapter '{name}' already registered")
- if existing.state in (AdapterState.RETIRING, AdapterState.CLEANUP):
- raise ValueError(f"Adapter '{name}' is still cleaning up; retry shortly")
- if (save_dir := getattr(config, "save", None)) is not None:
- for record in self.in_state(*LIVE_STATES).values():
- other_save = getattr(record.config, "save", None)
- if other_save is not None and Path(other_save).resolve() == Path(save_dir).resolve():
- raise ValueError(
- f"Adapter '{name}' save dir '{save_dir}' is already used by adapter '{record.name}'"
- )
- if not self.free_slots:
- raise RuntimeError(f"No free adapter slots (max {self.max_adapters})")
- slot = min(self.free_slots)
- self.free_slots.remove(slot)
- self.records.pop(name, None)
- self.records[name] = AdapterRecord(name=name, slot=slot, config=config)
- return {"name": name, "slot": slot}
-
- def deregister(self, name: str) -> None:
- record = self.records.get(name)
- if record is not None and record.state in (AdapterState.PENDING, AdapterState.ACTIVE):
- record.state = AdapterState.RETIRING
-
- def retire_adapters(self) -> list[str]:
- retired = sorted(self.in_state(AdapterState.RETIRING))
- for name in retired:
- self.records[name].state = AdapterState.CLEANUP
- return retired
-
- def free_slot(self, name: str) -> int:
- record = self.records.get(name)
- if record is None or record.state is not AdapterState.CLEANUP:
- return -1
- self.free_slots.add(record.slot)
- record.state = AdapterState.COMPLETED
- self.records[name] = self.records.pop(name)
- completed = self.in_state(AdapterState.COMPLETED)
- for oldest in list(completed)[: len(completed) - MAX_COMPLETED_RECORDS]:
- self.records.pop(oldest)
- return record.slot
-
- def adapter_state(self, name: str) -> AdapterState | None:
- record = self.records.get(name)
- if record is None:
- return None
- if record.state is AdapterState.COMPLETED:
- self.records[name] = self.records.pop(name)
- return record.state
-
- def record_weight_update(self, names: list[str]) -> None:
- """A weight push landed: bump slot versions, promote PENDING to ACTIVE."""
- for name in names:
- record = self.find(name)
- if record is None:
- continue
- self.slot_versions[record.slot] += 1
- if record.state is AdapterState.PENDING:
- record.state = AdapterState.ACTIVE
-
- def record_batch_adapters(self, rollout_id: int, groups: dict[str, int], step_names: list[str]) -> None:
- """Register what a train batch contains before it trains.
-
- ``groups`` maps adapter name -> prompt groups riding in this batch;
- ``step_names`` lists adapters whose adapter batch completes with
- this batch (decided by the collection loop, which caps per-adapter
- contributions at the adapter's remaining groups).
- """
- unknown = set(step_names) - set(groups)
- assert not unknown, f"step adapters {sorted(unknown)} not present in batch groups"
- self.batch_records[rollout_id] = {"groups": dict(groups), "step_names": list(step_names)}
- while len(self.batch_records) > MAX_BATCH_RECORDS:
- self.batch_records.pop(next(iter(self.batch_records)))
-
- def mark_batch_trained(self, rollout_id: int) -> list[str]:
- """Bank the batch's trained groups and fire steps; returns adapters that stepped. Only place
- accumulation/step state advances, so a failed/retried train call leaves the registry untouched."""
- record_entry = self.batch_records.pop(rollout_id, None)
- if record_entry is None:
- return []
- stepped = []
- reached_num_step = []
- for name, n_groups in record_entry["groups"].items():
- record = self.records.get(name)
- if record is None or record.state not in (
- AdapterState.ACTIVE,
- AdapterState.RETIRING,
- AdapterState.CLEANUP,
- ):
- continue
- record.accumulated_groups += n_groups
- if name in record_entry["step_names"]:
- target = record.config.rollout_batch_size
- if record.accumulated_groups != target:
- logger.warning(
- f"Adapter '{name}' stepped with accumulated_groups={record.accumulated_groups} "
- f"!= rollout_batch_size={target}; adapter batch accounting drifted"
- )
- record.step += 1
- record.accumulated_groups = 0
- stepped.append(name)
- if (
- getattr(record.config, "num_step", None) is not None
- and record.state is AdapterState.ACTIVE
- and (record.step - record.start_step) >= record.config.num_step
- ):
- reached_num_step.append(name)
- for name in reached_num_step:
- logger.info(
- f"Adapter '{name}' reached num_step={self.records[name].config.num_step} "
- f"(start_step={self.records[name].start_step}, step={self.records[name].step}), deregistering"
- )
- self.deregister(name)
- return stepped
-
- def resolve_num_step(self, name: str, dataset_rows: int) -> None:
- """Derive num_step from num_epoch once the data source knows the
- post-filter dataset length. No-op when num_step was set explicitly."""
- record = self.find(name)
- if record is None or not isinstance(record.config, AdapterRunConfig):
- return
- if record.config.num_step is not None:
- return
- num_epoch = record.config.num_epoch or 1
- num_step = max(1, num_epoch * dataset_rows // record.config.rollout_batch_size)
- record.config = replace(record.config, num_step=num_step)
- logger.info(f"Adapter '{name}': num_epoch={num_epoch} x {dataset_rows} rows -> num_step={num_step}")
-
- def set_step(self, name: str, step: int) -> None:
- if (record := self.find(name)) is not None:
- record.step = step
- record.start_step = step
-
- def step_count(self, name: str) -> int:
- record = self.find(name)
- return record.step if record is not None else 0
-
- def view(self, record: AdapterRecord) -> AdapterRun:
- return AdapterRun(
- name=record.name,
- config=record.config,
- slot=record.slot,
- version=self.slot_versions[record.slot],
- step=record.step,
- accumulated_groups=record.accumulated_groups,
- registration_id=record.registration_id,
- )
-
- def active_adapters(self) -> dict[str, AdapterRun]:
- """Sampleable view: RETIRING keeps serving until retired."""
- return {
- name: self.view(record)
- for name, record in self.in_state(AdapterState.ACTIVE, AdapterState.RETIRING).items()
- }
-
- def snapshot(self) -> dict:
- def views(state: AdapterState) -> dict[str, AdapterRun]:
- return {name: self.view(record) for name, record in self.in_state(state).items()}
-
- return {
- "pending": views(AdapterState.PENDING),
- "active": views(AdapterState.ACTIVE),
- "retiring": views(AdapterState.RETIRING),
- "cleanup": list(self.in_state(AdapterState.CLEANUP)),
- "completed": list(self.in_state(AdapterState.COMPLETED)),
- }
diff --git a/miles/ray/rollout/metrics.py b/miles/ray/rollout/metrics.py
index d32f1390a1..071172c0bf 100644
--- a/miles/ray/rollout/metrics.py
+++ b/miles/ray/rollout/metrics.py
@@ -15,7 +15,7 @@
has_repetition,
)
from miles.utils.tracking_utils import tracking
-from miles.utils.types import AdapterRef, Sample
+from miles.utils.types import Sample
logger = logging.getLogger(__name__)
@@ -136,13 +136,12 @@ def _compute_metrics_from_samples(args, samples):
return log_dict
-def _get_rollout_key(sample: Sample, position: int) -> tuple[AdapterRef | None, str, int | None, int]:
- # Adapter identity scopes the IDs because each Multi-LoRA data source numbers them independently.
+def _get_rollout_key(sample: Sample, position: int) -> tuple[str, int | None, int]:
if sample.rollout_id is not None:
- return (sample.adapter, "rollout", sample.group_index, sample.rollout_id)
+ return ("rollout", sample.group_index, sample.rollout_id)
if sample.index is not None:
- return (sample.adapter, "sample", sample.group_index, sample.index)
- return (sample.adapter, "position", sample.group_index, position)
+ return ("sample", sample.group_index, sample.index)
+ return ("position", sample.group_index, position)
def _compute_episode_response_length_metrics(samples: list[Sample]) -> dict[str, float]:
@@ -153,11 +152,8 @@ def _compute_episode_response_length_metrics(samples: list[Sample]) -> dict[str,
computing batch-level statistics. Effective lengths count only trainable
tokens; total lengths count both masked and unmasked tokens in every sample.
"""
- if any(sample.adapter is not None for sample in samples):
- return {}
-
- effective_lengths_by_rollout: dict[tuple[AdapterRef | None, str, int | None, int], int] = {}
- total_lengths_by_rollout: dict[tuple[AdapterRef | None, str, int | None, int], int] = {}
+ effective_lengths_by_rollout: dict[tuple[str, int | None, int], int] = {}
+ total_lengths_by_rollout: dict[tuple[str, int | None, int], int] = {}
for position, sample in enumerate(samples):
rollout_key = _get_rollout_key(sample, position)
effective_response_length = 0 if sample.remove_sample else sample.effective_response_length
@@ -183,10 +179,9 @@ def _compute_training_sample_metrics(args: Any, samples: list[Sample]) -> dict[s
Session compaction can turn one rollout into several training samples. The
sample count includes every resulting row, while the reward first averages
sibling rows that share a rollout ID so long rollouts do not receive more
- metric weight merely because they produced more samples. Adapter identity
- scopes these IDs because each Multi-LoRA data source numbers them independently.
+ metric weight merely because they produced more samples.
"""
- rewards_by_rollout: dict[tuple[AdapterRef | None, str, int | None, int], list[float]] = {}
+ rewards_by_rollout: dict[tuple[str, int | None, int], list[float]] = {}
use_metadata_reward = bool(samples and samples[0].metadata and "raw_reward" in samples[0].metadata)
for position, sample in enumerate(samples):
rollout_key = _get_rollout_key(sample, position)
diff --git a/miles/ray/rollout/rollout_data_conversion.py b/miles/ray/rollout/rollout_data_conversion.py
index b948856dd3..74b7178b13 100644
--- a/miles/ray/rollout/rollout_data_conversion.py
+++ b/miles/ray/rollout/rollout_data_conversion.py
@@ -1,7 +1,6 @@
import itertools
import logging
-from miles.utils.multi_lora import is_multi_lora_enabled
from miles.utils.types import Sample
logger = logging.getLogger(__name__)
@@ -12,15 +11,6 @@ def postprocess_rollout_data(args, data, train_parallel_config):
validate_compact_rollout_ids(data)
- # Multi-LoRA: record group boundaries (heterogeneous per-adapter group sizes)
- # and lift the collection loop's batch-level step decision out of sample metadata,
- # both before flattening.
- if is_multi_lora_enabled(args) and isinstance(data[0], list):
- metadata["prompt_group_sizes"] = [_nested_sample_count(group) for group in data]
- head = _first_sample(data[0])
- metadata["step_slots"] = list(head.metadata.pop("step_slots", []))
- metadata["step_adapter_names"] = list(head.metadata.pop("step_adapter_names", []))
-
# flatten the data if it is a list of lists
while isinstance(data[0], list):
data = list(itertools.chain.from_iterable(data))
@@ -91,17 +81,6 @@ def _compute_dynamic_global_batch_size(args, train_parallel_config, num_samples:
dp_size = train_parallel_config["dp_size"]
original_gbs = args.global_batch_size
- if is_multi_lora_enabled(args):
- # Batches take groups in multiples of each adapter's
- # min_groups_per_dp_split, so this holds by construction; a violation
- # means a generate fn's group shape broke the invariant.
- if num_samples % dp_size != 0:
- raise ValueError(
- f"Multi-LoRA batch of {num_samples} samples is not divisible by dp_size={dp_size}; "
- "the min_groups_per_dp_split invariant was violated (variable-size generate fn output?)"
- )
- return num_samples
-
# Round down to a multiple of dp_size to ensure only one training step
dynamic_gbs = (num_samples // dp_size) * dp_size
diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py
index 5f6048514c..a647bb9d97 100644
--- a/miles/ray/rollout/train_data_conversion.py
+++ b/miles/ray/rollout/train_data_conversion.py
@@ -15,6 +15,7 @@
ROLLOUT_DATA_TENSOR_DTYPES = {
"tokens": "int32",
+ "target_tokens": "int32",
"loss_masks": "int32",
"rollout_log_probs": "float32",
"rollout_sampling_mask_ids": "int32",
@@ -151,21 +152,6 @@ def convert_samples_to_train_data(
if samples[0].teacher_log_probs is not None:
train_data["teacher_log_probs"] = [sample.teacher_log_probs for sample in samples]
- if any(sample.adapter is not None for sample in samples):
- assert all(sample.adapter is not None for sample in samples), "Cannot mix adapter and adapter-less samples"
- train_data["adapter_slots"] = [sample.adapter.slot for sample in samples]
- # Slots whose adapter batch completes with this batch: the trainer scales their
- # accumulated gradients by 1/adapter-batch-size and advances the LR schedule.
- step_slots = sorted(metadata.get("step_slots", []))
- train_data["step_slots"] = step_slots
- train_data["step_adapter_names"] = sorted(metadata.get("step_adapter_names", []))
- step_slot_set = set(step_slots)
- train_data["step_adapter_batch_sizes"] = {
- sample.adapter.slot: sample.metadata["adapter_global_batch_size"]
- for sample in samples
- if sample.adapter.slot in step_slot_set
- }
-
if (prompt_group_sizes := metadata.get("prompt_group_sizes")) is not None:
train_data["prompt_group_sizes"] = prompt_group_sizes
@@ -397,6 +383,9 @@ def _package_shards(args, data: dict[str, Any], partitions) -> list[dict[str, An
"seq_witness_ids",
"weight_versions",
"adapter_slots",
+ "loss_weights",
+ "advantages",
+ "target_tokens",
]:
if key not in data:
continue
@@ -407,10 +396,9 @@ def _package_shards(args, data: dict[str, Any], partitions) -> list[dict[str, An
"raw_reward",
"total_lengths",
"dynamic_global_batch_size",
- "step_slots",
- "step_adapter_names",
- "step_adapter_batch_sizes",
"prompt_group_sizes",
+ "loss_fn",
+ "loss_fn_config",
]:
if key not in data:
continue
diff --git a/miles/ray/specs/train.py b/miles/ray/specs/train.py
index 188da4394b..f1528aed4d 100644
--- a/miles/ray/specs/train.py
+++ b/miles/ray/specs/train.py
@@ -7,6 +7,7 @@
from miles.utils.environ import default_fp8_block_scaling_fp32_scales
from miles.utils.ft_utils.indep_dp import create_tcp_store
from miles.utils.megatron_args_utils import compute_megatron_world_size_except_dp
+from miles.utils.multi_lora import is_multi_lora_enabled
from miles.utils.workers.worker_spec import PortInfo, SchedulingSpec, ServeWorkerSpec, WorkerLaunchContext
MASTER_PORT_NAME = "master"
@@ -95,7 +96,11 @@ def _compute_spec_trainer(
num_gpu_slots_per_worker=1,
pg_name="actor",
),
- worker_class=_TRAINER_ACTOR_CLASSES[args.train_backend],
+ worker_class=(
+ "miles.backends.megatron_utils.lora.actor.MultiLoRATrainRayActor"
+ if args.train_backend == "megatron" and role == "actor" and is_multi_lora_enabled(args)
+ else _TRAINER_ACTOR_CLASSES[args.train_backend]
+ ),
ctor_kwargs=lambda ctx: dict(
args=args,
world_size=gpus_per_cell,
diff --git a/miles/ray/train/group.py b/miles/ray/train/group.py
index 7b0802850d..81eb9ac3a9 100644
--- a/miles/ray/train/group.py
+++ b/miles/ray/train/group.py
@@ -92,6 +92,9 @@ def _expected_num_cells(self) -> int:
def _cells(self) -> list[TrainerCell]:
return sorted(self._cells_by_id.values(), key=lambda cell: cell.cell_index)
+ def has_errored_cell(self) -> bool:
+ return any(cell.is_errored for cell in self._cells)
+
@property
def cell_ids(self) -> list[str]:
return [cell.cell_id for cell in self._cells]
@@ -388,8 +391,39 @@ async def offload_grad_buffer(self):
# Catch *without* retry: cells w/ exceptions are auto marked errored, and will not be used
await self._execute_all_alive_and_catch("offload_grad_buffer")
- async def reconcile_adapters(self) -> None:
- await asyncio.gather(*[cell.execute("reconcile_adapters") for cell in self._cells])
+ # ------------------------ API :: multi-LoRA slot commands ------------------------
+
+ async def _execute_slots(self, fn_name: str, **kwargs) -> list:
+ (cell,) = self._cells
+ assert cell.is_alive, "the Tinker trainer cell is unavailable"
+ return await cell.execute(fn_name, **kwargs)
+
+ async def forward_backward(self, batch_id: int, data_ref) -> list:
+ return await self._execute_slots("forward_backward", batch_id=batch_id, rollout_data_ref=data_ref)
+
+ async def optim_step(self, adam_params_by_slot: dict[int, dict]) -> list:
+ return await self._execute_slots("optim_step", adam_params_by_slot=adam_params_by_slot)
+
+ async def forward_only(self, batch_id: int, data_ref) -> list:
+ return await self._execute_slots("forward_only", batch_id=batch_id, rollout_data_ref=data_ref)
+
+ async def load_slot(
+ self, slot: int, rank: int, alpha: float, ckpt_path: str | None = None, load_optimizer: bool = True
+ ) -> list:
+ return await self._execute_slots(
+ "load_slot", slot=slot, rank=rank, alpha=alpha, ckpt_path=ckpt_path, load_optimizer=load_optimizer
+ )
+
+ async def save_slot(self, slot: int, path: str, metadata: dict | None = None) -> list:
+ return await self._execute_slots("save_slot", slot=slot, path=path, metadata=metadata)
+
+ async def export_slot(self, slot: int, rank: int, alpha: float, path: str, metadata: dict | None = None) -> list:
+ return await self._execute_slots(
+ "export_slot", slot=slot, rank=rank, alpha=alpha, path=path, metadata=metadata
+ )
+
+ async def unload_slot(self, slot: int) -> list:
+ return await self._execute_slots("unload_slot", slot=slot)
async def set_rollout_executor(self):
await asyncio.gather(*[cell.set_rollout_executor() for cell in self._cells])
diff --git a/miles/rollout/generate_utils/prefill_logprobs.py b/miles/rollout/generate_utils/prefill_logprobs.py
index 436744020d..c12010fcb5 100644
--- a/miles/rollout/generate_utils/prefill_logprobs.py
+++ b/miles/rollout/generate_utils/prefill_logprobs.py
@@ -7,15 +7,12 @@
from miles.rollout.generate_utils.generate_endpoint_utils import compute_routing_headers, policy_uses_routing_key
from miles.utils.http_utils import post
from miles.utils.lora import LORA_ADAPTER_NAME, lora_rollout_enabled
-from miles.utils.multi_lora import slot_lora_name
from miles.utils.processing_utils import encode_image_for_rollout_engine
from miles.utils.types import Sample
def _lora_path_for_sample(args: Any, sample: Sample) -> str | None:
- """Adapter name to score under: the sample slot's __miles_slot_{N} name, the fixed single-LoRA name, or None."""
- if sample.adapter is not None:
- return slot_lora_name(sample.adapter.slot)
+ """Adapter name to score under: the fixed single-LoRA name, or None."""
if lora_rollout_enabled(args):
return LORA_ADAPTER_NAME
return None
diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py
index 502357802e..ff3c46c9ca 100644
--- a/miles/rollout/generate_utils/sample_utils.py
+++ b/miles/rollout/generate_utils/sample_utils.py
@@ -172,7 +172,6 @@ def _merge_metadata():
metadata=_merge_metadata(),
generate_function_path=_merge_equal_value("generate_function_path"),
train_metadata=_merge_equal_value("train_metadata"),
- adapter=_merge_equal_value("adapter"),
reward_spec=_merge_equal_value("reward_spec"),
routing_key=_merge_equal_value("routing_key"),
non_generation_time=_merge_equal_value("non_generation_time"),
diff --git a/miles/rollout/multi_lora/async_rollout.py b/miles/rollout/multi_lora/async_rollout.py
deleted file mode 100644
index 53d1b99a22..0000000000
--- a/miles/rollout/multi_lora/async_rollout.py
+++ /dev/null
@@ -1,584 +0,0 @@
-"""Fully-async multi-LoRA rollout: a background producer fills per-adapter buffers; batches are collected
-round-robin in ``min_groups_per_dp_split`` multiples without overshooting any adapter's remaining batch."""
-
-import asyncio
-import itertools
-import logging
-import threading
-import time
-from collections import defaultdict, deque
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Any
-
-from miles.ray.multi_lora.controller import AdaptersCache, get_multi_lora_controller
-from miles.rollout.base_types import RolloutFnTrainOutput
-from miles.rollout.filter_hub.common_filters import apply_preput_filters
-from miles.rollout.generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill
-from miles.rollout.sglang_rollout import GenerateState, generate_and_rm_group, get_model_url
-from miles.utils.async_utils import run
-from miles.utils.function_registry import load_function
-from miles.utils.metric_utils import compute_statistics, dict_add_prefix
-from miles.utils.multi_lora import EmptyBatchTimeoutError, min_groups_per_dp_split
-from miles.utils.tracking_utils import tracking
-from miles.utils.types import Sample
-
-logger = logging.getLogger(__name__)
-
-GenerateFn = Callable[..., Any]
-
-# Generate fns may return several samples per rollout; the manager flattens later.
-Group = list[Sample | list[Sample]]
-
-
-def iter_group_samples(group: Group):
- return itertools.chain.from_iterable(item if isinstance(item, list) else (item,) for item in group)
-
-
-def first_sample(group: Group) -> Sample:
- return group[0][0] if isinstance(group[0], list) else group[0]
-
-
-def group_adapter_name(group: Group) -> str | None:
- head = first_sample(group) if group else None
- return head.adapter.name if head is not None and head.adapter else None
-
-
-def group_sample_count(group: Group) -> int:
- return sum(1 for _ in iter_group_samples(group))
-
-
-# Safety valve, same convention as fully_async's queue.Queue(maxsize=1000):
-# never hit in practice, just bounds memory if training stalls entirely.
-MAX_BUFFERED_GROUPS = 1000
-EMPTY_BATCH_TIMEOUT_S = 30.0
-
-
-class GroupBuffer:
- """One adapter's FIFO of completed prompt groups; bounded — the oldest group is dropped when full."""
-
- def __init__(self) -> None:
- self._groups: deque[Group] = deque(maxlen=MAX_BUFFERED_GROUPS)
-
- def __len__(self) -> int:
- return len(self._groups)
-
- def put(self, group: Group) -> None:
- self._groups.append(group)
-
- def get(self, n_groups: int) -> list[Group]:
- """Remove and return the n oldest groups (queue.Queue-style API)."""
- return [self._groups.popleft() for _ in range(n_groups)]
-
- def drop_foreign(self, registration_id: str) -> int:
- """Drop groups stamped by a different registration of this adapter
- name: an in-flight generation of a retired tenant can land after the
- buffer was reset for a same-name re-registration. Unstamped groups
- (no adapter view at submission time) are kept. Returns the drop count."""
- if not self._groups:
- return 0
- kept: deque[Group] = deque(maxlen=MAX_BUFFERED_GROUPS)
- dropped = 0
- for group in self._groups:
- stamped = first_sample(group).metadata.get("registration_id")
- if stamped is not None and stamped != registration_id:
- dropped += 1
- else:
- kept.append(group)
- self._groups = kept
- return dropped
-
- def drop_stale(self, current_version: int, max_staleness: int | None) -> list[int]:
- """Drop groups generated too many weight versions ago; returns the
- staleness of each dropped group (for metrics)."""
- if max_staleness is None or not self._groups:
- return []
- kept: deque[Group] = deque(maxlen=MAX_BUFFERED_GROUPS)
- dropped: list[int] = []
- for group in self._groups:
- stamped = first_sample(group).metadata.get("slot_version")
- staleness = current_version - stamped if stamped is not None else 0
- if stamped is not None and staleness > max_staleness:
- for sample in iter_group_samples(group):
- sample.reset_for_retry()
- dropped.append(staleness)
- else:
- kept.append(group)
- self._groups = kept
- return dropped
-
-
-@dataclass
-class TrainBatch:
- """One train batch: the groups for one train call, with its per-adapter bookkeeping."""
-
- groups: list[Group]
- group_counts: dict[str, int] # prompt groups per adapter in this batch
- step_names: list[str] # adapters whose adapter batch completes -> they step
- step_slots: list[int]
-
-
-def remaining_groups(adapter) -> int:
- """Groups still needed to complete the adapter's batch."""
- remaining = adapter.config.rollout_batch_size - adapter.accumulated_groups
- assert remaining > 0, (
- f"adapter '{adapter.name}' accumulated_groups={adapter.accumulated_groups} >= "
- f"rollout_batch_size={adapter.config.rollout_batch_size}; batch accounting drifted"
- )
- return remaining
-
-
-async def process_group(
- args, group: list[Sample], sampling_params: dict, generate_fn: GenerateFn, data_source
-) -> Group | None:
- """Generate a group; returns None for aborted groups. The slot version is
- stamped at submission time (what the staleness filter compares against)."""
- adapter_name = group[0].adapter.name if group and group[0].adapter else None
- submission_version: int | None = None
- submission_registration: str | None = None
- if adapter_name is not None:
- adapter = await AdaptersCache().get(adapter_name)
- submission_version = adapter.version if adapter is not None else None
- submission_registration = adapter.registration_id if adapter is not None else None
-
- if submission_version is not None:
- for s in group:
- s.metadata["slot_version"] = submission_version
- s.metadata["registration_id"] = submission_registration
-
- result = await generate_fn(args, group, sampling_params)
-
- if submission_version is not None:
- for s in iter_group_samples(result):
- s.metadata["slot_version"] = submission_version
- s.metadata["registration_id"] = submission_registration
-
- if any(s.status == Sample.Status.ABORTED for s in iter_group_samples(result)):
- for s in iter_group_samples(result):
- s.reset_for_retry()
- # Re-queuing is not wired up (the per-adapter source is read-only).
- return None
- return result
-
-
-class MultiLoRAWorkerMetrics:
- """Cross-batch metric state; locked because the producer thread records while the trainer thread drains."""
-
- def __init__(self) -> None:
- self.lock = threading.Lock()
- self.dynamic_filter_drop_counts: dict[str, int] = defaultdict(int)
- # Staleness of dropped groups per adapter, drained every batch.
- self.staleness_values: dict[str, list[int]] = defaultdict(list)
- # Per-adapter shipped-sample values, flushed as step statistics when the adapter steps.
- self.step_rewards: dict[str, list[float]] = defaultdict(list)
- self.step_response_lens: dict[str, list[float]] = defaultdict(list)
- # Per-sample mean engine log prob (rough per-adapter entropy trend).
- self.step_log_prob_means: dict[str, list[float]] = defaultdict(list)
- # Group outcomes for zero-std rates: shipped group counts and each uniform-reward group's reward.
- self.step_group_counts: dict[str, int] = defaultdict(int)
- self.step_zero_std_rewards: dict[str, list[float]] = defaultdict(list)
-
- def record_dynamic_filter_drop(self, reason: str) -> None:
- with self.lock:
- self.dynamic_filter_drop_counts[reason] += 1
-
- def record_stale_drops(self, name: str, staleness_values: list[int]) -> None:
- with self.lock:
- self.staleness_values[name] += staleness_values
-
- def pop_stale_drops(self) -> dict[str, list[int]]:
- """Drain the staleness values of groups dropped since the last batch."""
- with self.lock:
- drained = dict(self.staleness_values)
- self.staleness_values.clear()
- return drained
-
- def record_shipped_samples(
- self, args, data: list[Group], step_names: list[str], adapters: dict
- ) -> dict[str, dict[str, float]]:
- """Accumulate shipped rewards/response lengths per adapter; flush whole-adapter-batch statistics
- for adapters stepping with this batch. Returns {adapter name: flushed metrics}."""
- with self.lock:
- for group in data:
- name = group_adapter_name(group)
- if name is None:
- continue
- group_rewards = []
- for sample in iter_group_samples(group):
- reward = sample.get_reward_value(args)
- group_rewards.append(reward)
- self.step_rewards[name].append(reward)
- self.step_response_lens[name].append(sample.effective_response_length)
- if sample.rollout_log_probs:
- self.step_log_prob_means[name].append(
- sum(sample.rollout_log_probs) / len(sample.rollout_log_probs)
- )
- self.step_group_counts[name] += 1
- if len(group_rewards) > 1 and all(reward == group_rewards[0] for reward in group_rewards):
- self.step_zero_std_rewards[name].append(round(group_rewards[0], 1))
-
- flushed: dict[str, dict[str, float]] = {}
- for name in step_names:
- rewards = self.step_rewards.pop(name, [])
- response_lens = self.step_response_lens.pop(name, [])
- log_prob_means = self.step_log_prob_means.pop(name, [])
- total_groups = self.step_group_counts.pop(name, 0)
- zero_std_rewards = self.step_zero_std_rewards.pop(name, [])
- if not rewards:
- continue
- expected = adapters[name].config.adapter_global_batch_size
- if len(rewards) != expected:
- logger.warning(
- f"Adapter '{name}' stepped with {len(rewards)} shipped samples, expected "
- f"adapter_global_batch_size={expected}; batch accounting drifted"
- )
- # Single-segment keys so "{name}/" matches the "{name}/*" glob (server globs one segment).
- flushed[name] = {
- **dict_add_prefix(compute_statistics(rewards), "raw_reward_"),
- **dict_add_prefix(compute_statistics(response_lens), "response_len_"),
- }
- if log_prob_means:
- flushed[name]["log_probs"] = sum(log_prob_means) / len(log_prob_means)
- if total_groups:
- zero = sum(1 for reward in zero_std_rewards if reward == 0.0)
- one = sum(1 for reward in zero_std_rewards if reward == 1.0)
- flushed[name]["zero_std_all_zero_percentage"] = zero / total_groups
- flushed[name]["zero_std_all_one_percentage"] = one / total_groups
- return flushed
-
- def discard_adapter(self, name: str) -> None:
- """Drop a retired adapter's partial step accumulation."""
- with self.lock:
- self.step_rewards.pop(name, None)
- self.step_response_lens.pop(name, None)
- self.step_log_prob_means.pop(name, None)
- self.step_group_counts.pop(name, None)
- self.step_zero_std_rewards.pop(name, None)
- self.staleness_values.pop(name, None)
-
- def pop_metrics(self) -> dict[str, float]:
- with self.lock:
- metrics = {
- f"rollout/dynamic_filter/drop_{reason}": count
- for reason, count in self.dynamic_filter_drop_counts.items()
- }
- self.dynamic_filter_drop_counts.clear()
- return metrics
-
-
-class AsyncMultiLoRAWorker:
- """Background producer filling bounded per-adapter completed-group buffers;
- the collection loop pops from them via ``get_groups``."""
-
- global_worker = None
- worker_lock = threading.Lock()
-
- def __init__(self, args, data_source, generate_fn: GenerateFn, concurrency: int = None) -> None:
- self.args = args
- self.data_source = data_source
- self.generate_fn = generate_fn
- self.concurrency = concurrency or args.rollout_batch_size
- self.running = True
- self.worker_thread: threading.Thread | None = None
- self.state = GenerateState(args)
- self.dynamic_filter = (
- load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path else None
- )
- # Guards the buffers: the producer thread puts while get_groups (trainer side) pops.
- self.buffer_lock = threading.Lock()
- self.buffers: dict[str, GroupBuffer] = defaultdict(GroupBuffer)
- # Round-robin cursor over adapters, persisting across get_groups calls and batches.
- self.rotation: deque[str] = deque()
- self.metrics = MultiLoRAWorkerMetrics()
- # Last seen registration id per adapter name; a change means re-registration -> drop inherited state.
- self.registrations: dict[str, str] = {}
- # Set when run_loop dies; collect_batch surfaces it instead of a misleading empty-batch timeout.
- self.failure: Exception | None = None
-
- @classmethod
- def get_or_create(cls, args, data_source, generate_fn: GenerateFn, concurrency: int = None):
- with cls.worker_lock:
- if cls.global_worker is None or not cls.global_worker.worker_thread.is_alive():
- cls.global_worker = cls(args, data_source, generate_fn, concurrency)
- cls.global_worker.start()
- return cls.global_worker
-
- def start(self) -> None:
- self.worker_thread = threading.Thread(target=self.thread_main, daemon=True)
- self.worker_thread.start()
-
- def stop(self) -> None:
- self.running = False
- if self.worker_thread and self.worker_thread.is_alive():
- self.worker_thread.join(timeout=5)
-
- @classmethod
- def stop_global(cls) -> None:
- with cls.worker_lock:
- if cls.global_worker is None:
- return
- cls.global_worker.stop()
- cls.global_worker = None
-
- def thread_main(self) -> None:
- asyncio.run(self.run_loop())
-
- async def run_loop(self) -> None:
- active: set[asyncio.Task] = set()
- max_concurrent = self.concurrency
- try:
- while self.running:
- done = {t for t in active if t.done()}
- for t in done:
- try:
- t.result()
- except Exception as e:
- logger.warning(f"generate task failed: {e}")
- active.discard(t)
-
- while len(active) < max_concurrent and self.running:
- samples = self.data_source.get_samples(1)
- if not samples:
- break
- active.add(asyncio.create_task(self.process_and_enqueue(samples[0])))
-
- await asyncio.sleep(0.01)
- except Exception as e:
- # Typically the data source: this stops production for EVERY
- # adapter, so record the cause for collect_batch to surface.
- self.failure = e
- logger.exception("multi-LoRA producer failed; generation is stopped")
- finally:
- for task in active:
- task.cancel()
- if active:
- await asyncio.gather(*active, return_exceptions=True)
-
- async def process_and_enqueue(self, group: list[Sample]) -> None:
- result = await process_group(self.args, group, self.state.sampling_params, self.generate_fn, self.data_source)
- if result is None:
- return
-
- filter_output = apply_preput_filters(self.args, self.dynamic_filter, result)
- if not filter_output.keep:
- if filter_output.reason:
- self.metrics.record_dynamic_filter_drop(filter_output.reason)
- return
-
- adapter_name = group_adapter_name(result)
- if adapter_name is None:
- return
- with self.buffer_lock:
- self.buffers[adapter_name].put(result)
-
- def queue_size(self) -> int:
- with self.buffer_lock:
- return sum(len(buffer) for buffer in self.buffers.values())
-
- def queue_sizes(self) -> dict[str, int]:
- """Buffered (completed, not yet shipped) prompt groups per adapter."""
- with self.buffer_lock:
- return {name: len(buffer) for name, buffer in self.buffers.items()}
-
- def get_groups(
- self, snapshot: dict, num_samples: int, group_counts: dict[str, int]
- ) -> tuple[list[Group], dict[str, int]]:
- """Pop groups round-robin in ``min_groups_per_dp_split`` multiples until ``num_samples`` is covered or
- nothing is poppable; returns them with an updated ``group_counts`` copy (prevents adapter overshoot)."""
- adapters = {**snapshot["active"], **snapshot["retiring"]}
- dp_size = self.args.multi_lora_dp_size
- max_staleness = getattr(self.args, "max_weight_staleness", None)
- group_counts = dict(group_counts) # updated copy; the argument is not modified
- popped: list[Group] = []
- popped_samples = 0
-
- with self.buffer_lock:
- # Retired adapters: discard their buffered tail and partial reward stats.
- for name in list(self.buffers):
- if name not in adapters:
- self.buffers.pop(name)
- self.metrics.discard_adapter(name)
- self.registrations.pop(name, None)
-
- # A re-registered name is a new tenant: drop buffered groups and
- # partial stats inherited from the old tenant.
- for name, adapter in adapters.items():
- previous = self.registrations.get(name)
- if previous is not None and previous != adapter.registration_id:
- self.buffers.pop(name, None)
- self.metrics.discard_adapter(name)
- logger.warning(f"Adapter '{name}' was re-registered; dropped the previous tenant's buffered state")
- self.registrations[name] = adapter.registration_id
-
- # Keep the rotation in sync with live adapters.
- self.rotation = deque(name for name in self.rotation if name in adapters)
- for name in sorted(set(adapters) - set(self.rotation)):
- self.rotation.append(name)
-
- while popped_samples < num_samples:
- made_progress = False
- for _ in range(len(self.rotation)):
- name = self.rotation[0]
- self.rotation.rotate(-1)
- adapter = adapters[name]
- buffer = self.buffers[name]
- if dropped := buffer.drop_stale(adapter.version, max_staleness):
- self.metrics.record_stale_drops(name, dropped)
- # In-flight stragglers of a retired same-name tenant that
- # landed after the re-registration sweep reset the buffer.
- if foreign := buffer.drop_foreign(adapter.registration_id):
- logger.warning(f"Dropped {foreign} buffered groups from a previous registration of '{name}'")
- min_groups_per_pop = min_groups_per_dp_split(adapter.config.n_samples_per_prompt, dp_size)
- trainable_groups = len(buffer) // min_groups_per_pop * min_groups_per_pop
- remaining_allowed_groups = max(0, remaining_groups(adapter) - group_counts.get(name, 0))
- groups_to_pop = min(min_groups_per_pop, trainable_groups, remaining_allowed_groups)
- if groups_to_pop <= 0:
- continue
- popped.extend(buffer.get(groups_to_pop))
- popped_samples += groups_to_pop * adapter.config.n_samples_per_prompt
- group_counts[name] = group_counts.get(name, 0) + groups_to_pop
- made_progress = True
- break
- if not made_progress:
- break # a full pass over rotation yielded nothing
- return popped, group_counts
-
-
-async def collect_batch(args, worker: AsyncMultiLoRAWorker, snapshot: dict) -> TrainBatch:
- """Pop group multiples until the batch reaches ``--global-batch-size`` samples, or it is non-empty and
- stalls for ``--multi-lora-max-coalesce-wait-s`` (the target can be unreachable; ship what there is)."""
- adapters = {**snapshot["active"], **snapshot["retiring"]}
- target_samples = args.global_batch_size
- wait_s = getattr(args, "multi_lora_max_coalesce_wait_s", 0.5)
- empty_wait_s = getattr(args, "multi_lora_max_empty_wait_s", EMPTY_BATCH_TIMEOUT_S)
-
- collected: list[Group] = []
- group_counts: dict[str, int] = {}
- total_samples = 0
- last_progress = time.time()
- last_warning = time.time()
-
- while total_samples < target_samples:
- if worker.failure is not None:
- raise RuntimeError(
- "multi-LoRA producer thread died; generation is stalled for every adapter"
- ) from worker.failure
- groups, group_counts = worker.get_groups(snapshot, target_samples - total_samples, group_counts)
- if groups:
- collected.extend(groups)
- total_samples += sum(adapters[group_adapter_name(g)].config.n_samples_per_prompt for g in groups)
- last_progress = time.time()
- continue
- stalled_s = time.time() - last_progress
- if collected and stalled_s > wait_s:
- break
- if not collected and stalled_s > empty_wait_s:
- raise EmptyBatchTimeoutError(
- "No poppable groups collected before empty timeout; this likely means every live adapter is "
- "below min_groups_per_dp_split (or sources are exhausted). "
- f"queue={worker.queue_size()} active={sorted(snapshot['active'])} retiring={sorted(snapshot['retiring'])}"
- )
- if not collected and time.time() - last_warning > 30:
- logger.warning(
- "No completed groups for 30s. "
- f"queue={worker.queue_size()} active={sorted(snapshot['active'])} "
- f"retiring={sorted(snapshot['retiring'])}"
- )
- last_warning = time.time()
- await asyncio.sleep(0.01)
-
- step_names = sorted(name for name, count in group_counts.items() if count == remaining_groups(adapters[name]))
- return TrainBatch(
- groups=collected,
- group_counts=group_counts,
- step_names=step_names,
- step_slots=sorted(adapters[name].slot for name in step_names),
- )
-
-
-async def generate_rollout_multi_lora_async(
- args, rollout_id: int, data_source, generate_fn: GenerateFn = generate_and_rm_group
-) -> RolloutFnTrainOutput:
- """Collect one train batch and record its contents on the controller."""
- assert args.rollout_global_dataset
-
- state = GenerateState(args)
- worker = AsyncMultiLoRAWorker.get_or_create(args, data_source, generate_fn)
- start_time = time.time()
- queue_sizes = worker.queue_sizes()
-
- # Driver contract: adapter state only changes between generate calls, so one snapshot serves the collection.
- snapshot = await get_multi_lora_controller().snapshot.remote()
- assert snapshot["active"] or snapshot["retiring"], "generate called with no live adapters"
-
- batch = await collect_batch(args, worker, snapshot)
-
- data = sorted(
- batch.groups,
- key=lambda group: (
- first_sample(group).adapter.slot if first_sample(group).adapter is not None else -1,
- first_sample(group).index,
- ),
- )
-
- # Per-sample adapter batch size (drives loss normalization) and batch-level step
- # decision (drives selective optimizer stepping), shipped via sample metadata.
- adapters = {**snapshot["active"], **snapshot["retiring"]}
- for group in data:
- adapter = adapters[group_adapter_name(group)]
- for sample in iter_group_samples(group):
- sample.metadata["adapter_global_batch_size"] = adapter.config.adapter_global_batch_size
- if data:
- head = first_sample(data[0])
- head.metadata["step_slots"] = list(batch.step_slots)
- head.metadata["step_adapter_names"] = list(batch.step_names)
-
- await get_multi_lora_controller().record_batch_adapters.remote(rollout_id, batch.group_counts, batch.step_names)
-
- if (x := args.rollout_sample_filter_path) is not None:
- load_function(x)(args, data)
-
- await recompute_samples_rollout_logprobs_via_prefill(
- args,
- [s for g in data for s in iter_group_samples(g)],
- url=get_model_url(args, "default"),
- sampling_params=state.sampling_params,
- )
-
- # Adapter metrics ride the adapter's own optimizer-step axis ({name}/step); this batch completes step + 1.
- for name, step_metrics in worker.metrics.record_shipped_samples(args, data, batch.step_names, adapters).items():
- step_key = f"{name}/step"
- log_dict = {step_key: adapters[name].step + 1}
- log_dict |= {f"{name}/{key}": value for key, value in step_metrics.items()}
- tracking.log(args, log_dict, step_key=step_key)
-
- stale_drops = worker.metrics.pop_stale_drops()
- all_staleness = [staleness for values in stale_drops.values() for staleness in values]
- metrics = {
- **worker.metrics.pop_metrics(),
- "perf/fully_async/queue_length": sum(queue_sizes.values()),
- "perf/fully_async/stale_dropped": len(all_staleness),
- # {name}/perf/* rides rollout/step; two segments under {name}/ keep these off the step axis.
- **{f"{name}/perf/queue_length": size for name, size in queue_sizes.items()},
- **{f"{name}/perf/stale_dropped": len(stale_drops.get(name, [])) for name in adapters},
- "perf/fully_async/batch_wait_time": time.time() - start_time,
- "perf/fully_async/batch_n_adapters": len(batch.group_counts),
- "perf/fully_async/batch_n_groups": len(data),
- "perf/fully_async/batch_n_samples": sum(group_sample_count(group) for group in data),
- "perf/fully_async/batch_n_adapters_to_step": len(batch.step_names),
- }
- if all_staleness:
- metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(all_staleness) / len(all_staleness)
- metrics["perf/fully_async/stale_dropped_max_staleness"] = max(all_staleness)
- for name, values in stale_drops.items():
- if values:
- metrics[f"{name}/perf/stale_dropped_avg_staleness"] = sum(values) / len(values)
- metrics[f"{name}/perf/stale_dropped_max_staleness"] = max(values)
-
- return RolloutFnTrainOutput(samples=data, metrics=metrics)
-
-
-def generate_rollout_multi_lora(args, rollout_id: int, data_source, evaluation: bool = False):
- if evaluation:
- raise ValueError("Evaluation not supported in multi-LoRA async rollout")
- return run(generate_rollout_multi_lora_async(args, rollout_id, data_source))
diff --git a/miles/rollout/multi_lora/data_source.py b/miles/rollout/multi_lora/data_source.py
deleted file mode 100644
index 426b7af245..0000000000
--- a/miles/rollout/multi_lora/data_source.py
+++ /dev/null
@@ -1,137 +0,0 @@
-"""Round-robin per-adapter data source. Deregistration is step-based and
-lives in the controller (``mark_batch_trained``); every adapter gets a
-``num_step`` at registration, explicit or derived from ``num_epoch``."""
-
-import copy
-import logging
-from argparse import Namespace
-from collections import deque
-from concurrent.futures import ThreadPoolExecutor
-
-import ray
-
-from miles.ray.multi_lora.controller import get_multi_lora_controller
-from miles.rollout.data_source import DataSource, RolloutDataSource
-from miles.utils.adapter_config import AdapterRun
-from miles.utils.types import AdapterRef, RewardSpec, Sample
-
-logger = logging.getLogger(__name__)
-
-MAX_RECONCILE_WORKERS = 16
-
-
-def fetch_snapshot() -> dict:
- return ray.get(get_multi_lora_controller().snapshot.remote())
-
-
-def sampleable(snapshot: dict) -> dict[str, AdapterRun]:
- return {**snapshot["active"], **snapshot["retiring"]}
-
-
-class MultiLoRAAsyncDataSource(DataSource):
- def __init__(self, args: Namespace):
- self.args = args
- self.sources: dict[str, RolloutDataSource] = {}
- self.source_queue: deque = deque()
-
- def reconcile(self, adapters: dict[str, AdapterRun]) -> None:
- for name in list(self.sources):
- if name not in adapters:
- del self.sources[name]
- logger.info(f"Removed data source for adapter '{name}'")
- pending = [(name, a) for name, a in adapters.items() if name not in self.sources]
- if pending:
- workers = min(MAX_RECONCILE_WORKERS, len(pending))
- if workers > 1:
- with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="mlora-ds") as ex:
- built = list(ex.map(lambda na: (na[0], self.create_source(na[1])), pending))
- else:
- built = [(name, self.create_source(a)) for name, a in pending]
- for name, source in built:
- self.sources[name] = source
- logger.info(f"Created data source for adapter '{name}'")
- # Post-filter dataset length; the controller derives num_step
- # from num_epoch for adapters that didn't set it.
- ray.get(get_multi_lora_controller().resolve_num_step.remote(name, len(source.dataset)))
- self.update_queue(set(adapters))
-
- def create_source(self, adapter: AdapterRun) -> RolloutDataSource:
- config = adapter.config
- adapter_args = copy.copy(self.args)
- adapter_args.prompt_data = config.data
- adapter_args.input_key = config.input_key or self.args.input_key
- adapter_args.label_key = config.label_key or self.args.label_key
- adapter_args.metadata_key = config.metadata_key or self.args.metadata_key
- adapter_args.save = config.save or self.args.save
- adapter_args.load = config.save or self.args.load
- adapter_args.n_samples_per_prompt = config.n_samples_per_prompt or self.args.n_samples_per_prompt
- adapter_args.start_rollout_id = 0
- return RolloutDataSource(adapter_args)
-
- def update_queue(self, active_names: set[str]) -> None:
- new_queue: deque = deque()
- in_queue: set[str] = set()
- while self.source_queue:
- if (name := self.source_queue.popleft()) in active_names:
- new_queue.append(name)
- in_queue.add(name)
- for name in active_names:
- if name not in in_queue:
- new_queue.append(name)
- self.source_queue = new_queue
-
- def get_samples(self, num_samples: int = 1) -> list[list[Sample]]:
- """Return the next prompt group, round-robined across adapters.
-
- One rotation of the queue: pull one group from the first adapter that
- yields, stamp it, and return. Empty list when no adapter can produce.
- """
- assert num_samples == 1, "the async producer dispatches one prompt group at a time"
- snapshot = fetch_snapshot()
- adapters = sampleable(snapshot)
- self.reconcile(adapters)
- self.update_queue(set(self.sources))
-
- for _ in range(len(self.source_queue)):
- name = self.source_queue.popleft()
- self.source_queue.append(name)
- source = self.sources[name]
- groups = source.get_samples(1)
- if not groups:
- continue
-
- adapter = adapters[name]
- config = adapter.config
- ref = AdapterRef(name=name, slot=adapter.slot)
- reward_spec = RewardSpec(rm_type=config.rm_type, custom_rm_path=config.custom_rm_path)
- for sample in groups[0]:
- sample.adapter = ref
- sample.reward_spec = reward_spec
- sample.metadata = {**config.metadata, **sample.metadata}
-
- return groups
-
- return []
-
- def add_samples(self, samples: list[list[Sample]]) -> None:
- """Recycle retried/aborted groups; drop groups for deregistered adapters."""
- adapters = sampleable(fetch_snapshot())
- self.reconcile(adapters)
- for group in samples:
- name = group[0].adapter.name if group and group[0].adapter else None
- if not name or name not in self.sources or name not in adapters:
- continue
- self.sources[name].add_samples([group])
-
- def save(self, rollout_id):
- for source in self.sources.values():
- source.save(rollout_id)
-
- def load(self, rollout_id=None):
- for source in self.sources.values():
- source.load(rollout_id)
-
- def close(self) -> None:
- from miles.rollout.multi_lora.async_rollout import AsyncMultiLoRAWorker
-
- AsyncMultiLoRAWorker.stop_global()
diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py
index 2d59d26389..61fbb935fd 100644
--- a/miles/rollout/rm_hub/__init__.py
+++ b/miles/rollout/rm_hub/__init__.py
@@ -5,7 +5,6 @@
import aiohttp
from miles.utils.function_registry import load_function
-from miles.utils.multi_lora import is_multi_lora_enabled
from miles.utils.types import Sample
from .deepscaler import get_deepscaler_rule_based_reward, get_gemma_math_reward
@@ -101,7 +100,7 @@ async def batched_async_rm(
sample.reward = reward
return None
- if args.custom_rm_path is not None and not is_multi_lora_enabled(args):
+ if args.custom_rm_path is not None:
rm_function = load_function(args.custom_rm_path)
return await rm_function(args, samples, **kwargs)
tasks = [async_rm(args, sample, **kwargs) for sample in samples]
diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py
index 171eabd260..40935a8b63 100644
--- a/miles/rollout/sglang_rollout.py
+++ b/miles/rollout/sglang_rollout.py
@@ -26,7 +26,6 @@
from miles.utils.lifecycle import TrajectoryLifecycle
from miles.utils.lora import LORA_ADAPTER_NAME, lora_rollout_enabled
from miles.utils.misc import SingletonMeta, call_agent_abort_hook
-from miles.utils.multi_lora import make_rid, slot_lora_name
from miles.utils.processing_utils import (
call_processor,
encode_image_for_rollout_engine,
@@ -183,22 +182,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
if getattr(args, "use_opd", False) and opd_top_k > 0 and opd_top_k_strategy != "only-teacher":
payload["top_logprobs_num"] = opd_top_k
- if sample.adapter is not None:
- from miles.ray.multi_lora.controller import AdaptersCache
-
- if (adapter := await AdaptersCache().get(sample.adapter.name)) is None:
- # Adapter deregistered: don't POST, or an orphan the abort round can't see
- # would keep decoding under the slot's next tenant and pollute its group.
- logger.warning(
- f"Dropping generation for adapter '{sample.adapter.name}' (slot {sample.adapter.slot}): "
- "adapter is no longer sampleable"
- )
- sample.status = Sample.Status.ABORTED
- return sample
- payload["lora_path"] = slot_lora_name(sample.adapter.slot)
- payload["rid"] = make_rid(sample.adapter.name)
- payload["extra_key"] = f"{sample.adapter.name}:v{adapter.version}"
- elif lora_rollout_enabled(args):
+ if lora_rollout_enabled(args):
payload["lora_path"] = LORA_ADAPTER_NAME
if args.use_rollout_routing_replay:
diff --git a/miles/rollout/multi_lora/__init__.py b/miles/tinker/__init__.py
similarity index 100%
rename from miles/rollout/multi_lora/__init__.py
rename to miles/tinker/__init__.py
diff --git a/miles/tinker/arguments.py b/miles/tinker/arguments.py
new file mode 100644
index 0000000000..ffc930e7f4
--- /dev/null
+++ b/miles/tinker/arguments.py
@@ -0,0 +1,57 @@
+import argparse
+
+from miles.utils.hf_config import load_hf_config
+
+
+def add_tinker_arguments(parser):
+ group = parser.add_argument_group("Tinker")
+
+ def add_argument(name, **kwargs):
+ return group.add_argument(f"--tinker-{name}", **kwargs)
+
+ add_argument("server-host", default="0.0.0.0")
+ add_argument("server-port", type=int, default=10613)
+ add_argument(
+ "base-model",
+ help="Model name advertised by the gateway (default: --hf-checkpoint)",
+ )
+ add_argument(
+ "checkpoint-root",
+ help="Directory for tinker:// checkpoints (default: /tinker)",
+ )
+ add_argument("train-attn", action=argparse.BooleanOptionalAction, default=True)
+ add_argument("train-mlp", action=argparse.BooleanOptionalAction, default=True)
+ add_argument("train-unembed", action=argparse.BooleanOptionalAction, default=True)
+ return parser
+
+
+def configure_tinker_args(args):
+ assert args.train_backend == "megatron", "Tinker requires the Megatron backend"
+ assert (
+ args.target_modules is None and args.exclude_modules is None
+ ), "Tinker uses --tinker-train-attn/mlp/unembed; --target-modules and --exclude-modules are not supported"
+ modules = _resolve_target_modules(
+ load_hf_config(args.hf_checkpoint),
+ train_attn=args.tinker_train_attn,
+ train_mlp=args.tinker_train_mlp,
+ train_unembed=args.tinker_train_unembed,
+ )
+ # The common LoRA validator parses and validates this before trainer/engine initialization.
+ args.target_modules = ",".join(modules)
+
+
+def _resolve_target_modules(hf_config, *, train_attn, train_mlp, train_unembed):
+ # Other architectures need their own complete attention/MLP mapping.
+ assert hf_config.model_type in (
+ "qwen3",
+ "qwen3_moe",
+ ), f"Tinker target layout is not defined for model_type={hf_config.model_type!r}"
+ modules = []
+ if train_attn:
+ modules.extend(("q_proj", "k_proj", "v_proj", "o_proj"))
+ if train_mlp:
+ modules.extend(("gate_proj", "up_proj", "down_proj"))
+ if train_unembed:
+ modules.append("lm_head")
+ assert modules, "Tinker requires at least one trainable LoRA module group"
+ return modules
diff --git a/tests/fast/ray/multi_lora/__init__.py b/miles/tinker/core/__init__.py
similarity index 100%
rename from tests/fast/ray/multi_lora/__init__.py
rename to miles/tinker/core/__init__.py
diff --git a/miles/tinker/core/future.py b/miles/tinker/core/future.py
new file mode 100644
index 0000000000..38c94b48d5
--- /dev/null
+++ b/miles/tinker/core/future.py
@@ -0,0 +1,84 @@
+"""In-memory future results. Request deduplication lasts only for the current process."""
+
+import asyncio
+import time
+import uuid
+from dataclasses import dataclass, field
+
+from miles.tinker.core.types import OwnershipError
+
+PENDING = "pending"
+DONE = "done"
+FAILED = "failed"
+
+_FINISHED_TTL_S = 3600.0
+
+
+@dataclass
+class RequestFuture:
+ request_id: str
+ model_id: str
+ tenant: str
+ state: str = PENDING
+ result: dict | None = None
+ error: str | None = None
+ error_category: str | None = None
+ finished_at: float | None = None
+ # long-poll wakeup: set when the future settles
+ settled: asyncio.Event = field(default_factory=asyncio.Event)
+
+
+class RequestFutureStore:
+ def __init__(self) -> None:
+ self._futures: dict[str, RequestFuture] = {}
+
+ def create(self, model_id: str, tenant: str) -> RequestFuture:
+ future = RequestFuture(request_id=f"req-{uuid.uuid4().hex}", model_id=model_id, tenant=tenant)
+ self._futures[future.request_id] = future
+ return future
+
+ def resolve(self, request_id: str, result: dict) -> None:
+ future = self._futures[request_id]
+ if future.state != PENDING: # e.g. already failed by lease expiry
+ return
+ future.state = DONE
+ future.result = result
+ future.finished_at = time.monotonic()
+ future.settled.set()
+
+ def fail(self, request_id: str, error: str, category: str) -> None:
+ future = self._futures[request_id]
+ if future.state != PENDING:
+ return
+ future.state = FAILED
+ future.error = error
+ future.error_category = category
+ future.finished_at = time.monotonic()
+ future.settled.set()
+
+ def get(self, request_id: str, tenant: str) -> RequestFuture | None:
+ """Return None for unknown or expired futures; enforce tenant ownership otherwise."""
+ self._sweep()
+ future = self._futures.get(request_id)
+ if future is None:
+ return None
+ if future.tenant != tenant:
+ raise OwnershipError(f"request {request_id} does not belong to this tenant")
+ return future
+
+ def request_id_for_retry(self, request_id: str, model_id: str, tenant: str) -> str:
+ if self.get(request_id, tenant) is not None:
+ return request_id
+ replacement = self.create(model_id, tenant)
+ self.fail(replacement.request_id, "result expired after retention", "user")
+ return replacement.request_id
+
+ def _sweep(self) -> None:
+ now = time.monotonic()
+ expired = [
+ request_id
+ for request_id, future in self._futures.items()
+ if future.finished_at is not None and now - future.finished_at > _FINISHED_TTL_S
+ ]
+ for request_id in expired:
+ del self._futures[request_id]
diff --git a/miles/tinker/core/input_validation.py b/miles/tinker/core/input_validation.py
new file mode 100644
index 0000000000..63b4f49fc5
--- /dev/null
+++ b/miles/tinker/core/input_validation.py
@@ -0,0 +1,156 @@
+import math
+import re
+
+from miles.tinker.core.types import (
+ LOSS_FN_INPUTS,
+ LOSS_INPUT_KEYS,
+ CommandOp,
+ GatewayConfig,
+ ModelRecord,
+ UserInputError,
+)
+
+
+def validate_model_config(lora_config: dict, config: GatewayConfig) -> None:
+ """Reject per-model settings that conflict with the fixed server adapter layout."""
+ if lora_config.get("seed") is not None:
+ raise UserInputError("lora_config.seed is not supported: adapter initialization is not per-model seedable")
+ layout = {
+ "train_attn": config.trains_attn,
+ "train_mlp": config.trains_mlp,
+ "train_unembed": config.trains_unembed,
+ }
+ for field, layout_trains in layout.items():
+ requested = lora_config.get(field)
+ if requested is not None and requested != layout_trains:
+ raise UserInputError(
+ f"lora_config.{field}={requested} conflicts with this gateway's adapter layout "
+ f"({field}={layout_trains}); the layout is fixed by --tinker-train-attn/mlp/unembed at server start"
+ )
+ rank = lora_config.get("rank", 32)
+ if type(rank) is not int or not 1 <= rank <= config.max_lora_rank:
+ raise UserInputError(f"lora_config.rank must be an integer in [1, {config.max_lora_rank}], got {rank!r}")
+
+
+def validate_batch_payload(op: CommandOp, payload: dict, config: GatewayConfig) -> None:
+ if not op.is_batch():
+ return
+ datums = payload["datums"]
+ if not datums:
+ raise UserInputError("forward_backward with no data")
+ if len(datums) > config.max_datums_per_request:
+ raise UserInputError(f"{len(datums)} datums exceeds max_datums_per_request={config.max_datums_per_request}")
+ required_inputs = LOSS_FN_INPUTS.get(payload["loss_fn"])
+ if required_inputs is None:
+ raise UserInputError(f"unknown loss_fn {payload['loss_fn']!r}; known: {sorted(LOSS_FN_INPUTS)}")
+ loss_fn_config = payload.get("loss_fn_config")
+ if loss_fn_config is not None:
+ if not isinstance(loss_fn_config, dict):
+ raise UserInputError("loss_fn_config must be an object")
+ config_keys = {
+ "ppo": ("clip_low_threshold", "clip_high_threshold"),
+ "cispo": ("clip_low_threshold", "clip_high_threshold"),
+ "dro": ("beta",),
+ }.get(payload["loss_fn"], ())
+ for key in config_keys:
+ if key in loss_fn_config:
+ value = loss_fn_config[key]
+ if type(value) not in (int, float) or not math.isfinite(value):
+ raise UserInputError(f"loss_fn_config[{key!r}] must be a finite number")
+ total_tokens = 0
+ for index, datum in enumerate(datums):
+ validate_token_ids(datum["tokens"], config.vocab_size, f"datum {index}: model_input")
+ validate_token_ids(datum["target_tokens"], config.vocab_size, f"datum {index}: target_tokens")
+ if len(datum["tokens"]) > config.max_tokens_per_datum:
+ raise UserInputError(f"datum {index}: {len(datum['tokens'])} tokens exceeds {config.max_tokens_per_datum}")
+ total_tokens += len(datum["tokens"])
+ for wire_key in required_inputs:
+ values = datum.get(LOSS_INPUT_KEYS[wire_key])
+ if values is None:
+ raise UserInputError(
+ f"datum {index}: loss_fn {payload['loss_fn']!r} needs loss_fn_inputs[{wire_key!r}]"
+ )
+ if len(values) != datum["target_len"]:
+ raise UserInputError(
+ f"datum {index}: loss_fn_inputs[{wire_key!r}] has {len(values)} values "
+ f"for {datum['target_len']} target tokens"
+ )
+ unread = [
+ wire_key
+ for wire_key, datum_key in LOSS_INPUT_KEYS.items()
+ if wire_key not in required_inputs and datum_key in datum
+ ]
+ if unread:
+ raise UserInputError(
+ f"datum {index}: loss_fn {payload['loss_fn']!r} does not read loss_fn_inputs {unread}"
+ )
+ if total_tokens > config.max_tokens_per_request:
+ raise UserInputError(f"{total_tokens} tokens exceeds max_tokens_per_request={config.max_tokens_per_request}")
+
+
+def validate_sample_payload(payload: dict, config: GatewayConfig) -> None:
+ validate_token_ids(payload["prompt_tokens"], config.vocab_size, "prompt")
+ num_samples = payload.get("num_samples", 1)
+ if type(num_samples) is not int or not 1 <= num_samples <= config.max_samples_per_request:
+ raise UserInputError(f"num_samples must be an integer in [1, {config.max_samples_per_request}]")
+ topk = payload.get("topk_prompt_logprobs", 0)
+ if type(topk) is not int or topk < 0:
+ raise UserInputError("topk_prompt_logprobs must be a nonnegative integer")
+
+
+def validate_seq_id(value, name: str, minimum: int = 1) -> int:
+ # request seq_ids are 1-based (last_enqueued_seq_id starts at 0); idempotency keys are 0-based
+ if not isinstance(value, int) or value < minimum:
+ raise UserInputError(f"{name} must be an integer >= {minimum}, got {value!r}")
+ return value
+
+
+def validate_checkpoint_segment(segment: str) -> None:
+ """Reject client path segments that could escape the checkpoint root."""
+ if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}", segment) is None:
+ raise UserInputError(f"invalid checkpoint path segment {segment!r}")
+
+
+def validate_checkpoint_compatibility(meta: dict, record: ModelRecord, config: GatewayConfig, shown_path: str) -> None:
+ """Reject settings that would change the saved tensors' meaning."""
+ expected = {
+ "base_model": record.base_model,
+ "lora_rank": record.lora_rank,
+ "lora_alpha": record.lora_alpha,
+ "train_attn": config.trains_attn,
+ "train_mlp": config.trains_mlp,
+ "train_unembed": config.trains_unembed,
+ }
+ for key, value in expected.items():
+ if meta[key] != value:
+ raise UserInputError(
+ f"checkpoint {shown_path!r} was saved with {key}={meta[key]!r}; this model expects {key}={value!r}"
+ )
+
+
+def validate_save_options(payload: dict) -> None:
+ if payload.get("ttl_seconds") is not None:
+ raise UserInputError("ttl_seconds is not supported: checkpoints on this gateway do not expire")
+ if payload.get("user_metadata") is not None:
+ raise UserInputError("user_metadata is not supported by this gateway")
+
+
+def validate_token_ids(tokens: list, vocab_size: int, field: str) -> None:
+ if any(type(token) is not int or not 0 <= token < vocab_size for token in tokens):
+ raise UserInputError(f"{field} must contain integer token IDs in [0, {vocab_size})")
+
+
+def validate_checkpoint_metadata(meta, shown_path: str) -> None:
+ fields = {
+ "tenant_digest": (str,),
+ "base_model": (str,),
+ "lora_rank": (int,),
+ "lora_alpha": (int, float),
+ "train_attn": (bool,),
+ "train_mlp": (bool,),
+ "train_unembed": (bool,),
+ }
+ if not isinstance(meta, dict) or any(type(meta.get(key)) not in types for key, types in fields.items()):
+ raise UserInputError(f"checkpoint {shown_path!r} has invalid or unsupported metadata")
+ if meta["lora_rank"] <= 0 or not math.isfinite(meta["lora_alpha"]):
+ raise UserInputError(f"checkpoint {shown_path!r} has invalid LoRA metadata")
diff --git a/miles/tinker/core/model_queue.py b/miles/tinker/core/model_queue.py
new file mode 100644
index 0000000000..15ac077aee
--- /dev/null
+++ b/miles/tinker/core/model_queue.py
@@ -0,0 +1,77 @@
+"""Per-model command ordering and batch runs separated by barriers."""
+
+from collections import deque
+from dataclasses import dataclass, field
+
+from miles.tinker.core.types import Command
+
+
+@dataclass
+class PendingRequest:
+ """One submitted command and its completion accounting."""
+
+ command: Command
+ datums: list[dict] = field(default_factory=list) # batch ops only
+ num_issued_datums: int = 0
+ outputs: list[dict | None] = field(default_factory=list)
+ remaining: int = 0
+
+ @property
+ def is_batch_op(self) -> bool:
+ return self.command.op.is_batch()
+
+ def record_output(self, local_index: int, output: dict) -> bool:
+ """Store one datum's result; True once every datum has reported."""
+ self.outputs[local_index] = output
+ self.remaining -= 1
+ return self.remaining == 0
+
+ def pack_key(self) -> tuple:
+ """Datums pack into one BatchUnit only within the same (op, loss_fn, config)."""
+ config = self.command.payload.get("loss_fn_config") or {}
+ return (self.command.op, self.command.payload["loss_fn"], tuple(sorted(config.items())))
+
+
+class ModelRequestQueue:
+ def __init__(self, model_id: str, tenant: str, slot: int) -> None:
+ self.model_id = model_id
+ self.tenant = tenant
+ self.slot = slot
+ # seq_ids are 1-based
+ self.last_enqueued_seq_id = 0
+ self.pending_by_seq: dict[int, Command] = {}
+ self.queue: deque[PendingRequest] = deque()
+
+ def submit(self, command: Command) -> None:
+ """Accept one deduplicated command; feed the queue in seq order."""
+ self.pending_by_seq[command.seq_id] = command
+ self._enqueue_contiguous_commands()
+
+ def _enqueue_contiguous_commands(self) -> None:
+ while self.last_enqueued_seq_id + 1 in self.pending_by_seq:
+ self.last_enqueued_seq_id += 1
+ next_command = self.pending_by_seq.pop(self.last_enqueued_seq_id)
+ pending = PendingRequest(command=next_command)
+ if pending.is_batch_op and pending.command.validation_error is None:
+ pending.datums = next_command.payload["datums"]
+ pending.remaining = len(pending.datums)
+ pending.outputs = [None] * len(pending.datums)
+ self.queue.append(pending)
+
+ def open_batch_run(self) -> list[PendingRequest]:
+ """The leading run of batch-op commands; their datums are all issuable."""
+ run = []
+ for pending in self.queue:
+ if not pending.is_batch_op or pending.command.validation_error is not None:
+ break
+ run.append(pending)
+ return run
+
+ def ready_barrier(self) -> PendingRequest | None:
+ """The head barrier, executable once every batch op ahead of it completed."""
+ if self.queue and not self.queue[0].is_batch_op and self.queue[0].command.validation_error is None:
+ return self.queue[0]
+ return None
+
+ def finish(self, pending: PendingRequest) -> None:
+ self.queue.remove(pending)
diff --git a/miles/tinker/core/scheduler.py b/miles/tinker/core/scheduler.py
new file mode 100644
index 0000000000..14ff0c3a07
--- /dev/null
+++ b/miles/tinker/core/scheduler.py
@@ -0,0 +1,130 @@
+"""Choose the earliest ready seed, then pack compatible work across model queues."""
+
+from dataclasses import dataclass
+
+from miles.tinker.core.model_queue import ModelRequestQueue, PendingRequest
+from miles.tinker.core.types import CommandOp
+
+
+@dataclass
+class DatumRef:
+ """Pointer to ``request.datums[local_index]``; the datum's output is written back through it."""
+
+ model_queue: ModelRequestQueue
+ request: PendingRequest
+ local_index: int
+
+ @property
+ def datum(self) -> dict:
+ return self.request.datums[self.local_index]
+
+ @property
+ def arrival(self) -> int:
+ return self.request.command.arrival
+
+
+@dataclass
+class BatchUnit:
+ """One forward pass on the trainer: datums packed from any number of requests."""
+
+ op: CommandOp # FORWARD_BACKWARD | FORWARD_ONLY
+ loss_fn: str | None
+ loss_fn_config: dict | None
+ datums: list[DatumRef]
+
+
+@dataclass
+class BarrierUnit:
+ """One barrier call after all preceding batch requests in its model queue finish."""
+
+ op: CommandOp # OPTIM_STEP | SAVE_STATE | LOAD_STATE | SAVE_WEIGHTS_FOR_SAMPLER
+ entries: list[tuple[ModelRequestQueue, PendingRequest]]
+
+
+class RequestScheduler:
+ def __init__(self, batch_token_budget: int) -> None:
+ self.batch_token_budget = batch_token_budget
+ self._model_queues: dict[str, ModelRequestQueue] = {}
+
+ def add_model_queue(self, model_queue: ModelRequestQueue) -> None:
+ self._model_queues[model_queue.model_id] = model_queue
+
+ def remove_model_queue(self, model_id: str) -> None:
+ del self._model_queues[model_id]
+
+ def model_queue(self, model_id: str) -> ModelRequestQueue:
+ return self._model_queues[model_id]
+
+ def ready_rejections(self) -> list[tuple[ModelRequestQueue, PendingRequest]]:
+ return [
+ (model_queue, model_queue.queue[0])
+ for model_queue in self._model_queues.values()
+ if model_queue.queue and model_queue.queue[0].command.validation_error is not None
+ ]
+
+ def schedule_next(self) -> BatchUnit | BarrierUnit | None:
+ """Arrival order selects the seed; compatible work may overtake intervening requests."""
+ datums = self._ready_datums()
+ barriers = self._ready_barriers()
+
+ datum_seed = min(datums, key=lambda ref: ref.arrival) if datums else None
+ barrier_seed = min(barriers, key=lambda e: e[1].command.arrival) if barriers else None
+ if datum_seed is None and barrier_seed is None:
+ return None
+ if barrier_seed is not None and (datum_seed is None or barrier_seed[1].command.arrival < datum_seed.arrival):
+ return self._merge_barriers(barrier_seed, barriers)
+ return self._pack_batch(datum_seed, datums)
+
+ def _ready_datums(self) -> list[DatumRef]:
+ datums = []
+ for model_queue in self._model_queues.values():
+ for request in model_queue.open_batch_run():
+ datums.extend(
+ DatumRef(model_queue, request, index)
+ for index in range(request.num_issued_datums, len(request.datums))
+ )
+ return datums
+
+ def _ready_barriers(self) -> list[tuple[ModelRequestQueue, PendingRequest]]:
+ return [
+ (model_queue, barrier)
+ for model_queue in self._model_queues.values()
+ if (barrier := model_queue.ready_barrier()) is not None
+ ]
+
+ def _pack_batch(self, seed: DatumRef, datums: list[DatumRef]) -> BatchUnit:
+ pack_key = seed.request.pack_key()
+ compatible = sorted(
+ (ref for ref in datums if ref.request.pack_key() == pack_key),
+ key=lambda ref: (ref.arrival, ref.local_index),
+ )
+ picked: list[DatumRef] = []
+ tokens = 0
+ for ref in compatible:
+ datum_tokens = len(ref.datum["tokens"])
+ if picked and tokens + datum_tokens > self.batch_token_budget:
+ break
+ picked.append(ref)
+ tokens += datum_tokens
+ for ref in picked:
+ ref.request.num_issued_datums += 1
+ command = seed.request.command
+ return BatchUnit(
+ op=command.op,
+ loss_fn=command.payload.get("loss_fn"),
+ loss_fn_config=command.payload.get("loss_fn_config"),
+ datums=picked,
+ )
+
+ def _merge_barriers(
+ self,
+ seed: tuple[ModelRequestQueue, PendingRequest],
+ barriers: list[tuple[ModelRequestQueue, PendingRequest]],
+ ) -> BarrierUnit:
+ op = seed[1].command.op
+ if op == CommandOp.OPTIM_STEP:
+ # optim barriers of different models step in one trainer call
+ entries = [(model_queue, barrier) for model_queue, barrier in barriers if barrier.command.op == op]
+ else:
+ entries = [seed]
+ return BarrierUnit(op=op, entries=entries)
diff --git a/miles/tinker/core/service.py b/miles/tinker/core/service.py
new file mode 100644
index 0000000000..b89b8cdac5
--- /dev/null
+++ b/miles/tinker/core/service.py
@@ -0,0 +1,548 @@
+"""Sessions, models, futures, and ordered trainer dispatch.
+
+The trainer lock serializes trainer calls across dispatch, model creation, and lease expiry."""
+
+import asyncio
+import logging
+import os
+import time
+import uuid
+
+from miles.tinker.core.future import RequestFuture, RequestFutureStore
+from miles.tinker.core.input_validation import (
+ validate_batch_payload,
+ validate_checkpoint_compatibility,
+ validate_checkpoint_segment,
+ validate_model_config,
+ validate_sample_payload,
+ validate_seq_id,
+)
+from miles.tinker.core.model_queue import ModelRequestQueue
+from miles.tinker.core.scheduler import BarrierUnit, BatchUnit, RequestScheduler
+from miles.tinker.core.types import (
+ Command,
+ CommandOp,
+ GatewayConfig,
+ ModelRecord,
+ OwnershipError,
+ SamplingSessionRecord,
+ SessionRecord,
+ UserInputError,
+)
+from miles.tinker.core.utils import (
+ build_checkpoint_metadata,
+ parse_tinker_path,
+ read_checkpoint_metadata,
+ resolve_checkpoint_dir,
+ resolve_sampler_checkpoint,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class TinkerService:
+ def __init__(self, backend, config: GatewayConfig) -> None:
+ self.backend = backend
+ self.config = config
+ self.futures = RequestFutureStore()
+ self.scheduler = RequestScheduler(config.batch_token_budget)
+ self.models: dict[str, ModelRecord] = {}
+ self.sessions: dict[str, SessionRecord] = {}
+ self.sampling_sessions: dict[str, SamplingSessionRecord] = {}
+ self.free_slots = set(range(config.n_slots))
+ self._wake = asyncio.Event()
+ self._trainer_lock = asyncio.Lock()
+ self._sample_tasks: dict[str, tuple] = {} # request_id -> (task, session_id)
+ self._create_tasks: set = set()
+ self._arrival_counter = 0
+ self._batch_counter = 0
+ self._background_error: BaseException | None = None
+ # why each model closed, so later requests get the reason instead of "unknown model"
+ self._close_reasons: dict[str, str] = {}
+
+ async def run(self) -> None:
+ sweep_task = asyncio.create_task(self._run_lease_sweeper())
+ sweep_task.add_done_callback(self._observe_background_task)
+ try:
+ while True:
+ if self._background_error is not None:
+ raise self._background_error
+ # unit selection shares the critical section with execution, so
+ # lease expiry cannot reclaim a model queue between the two
+ async with self._trainer_lock:
+ rejections = self.scheduler.ready_rejections()
+ if rejections:
+ for model_queue, pending in rejections:
+ await self._finish_request(
+ model_queue,
+ pending,
+ {"error": pending.command.validation_error, "error_category": "user"},
+ )
+ continue
+ unit = self.scheduler.schedule_next()
+ if unit is not None:
+ if isinstance(unit, BatchUnit):
+ await self._run_batch(unit)
+ else:
+ await self._run_barrier(unit)
+ if self.backend.trainer_dead():
+ raise RuntimeError("the trainer workers died; exiting so clients get refused connections")
+ continue
+ await self._wake.wait()
+ self._wake.clear()
+ finally:
+ tasks = [sweep_task, *self._create_tasks, *(task for task, _ in self._sample_tasks.values())]
+ for task in tasks:
+ task.cancel()
+ # Cleanup failures must not replace the error that stopped the dispatcher.
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+ def _observe_background_task(self, task: asyncio.Task) -> None:
+ if not task.cancelled() and (error := task.exception()) is not None:
+ if self._background_error is None:
+ self._background_error = error
+ self._wake.set()
+
+ def create_session(self, tenant: str) -> str:
+ session_id = f"session-{uuid.uuid4().hex}"
+ self.sessions[session_id] = SessionRecord(tenant=tenant, last_heartbeat=time.monotonic())
+ return session_id
+
+ def _session_for(self, tenant: str, session_id: str) -> SessionRecord:
+ session = self.sessions.get(session_id)
+ if session is None or session.tenant != tenant:
+ raise UserInputError(f"unknown session {session_id!r}; create a session first")
+ return session
+
+ def heartbeat(self, tenant: str, session_id: str) -> bool:
+ session = self.sessions.get(session_id)
+ if session is None or session.tenant != tenant:
+ return False
+ session.last_heartbeat = time.monotonic()
+ return True
+
+ def create_model(self, tenant: str, payload: dict) -> tuple[str, str]:
+ """Two-phase like every command: allocate now, initialize the slot behind the future."""
+ session = self._session_for(tenant, payload["session_id"])
+ model_seq_id = validate_seq_id(payload["model_seq_id"], "model_seq_id", minimum=0)
+ if (previous := session.models_by_seq.get(model_seq_id)) is not None:
+ request_id, model_id = previous
+ request_id = self.futures.request_id_for_retry(request_id, model_id, tenant)
+ session.models_by_seq[model_seq_id] = (request_id, model_id)
+ return request_id, model_id
+ base_model = payload["base_model"]
+ if base_model != self.config.base_model:
+ raise UserInputError(f"this gateway serves {self.config.base_model!r}, not {base_model!r}")
+ lora_config = payload.get("lora_config") or {}
+ validate_model_config(lora_config, self.config)
+ rank = lora_config.get("rank", 32)
+ alpha = self.config.lora_alpha if self.config.lora_alpha is not None else float(2 * rank)
+ if not self.free_slots:
+ raise UserInputError(f"no free adapter slots (capacity {self.config.n_slots})")
+ slot = min(self.free_slots)
+ self.free_slots.remove(slot)
+
+ model_id = f"model-{uuid.uuid4().hex[:12]}"
+ future = self.futures.create(model_id, tenant)
+ record = ModelRecord(
+ model_id=model_id,
+ session_id=payload["session_id"],
+ tenant=tenant,
+ slot=slot,
+ base_model=base_model,
+ lora_rank=rank,
+ lora_alpha=alpha,
+ create_request_id=future.request_id,
+ )
+ self.models[model_id] = record
+ self.scheduler.add_model_queue(ModelRequestQueue(model_id, tenant, slot))
+ task = asyncio.create_task(self._run_create_model(record))
+ self._create_tasks.add(task)
+ task.add_done_callback(self._create_tasks.discard)
+ task.add_done_callback(self._observe_background_task)
+ session.models_by_seq[model_seq_id] = (future.request_id, model_id)
+ return future.request_id, model_id
+
+ async def _run_create_model(self, record: ModelRecord) -> None:
+ async with self._trainer_lock:
+ if self.models.get(record.model_id) is not record:
+ return
+ failure = await self.backend.load_slot(record.slot, record.lora_rank, record.lora_alpha)
+ record.slot_initialized = True
+ if failure is not None:
+ await self._close_model(record.model_id, failure["error"], "server")
+ return
+ self.futures.resolve(record.create_request_id, {"op": "create_model", "model_id": record.model_id})
+
+ def get_model(self, tenant: str, model_id: str) -> ModelRecord:
+ record = self.models.get(model_id)
+ if record is None:
+ reason = self._close_reasons.get(model_id)
+ if reason is not None:
+ raise UserInputError(f"model {model_id!r} was unloaded: {reason}")
+ raise UserInputError(f"unknown model {model_id!r}")
+ if record.tenant != tenant:
+ raise OwnershipError(f"model {model_id} does not belong to this tenant")
+ return record
+
+ async def _close_model(self, model_id: str, error: str, category: str) -> None:
+ """Free a model's slot and fail its pending requests; requires the trainer lock. Idempotent."""
+ record = self.models.pop(model_id, None)
+ if record is None:
+ return
+ self._close_reasons[model_id] = error
+ while len(self._close_reasons) > 4 * self.config.n_slots:
+ self._close_reasons.pop(next(iter(self._close_reasons)))
+ self.scheduler.remove_model_queue(model_id)
+ for request_id in [record.create_request_id, *record.request_id_by_seq.values()]:
+ if self.futures.get(request_id, record.tenant) is not None:
+ self.futures.fail(request_id, error, category)
+ if self.backend.trainer_dead():
+ return
+ if record.slot_initialized:
+ failure = await self.backend.unload_slot(record.slot)
+ if failure is not None:
+ logger.error("slot %s remains unavailable after unload failed: %s", record.slot, failure["error"])
+ return
+ self.free_slots.add(record.slot)
+
+ def submit(self, tenant: str, op: str, payload: dict) -> str:
+ """Admit a decoded command; content errors settle its future as a user failure."""
+ try:
+ op = CommandOp(op)
+ except ValueError:
+ raise UserInputError(f"unknown command op {op!r}") from None
+ model_id = payload["model_id"]
+ record = self.get_model(tenant, model_id)
+ seq_id = validate_seq_id(payload["seq_id"], "seq_id")
+ model_queue = self.scheduler.model_queue(model_id)
+
+ # retries must not accumulate gradients twice
+ if seq_id in record.request_id_by_seq:
+ request_id = self.futures.request_id_for_retry(record.request_id_by_seq[seq_id], model_id, tenant)
+ record.request_id_by_seq[seq_id] = request_id
+ return request_id
+
+ future = self.futures.create(model_id, tenant)
+ record.request_id_by_seq[seq_id] = future.request_id
+ self._arrival_counter += 1
+ validation_error = payload.get("validation_error")
+ if validation_error is None:
+ try:
+ validate_batch_payload(op, payload, self.config)
+ except UserInputError as error:
+ validation_error = str(error)
+ model_queue.submit(
+ Command(
+ model_id=model_id,
+ seq_id=seq_id,
+ op=op,
+ payload={key: value for key, value in payload.items() if key != "validation_error"},
+ request_id=future.request_id,
+ arrival=self._arrival_counter,
+ validation_error=validation_error,
+ )
+ )
+ self._wake.set()
+ return future.request_id
+
+ def retrieve_future(self, tenant: str, request_id: str) -> RequestFuture | None:
+ return self.futures.get(request_id, tenant)
+
+ async def _run_batch(self, batch: BatchUnit) -> None:
+ # slot-contiguous order; outputs come back aligned to it
+ refs = sorted(batch.datums, key=lambda ref: ref.model_queue.slot)
+ slot_datums = [(ref.model_queue.slot, ref.datum) for ref in refs]
+ self._batch_counter += 1
+ forward = (
+ self.backend.forward_backward if batch.op == CommandOp.FORWARD_BACKWARD else self.backend.forward_only
+ )
+ try:
+ outputs = await forward(self._batch_counter, slot_datums, batch.loss_fn, batch.loss_fn_config)
+ except UserInputError as error:
+ await self._fail_batch(batch, str(error), "user")
+ return
+ if isinstance(outputs, dict) and "error" in outputs:
+ await self._fail_batch(batch, outputs["error"], "server")
+ return
+
+ assert len(outputs) == len(refs), f"unit returned {len(outputs)} outputs for {len(refs)} datums"
+ for ref, output in zip(refs, outputs, strict=True):
+ request = ref.request
+ if request.record_output(ref.local_index, output):
+ await self._finish_request(
+ ref.model_queue, request, {"op": request.command.op, "outputs": request.outputs}
+ )
+
+ async def _fail_batch(self, batch: BatchUnit, error: str, category: str) -> None:
+ requests = {ref.request.command.request_id: (ref.model_queue, ref.request) for ref in batch.datums}
+ for model_queue, pending in requests.values():
+ await self._finish_request(model_queue, pending, {"error": error, "error_category": category})
+
+ async def _run_barrier(self, barrier: BarrierUnit) -> None:
+ try:
+ outcomes = await self._dispatch_barrier_op(barrier)
+ except (UserInputError, OwnershipError) as error:
+ outcomes = [{"error": str(error), "error_category": "user"} for _ in barrier.entries]
+ for (model_queue, pending), outcome in zip(barrier.entries, outcomes, strict=True):
+ await self._finish_request(model_queue, pending, outcome)
+
+ async def _dispatch_barrier_op(self, barrier: BarrierUnit) -> list[dict]:
+ """Return one result or error per entry; only the caller settles futures and retires models."""
+ if barrier.op == CommandOp.OPTIM_STEP:
+ return await self._step_optimizers(barrier.entries)
+ ((model_queue, pending),) = barrier.entries # every other barrier is single-entry
+ record = self.models[model_queue.model_id]
+ payload = pending.command.payload
+ if barrier.op == CommandOp.SAVE_STATE:
+ return [await self._save_state(record, payload)]
+ if barrier.op == CommandOp.LOAD_STATE:
+ return [await self._load_state(record, payload)]
+ if barrier.op == CommandOp.SAVE_WEIGHTS_FOR_SAMPLER:
+ return [await self._save_weights_for_sampler(record, payload)]
+ raise UserInputError(f"unknown barrier op {barrier.op!r}")
+
+ async def _step_optimizers(self, entries: list) -> list[dict]:
+ adam_params_by_slot = {
+ model_queue.slot: pending.command.payload["adam_params"] for model_queue, pending in entries
+ }
+ slot_outcomes = await self.backend.optim_step(adam_params_by_slot)
+ outcomes = []
+ for model_queue, _ in entries:
+ outcome = slot_outcomes[model_queue.slot]
+ if "error" in outcome:
+ outcomes.append(outcome)
+ else:
+ outcomes.append({"op": "optim_step", "metrics": {key: float(value) for key, value in outcome.items()}})
+ return outcomes
+
+ async def _finish_request(self, model_queue, pending, outcome: dict) -> None:
+ if "error" in outcome:
+ category = outcome.get("error_category", "server")
+ if pending.command.op.requires_model_close_on_failure():
+ await self._close_model(
+ model_queue.model_id,
+ f"model training failed ({outcome['error']}); create a new model and restore from a checkpoint",
+ category,
+ )
+ return
+ self.futures.fail(pending.command.request_id, outcome["error"], category)
+ else:
+ self.futures.resolve(pending.command.request_id, outcome)
+ model_queue.finish(pending)
+
+ async def _save_state(self, record: ModelRecord, payload: dict) -> dict:
+ """Save parameters and optimizer state; call after optim_step to persist accumulated training work."""
+ name = payload["name"] or f"checkpoint-{payload['seq_id']:06d}"
+ validate_checkpoint_segment(name)
+ checkpoint_dir = resolve_checkpoint_dir(self.config.checkpoint_root, record.model_id, "weights", name)
+ if not payload["overwrite"] and os.path.exists(checkpoint_dir):
+ raise UserInputError(f"checkpoint {name!r} already exists; pass overwrite=True to replace it")
+ if os.path.isdir(checkpoint_dir) and not os.path.islink(checkpoint_dir):
+ raise UserInputError(f"cannot overwrite legacy checkpoint {name!r}; save under a new name")
+ failure = await self.backend.save_slot(
+ record.slot, checkpoint_dir, metadata=build_checkpoint_metadata(record, self.config)
+ )
+ if failure is not None:
+ return failure
+ return {"op": "save_state", "path": f"tinker://{record.model_id}/weights/{name}"}
+
+ async def _load_state(self, record: ModelRecord, payload: dict) -> dict:
+ source_id, kind, name = parse_tinker_path(payload["path"])
+ if kind != "weights":
+ raise UserInputError("cannot load sampler weights into a training model; use a save_state checkpoint")
+ checkpoint_dir = os.path.realpath(resolve_checkpoint_dir(self.config.checkpoint_root, source_id, kind, name))
+ source_tenant = payload.get("weights_access_token")
+ if source_tenant is None:
+ source_tenant = record.tenant
+ meta = read_checkpoint_metadata(checkpoint_dir, source_tenant, payload["path"])
+ validate_checkpoint_compatibility(meta, record, self.config, payload["path"])
+ failure = await self.backend.load_slot(
+ record.slot,
+ record.lora_rank,
+ record.lora_alpha,
+ ckpt_path=checkpoint_dir,
+ load_optimizer=payload["optimizer"],
+ )
+ if failure is not None:
+ return failure
+ return {"op": "load_state"}
+
+ async def _save_weights_for_sampler(self, record: ModelRecord, payload: dict) -> dict:
+ version = payload.get("sampler_path")
+ if version is None:
+ version = str(record.next_sampler_version)
+ record.next_sampler_version += 1
+ else:
+ validate_checkpoint_segment(version)
+ path = resolve_checkpoint_dir(self.config.checkpoint_root, record.model_id, "sampler_weights", version)
+ if os.path.exists(path):
+ raise UserInputError(f"sampler weights {version!r} already exist; save under a new name")
+ failure = await self.backend.export_slot(
+ record.slot,
+ record.lora_rank,
+ record.lora_alpha,
+ path,
+ metadata=build_checkpoint_metadata(record, self.config),
+ )
+ if failure is not None:
+ return failure
+ result = {
+ "op": "save_weights_for_sampler",
+ "path": f"tinker://{record.model_id}/sampler_weights/{version}",
+ }
+ if payload.get("sampler_path") is None:
+ # unnamed saves return a sampling session bound to the new version
+ result["sampling_session_id"] = self._new_sampling_session(
+ record.tenant, record.session_id, result["path"]
+ )
+ return result
+
+ def weights_info(self, tenant: str, tinker_path: str) -> dict:
+ """What the SDK needs to rebuild a training client from a checkpoint."""
+ model_id, kind, name = parse_tinker_path(tinker_path)
+ meta = read_checkpoint_metadata(
+ resolve_checkpoint_dir(self.config.checkpoint_root, model_id, kind, name), tenant, tinker_path
+ )
+ return {
+ "base_model": meta["base_model"],
+ "is_lora": True,
+ "lora_rank": meta["lora_rank"],
+ "train_attn": meta["train_attn"],
+ "train_mlp": meta["train_mlp"],
+ "train_unembed": meta["train_unembed"],
+ }
+
+ def create_sampling_session(self, tenant: str, payload: dict) -> str:
+ session = self._session_for(tenant, payload["session_id"])
+ base_model = payload.get("base_model")
+ if base_model is not None and base_model != self.config.base_model:
+ raise UserInputError(f"this gateway serves {self.config.base_model!r}, not {base_model!r}")
+ seq_id = validate_seq_id(payload["sampling_session_seq_id"], "sampling_session_seq_id", minimum=0)
+ if (previous := session.sampling_sessions_by_seq.get(seq_id)) is not None:
+ return previous
+ sampling_session_id = self._new_sampling_session(tenant, payload["session_id"], payload.get("model_path"))
+ session.sampling_sessions_by_seq[seq_id] = sampling_session_id
+ return sampling_session_id
+
+ def _new_sampling_session(self, tenant: str, session_id: str, model_path: str | None) -> str:
+ if model_path is not None:
+ resolve_sampler_checkpoint(self.config.checkpoint_root, tenant, model_path, self.config.base_model)
+ sampling_session_id = f"sampling-{uuid.uuid4().hex}"
+ self.sampling_sessions[sampling_session_id] = SamplingSessionRecord(
+ tenant=tenant, model_path=model_path, session_id=session_id
+ )
+ return sampling_session_id
+
+ def get_sampler(self, tenant: str, sampling_session_id: str) -> dict:
+ session = self.sampling_sessions.get(sampling_session_id)
+ if session is None:
+ raise UserInputError(f"unknown sampling session {sampling_session_id!r}")
+ if session.tenant != tenant:
+ raise OwnershipError("sampling session does not belong to this tenant")
+ return {
+ "sampler_id": sampling_session_id,
+ "base_model": self.config.base_model,
+ "model_path": session.model_path,
+ }
+
+ def submit_sample(self, tenant: str, payload: dict) -> tuple[str, list[str]]:
+ base_model = payload.get("base_model")
+ if base_model is not None and base_model != self.config.base_model:
+ raise UserInputError(f"this gateway serves {self.config.base_model!r}, not {base_model!r}")
+ model_path = payload.get("model_path")
+ sampling_session = None
+ if payload.get("sampling_session_id"):
+ sampling_session_id = payload["sampling_session_id"]
+ sampling_session = self.sampling_sessions.get(sampling_session_id)
+ if sampling_session is None:
+ raise UserInputError(
+ f"unknown sampling session {sampling_session_id!r}; create a sampling session first"
+ )
+ if sampling_session.tenant != tenant:
+ raise OwnershipError("sampling session does not belong to this tenant")
+ model_path = model_path or sampling_session.model_path
+ seq_id = validate_seq_id(payload["seq_id"], "seq_id", minimum=0)
+ if (previous := sampling_session.samples_by_seq.get(seq_id)) is not None:
+ request_id, sequence_ids = previous
+ request_id = self.futures.request_id_for_retry(request_id, model_path or "base", tenant)
+ sampling_session.samples_by_seq[seq_id] = (request_id, sequence_ids)
+ return request_id, sequence_ids
+ validate_sample_payload(payload, self.config)
+ lora_name, lora_path = (
+ resolve_sampler_checkpoint(self.config.checkpoint_root, tenant, model_path, self.config.base_model)
+ if model_path
+ else (None, None)
+ )
+ future = self.futures.create(model_path or "base", tenant)
+ sequence_ids = [f"seq-{uuid.uuid4().hex}" for _ in range(payload.get("num_samples", 1))]
+ task = asyncio.create_task(self._run_sample(future.request_id, sequence_ids, payload, lora_name, lora_path))
+ session_id = sampling_session.session_id if sampling_session is not None else None
+ self._sample_tasks[future.request_id] = (task, session_id)
+ task.add_done_callback(lambda _t, rid=future.request_id: self._sample_tasks.pop(rid, None))
+ task.add_done_callback(self._observe_background_task)
+ if sampling_session is not None:
+ sampling_session.samples_by_seq[seq_id] = (future.request_id, sequence_ids)
+ return future.request_id, sequence_ids
+
+ async def _run_sample(
+ self,
+ request_id: str,
+ sequence_ids: list[str],
+ payload: dict,
+ lora_name: str | None,
+ lora_path: str | None = None,
+ ) -> None:
+ try:
+ result = await self.backend.sample(payload, lora_name, lora_path)
+ except asyncio.CancelledError:
+ self.futures.fail(request_id, "cancelled", "user")
+ except UserInputError as error:
+ self.futures.fail(request_id, str(error), "user")
+ else:
+ if "error" in result:
+ self.futures.fail(request_id, result["error"], "server")
+ else:
+ for sequence_id, sequence in zip(sequence_ids, result["sequences"], strict=True):
+ sequence["sequence_id"] = sequence_id
+ self.futures.resolve(request_id, {"op": "sample", **result})
+
+ def cancel(self, tenant: str, request_id: str) -> None:
+ """Cancel an in-flight sampling future; training commands have no
+ cancel in the protocol and are ignored."""
+ if self.futures.get(request_id, tenant) is None:
+ return
+ entry = self._sample_tasks.get(request_id)
+ if entry is not None:
+ entry[0].cancel()
+
+ async def _run_lease_sweeper(self) -> None:
+ """Reclaim from stale sessions: cancel sampling, unload models, free
+ slots. Training state dies with the lease; only checkpoints survive."""
+ while True:
+ await asyncio.sleep(30)
+ await self._expire_sessions()
+
+ async def _expire_sessions(self) -> None:
+ async with self._trainer_lock:
+ now = time.monotonic()
+ expired_sessions = {
+ session_id
+ for session_id, session in self.sessions.items()
+ if now - session.last_heartbeat >= self.config.lease_timeout_s
+ }
+ for session_id in expired_sessions:
+ del self.sessions[session_id]
+ for sampling_session_id, record in list(self.sampling_sessions.items()):
+ if record.session_id in expired_sessions:
+ del self.sampling_sessions[sampling_session_id]
+
+ for request_id, (task, session_id) in list(self._sample_tasks.items()):
+ if session_id in expired_sessions:
+ logger.warning(f"lease expired for session of sample {request_id}; cancelling")
+ self.futures.fail(request_id, "lease expired", "user")
+ task.cancel()
+ for model_id, record in list(self.models.items()):
+ if record.session_id in expired_sessions:
+ logger.warning(f"lease expired for {model_id}; freeing slot {record.slot}")
+ await self._close_model(model_id, "lease expired", "user")
diff --git a/miles/tinker/core/types.py b/miles/tinker/core/types.py
new file mode 100644
index 0000000000..df24db93c0
--- /dev/null
+++ b/miles/tinker/core/types.py
@@ -0,0 +1,108 @@
+"""Shared commands, datums, and configuration for the Tinker gateway.
+
+An internal datum has tokens = model_input + target_tokens[-1:] and explicit target_tokens.
+Its target_len counts every label position, including prompt positions; loss inputs align to it.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class CommandOp(str, Enum):
+ FORWARD_BACKWARD = "forward_backward"
+ FORWARD_ONLY = "forward_only"
+ OPTIM_STEP = "optim_step"
+ SAVE_STATE = "save_state"
+ LOAD_STATE = "load_state"
+ SAVE_WEIGHTS_FOR_SAMPLER = "save_weights_for_sampler"
+
+ def is_batch(self) -> bool:
+ """Batch ops pack into BatchUnits; every other op is a barrier."""
+ return self in (CommandOp.FORWARD_BACKWARD, CommandOp.FORWARD_ONLY)
+
+ def requires_model_close_on_failure(self) -> bool:
+ return self in (CommandOp.FORWARD_BACKWARD, CommandOp.OPTIM_STEP, CommandOp.LOAD_STATE)
+
+
+# wire loss_fn_inputs key -> internal datum key
+LOSS_INPUT_KEYS = {"weights": "weights", "advantages": "advantages", "logprobs": "sampling_logprobs"}
+
+# reject missing loss inputs before they enter a shared batch
+LOSS_FN_INPUTS = {
+ "cross_entropy": ("weights",),
+ "importance_sampling": ("logprobs", "advantages"),
+ "ppo": ("logprobs", "advantages"),
+ "cispo": ("logprobs", "advantages"),
+ "dro": ("logprobs", "advantages"),
+}
+
+
+class UserInputError(Exception):
+ """Rejected request content; fails the future with category user."""
+
+
+class OwnershipError(Exception):
+ """model/checkpoint does not belong to the caller's tenant."""
+
+
+@dataclass
+class GatewayConfig:
+ base_model: str
+ n_slots: int
+ checkpoint_root: str
+ vocab_size: int
+ max_datums_per_request: int = 1024
+ max_tokens_per_datum: int = 32768
+ max_tokens_per_request: int = 4_000_000
+ max_samples_per_request: int = 64
+ max_lora_rank: int = 32 # the slot capacity the server was built with (--lora-rank)
+ lora_alpha: float | None = None # None: 2 * rank
+ # what the server-wide adapter layout trains; create_model rejects deviations
+ trains_attn: bool = True
+ trains_mlp: bool = True
+ trains_unembed: bool = False
+ lease_timeout_s: float = 300.0 # sessions stale beyond this lose their sampling, models, and slots
+ batch_token_budget: int = 262_144 # packing bound per BatchUnit
+
+
+@dataclass
+class Command:
+ model_id: str
+ seq_id: int
+ op: CommandOp
+ payload: dict
+ request_id: str
+ arrival: int # global submit order for selecting the scheduler's seed
+ validation_error: str | None = None
+
+
+@dataclass
+class ModelRecord:
+ model_id: str
+ session_id: str
+ tenant: str
+ slot: int
+ base_model: str
+ lora_rank: int
+ lora_alpha: float
+ create_request_id: str
+ request_id_by_seq: dict[int, str] = field(default_factory=dict)
+ slot_initialized: bool = False
+ # failed publications burn their version number, leaving gaps
+ next_sampler_version: int = 1
+
+
+@dataclass
+class SessionRecord:
+ tenant: str
+ last_heartbeat: float
+ models_by_seq: dict[int, tuple[str, str]] = field(default_factory=dict)
+ sampling_sessions_by_seq: dict[int, str] = field(default_factory=dict)
+
+
+@dataclass
+class SamplingSessionRecord:
+ tenant: str
+ model_path: str | None
+ session_id: str
+ samples_by_seq: dict[int, tuple[str, list[str]]] = field(default_factory=dict)
diff --git a/miles/tinker/core/utils.py b/miles/tinker/core/utils.py
new file mode 100644
index 0000000000..b73c99706c
--- /dev/null
+++ b/miles/tinker/core/utils.py
@@ -0,0 +1,74 @@
+import hashlib
+import json
+import os
+from pathlib import Path
+
+from miles.tinker.core.input_validation import validate_checkpoint_metadata, validate_checkpoint_segment
+from miles.tinker.core.types import GatewayConfig, ModelRecord, OwnershipError, UserInputError
+
+
+def parse_tinker_path(path: str) -> tuple[str, str, str]:
+ if not path.startswith("tinker://"):
+ raise UserInputError(f"not a tinker path: {path!r}")
+ parts = path.removeprefix("tinker://").split("/")
+ if len(parts) != 3 or parts[1] not in ("weights", "sampler_weights"):
+ raise UserInputError(f"malformed tinker path: {path!r}")
+ for segment in parts:
+ validate_checkpoint_segment(segment)
+ return parts[0], parts[1], parts[2]
+
+
+def resolve_checkpoint_dir(checkpoint_root: str, model_id: str, kind: str, name: str) -> str:
+ root = os.path.realpath(checkpoint_root)
+ path = f"{root}/{model_id}/{kind}/{name}"
+ resolved = os.path.realpath(path)
+ assert resolved.startswith(root + os.sep), f"checkpoint path {resolved!r} escapes {root!r}"
+ # saves replace the public link, not the version it currently points to
+ return path
+
+
+def build_checkpoint_metadata(record: ModelRecord, config: GatewayConfig) -> dict:
+ return {
+ # the digest proves ownership without persisting the bearer credential itself
+ "tenant_digest": _tenant_digest(record.tenant),
+ "base_model": record.base_model,
+ "lora_rank": record.lora_rank,
+ "lora_alpha": record.lora_alpha,
+ "train_attn": config.trains_attn,
+ "train_mlp": config.trains_mlp,
+ "train_unembed": config.trains_unembed,
+ }
+
+
+def read_checkpoint_metadata(checkpoint_dir: str, tenant: str, shown_path: str) -> dict:
+ meta_file = Path(checkpoint_dir) / "META.json"
+ if not meta_file.exists():
+ raise UserInputError(f"unknown checkpoint {shown_path!r}")
+ try:
+ meta = json.loads(meta_file.read_text())
+ except (OSError, UnicodeError, json.JSONDecodeError) as error:
+ raise UserInputError(f"cannot read checkpoint {shown_path!r}: {error}") from error
+ validate_checkpoint_metadata(meta, shown_path)
+ if meta["tenant_digest"] != _tenant_digest(tenant):
+ raise OwnershipError(f"checkpoint {shown_path!r} does not belong to this tenant")
+ return meta
+
+
+def resolve_sampler_checkpoint(checkpoint_root: str, tenant: str, model_path: str, base_model: str) -> tuple[str, str]:
+ """Return the adapter name and directory so engines can reload evicted snapshots."""
+ model_id, kind, name = parse_tinker_path(model_path)
+ if kind != "sampler_weights":
+ raise UserInputError(
+ f"sampling from training checkpoints is not supported: {model_path!r}; use save_weights_for_sampler()"
+ )
+ checkpoint_dir = resolve_checkpoint_dir(checkpoint_root, model_id, "sampler_weights", name)
+ meta = read_checkpoint_metadata(checkpoint_dir, tenant, model_path)
+ if meta["base_model"] != base_model:
+ raise UserInputError(
+ f"checkpoint {model_path!r} uses base_model={meta['base_model']!r}; this server serves {base_model!r}"
+ )
+ return f"{model_id}@{name}", checkpoint_dir
+
+
+def _tenant_digest(tenant: str) -> str:
+ return hashlib.sha256(tenant.encode()).hexdigest()
diff --git a/miles/tinker/runtime.py b/miles/tinker/runtime.py
new file mode 100644
index 0000000000..b584965a8f
--- /dev/null
+++ b/miles/tinker/runtime.py
@@ -0,0 +1,216 @@
+"""Translate gateway datums to trainer batches and sampling requests to SGLang."""
+
+import asyncio
+
+import httpx
+
+from miles.ray.rollout.train_data_conversion import ROLLOUT_DATA_VALUE_SPEC
+from miles.tinker.core.types import UserInputError
+from miles.utils import object_store
+from miles.utils.http_utils import post
+from tinker.types.sample_response import MASK_LOGPROB
+
+# internal datum key -> trainer batch key
+DATUM_TO_BATCH_KEYS = {"weights": "loss_weights", "advantages": "advantages", "sampling_logprobs": "rollout_log_probs"}
+
+
+def _pad_to_dp_multiple(slot_datums: list, dp_size: int) -> list:
+ """Equalize DP shares with zero-mask datums; their outputs are dropped after the loss pass."""
+ remainder = len(slot_datums) % dp_size
+ if remainder == 0:
+ return slot_datums
+ slot, datum = slot_datums[-1]
+ filler = dict(datum) | {"padding": True}
+ return slot_datums + [(slot, filler)] * (dp_size - remainder)
+
+
+def _build_train_data(slot_datums: list) -> dict:
+ """Miles response_lengths select label positions here, including prompt targets."""
+ datums = [datum for _, datum in slot_datums]
+ train_data = {
+ "tokens": [datum["tokens"] for datum in datums],
+ "target_tokens": [datum["target_tokens"] for datum in datums],
+ "loss_masks": [[0 if datum.get("padding") else 1] * datum["target_len"] for datum in datums],
+ "response_lengths": [datum["target_len"] for datum in datums],
+ "total_lengths": [len(datum["tokens"]) for datum in datums],
+ "sample_indices": list(range(len(datums))),
+ "adapter_slots": [slot for slot, _ in slot_datums],
+ "dynamic_global_batch_size": len(datums),
+ }
+ for datum_key, batch_key in DATUM_TO_BATCH_KEYS.items():
+ if datum_key in datums[0]:
+ train_data[batch_key] = [datum[datum_key] for datum in datums]
+ return train_data
+
+
+class MilesBackend:
+ def __init__(self, trainer, router_url: str, dp_size: int = 1) -> None:
+ self.trainer = trainer
+ self.router_url = router_url
+ self.dp_size = dp_size
+
+ def trainer_dead(self) -> bool:
+ return self.trainer.has_errored_cell()
+
+ async def load_slot(
+ self, slot: int, rank: int, alpha: float, ckpt_path: str | None = None, load_optimizer: bool = True
+ ) -> dict | None:
+ return _slot_failure(
+ await self.trainer.load_slot(slot, rank, alpha, ckpt_path=ckpt_path, load_optimizer=load_optimizer)
+ )
+
+ async def unload_slot(self, slot: int) -> dict | None:
+ return _slot_failure(await self.trainer.unload_slot(slot))
+
+ async def forward_backward(
+ self, batch_id: int, slot_datums: list, loss_fn: str, loss_fn_config: dict
+ ) -> list[dict] | dict:
+ return await self._execute_batch("forward_backward", batch_id, slot_datums, loss_fn, loss_fn_config)
+
+ async def forward_only(
+ self, batch_id: int, slot_datums: list, loss_fn: str, loss_fn_config: dict
+ ) -> list[dict] | dict:
+ return await self._execute_batch("forward_only", batch_id, slot_datums, loss_fn, loss_fn_config)
+
+ async def _execute_batch(
+ self, method: str, batch_id: int, slot_datums: list, loss_fn: str, loss_fn_config: dict
+ ) -> list[dict] | dict:
+ train_data = _build_train_data(_pad_to_dp_multiple(slot_datums, self.dp_size))
+ train_data["loss_fn"] = loss_fn
+ train_data["loss_fn_config"] = loss_fn_config
+ worker_results = await self._call_trainer(method, batch_id, train_data)
+ by_index: dict[int, dict] = {}
+ for worker_result in worker_results:
+ if "error" in worker_result:
+ return worker_result
+ for datum_output in worker_result["per_datum"]:
+ index = int(datum_output["sample_index"])
+ if index not in by_index:
+ by_index[index] = {
+ "loss": float(datum_output["loss"]),
+ "logprobs": datum_output["logprobs"].tolist(),
+ }
+ return [by_index[index] for index in range(len(slot_datums))]
+
+ async def _call_trainer(self, method: str, batch_id: int, train_data: dict) -> list:
+ store = object_store.get_instance()
+ data_ref = store.put(value=train_data, value_spec=ROLLOUT_DATA_VALUE_SPEC)
+ try:
+ return await getattr(self.trainer, method)(batch_id=batch_id, data_ref=data_ref)
+ finally:
+ store.remove(data_ref)
+
+ async def optim_step(self, adam_params_by_slot: dict[int, dict]) -> dict[int, dict]:
+ worker_results = await self.trainer.optim_step(adam_params_by_slot=adam_params_by_slot)
+ return worker_results[0]
+
+ async def save_slot(self, slot: int, path: str, metadata: dict | None = None) -> dict | None:
+ return _slot_failure(await self.trainer.save_slot(slot=slot, path=path, metadata=metadata))
+
+ async def export_slot(
+ self, slot: int, rank: int, alpha: float, path: str, metadata: dict | None = None
+ ) -> dict | None:
+ return _slot_failure(
+ await self.trainer.export_slot(slot=slot, rank=rank, alpha=alpha, path=path, metadata=metadata)
+ )
+
+ # -------- sampling --------
+
+ async def sample(self, payload: dict, lora_name: str | None, lora_path: str | None = None) -> dict:
+ request = self._generate_request(payload, lora_name, lora_path)
+ try:
+ responses = await asyncio.gather(
+ *[
+ post(f"{self.router_url}/generate", _with_sample_seed(request, index))
+ for index in range(payload["num_samples"])
+ ]
+ )
+ except httpx.HTTPError as error:
+ return {"error": str(error)}
+ sequences = [_to_sequence(response) for response in responses]
+ for sequence in sequences:
+ if "error" in sequence:
+ return sequence
+ result = {"sequences": sequences}
+ if payload["prompt_logprobs"]:
+ result["prompt_logprobs"] = _prompt_logprobs(responses[0])
+ if payload["topk_prompt_logprobs"]:
+ result["topk_prompt_logprobs"] = _topk_prompt_logprobs(responses[0], payload["topk_prompt_logprobs"])
+ return result
+
+ def _generate_request(self, payload: dict, lora_name: str | None, lora_path: str | None = None) -> dict:
+ params = payload["sampling_params"]
+ max_tokens = params.get("max_tokens")
+ if max_tokens is None:
+ raise UserInputError("sampling_params.max_tokens is required")
+ sampling_params = {
+ "max_new_tokens": max_tokens,
+ "temperature": params.get("temperature", 1.0),
+ "top_p": params.get("top_p", 1.0),
+ "top_k": params.get("top_k", -1),
+ }
+ if params.get("seed") is not None:
+ sampling_params["sampling_seed"] = params["seed"]
+ stop = params.get("stop")
+ if stop is not None:
+ if stop == []:
+ # tinker defines stop=[] as disabling every stop token, EOS included
+ sampling_params["ignore_eos"] = True
+ elif isinstance(stop, list) and isinstance(stop[0], int):
+ sampling_params["stop_token_ids"] = stop
+ else:
+ sampling_params["stop"] = stop
+ request = {"input_ids": payload["prompt_tokens"], "sampling_params": sampling_params, "return_logprob": True}
+ if payload["prompt_logprobs"] or payload["topk_prompt_logprobs"]:
+ request["logprob_start_len"] = 0
+ if payload["topk_prompt_logprobs"]:
+ request["top_logprobs_num"] = payload["topk_prompt_logprobs"]
+ if lora_name is not None:
+ request["lora_path"] = lora_name
+ if lora_path is not None:
+ # request-carried backfill source: the engine refills an evicted version itself
+ request["lora_backfill_paths"] = {lora_name: lora_path}
+ return request
+
+
+def _slot_failure(worker_results: list) -> dict | None:
+ return next((result for result in worker_results if result is not None), None)
+
+
+def _with_sample_seed(request: dict, index: int) -> dict:
+ """Give each sample a distinct seed when the caller pins the request seed."""
+ request = dict(request)
+ params = dict(request["sampling_params"])
+ if (seed := params.get("sampling_seed")) is not None:
+ params["sampling_seed"] = seed + index
+ request["sampling_params"] = params
+ return request
+
+
+def _prompt_logprobs(response: dict) -> list[float]:
+ entries = response["meta_info"]["input_token_logprobs"]
+ return [float("nan") if entry[0] is None else float(entry[0]) for entry in entries]
+
+
+def _topk_prompt_logprobs(response: dict, k: int) -> dict:
+ token_ids, logprobs = [], []
+ for position in response["meta_info"]["input_top_logprobs"]:
+ candidates = position or []
+ candidate_token_ids = [entry[1] for entry in candidates][:k]
+ candidate_logprobs = [float(entry[0]) for entry in candidates][:k]
+ token_ids.append(candidate_token_ids + [0] * (k - len(candidate_token_ids)))
+ logprobs.append(candidate_logprobs + [MASK_LOGPROB] * (k - len(candidate_logprobs)))
+ return {"token_ids": token_ids, "logprobs": logprobs}
+
+
+def _to_sequence(response: dict) -> dict:
+ output_token_logprobs = response["meta_info"]["output_token_logprobs"]
+ finish = response["meta_info"]["finish_reason"]["type"]
+ if finish == "abort":
+ # a truncated sequence must fail the request, not pass as a completed sample
+ return {"error": "the engine aborted this sample; resubmit the request"}
+ return {
+ "tokens": [entry[1] for entry in output_token_logprobs],
+ "logprobs": [entry[0] for entry in output_token_logprobs],
+ "stop_reason": "length" if finish == "length" else "stop",
+ }
diff --git a/miles/tinker/server/__init__.py b/miles/tinker/server/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/miles/tinker/server/app.py b/miles/tinker/server/app.py
new file mode 100644
index 0000000000..5012e60c0e
--- /dev/null
+++ b/miles/tinker/server/app.py
@@ -0,0 +1,184 @@
+"""Tinker SDK endpoints with JSON/protobuf translation and tenant authentication."""
+
+import asyncio
+from contextlib import suppress
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response
+
+from miles.tinker.core.future import FAILED, PENDING
+from miles.tinker.core.service import TinkerService
+from miles.tinker.core.types import OwnershipError, UserInputError
+from miles.tinker.server.encoding import (
+ decode_command,
+ decode_sample_request,
+ render_result,
+ validate_create_model,
+ validate_create_sampling_session,
+)
+from miles.tinker.server.proto_codec import (
+ PROTO_CONTENT_TYPE,
+ PROTO_ENCODERS,
+ decode_forward_backward_request,
+ maybe_decompress,
+)
+
+RETRIEVE_LONG_POLL_S = 30.0
+
+COMMAND_ROUTES = {
+ "/api/v1/optim_step": "optim_step",
+ "/api/v1/save_weights": "save_state",
+ "/api/v1/load_weights": "load_state",
+ "/api/v1/save_weights_for_sampler": "save_weights_for_sampler",
+}
+
+
+def _tenant(request: Request) -> str:
+ """Require X-API-Key or a bearer token to establish the tenant identity."""
+ key = (
+ request.headers.get("x-api-key")
+ or (request.headers.get("authorization") or "").removeprefix("Bearer ").strip()
+ )
+ if not key:
+ raise UserInputError("missing API key: send it in the X-API-Key header")
+ return key
+
+
+def build_app(service: TinkerService) -> FastAPI:
+ app = FastAPI()
+
+ @app.exception_handler(UserInputError)
+ async def _user_error(request: Request, error: UserInputError):
+ return JSONResponse(status_code=400, content={"error": str(error)})
+
+ @app.exception_handler(OwnershipError)
+ async def _ownership_error(request: Request, error: OwnershipError):
+ return JSONResponse(status_code=403, content={"error": str(error)})
+
+ @app.get("/api/v1/healthz")
+ async def healthz():
+ return {"status": "ok"}
+
+ @app.post("/api/v1/client/config")
+ @app.get("/api/v1/client/config")
+ async def client_config():
+ return {}
+
+ @app.post("/api/v1/client/dynamic_config")
+ @app.get("/api/v1/client/dynamic_config")
+ async def client_dynamic_config():
+ return {}
+
+ @app.post("/api/v1/telemetry")
+ async def telemetry():
+ return {"status": "accepted"}
+
+ @app.post("/api/v1/create_session")
+ async def create_session(request: Request):
+ session_id = service.create_session(_tenant(request))
+ return {"type": "create_session", "session_id": session_id}
+
+ @app.post("/api/v1/session_heartbeat")
+ async def session_heartbeat(request: Request):
+ tenant = _tenant(request)
+ payload = await request.json()
+ if not service.heartbeat(tenant, payload["session_id"]):
+ return JSONResponse(status_code=410, content={"error": "unknown or expired session"})
+ return {"type": "session_heartbeat"}
+
+ @app.post("/api/v1/get_server_capabilities")
+ @app.get("/api/v1/get_server_capabilities")
+ async def get_server_capabilities():
+ return {"supported_models": [{"model_name": service.config.base_model, "trainable": True, "sampleable": True}]}
+
+ @app.post("/api/v1/create_model")
+ async def create_model(request: Request):
+ payload = await request.json()
+ validate_create_model(payload)
+ request_id, model_id = service.create_model(_tenant(request), payload)
+ return {"request_id": request_id, "model_id": model_id}
+
+ @app.post("/api/v1/get_info")
+ async def get_info(request: Request):
+ payload = await request.json()
+ record = service.get_model(_tenant(request), payload["model_id"])
+ return {
+ "type": "get_info",
+ "model_id": record.model_id,
+ "model_name": record.base_model,
+ "is_lora": True,
+ "lora_rank": record.lora_rank,
+ "model_data": {"model_name": record.base_model},
+ }
+
+ @app.post("/api/v1/forward_backward")
+ async def forward_backward(request: Request):
+ if PROTO_CONTENT_TYPE in request.headers.get("content-type", ""):
+ body = maybe_decompress(await request.body(), request.headers.get("content-encoding"))
+ op, payload = decode_forward_backward_request(body)
+ else:
+ op, payload = decode_command("forward_backward", await request.json())
+ request_id = service.submit(_tenant(request), op, payload)
+ return {"request_id": request_id, "model_id": payload["model_id"]}
+
+ for route, route_op in COMMAND_ROUTES.items():
+
+ def _command_handler(route_op: str):
+ async def command(request: Request):
+ op, payload = decode_command(route_op, await request.json())
+ request_id = service.submit(_tenant(request), op, payload)
+ return {"request_id": request_id, "model_id": payload["model_id"]}
+
+ return command
+
+ app.post(route)(_command_handler(route_op))
+
+ @app.post("/api/v1/weights_info")
+ async def weights_info(request: Request):
+ payload = await request.json()
+ return service.weights_info(_tenant(request), payload["tinker_path"])
+
+ @app.post("/api/v1/retrieve_future")
+ async def retrieve_future(request: Request):
+ payload = await request.json()
+ future = service.retrieve_future(_tenant(request), payload["request_id"])
+ if future is None:
+ return JSONResponse(status_code=410, content={"error": "unknown or expired request"})
+ if future.state == PENDING:
+ # long-poll: the SDK retries try_again with no backoff, so answer on
+ # settlement instead of turning every pending future into a busy-poll
+ with suppress(asyncio.TimeoutError):
+ await asyncio.wait_for(future.settled.wait(), timeout=RETRIEVE_LONG_POLL_S)
+ if future.state == PENDING:
+ return {"type": "try_again", "queue_state": "active"}
+ if future.state == FAILED:
+ return {"error": future.error, "category": future.error_category}
+ encoder = PROTO_ENCODERS.get(future.result["op"])
+ if encoder is not None and PROTO_CONTENT_TYPE in request.headers.get("accept", ""):
+ return Response(content=encoder(future.result), media_type=PROTO_CONTENT_TYPE)
+ return render_result(future.result)
+
+ @app.post("/api/v1/cancel_future")
+ async def cancel_future(request: Request):
+ payload = await request.json()
+ service.cancel(_tenant(request), payload["request_id"])
+ return {"status": "ok"}
+
+ @app.post("/api/v1/create_sampling_session")
+ async def create_sampling_session(request: Request):
+ payload = await request.json()
+ validate_create_sampling_session(payload)
+ sampling_session_id = service.create_sampling_session(_tenant(request), payload)
+ return {"type": "create_sampling_session", "sampling_session_id": sampling_session_id}
+
+ @app.get("/api/v1/samplers/{sampling_session_id}")
+ async def get_sampler(sampling_session_id: str, request: Request):
+ return service.get_sampler(_tenant(request), sampling_session_id)
+
+ @app.post("/api/v1/asample")
+ async def asample(request: Request):
+ payload = decode_sample_request(await request.json())
+ request_id, sequence_ids = service.submit_sample(_tenant(request), payload)
+ return {"request_id": request_id, "sample_sequence_ids": sequence_ids}
+
+ return app
diff --git a/miles/tinker/server/encoding.py b/miles/tinker/server/encoding.py
new file mode 100644
index 0000000000..6abb1b34fb
--- /dev/null
+++ b/miles/tinker/server/encoding.py
@@ -0,0 +1,218 @@
+"""Translate SDK JSON requests and results.
+
+Datums encode input x and explicit target labels t as x + [t[-1]],
+so each output scores logprob(t[i] | x[0..i])."""
+
+import math
+
+import pydantic
+
+from miles.tinker.core.input_validation import validate_save_options
+from miles.tinker.core.types import LOSS_INPUT_KEYS, UserInputError
+from tinker import types as tinker_types
+from tinker.types.sample_response import MASK_LOGPROB
+
+# materialized at the boundary so core and the executor can require every key
+ADAM_PARAM_DEFAULTS = tinker_types.AdamParams().model_dump()
+
+# JSON-wire request models; forward_backward travels as protobuf and its datums
+# are checked in build_datum instead
+REQUEST_TYPES = {
+ "optim_step": tinker_types.OptimStepRequest,
+ "save_state": tinker_types.SaveWeightsRequest,
+ "load_state": tinker_types.LoadWeightsRequest,
+ "save_weights_for_sampler": tinker_types.SaveWeightsForSamplerRequest,
+}
+
+
+def validate_against_sdk(request_type, payload: dict) -> None:
+ """The pinned SDK's request models are the wire contract; reject what they reject."""
+ try:
+ request_type.model_validate(payload)
+ except pydantic.ValidationError as error:
+ first = error.errors()[0]
+ location = ".".join(str(part) for part in first["loc"])
+ raise UserInputError(f"{location or 'request'}: {first['msg']}") from None
+
+
+def validate_create_model(payload: dict) -> None:
+ validate_against_sdk(tinker_types.CreateModelRequest, payload)
+
+
+def validate_create_sampling_session(payload: dict) -> None:
+ validate_against_sdk(tinker_types.CreateSamplingSessionRequest, payload)
+
+
+def decode_command(op: str, payload: dict) -> tuple[str, dict]:
+ """Decode content errors into ordered failures when the envelope is identifiable."""
+ try:
+ envelope = {"model_id": payload["model_id"], "seq_id": payload["seq_id"]}
+ except KeyError as error:
+ raise UserInputError(f"missing command envelope field {error.args[0]!r}") from None
+ try:
+ return _decode_command(op, payload, envelope)
+ except (UserInputError, KeyError, TypeError, ValueError, IndexError) as error:
+ if op == "forward_backward" and payload.get("forward_only"):
+ op = "forward_only"
+ return op, envelope | {"validation_error": str(error)}
+
+
+def _decode_command(op: str, payload: dict, decoded: dict) -> tuple[str, dict]:
+ if (request_type := REQUEST_TYPES.get(op)) is not None:
+ validate_against_sdk(request_type, payload)
+ if op == "forward_backward":
+ fb_input = payload["forward_backward_input"]
+ datums = [
+ (model_input_tokens(datum["model_input"]), _decode_inputs(datum["loss_fn_inputs"]))
+ for datum in fb_input["data"]
+ ]
+ decoded |= {
+ "datums": [build_datum(tokens, inputs, i) for i, (tokens, inputs) in enumerate(datums)],
+ "loss_fn": fb_input["loss_fn"],
+ "loss_fn_config": fb_input.get("loss_fn_config") or {},
+ }
+ return ("forward_only" if payload.get("forward_only") else op), decoded
+ if op == "optim_step":
+ return op, decoded | {"adam_params": {**ADAM_PARAM_DEFAULTS, **payload["adam_params"]}}
+ if op == "save_state":
+ validate_save_options(payload)
+ return op, decoded | {"name": payload.get("path"), "overwrite": bool(payload.get("overwrite", False))}
+ if op == "load_state":
+ return op, decoded | {
+ "path": payload["path"],
+ "optimizer": payload["optimizer"],
+ "weights_access_token": payload.get("weights_access_token"),
+ }
+ if op == "save_weights_for_sampler":
+ validate_save_options(payload)
+ return op, decoded | {"sampler_path": payload.get("path")}
+ raise UserInputError(f"unknown command op {op!r}")
+
+
+def model_input_tokens(model_input: dict) -> list[int]:
+ tokens: list[int] = []
+ for chunk in model_input["chunks"]:
+ if chunk.get("type") != "encoded_text":
+ raise UserInputError(f"unsupported model_input chunk type: {chunk.get('type')}")
+ tokens.extend(chunk["tokens"])
+ return tokens
+
+
+def build_datum(input_tokens: list[int], inputs: dict[str, list], index: int) -> dict:
+ """One decoded datum (token list + loss_fn_inputs lists) -> internal datum."""
+ unknown = set(inputs) - set(LOSS_INPUT_KEYS) - {"target_tokens"}
+ if unknown:
+ raise UserInputError(f"datum {index}: unknown loss_fn_inputs {sorted(unknown)}")
+ for name, values in inputs.items():
+ if any(isinstance(value, (list, tuple)) for value in values):
+ raise UserInputError(
+ f"datum {index}: loss_fn_inputs[{name!r}] must be 1-D; multi-target inputs are not supported"
+ )
+ targets = list(inputs["target_tokens"])
+ if len(targets) != len(input_tokens):
+ raise UserInputError(
+ f"datum {index}: target_tokens length {len(targets)} != model_input length {len(input_tokens)}"
+ )
+ datum = {"tokens": input_tokens + targets[-1:], "target_len": len(targets), "target_tokens": targets}
+ for wire_key, datum_key in LOSS_INPUT_KEYS.items():
+ if wire_key in inputs:
+ datum[datum_key] = [float(value) for value in inputs[wire_key]]
+ return datum
+
+
+def _decode_inputs(loss_fn_inputs: dict) -> dict[str, list]:
+ return {name: tensor_data_to_list(value) for name, value in loss_fn_inputs.items()}
+
+
+def tensor_data_to_list(tensor_data) -> list:
+ if isinstance(tensor_data, list):
+ return tensor_data
+ if not isinstance(tensor_data, dict):
+ raise UserInputError(f"expected TensorData, got {type(tensor_data).__name__}")
+ if len(tensor_data.get("shape") or []) > 1:
+ raise UserInputError("multi-target inputs are not supported; loss_fn_inputs must be 1-D")
+ if tensor_data.get("sparse_crow_indices") is not None:
+ return _dense_from_csr(tensor_data)
+ data = tensor_data.get("data")
+ if data is None:
+ raise UserInputError("TensorData without data")
+ return list(data)
+
+
+def _dense_from_csr(tensor_data: dict) -> list:
+ (length,) = tensor_data["shape"]
+ if len(tensor_data["sparse_crow_indices"]) != 2:
+ raise UserInputError("1-D CSR expected")
+ dense = [0] * length
+ for col, value in zip(tensor_data["sparse_col_indices"], tensor_data["data"], strict=True):
+ dense[col] = value
+ return dense
+
+
+def decode_sample_request(payload: dict) -> dict:
+ validate_against_sdk(tinker_types.SampleRequest, payload)
+ return {
+ "model_path": payload.get("model_path"),
+ "base_model": payload.get("base_model"),
+ "sampling_session_id": payload.get("sampling_session_id"),
+ "seq_id": payload.get("seq_id"),
+ "num_samples": payload.get("num_samples", 1),
+ "prompt_tokens": model_input_tokens(payload["prompt"]),
+ "sampling_params": payload.get("sampling_params") or {},
+ "prompt_logprobs": bool(payload.get("prompt_logprobs")),
+ "topk_prompt_logprobs": payload.get("topk_prompt_logprobs", 0) or 0,
+ }
+
+
+# -------- result rendering (JSON; proto_codec renders the binary forms) --------
+
+
+def render_result(result: dict) -> dict:
+ op = result["op"]
+ if op in ("forward_backward", "forward_only"):
+ outputs = result["outputs"]
+ return {
+ "type": "forward_backward",
+ "loss_fn_output_type": "ArrayRecord",
+ "loss_fn_outputs": [
+ {"loss:sum": _tensor_json([output["loss"]]), "logprobs": _tensor_json(output["logprobs"])}
+ for output in outputs
+ ],
+ "metrics": {"loss:sum": float(sum(output["loss"] for output in outputs))},
+ }
+ if op == "sample":
+ rendered = {"type": "sample", "sequences": result["sequences"]}
+ if result.get("prompt_logprobs") is not None:
+ rendered["prompt_logprobs"] = [
+ None if math.isnan(logprob) else logprob for logprob in result["prompt_logprobs"]
+ ]
+ if result.get("topk_prompt_logprobs") is not None:
+ topk = result["topk_prompt_logprobs"]
+ rendered["topk_prompt_logprobs"] = [
+ [
+ (token_id, logprob)
+ for token_id, logprob in zip(ids, probs, strict=True)
+ if (token_id, logprob) != (0, MASK_LOGPROB)
+ ]
+ or None
+ for ids, probs in zip(topk["token_ids"], topk["logprobs"], strict=True)
+ ]
+ return rendered
+ if op == "create_model":
+ return {"type": "create_model", "model_id": result["model_id"]}
+ if op == "save_state":
+ return {"type": "save_weights", "path": result["path"]}
+ if op == "save_weights_for_sampler":
+ rendered = {"type": "save_weights_for_sampler", "path": result["path"]}
+ if "sampling_session_id" in result:
+ rendered["sampling_session_id"] = result["sampling_session_id"]
+ return rendered
+ if op == "load_state":
+ return {"type": "load_weights"}
+ if op == "optim_step":
+ return {"type": "optim_step", "metrics": result["metrics"]}
+ raise AssertionError(f"unrenderable result op {op!r}")
+
+
+def _tensor_json(values: list[float]) -> dict:
+ return {"dtype": "float32", "shape": [len(values)], "data": values}
diff --git a/miles/tinker/server/proto_codec.py b/miles/tinker/server/proto_codec.py
new file mode 100644
index 0000000000..3dec23f645
--- /dev/null
+++ b/miles/tinker/server/proto_codec.py
@@ -0,0 +1,138 @@
+"""Translate Tinker SDK protobuf requests and results, including compressed requests."""
+
+import numpy as np
+
+from miles.tinker.core.types import UserInputError
+from miles.tinker.server.encoding import build_datum
+from tinker.proto import tinker_public_pb2 as public_pb
+
+PROTO_CONTENT_TYPE = "application/x-protobuf"
+
+_PROTO_DTYPE_TO_NUMPY = {
+ public_pb.DTYPE_FLOAT32: np.dtype(np.float32),
+ public_pb.DTYPE_INT64: np.dtype(np.int64),
+ public_pb.DTYPE_INT32: np.dtype(np.int32),
+}
+_STOP_REASON_TO_PROTO = {
+ "stop": public_pb.STOP_REASON_STOP,
+ "length": public_pb.STOP_REASON_LENGTH,
+}
+
+
+def maybe_decompress(body: bytes, content_encoding: str | None) -> bytes:
+ if content_encoding == "zstd":
+ import zstandard
+
+ return zstandard.ZstdDecompressor().decompress(body)
+ return body
+
+
+def decode_forward_backward_request(body: bytes) -> tuple[str, dict]:
+ """ForwardBackwardRequest proto -> (op, internal payload)."""
+ message = public_pb.ForwardBackwardRequest()
+ message.ParseFromString(body)
+ op = "forward_only" if message.forward_only else "forward_backward"
+ try:
+ return op, _decode_forward_backward(message)
+ except (UserInputError, KeyError, TypeError, ValueError, IndexError) as error:
+ return op, {"model_id": message.model_id, "seq_id": message.seq_id, "validation_error": str(error)}
+
+
+def _decode_forward_backward(message) -> dict:
+ datums = []
+ for index, datum in enumerate(message.data):
+ tokens: list[int] = []
+ for chunk in datum.model_input:
+ if chunk.WhichOneof("chunk") != "encoded_text":
+ raise UserInputError(f"unsupported model_input chunk type: {chunk.WhichOneof('chunk')}")
+ tokens.extend(np.frombuffer(chunk.encoded_text.tokens, dtype=np.int32).tolist())
+ inputs = {name: _decode_tensor(name, tensor) for name, tensor in datum.loss_fn_inputs.items()}
+ datums.append(build_datum(tokens, inputs, index))
+
+ loss_fn_config = dict(message.loss_fn_config)
+ # Tinker SDK's v2 protobuf config supports both numeric and string values.
+ if message.loss_fn_config_v2:
+ loss_fn_config = {}
+ for name, value in message.loss_fn_config_v2.items():
+ kind = value.WhichOneof("value")
+ if kind is None:
+ raise UserInputError(f"loss_fn_config[{name!r}]: missing number or text value")
+ loss_fn_config[name] = getattr(value, kind)
+
+ decoded = {
+ "model_id": message.model_id,
+ "seq_id": message.seq_id,
+ "datums": datums,
+ "loss_fn": message.loss_fn,
+ "loss_fn_config": loss_fn_config,
+ }
+ return decoded
+
+
+def _decode_tensor(name: str, tensor) -> list:
+ np_dtype = _PROTO_DTYPE_TO_NUMPY.get(tensor.dtype)
+ if np_dtype is None:
+ raise UserInputError(f"loss_fn_inputs[{name!r}]: unsupported tensor dtype {tensor.dtype}")
+ if len(tensor.shape) > 1:
+ raise UserInputError(f"loss_fn_inputs[{name!r}]: shape {list(tensor.shape)} is not supported (1-D only)")
+ encoding = tensor.WhichOneof("encoding")
+ if encoding == "dense":
+ return np.frombuffer(tensor.dense, dtype=np_dtype).tolist()
+ if encoding == "sparse_csr":
+ (length,) = tensor.shape
+ dense = np.zeros(length, dtype=np_dtype)
+ cols = np.frombuffer(tensor.sparse_csr.col_indices, dtype=np.int64)
+ dense[cols] = np.frombuffer(tensor.sparse_csr.values, dtype=np_dtype)
+ return dense.tolist()
+ raise UserInputError(f"loss_fn_inputs[{name!r}]: tensor without data")
+
+
+# -------- result encoding --------
+
+
+def encode_forward_backward_output(result: dict) -> bytes:
+ message = public_pb.ForwardBackwardOutput()
+ message.loss_fn_output_type = "ArrayRecord"
+ outputs = result["outputs"]
+ message.metrics["loss:sum"] = float(sum(output["loss"] for output in outputs))
+
+ record = message.loss_fn_outputs.add()
+ record.num_datums = len(outputs)
+ for field, arrays in (
+ ("loss:sum", [np.asarray([output["loss"]], dtype=np.float32) for output in outputs]),
+ ("logprobs", [np.asarray(output["logprobs"], dtype=np.float32) for output in outputs]),
+ ):
+ batched = record.fields[field]
+ batched.dtype = public_pb.DTYPE_FLOAT32
+ offsets = np.zeros(len(arrays) + 1, dtype=np.int64)
+ np.cumsum([array.nbytes for array in arrays], out=offsets[1:])
+ batched.offsets = offsets.tobytes()
+ batched.data = b"".join(array.tobytes() for array in arrays)
+ # flat per-datum arrays: empty trailing_shape decodes as shape [n]
+ return message.SerializeToString()
+
+
+def encode_sample_response(result: dict) -> bytes:
+ message = public_pb.SampleResponse()
+ for sequence in result["sequences"]:
+ out = message.sequences.add()
+ out.stop_reason = _STOP_REASON_TO_PROTO[sequence["stop_reason"]]
+ out.tokens = np.asarray(sequence["tokens"], dtype=np.int32).tobytes()
+ if sequence.get("logprobs") is not None:
+ out.logprobs = np.asarray(sequence["logprobs"], dtype=np.float32).tobytes()
+ if result.get("prompt_logprobs") is not None:
+ message.prompt_logprobs = np.asarray(result["prompt_logprobs"], dtype=np.float32).tobytes()
+ if result.get("topk_prompt_logprobs") is not None:
+ topk = result["topk_prompt_logprobs"]
+ token_ids = np.asarray(topk["token_ids"], dtype=np.int32)
+ message.topk_prompt_logprobs.prompt_length, message.topk_prompt_logprobs.k = token_ids.shape
+ message.topk_prompt_logprobs.token_ids = token_ids.tobytes()
+ message.topk_prompt_logprobs.logprobs = np.asarray(topk["logprobs"], dtype=np.float32).tobytes()
+ return message.SerializeToString()
+
+
+PROTO_ENCODERS = {
+ "forward_backward": encode_forward_backward_output,
+ "forward_only": encode_forward_backward_output,
+ "sample": encode_sample_response,
+}
diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py
deleted file mode 100644
index 6c4d224a86..0000000000
--- a/miles/utils/adapter_config.py
+++ /dev/null
@@ -1,93 +0,0 @@
-"""Adapter config parsing for multi-LoRA training.
-
-``AdapterRunConfig`` carries only static, YAML-sourced configuration; the
-mutable slot is owned by the controller and exposed through ``AdapterRun``
-views.
-"""
-
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-import yaml
-
-
-@dataclass(frozen=True)
-class AdapterRunConfig:
-
- data: str
-
- # resolves them to CLI defaults if None (--lora-rank / --lora-alpha) on register.
- rank: int | None = None
- alpha: int | None = None
-
- # Prompt groups consumed per optimizer step for this adapter (group units,
- # like --rollout-batch-size, which it defaults to). The samples-per-step
- # analog of --global-batch-size is derived: adapter_global_batch_size =
- # rollout_batch_size * n_samples_per_prompt.
- rollout_batch_size: int | None = None
- n_samples_per_prompt: int | None = None
-
- save: str | Path | None = None
-
- input_key: str = "text"
- label_key: str | None = None
- metadata_key: str | None = None
-
- rm_type: str | None = None
- custom_rm_path: str | None = None
-
- # Stop after N optimizer steps; derived from num_epoch (default 1) when absent.
- num_step: int | None = None
- num_epoch: int | None = None
-
- metadata: dict[str, Any] = field(default_factory=dict)
-
- @property
- def adapter_global_batch_size(self) -> int:
- """Samples per optimizer step (per-adapter analog of --global-batch-size)."""
- assert self.rollout_batch_size is not None and self.n_samples_per_prompt is not None
- return self.rollout_batch_size * self.n_samples_per_prompt
-
-
-@dataclass(frozen=True)
-class AdapterRun:
- """Read-only join view of a run's static config and current slot."""
-
- name: str
- config: AdapterRunConfig
- slot: int
- version: int = 0
- step: int = 0
- # Committed prompt groups accumulated toward the current optimizer step.
- accumulated_groups: int = 0
- # Unique per registration (see AdapterRecord.registration_id): lets the
- # rollout worker tell a re-registered name apart from the previous tenant.
- registration_id: str = ""
-
-
-def parse_adapter_run_yaml(path: Path) -> AdapterRunConfig:
- """Parse a single adapter.yaml file.
-
- ``rank``, ``alpha`` and ``save`` are optional in the YAML; when absent the
- caller (e.g. the multi-LoRA controller) is responsible for resolving them.
- """
- with open(path) as f:
- raw = yaml.safe_load(f)
-
- return AdapterRunConfig(
- rank=raw.get("rank"),
- alpha=raw.get("alpha"),
- data=raw["data"],
- rollout_batch_size=raw.get("rollout_batch_size"),
- n_samples_per_prompt=raw.get("n_samples_per_prompt"),
- save=Path(raw["save"]) if raw.get("save", None) else None,
- input_key=raw.get("input_key", "text"),
- label_key=raw.get("label_key"),
- metadata_key=raw.get("metadata_key"),
- rm_type=raw.get("rm_type"),
- custom_rm_path=raw.get("custom_rm_path"),
- num_step=raw.get("num_step"),
- num_epoch=raw.get("num_epoch"),
- metadata=raw.get("metadata") or {},
- )
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index f5045d29b9..787e6a9c39 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -104,6 +104,8 @@ def reset_arg(parser, name, **kwargs):
def get_miles_extra_args_provider(add_custom_arguments=None):
def add_miles_arguments(parser):
+ parser.set_defaults(entry="train")
+
def add_run_uuid_arguments(parser):
parser.add_argument(
"--run-uuid",
@@ -1127,10 +1129,11 @@ def add_data_arguments(parser):
parser.add_argument(
"--rollout-batch-size",
type=int,
- required=True,
+ default=None,
help=(
"The number of prompts in each rollout step. "
"The total data returned should be rollout_batch_size * n_samples_per_prompt. "
+ "Required for the train entry. "
),
)
parser.add_argument(
@@ -1841,78 +1844,6 @@ def add_lora_arguments(parser):
default=0,
help="Maximum number of concurrent adapter slots for multi-LoRA. Set to 0 to disable multi-LoRA (default: 0)",
)
- parser.add_argument(
- "--multi-lora-adapter",
- nargs=2,
- action="append",
- type=str,
- dest="multi_lora_adapters",
- default=[],
- )
- parser.add_argument(
- "--multi-lora-idle-poll-s",
- type=float,
- default=5.0,
- help="When no adapter is RUNNING, the trainer polls for new registrations every this many seconds (default: 5.0)",
- )
- parser.add_argument(
- "--multi-lora-http-server-path",
- type=str,
- default=None,
- help=(
- "Dotted path to a MultiLoRAHTTPServer subclass to use for the multi-LoRA "
- "controller's HTTP server (default: MultiLoRAHTTPServer)"
- ),
- )
- parser.add_argument(
- "--multi-lora-backend-path",
- type=str,
- default=None,
- help=(
- "Dotted path to a MultiLoRABackend subclass for the multi-LoRA controller, "
- "e.g. to add custom adapter validation via validate_adapter (default: MultiLoRABackend)"
- ),
- )
- parser.add_argument(
- "--multi-lora-api-port",
- type=int,
- default=8068,
- help="Port for the multi-LoRA controller's control-plane API, served from the head node (default: 8068)",
- )
- parser.add_argument(
- "--multi-lora-disable-service-mode",
- action="store_false",
- dest="multi_lora_service_mode",
- help="Disable service mode. By default, the trainer waits indefinitely for new adapters. With this flag, it exits after all adapters have been processed.",
- )
- parser.add_argument(
- "--multi-lora-max-adapter-global-batch-size",
- type=int,
- default=None,
- help=(
- "Registration-time upper bound on an adapter's samples per optimizer "
- "step (rollout_batch_size x n_samples_per_prompt). Defaults to 4x "
- "--global-batch-size."
- ),
- )
- parser.add_argument(
- "--multi-lora-max-coalesce-wait-s",
- type=float,
- default=0.5,
- help=(
- "Maximum time ready groups wait for the batch to fill toward "
- "--global-batch-size before training starts on what is ready (default: 0.5)."
- ),
- )
- parser.add_argument(
- "--multi-lora-max-empty-wait-s",
- type=float,
- default=30.0,
- help=(
- "How long a generate call waits for the first poppable group before "
- "failing with an empty-batch timeout (default: 30)."
- ),
- )
return parser
def add_router_arguments(parser):
@@ -2722,7 +2653,8 @@ def add_user_provided_function_arguments(parser):
return add_miles_arguments
-def parse_args(add_custom_arguments=None):
+def parse_args(add_custom_arguments=None, entry="train", preprocess_args=None):
+ assert entry in ("train", "serve"), f"unknown entry {entry!r}"
# Users may call `parse_args` very early, thus we ensure logger is configured here
configure_logger_raw("main")
@@ -2766,6 +2698,9 @@ def parse_args(add_custom_arguments=None):
# locates the per-test record). No CLI flag: non-CI runs always stay False.
args.ci_enable_metrics_capture = bool(os.environ.get(RECORD_DIR_ENV))
+ args.entry = entry
+ if preprocess_args is not None:
+ preprocess_args(args)
miles_validate_args(args)
if backend == "megatron":
@@ -2871,7 +2806,7 @@ def _validate_rematerialize_param_from_master_weight(args):
assert (
args.train_backend == "megatron"
), "--rematerialize-param-from-master-weight reads Megatron's distributed-optimizer main params"
- from miles.backends.megatron_utils.lora_utils import is_lora_enabled
+ from miles.backends.megatron_utils.lora.utils import is_lora_enabled
assert not is_lora_enabled(args), "--rematerialize-param-from-master-weight does not support LoRA"
assert not args.debug_disable_optimizer, "--debug-disable-optimizer leaves no main params to rematerialize from"
@@ -3558,27 +3493,30 @@ def miles_validate_args(args):
args.grpo_std_normalization = False
logger.info("n_samples_per_prompt is set to 1, grpo_std_normalization will be set to False.")
- if args.over_sampling_batch_size is None:
- args.over_sampling_batch_size = args.rollout_batch_size
+ if args.entry == "train":
+ assert args.rollout_batch_size is not None, "please set --rollout-batch-size"
- assert args.over_sampling_batch_size >= args.rollout_batch_size, (
- f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than or equal to "
- f"rollout_batch_size {args.rollout_batch_size}"
- )
+ if args.over_sampling_batch_size is None:
+ args.over_sampling_batch_size = args.rollout_batch_size
+
+ assert args.over_sampling_batch_size >= args.rollout_batch_size, (
+ f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than or equal to "
+ f"rollout_batch_size {args.rollout_batch_size}"
+ )
- if args.num_epoch is not None:
- if args.num_rollout is not None:
- logger.info("Both num_epoch and num_rollout are set, num_epoch will be ignored.")
+ if args.num_epoch is not None:
+ if args.num_rollout is not None:
+ logger.info("Both num_epoch and num_rollout are set, num_epoch will be ignored.")
+ else:
+ assert args.rollout_global_dataset, (
+ "num_epoch is set, but rollout_global_dataset is not set, "
+ "please remove --disable-rollout-global-dataset to use num_epoch"
+ )
else:
- assert args.rollout_global_dataset, (
- "num_epoch is set, but rollout_global_dataset is not set, "
- "please remove --disable-rollout-global-dataset to use num_epoch"
+ # if num_epoch is not set, we should set num_rollout
+ assert args.num_rollout is not None, (
+ "num_epoch is not set, but num_rollout is not set, " "please set --num-rollout or --num-epoch"
)
- else:
- # if num_epoch is not set, we should set num_rollout
- assert args.num_rollout is not None, (
- "num_epoch is not set, but num_rollout is not set, " "please set --num-rollout or --num-epoch"
- )
if args.enable_mtp_training:
assert args.mtp_num_layers, "mtp_num_layers must be set when enable_mtp_training is set"
diff --git a/miles/utils/lora.py b/miles/utils/lora.py
index 7a0984e459..ce97106041 100644
--- a/miles/utils/lora.py
+++ b/miles/utils/lora.py
@@ -1,4 +1,6 @@
+import json
from argparse import Namespace
+from pathlib import Path
LORA_ADAPTER_NAME = "miles_lora"
@@ -25,3 +27,13 @@ def lora_rollout_enabled(args: Namespace) -> bool:
def lora_base_cpu_backup_enabled(args: Namespace) -> bool:
"""LoRA + --colocate + --lora-base-cpu-backup all set."""
return is_lora_enabled(args) and getattr(args, "colocate", False) and getattr(args, "lora_base_cpu_backup", False)
+
+
+def save_adapter_to_disk(out_dir, config: dict, tensors: dict) -> None:
+ """Write a LoRA adapter dir (adapter_config.json + adapter_model.safetensors)."""
+ import safetensors.torch # lazy: this module is imported on paths that never touch weights
+
+ out = Path(out_dir)
+ out.mkdir(parents=True, exist_ok=True)
+ (out / "adapter_config.json").write_text(json.dumps(config, indent=2))
+ safetensors.torch.save_file(tensors, str(out / "adapter_model.safetensors"))
diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py
index 775064f2ce..17a181ab08 100644
--- a/miles/utils/multi_lora.py
+++ b/miles/utils/multi_lora.py
@@ -1,51 +1,30 @@
-"""Small multi-LoRA helpers shared across the rollout, trainer, and controller.
+"""Multi-LoRA arguments and adapter configuration. Slot mechanics live in
+``miles/backends/megatron_utils/lora/slots.py``."""
-The controller-side machinery (AdapterRegistry, MultiLoRABackend,
-MultiLoRAHTTPServer) lives in ``miles/ray/multi_lora/``.
-"""
-
-import logging
-import uuid
+from dataclasses import dataclass
from typing import Any
-logger = logging.getLogger(__name__)
-
__all__ = [
- "EmptyBatchTimeoutError",
- "RID_SEPARATOR",
- "define_new_adapter_metrics",
+ "AdapterSpec",
"is_multi_lora_enabled",
- "make_rid",
- "min_groups_per_dp_split",
- "parse_adapter",
- "slot_lora_name",
+ "targets_expert_leaves",
"validate_multi_lora_args",
]
-# Must not appear in adapter names so rid prefix aborts can't cross adapters.
-RID_SEPARATOR = "::"
-
+@dataclass(frozen=True)
+class AdapterSpec:
+ """The slot and scaling needed to export one adapter."""
-class EmptyBatchTimeoutError(RuntimeError):
- """No trainable groups arrived before empty-wait timeout."""
+ slot: int
+ rank: int
+ alpha: float
def is_multi_lora_enabled(args: Any) -> bool:
return getattr(args, "multi_lora", False)
-def define_new_adapter_metrics(snapshot: dict) -> None:
- """Declare metric axes for new adapters ({name}/* -> {name}/step, {name}/perf/* -> rollout/step); must run
- in the primary tracking writer. Already-declared adapters are skipped, so calling every snapshot is free."""
- # lazy import tracking deps
- from miles.utils.tracking_utils.tracking import define_step_key_metric_group
-
- for name in {**snapshot["pending"], **snapshot["active"], **snapshot["retiring"]}:
- define_step_key_metric_group(prefix=name, step_key=f"{name}/step")
- define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step")
-
-
# Leaf module names that can live inside MoE experts (they also name the dense MLP
# projections); the bulk aliases expand to them during target-module resolution.
_EXPERT_LEAF_NAMES = frozenset({"linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj"})
@@ -64,27 +43,22 @@ def targets_expert_leaves(target_modules: Any) -> bool:
def validate_multi_lora_args(args: Any) -> None:
- """Set ``args.multi_lora``, then validate and default the multi-LoRA arg
- surface. Called from ``miles_validate_args``; a no-op for normal runs."""
+ """Set ``args.multi_lora``, then validate the trainer-side constraints of
+ the slot machinery. A no-op for normal runs."""
args.multi_lora = getattr(args, "multi_lora_n_adapters", 0) > 0
if not args.multi_lora:
return
- # Swap in the multi-LoRA rollout fn and data source unless the user pointed these flags elsewhere.
- if args.rollout_function_path is None:
- args.rollout_function_path = "miles.rollout.multi_lora.async_rollout.generate_rollout_multi_lora"
- if args.data_source_path == "miles.rollout.data_source.RolloutDataSourceWithBuffer":
- args.data_source_path = "miles.rollout.multi_lora.data_source.MultiLoRAAsyncDataSource"
- # The per-adapter data source is inherently global (the controller owns
- # what is sampleable); rollout workers must not shard it.
- args.rollout_global_dataset = True
assert args.lora_rank > 0, "--lora-rank must be set when --multi-lora-n-adapters > 0"
assert args.target_modules is not None, "--target-modules must be set when --multi-lora-n-adapters > 0"
assert args.train_backend == "megatron", "Multi-LoRA currently requires --train-backend megatron"
# Adapter routing is only recompute-safe without pipelining; enforce at launch.
+ assert getattr(args, "context_parallel_size", 1) == 1, (
+ "multi-LoRA requires --context-parallel-size 1: the Tinker losses zip "
+ "full-length per-datum vectors against log_probs, which CP would shard"
+ )
assert getattr(args, "pipeline_model_parallel_size", 1) == 1, (
- "Multi-LoRA requires --pipeline-model-parallel-size 1: no single rank holds a "
- "complete adapter to push to the rollout engines, and a pipelined schedule would "
+ "Multi-LoRA requires --pipeline-model-parallel-size 1: a pipelined schedule would "
"recompute activations against a later micro-batch's adapter routing."
)
# Per-slot token spans assume sequence-major contiguous sample packing, which only 'thd' provides.
@@ -96,16 +70,11 @@ def validate_multi_lora_args(args: Any) -> None:
"Multi-LoRA does not support --experts-shared-outer-loras; MoE expert adapters "
"use the per-expert layout. Drop the flag (and --sglang-experts-shared-outer-loras)."
)
- # Expert-parallel sizes are checked post-finalize in _validate_multi_lora_moe_support:
- # --expert-tensor-parallel-size stays None until Megatron's own validate_args resolves it.
assert "muon" not in str(getattr(args, "optimizer", "")).lower(), (
"Multi-LoRA does not support Muon: per-adapter decoupled stepping is only "
"implemented for Adam-family per-slot optimizers"
)
- assert not args.colocate, (
- "Multi-LoRA requires disaggregated rollout engines: weight sync is only "
- "implemented for the distributed path, not the colocated tensor path."
- )
+ assert not args.colocate, "Multi-LoRA requires separate training and sampling GPUs to retain accumulated gradients"
assert (
not getattr(args, "indep_dp", False) and "train" not in args.ft_components
), "Multi-LoRA does not support independent-DP training; remove 'train' from --ft-components"
@@ -118,91 +87,16 @@ def validate_multi_lora_args(args: Any) -> None:
"optimizers); the witness module assumes use_distributed_optimizer"
)
assert getattr(args, "sglang_tokenizer_worker_num", 1) == 1, (
- "Multi-LoRA requires --sglang-tokenizer-worker-num 1: each tokenizer "
- "worker process holds its own LoRA registry, so per-step adapter "
- "upserts resolve against whichever worker the router picks and fail "
- "non-deterministically. sglang rejects the upsert at runtime anyway; "
- "fail at launch instead of burning GPU time until the first weight push."
+ "Multi-LoRA requires --sglang-tokenizer-worker-num 1: dynamic adapter loading "
+ "requires a single tokenizer-side LoRA registry."
)
assert not args.calculate_per_token_loss, (
"Multi-LoRA normalizes each sample by its adapter batch "
"(sample-mean); per-token loss normalization would make adapter batch weights "
"depend on batch contents. Drop --calculate-per-token-loss."
)
- assert args.multi_lora_max_coalesce_wait_s >= 0, "--multi-lora-max-coalesce-wait-s must be non-negative"
assert (getattr(args, "optimizer", "adam") or "adam").lower() == "adam", (
- "Multi-LoRA requires --optimizer adam: the per-slot optimizer isolation "
- "(build_multi_lora_optimizer, slot retirement state cleanup) only implements "
- f"Adam semantics; got --optimizer {args.optimizer}"
- )
- # --global-batch-size may legitimately be unset (Megatron derives it later);
- # leave the adapter cap unset too rather than multiplying None.
- if args.multi_lora_max_adapter_global_batch_size is None and getattr(args, "global_batch_size", None) is not None:
- args.multi_lora_max_adapter_global_batch_size = 4 * args.global_batch_size
- if args.multi_lora_max_adapter_global_batch_size is not None:
- assert (
- args.multi_lora_max_adapter_global_batch_size > 0
- ), "--multi-lora-max-adapter-global-batch-size must be positive"
-
- # Trainer DP size, used to validate adapter batch shapes; guarded for harnesses without megatron args set.
- if all(
- hasattr(args, name)
- for name in (
- "world_size",
- "tensor_model_parallel_size",
- "pipeline_model_parallel_size",
- "context_parallel_size",
- )
- ):
- from miles.utils.megatron_args_utils import compute_megatron_world_size_except_dp
-
- model_parallel = compute_megatron_world_size_except_dp(args)
- assert (
- args.world_size % model_parallel == 0
- ), f"actor world size {args.world_size} is not divisible by tp*pp*cp {model_parallel}"
- args.multi_lora_dp_size = args.world_size // model_parallel
- else:
- args.multi_lora_dp_size = None
-
- # Batches are variable-sized; carry the exact sample
- # count through rollout conversion instead of trimming to --global-batch-size.
- assert not args.disable_rollout_trim_samples, (
- "Multi-LoRA computes the exact dynamic batch size in rollout postprocessing; "
- "do not pass --disable-rollout-trim-samples"
+ "Multi-LoRA requires --optimizer adam: the per-slot SlotOptimizer only "
+ f"implements Adam semantics; got --optimizer {args.optimizer}"
)
- args.use_dynamic_global_batch_size = True
args.megatron_to_hf_mode = "bridge"
-
-
-def make_rid(adapter_name: str) -> str:
- return f"{adapter_name}{RID_SEPARATOR}{uuid.uuid4().hex}"
-
-
-def parse_adapter(rid: str) -> str:
- return rid.rsplit(RID_SEPARATOR, 1)[0]
-
-
-def slot_lora_name(slot: int) -> str:
- """Engine-side LoRA adapter name for a controller slot. Weight pushes and
- every inference request (rollout and prefill scoring) must agree on this."""
- return f"__miles_slot_{slot}"
-
-
-def min_groups_per_dp_split(n_samples_per_prompt: int, dp_size: int) -> int:
- """Minimum prompt-group count that splits cleanly across data-parallel
- ranks.
-
- Train batches only pop groups in multiples of this value, so each popped
- slice has a sample count divisible by ``dp_size`` with no trimming.
-
- Requires ``n_samples_per_prompt`` and ``dp_size`` to divide each other
- (one must be a multiple of the other).
- """
- larger = max(dp_size, n_samples_per_prompt)
- smaller = min(dp_size, n_samples_per_prompt)
- if larger % smaller == 0:
- return larger // n_samples_per_prompt
- raise ValueError(
- f"n_samples_per_prompt={n_samples_per_prompt} must be a divisor or a multiple of "
- f"the data-parallel size {dp_size} so whole prompt groups can split evenly across ranks"
- )
diff --git a/miles/utils/types.py b/miles/utils/types.py
index daa514a055..8f9e4cc550 100644
--- a/miles/utils/types.py
+++ b/miles/utils/types.py
@@ -57,14 +57,6 @@ def from_meta_info(meta_info: dict, output_end: int) -> "WeightVersionsPerCall":
)
-@dataclass(frozen=True)
-class AdapterRef:
- """Which LoRA adapter a sample is bound to (training slot routing, inference lora_path); ``None`` = no adapter."""
-
- name: str
- slot: int
-
-
@dataclass(frozen=True)
class RewardSpec:
"""Per-sample spec of how the response is scored; intentionally decoupled from adapter routing."""
@@ -123,9 +115,7 @@ class Status(Enum):
# metadata used during training, e.g., what loss to use for this sample.
train_metadata: dict | None = None
- # MultiLoRA: which adapter this sample trains/infers with
- adapter: AdapterRef | None = None
- # Per-sample reward dispatch override (e.g., per-adapter RM in multi-LoRA)
+ # Per-sample reward dispatch override
reward_spec: RewardSpec | None = None
# Per-sample routing key for the router's consistent_hashing policy (sent as X-SMG-Routing-Key)
diff --git a/miles_plugins/models/inkling/lora.py b/miles_plugins/models/inkling/lora.py
index 7052437934..9eee20da7f 100644
--- a/miles_plugins/models/inkling/lora.py
+++ b/miles_plugins/models/inkling/lora.py
@@ -373,7 +373,7 @@ def lm_head_forward(inputs, *forward_args, **forward_kwargs):
def apply_inkling_lora(model, args):
"""Attach Inkling LoRA to ONE built model chunk (before Float16Module / DDP wrapping)."""
- from miles.backends.megatron_utils.lora_utils import patch_param_grad_buffer_for_colocate_mode_lora
+ from miles.backends.megatron_utils.lora.utils import patch_param_grad_buffer_for_colocate_mode_lora
from miles_plugins.models.inkling.layers import InklingDenseMLP, InklingSelfAttention, InklingSharedExperts
diff --git a/serve_tinker.py b/serve_tinker.py
new file mode 100644
index 0000000000..563e589917
--- /dev/null
+++ b/serve_tinker.py
@@ -0,0 +1,106 @@
+import asyncio
+import logging
+from contextlib import suppress
+
+import uvicorn
+
+from miles.ray.rollout.inference_controller import InferenceController
+from miles.ray.train.group import TrainerController
+from miles.ray.wiring import launch_worker_manager
+from miles.tinker.arguments import add_tinker_arguments, configure_tinker_args
+from miles.tinker.core.service import TinkerService
+from miles.tinker.core.types import GatewayConfig
+from miles.tinker.runtime import MilesBackend
+from miles.tinker.server.app import build_app
+from miles.utils import object_store
+from miles.utils.arguments import parse_args
+from miles.utils.audit_utils.process_identity import MainProcessIdentity
+from miles.utils.hf_config import load_hf_config
+from miles.utils.http_utils import init_http_client
+from miles.utils.logging_utils import configure_logger
+
+logger = logging.getLogger(__name__)
+
+
+async def serve(args):
+ assert args.multi_lora, "serve_tinker requires --multi-lora-n-adapters > 0"
+ assert args.load == args.hf_checkpoint, "Tinker trainers and engines must load the same frozen HF base"
+ checkpoint_root = args.tinker_checkpoint_root or (args.save and f"{args.save}/tinker")
+ assert checkpoint_root, "set --tinker-checkpoint-root (or --save to derive /tinker)"
+ hf_config = load_hf_config(args.hf_checkpoint)
+ max_tokens_per_datum = hf_config.max_position_embeddings
+ if args.max_tokens_per_gpu is not None:
+ # The trainer pads each packed microbatch to this multiple.
+ pad_size = args.tensor_model_parallel_size * args.data_pad_size_multiplier
+ trainer_token_limit = args.max_tokens_per_gpu // pad_size * pad_size
+ max_tokens_per_datum = min(max_tokens_per_datum, trainer_token_limit)
+ assert max_tokens_per_datum > 0, "trainer token budget must fit at least one padding block"
+ configure_logger(args, source=MainProcessIdentity())
+
+ init_http_client(args)
+
+ _worker_manager = launch_worker_manager(args)
+ object_store.init_instance(args, contribute_segment=False)
+
+ inference_controller = InferenceController(args)
+ await inference_controller.init()
+
+ trainer = TrainerController(
+ args=args,
+ role="actor",
+ with_ref=False,
+ with_opd_teacher=False,
+ inference_controller=None,
+ rollout_executor=None,
+ )
+ await trainer.init()
+
+ config = GatewayConfig(
+ base_model=args.tinker_base_model or args.hf_checkpoint,
+ n_slots=args.multi_lora_n_adapters,
+ checkpoint_root=checkpoint_root,
+ vocab_size=hf_config.vocab_size,
+ max_tokens_per_datum=max_tokens_per_datum,
+ lora_alpha=args.lora_alpha,
+ max_lora_rank=args.lora_rank,
+ trains_attn=args.tinker_train_attn,
+ trains_mlp=args.tinker_train_mlp,
+ trains_unembed=args.tinker_train_unembed,
+ )
+ router_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}"
+ actor_world_size = args.actor_num_nodes * args.actor_num_gpus_per_node
+ dp_size = actor_world_size // (
+ args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size
+ )
+ service = TinkerService(MilesBackend(trainer, router_url, dp_size=dp_size), config)
+
+ server = uvicorn.Server(
+ uvicorn.Config(
+ build_app(service), host=args.tinker_server_host, port=args.tinker_server_port, log_level="info"
+ )
+ )
+ logger.info(f"tinker gateway serving {config.base_model} on :{args.tinker_server_port}")
+ # supervise both: a crashed dispatcher must take the HTTP server down with it,
+ # not keep answering /healthz while every training future pends forever
+ service_task = asyncio.create_task(service.run())
+ server_task = asyncio.create_task(server.serve())
+ try:
+ done, _ = await asyncio.wait({service_task, server_task}, return_when=asyncio.FIRST_COMPLETED)
+ for task in done:
+ task.result()
+ finally:
+ for task in (service_task, server_task):
+ task.cancel()
+ with suppress(asyncio.CancelledError):
+ await task
+
+ await inference_controller.dispose()
+ await trainer.dispose()
+
+
+if __name__ == "__main__":
+ args = parse_args(add_tinker_arguments, entry="serve", preprocess_args=configure_tinker_args)
+ # commands ship one work unit at a time; its size is the batch size
+ args.use_dynamic_global_batch_size = True
+ args.delay_split_train_data_by_dp = True
+ asyncio.run(serve(args))
diff --git a/tests/ci/labels.py b/tests/ci/labels.py
index 5d9c8457d4..d6ec09268e 100644
--- a/tests/ci/labels.py
+++ b/tests/ci/labels.py
@@ -26,6 +26,7 @@
"long": "Long-running training tests",
"ckpt": "Checkpoint save / load tests",
"lora": "LoRA training tests",
+ "multi-lora": "Tinker cookbook and multi-tenant gateway tests",
"eval": "Evaluation machinery tests (shared-engine / fleet / external postures)",
"precision": "Numerical precision parity tests",
"ft-short": "Fault-tolerance trainer comparison tests (no_failure / deterministic / with_failure)",
diff --git a/tests/ci/requirements-ci-cpu.txt b/tests/ci/requirements-ci-cpu.txt
index e812c695af..df5238d158 100644
--- a/tests/ci/requirements-ci-cpu.txt
+++ b/tests/ci/requirements-ci-cpu.txt
@@ -9,6 +9,7 @@ partial_json_parser==0.2.1.1.post7
pyzmq==27.1.0
sentencepiece==0.2.1
tiktoken==0.13.0
+tinker==0.26.2 # SDK wire schema for the tinker gateway proto round-trip tests
torch==2.11.0
torchvision==0.26.0
xgrammar==0.2.1
diff --git a/tests/e2e/lora/test_slot_checkpoint_resume.py b/tests/e2e/lora/test_slot_checkpoint_resume.py
new file mode 100644
index 0000000000..f5015197e0
--- /dev/null
+++ b/tests/e2e/lora/test_slot_checkpoint_resume.py
@@ -0,0 +1,180 @@
+"""A slot checkpoint preserves FP32 Adam state and resumes into another slot exactly."""
+
+import argparse
+import os
+import subprocess
+import sys
+import tempfile
+from argparse import Namespace
+from datetime import timedelta
+from pathlib import Path
+
+import pytest
+import torch
+import torch.distributed as dist
+from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear, init_adapter_slot
+from megatron.core import parallel_state
+from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from tests.ci.ci_register import register_cuda_ci
+from torch.utils._pytree import tree_flatten, tree_map
+
+from miles.backends.megatron_utils.lora.checkpoint import load_slot, save_slot
+from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer, adapter_slot_parameters, step_slot_optimizers
+from miles.utils.distributed_utils import init_gloo_group
+
+register_cuda_ci(est_time=120, suite="stage-b-2-gpu-h200", labels=["lora"], hardware=["hopper"])
+
+ADAM = dict(learning_rate=3e-4, beta1=0.9, beta2=0.95, eps=1e-8, weight_decay=0.01, grad_clip_norm=1.0)
+
+
+class CheckpointModel(MegatronModule):
+ def __init__(self, config):
+ super().__init__(config)
+ self.layers = torch.nn.ModuleList()
+ for index, (input_size, output_size) in enumerate(((16, 32), (32, 16))):
+ linear = ColumnParallelLinear(
+ input_size,
+ output_size,
+ config=config,
+ init_method=config.init_method,
+ bias=False,
+ gather_output=False,
+ )
+ linear.requires_grad_(False)
+ self.layers.append(MultiLoRALinear(linear, n_adapters=2, dim=4, alpha=8, full_name=f"layers.{index}"))
+
+
+def snapshot(model, optimizer):
+ state = {
+ "weights": [param for param in adapter_slot_parameters(model, optimizer.slot)],
+ "optimizer": optimizer._inner.state_dict(),
+ }
+ return tree_map(lambda value: value.detach().cpu().clone() if isinstance(value, torch.Tensor) else value, state)
+
+
+def assert_same_state(actual, expected):
+ actual_values, actual_spec = tree_flatten(actual)
+ expected_values, expected_spec = tree_flatten(expected)
+ assert actual_spec == expected_spec
+ for actual_value, expected_value in zip(actual_values, expected_values, strict=True):
+ if isinstance(expected_value, torch.Tensor):
+ torch.testing.assert_close(actual_value, expected_value, rtol=0, atol=0)
+ else:
+ assert actual_value == expected_value
+
+
+def apply_step(model, optimizer, step):
+ # Identical post-all-reduce gradients isolate checkpoint/Adam correctness from forward kernels.
+ for index, param in enumerate(adapter_slot_parameters(model, optimizer.slot)):
+ gradient = torch.arange(param.numel(), device=param.device, dtype=torch.float32).reshape(param.shape)
+ param.main_grad.copy_(torch.sin(gradient + index + step) * 0.1)
+ outcome = step_slot_optimizers({optimizer.slot: optimizer}, {optimizer.slot: ADAM})
+ assert outcome[optimizer.slot]["grad_norm"] > 0
+
+
+def run_worker(checkpoint_dir):
+ torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+ dist.init_process_group("nccl", timeout=timedelta(seconds=120))
+ assert dist.get_world_size() == 2
+ parallel_state.initialize_model_parallel()
+ init_gloo_group()
+ model_parallel_cuda_manual_seed(1234)
+ config = TransformerConfig(
+ num_layers=1,
+ hidden_size=16,
+ num_attention_heads=2,
+ params_dtype=torch.bfloat16,
+ bf16=True,
+ use_cpu_initialization=False,
+ )
+ model = [
+ DistributedDataParallel(
+ config,
+ DistributedDataParallelConfig(grad_reduce_in_fp32=True, overlap_grad_reduce=False),
+ CheckpointModel(config).cuda(),
+ )
+ ]
+ args = Namespace(optimizer="adam", bf16=True, use_gloo_process_groups=True, lr=ADAM["learning_rate"])
+ for slot in (0, 1):
+ init_adapter_slot(model, slot, rank=4, alpha=8)
+ source = SlotOptimizer(args, model, slot=0)
+ for step in (1, 2, 3):
+ apply_step(model, source, step)
+ saved = snapshot(model, source)
+ masters = [
+ param
+ for child in source._inner.chained_optimizers
+ for group in child.fp32_from_float16_groups
+ for param in group
+ ]
+ assert sum(param.numel() for param in masters) == 192, "DP=2 must own half of the slot's 384 parameters"
+ assert any(
+ not torch.equal(param, param.bfloat16().float()) for param in masters
+ ), "warmup must create sub-BF16 precision"
+ assert all(param.dtype == torch.float32 for param in masters)
+ for child in source._inner.chained_optimizers:
+ assert child.optimizer.state
+ assert all(group["step"] == 3 for group in child.param_groups)
+ for state in child.optimizer.state.values():
+ for name in ("exp_avg", "exp_avg_sq"):
+ assert state[name].dtype == torch.float32 and torch.count_nonzero(state[name]) > 0
+ print(f"rank={dist.get_rank()} warmup complete; owned_masters=192; sub_bf16_precision=True", flush=True)
+
+ save_slot(model, source, str(checkpoint_dir))
+ resumed = SlotOptimizer(args, model, slot=1)
+ load_slot(model, resumed, str(checkpoint_dir), load_optimizer=True)
+ assert_same_state(snapshot(model, resumed), saved)
+ assert_same_state(snapshot(model, source), saved)
+ print(f"rank={dist.get_rank()} cross-slot restore is exact", flush=True)
+
+ apply_step(model, source, 4)
+ uninterrupted = snapshot(model, source)
+ apply_step(model, resumed, 4)
+ assert_same_state(snapshot(model, resumed), uninterrupted)
+ assert_same_state(snapshot(model, source), uninterrupted)
+ assert all(group["step"] == 4 for child in resumed._inner.chained_optimizers for group in child.param_groups)
+ print(f"rank={dist.get_rank()} resumed Adam step is exact", flush=True)
+
+ load_slot(model, resumed, str(checkpoint_dir), load_optimizer=True)
+ resumed.reload_masters()
+ with pytest.raises(AssertionError):
+ assert_same_state(snapshot(model, resumed), saved)
+ print(f"rank={dist.get_rank()} negative control: rounding masters is detected", flush=True)
+
+ resumed = SlotOptimizer(args, model, slot=1)
+ load_slot(model, resumed, str(checkpoint_dir), load_optimizer=False)
+ apply_step(model, resumed, 4)
+ with pytest.raises(AssertionError):
+ assert_same_state(snapshot(model, resumed), uninterrupted)
+ print(f"rank={dist.get_rank()} negative control: weights-only restart diverges", flush=True)
+ dist.barrier()
+ parallel_state.destroy_model_parallel()
+ dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--checkpoint-root", type=Path)
+ parser.add_argument("--worker-dir", type=Path)
+ options = parser.parse_args()
+ if options.worker_dir is not None:
+ run_worker(options.worker_dir)
+ else:
+ with tempfile.TemporaryDirectory(prefix="slot-checkpoint-", dir=options.checkpoint_root) as directory:
+ subprocess.run(
+ [
+ sys.executable,
+ "-m",
+ "torch.distributed.run",
+ "--standalone",
+ "--nproc_per_node=2",
+ __file__,
+ "--worker-dir",
+ directory + "/checkpoint",
+ ],
+ check=True,
+ )
diff --git a/tests/e2e/lora/test_tinker_cookbook_gateway.py b/tests/e2e/lora/test_tinker_cookbook_gateway.py
new file mode 100644
index 0000000000..f16c251e2e
--- /dev/null
+++ b/tests/e2e/lora/test_tinker_cookbook_gateway.py
@@ -0,0 +1,33 @@
+"""Run the official cookbook SFT and RL recipes against a real Tinker gateway."""
+
+from tests.ci.ci_register import register_cuda_ci
+from tests.e2e.lora.tinker_gateway import BASE_MODEL, prepare_gateway, running_gateway
+
+import miles.utils.external_utils.command_utils as U
+
+register_cuda_ci(
+ est_time=2400,
+ suite="stage-c-8-gpu-h200",
+ labels=["lora", "weight-update", "multi-lora"],
+ hardware=["hopper"],
+)
+
+COOKBOOK_PIN = "git+https://github.com/thinking-machines-lab/tinker-cookbook@1f962eda3a2c"
+
+
+def prepare():
+ prepare_gateway()
+ U.exec_command_cpu(f"pip install tinker==0.26.2 {COOKBOOK_PIN}")
+
+
+def execute():
+ with running_gateway() as base_url:
+ U.exec_command_cpu(
+ "python examples/multi_lora/run_client_recipes.py "
+ f"--base-url {base_url} --base-model {BASE_MODEL} --mode both --steps 2"
+ )
+
+
+if __name__ == "__main__":
+ prepare()
+ execute()
diff --git a/tests/e2e/lora/test_tinker_multi_tenant_gateway.py b/tests/e2e/lora/test_tinker_multi_tenant_gateway.py
new file mode 100644
index 0000000000..5b497c4571
--- /dev/null
+++ b/tests/e2e/lora/test_tinker_multi_tenant_gateway.py
@@ -0,0 +1,27 @@
+"""Concurrent tenants learn distinct markers and sample their own adapter snapshots."""
+
+from tests.ci.ci_register import register_cuda_ci
+from tests.e2e.lora.tinker_gateway import BASE_MODEL, prepare_gateway, running_gateway
+
+import miles.utils.external_utils.command_utils as U
+
+register_cuda_ci(
+ est_time=2400,
+ suite="stage-c-8-gpu-h200",
+ labels=["multi-lora"],
+ hardware=["hopper"],
+)
+
+
+def execute():
+ with running_gateway() as base_url:
+ U.exec_command_cpu(
+ "python examples/multi_lora/run_multi_tenant_example.py "
+ f"--base-url {base_url} --base-model {BASE_MODEL} "
+ "--mode multi --clients 4 --lora-rank 8 --steps 12 --lr 1e-3 --max-tokens 24"
+ )
+
+
+if __name__ == "__main__":
+ prepare_gateway()
+ execute()
diff --git a/tests/e2e/lora/tinker_gateway.py b/tests/e2e/lora/tinker_gateway.py
new file mode 100644
index 0000000000..04868fa2ce
--- /dev/null
+++ b/tests/e2e/lora/tinker_gateway.py
@@ -0,0 +1,60 @@
+"""Shared gateway setup for Tinker GPU acceptance tests."""
+
+import os
+import signal
+import subprocess
+import time
+import urllib.request
+from contextlib import contextmanager, suppress
+
+import miles.utils.external_utils.command_utils as U
+
+MODEL_NAME = "Qwen3-4B-Instruct-2507"
+BASE_MODEL = f"Qwen/{MODEL_NAME}"
+GATEWAY_PORT = 10613
+SERVE_TIMEOUT_S = 1200
+
+
+def prepare_gateway():
+ U.exec_command_cpu("mkdir -p /root/models")
+ U.exec_command_cpu(f"hf download {BASE_MODEL} --local-dir /root/models/{MODEL_NAME}")
+ U.exec_command_cpu("pip install tinker==0.26.2")
+
+
+def _wait_for_gateway(server: subprocess.Popen) -> None:
+ deadline = time.time() + SERVE_TIMEOUT_S
+ url = f"http://127.0.0.1:{GATEWAY_PORT}/api/v1/healthz"
+ while time.time() < deadline:
+ if server.poll() is not None:
+ raise RuntimeError(f"gateway exited during startup with code {server.returncode}")
+ try:
+ with urllib.request.urlopen(url, timeout=2):
+ return
+ except OSError:
+ time.sleep(5)
+ raise TimeoutError(f"gateway not serving after {SERVE_TIMEOUT_S}s")
+
+
+@contextmanager
+def running_gateway():
+ serve_cmd = (
+ "python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py serve "
+ f"--hf-checkpoint /root/models/{MODEL_NAME} "
+ "--model-type qwen3-4B-Instruct-2507 --tp 2 --ep 1 --lora-rank 8 --lora-alpha 16 "
+ f'--extra-args "--tinker-base-model {BASE_MODEL}"'
+ )
+ server = subprocess.Popen(["bash", "-c", serve_cmd], start_new_session=True)
+ try:
+ _wait_for_gateway(server)
+ yield f"http://127.0.0.1:{GATEWAY_PORT}"
+ finally:
+ try:
+ with suppress(ProcessLookupError):
+ os.killpg(server.pid, signal.SIGTERM)
+ server.wait(timeout=30)
+ finally:
+ # Descendants can retain CI stdout after the launcher has exited.
+ with suppress(ProcessLookupError):
+ os.killpg(server.pid, signal.SIGKILL)
+ server.wait(timeout=30)
+ subprocess.run(["ray", "stop", "--force"], check=True, timeout=120)
diff --git a/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py b/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py
index bbecc98209..e835995781 100644
--- a/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py
+++ b/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py
@@ -1,4 +1,4 @@
-"""Tests for the bridge value-model hook in bridge_lora_helpers.
+"""Tests for the bridge value-model hook in lora/bridge.py.
``_make_value_model_hook`` swaps a bridge-built model's language-model head for
a scalar value head. Its body had no coverage, which let a constructor mismatch
@@ -74,7 +74,7 @@ def stub(name: str, attrs: dict | None = None, is_package: bool = False):
stub("megatron.core.utils", {"get_attr_wrapped_model": lambda *a, **k: None})
install("megatron.core.parallel_state", parallel_state)
- from miles.backends.megatron_utils import bridge_lora_helpers
+ from miles.backends.megatron_utils.lora import bridge as bridge_lora_helpers
# The hook imports LinearForLastLayer lazily from the sibling module, so
# swapping the module in sys.modules is enough to intercept it.
diff --git a/tests/fast/backends/megatron_utils/test_lora_utils.py b/tests/fast/backends/megatron_utils/test_lora_utils.py
index 42f569bf84..1cade45a81 100644
--- a/tests/fast/backends/megatron_utils/test_lora_utils.py
+++ b/tests/fast/backends/megatron_utils/test_lora_utils.py
@@ -1,4 +1,4 @@
-"""Unit tests for miles.backends.megatron_utils.lora_utils.
+"""Unit tests for miles.backends.megatron_utils.lora.utils.
Tests cover module name conversion, LoRA detection helpers, parameter identification,
exclude-module parsing, and LoRA sync config building — all without GPU.
@@ -13,8 +13,8 @@
import pytest
import torch
-import miles.backends.megatron_utils.lora_utils as lora_utils
-from miles.backends.megatron_utils.lora_utils import (
+import miles.backends.megatron_utils.lora.utils as lora_utils
+from miles.backends.megatron_utils.lora.utils import (
_get_lora_class_name,
_is_adapter_param_name,
build_lora_sync_config,
@@ -380,7 +380,7 @@ def _args(**overrides):
def test_a_single_adapter_on_an_inkling_checkpoint_publishes_the_shorthand(self, monkeypatch):
"""The engine was launched to auto-detect its targets, so the adapter must say the same."""
monkeypatch.setattr(
- "miles.backends.megatron_utils.lora_utils.sglang_lora_target_all_sentinel", lambda _a: True
+ "miles.backends.megatron_utils.lora.utils.sglang_lora_target_all_sentinel", lambda _a: True
)
assert build_lora_sync_config(self._args())["target_modules"] == "all-linear"
diff --git a/tests/fast/backends/megatron_utils/test_model_initialize.py b/tests/fast/backends/megatron_utils/test_model_initialize.py
index 6ac20b096d..6654ff013b 100644
--- a/tests/fast/backends/megatron_utils/test_model_initialize.py
+++ b/tests/fast/backends/megatron_utils/test_model_initialize.py
@@ -127,7 +127,7 @@ def _mock_megatron_environment():
},
)
_stub_module(
- "miles.backends.megatron_utils.bridge_lora_helpers",
+ "miles.backends.megatron_utils.lora.bridge",
{
"_ensure_model_list": MagicMock(),
"_setup_lora_model_via_bridge": MagicMock(),
diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
deleted file mode 100644
index 97ab9f3c5e..0000000000
--- a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
+++ /dev/null
@@ -1,46 +0,0 @@
-"""Adapter shards are keyed by (tp, pp, ep): EP ranks hold different local experts, and
-the realized coordinates are not the tp x pp x ep cross product when ETP < TP."""
-
-from miles.backends.megatron_utils.multi_lora_utils import all_megatron_checkpoints_exist, megatron_shard_name
-
-
-def _names(coords, ep_size):
- return {megatron_shard_name(*coord, ep_size) for coord in coords}
-
-
-def test_shard_name_omits_ep_suffix_without_expert_parallelism():
- # Checkpoints written before expert adapters existed must stay loadable.
- assert megatron_shard_name(0, 0, 0, ep_size=1) == "adapter_megatron_tp0_pp0.pt"
- assert megatron_shard_name(1, 2, 0, ep_size=1) == "adapter_megatron_tp1_pp2.pt"
-
-
-def test_shard_name_is_unique_per_expert_parallel_rank():
- names = {megatron_shard_name(0, 0, ep, ep_size=4) for ep in range(4)}
- assert len(names) == 4
- assert megatron_shard_name(0, 0, 2, ep_size=4) == "adapter_megatron_tp0_pp0_ep2.pt"
-
-
-def test_completeness_check_requires_every_realized_shard(tmp_path):
- coords = [(0, 0, 0), (0, 0, 1), (0, 0, 2)]
- for coord in coords[:2]:
- (tmp_path / megatron_shard_name(*coord, 3)).touch()
-
- assert not all_megatron_checkpoints_exist(tmp_path, _names(coords, 3))
-
- (tmp_path / megatron_shard_name(*coords[2], 3)).touch()
- assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 3))
-
-
-def test_completeness_ignores_unrealized_coordinates(tmp_path):
- # TP=2, EP=2, ETP=1: only (0,0,0) and (1,0,1) exist; a cross-product check
- # would demand four shards and never resume.
- coords = [(0, 0, 0), (1, 0, 1)]
- for coord in coords:
- (tmp_path / megatron_shard_name(*coord, 2)).touch()
-
- assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 2))
-
-
-def test_completeness_check_with_single_shard(tmp_path):
- (tmp_path / "adapter_megatron_tp0_pp0.pt").touch()
- assert all_megatron_checkpoints_exist(tmp_path, _names([(0, 0, 0)], 1))
diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_scheduler.py b/tests/fast/backends/megatron_utils/test_multi_lora_scheduler.py
deleted file mode 100644
index be48d7118e..0000000000
--- a/tests/fast/backends/megatron_utils/test_multi_lora_scheduler.py
+++ /dev/null
@@ -1,114 +0,0 @@
-"""Per-adapter LR schedules: parameters come from the global args, position is per adapter.
-Pins two fixes: late loads don't inherit the decayed position; resume rebuilds position from committed steps."""
-
-from types import SimpleNamespace
-
-import pytest
-
-from miles.backends.megatron_utils.multi_lora_scheduler import install_slot_scheduler, step_slot_schedulers
-
-LR = 2e-5
-
-
-def make_args(**overrides) -> SimpleNamespace:
- args = SimpleNamespace(
- lr=LR,
- min_lr=0.0,
- lr_warmup_init=0.0,
- lr_warmup_fraction=None,
- lr_warmup_iters=0,
- lr_decay_style="cosine",
- start_weight_decay=0.1,
- end_weight_decay=0.1,
- weight_decay_incr_style="constant",
- lr_wsd_decay_iters=None,
- lr_wsd_decay_style=None,
- )
- for key, value in overrides.items():
- setattr(args, key, value)
- return args
-
-
-def make_optimizer(n_slots: int = 2) -> SimpleNamespace:
- children = [SimpleNamespace(param_groups=[{"lr": 0.0, "weight_decay": 0.0}]) for _ in range(n_slots)]
- return SimpleNamespace(
- chained_optimizers=children,
- miles_slot_child_indices={slot: [slot] for slot in range(n_slots)},
- )
-
-
-def make_adapter(slot: int, num_step: int | None, samples_per_step: int = 64) -> SimpleNamespace:
- config = SimpleNamespace(num_step=num_step, adapter_global_batch_size=samples_per_step)
- return SimpleNamespace(slot=slot, name=f"a{slot}", config=config)
-
-
-def slot_lr(optimizer, slot: int) -> float:
- return optimizer.chained_optimizers[slot].param_groups[0]["lr"]
-
-
-def test_decaying_adapter_walks_its_own_cosine_schedule():
- optimizer = make_optimizer()
- adapter = make_adapter(slot=0, num_step=10)
- install_slot_scheduler(make_args(), optimizer, adapter, resume_step=0)
-
- assert slot_lr(optimizer, 0) == pytest.approx(LR) # fresh: top of the schedule
-
- step_slot_schedulers(optimizer, {0: 5 * 64}) # half the horizon
- assert slot_lr(optimizer, 0) == pytest.approx(LR / 2)
-
- step_slot_schedulers(optimizer, {0: 100 * 64}) # far past the horizon
- assert slot_lr(optimizer, 0) == pytest.approx(0.0) # clamped at min_lr
-
-
-def test_adapter_without_num_step_holds_constant():
- optimizer = make_optimizer()
- install_slot_scheduler(make_args(), optimizer, make_adapter(slot=0, num_step=None), resume_step=0)
-
- step_slot_schedulers(optimizer, {0: 12345 * 64})
- assert slot_lr(optimizer, 0) == pytest.approx(LR) # no horizon: never decays
-
-
-def test_resume_position_is_deterministic_from_committed_steps():
- stepped = make_optimizer()
- install_slot_scheduler(make_args(), stepped, make_adapter(slot=0, num_step=10), resume_step=0)
- step_slot_schedulers(stepped, {0: 5 * 64})
-
- resumed = make_optimizer()
- install_slot_scheduler(make_args(), resumed, make_adapter(slot=0, num_step=10), resume_step=5)
-
- assert slot_lr(resumed, 0) == pytest.approx(slot_lr(stepped, 0))
-
-
-def test_only_stepped_slots_advance():
- optimizer = make_optimizer()
- args = make_args()
- install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=10), resume_step=0)
- install_slot_scheduler(args, optimizer, make_adapter(slot=1, num_step=10), resume_step=0)
-
- lr_by_slot = step_slot_schedulers(optimizer, {0: 5 * 64})
-
- assert set(lr_by_slot) == {0}
- assert slot_lr(optimizer, 0) == pytest.approx(LR / 2)
- assert slot_lr(optimizer, 1) == pytest.approx(LR) # co-tenant untouched
-
-
-def test_slot_reuse_installs_a_fresh_schedule():
- optimizer = make_optimizer()
- args = make_args()
- install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=10), resume_step=0)
- step_slot_schedulers(optimizer, {0: 5 * 64})
-
- install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=20), resume_step=0)
- assert slot_lr(optimizer, 0) == pytest.approx(LR) # next tenant starts at the top
-
-
-def test_warmup_ramps_from_init_lr():
- optimizer = make_optimizer()
- args = make_args(lr_warmup_iters=2) # 2 adapter steps of warmup
- install_slot_scheduler(args, optimizer, make_adapter(slot=0, num_step=10), resume_step=0)
-
- assert slot_lr(optimizer, 0) == pytest.approx(0.0) # init_lr
- step_slot_schedulers(optimizer, {0: 64})
- assert slot_lr(optimizer, 0) == pytest.approx(LR / 2) # mid-warmup
- step_slot_schedulers(optimizer, {0: 64})
- assert slot_lr(optimizer, 0) == pytest.approx(LR) # warmed up
diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py b/tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py
deleted file mode 100644
index 93fcd3bd93..0000000000
--- a/tests/fast/backends/megatron_utils/test_multi_lora_slot_cleanup.py
+++ /dev/null
@@ -1,91 +0,0 @@
-"""zero_optimizer_state_for_adapter must reset a retired slot's Adam moments and step clock
-(group-level FusedAdam or per-param torch AdamW) while leaving co-tenant slots untouched."""
-
-import sys
-import types
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from miles.backends.megatron_utils.multi_lora_utils import zero_optimizer_state_for_adapter
-
-MLL_MODULE = "megatron.bridge.peft.multi_lora_layers"
-
-
-class FakeAdapter:
- def __init__(self, params):
- self._params = list(params)
-
- def parameters(self):
- return self._params
-
-
-class FakeMultiLoRALinear:
- def __init__(self, adapters):
- self.adapters = adapters
-
-
-@pytest.fixture()
-def rig(monkeypatch):
- # Stub the lazily imported bridge module so the test needs no bridge build that ships multi-LoRA.
- p0 = torch.nn.Parameter(torch.ones(4))
- p1 = torch.nn.Parameter(torch.ones(4))
- module = FakeMultiLoRALinear({0: FakeAdapter([p0]), 1: FakeAdapter([p1])})
- stub = types.ModuleType(MLL_MODULE)
- stub.MultiLoRALinear = FakeMultiLoRALinear
- stub._iter_multi_lora_modules = lambda model: [module]
- monkeypatch.setitem(sys.modules, MLL_MODULE, stub)
- return SimpleNamespace(p0=p0, p1=p1, model=object())
-
-
-def make_optimizer(groups, state):
- inner = SimpleNamespace(param_groups=groups, state=state)
- return inner, SimpleNamespace(chained_optimizers=[SimpleNamespace(optimizer=inner)])
-
-
-def test_group_level_fused_adam_clock_resets_only_for_the_retired_slot(rig):
- inner, optimizer = make_optimizer(
- groups=[
- {"params": [rig.p0], "miles_multi_lora_slot": 0, "step": 50},
- {"params": [rig.p1], "miles_multi_lora_slot": 1, "step": 50},
- ],
- state={
- rig.p0: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)},
- rig.p1: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)},
- },
- )
-
- zero_optimizer_state_for_adapter(optimizer, rig.model, 0)
-
- assert inner.param_groups[0]["step"] == 0
- assert inner.param_groups[1]["step"] == 50 # co-tenant slot untouched
- assert float(inner.state[rig.p0]["exp_avg"].abs().sum()) == 0.0
- assert float(inner.state[rig.p0]["exp_avg_sq"].abs().sum()) == 0.0
- assert float(inner.state[rig.p1]["exp_avg"].abs().sum()) == 4.0
-
-
-def test_tensor_valued_group_clock_resets_in_place(rig):
- step = torch.tensor(50)
- inner, optimizer = make_optimizer(
- groups=[{"params": [rig.p0], "miles_multi_lora_slot": 0, "step": step}],
- state={rig.p0: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)}},
- )
-
- zero_optimizer_state_for_adapter(optimizer, rig.model, 0)
-
- assert int(step) == 0 # zeroed in place, no rebinding needed
-
-
-def test_per_param_adamw_fallback_clock_resets(rig):
- # torch.optim.AdamW keeps the clock per param; groups carry no "step".
- inner, optimizer = make_optimizer(
- groups=[{"params": [rig.p0], "miles_multi_lora_slot": 0}],
- state={
- rig.p0: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4), "step": torch.tensor(50.0)},
- },
- )
-
- zero_optimizer_state_for_adapter(optimizer, rig.model, 0)
-
- assert float(inner.state[rig.p0]["step"]) == 0.0
diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_utils.py b/tests/fast/backends/megatron_utils/test_multi_lora_utils.py
deleted file mode 100644
index 7b4add13d5..0000000000
--- a/tests/fast/backends/megatron_utils/test_multi_lora_utils.py
+++ /dev/null
@@ -1,45 +0,0 @@
-import logging
-from types import SimpleNamespace
-
-from miles.backends.megatron_utils import multi_lora_optimizer, multi_lora_utils
-
-_MULTI_LORA_UTILS_LOGGER = "miles.backends.megatron_utils.multi_lora_utils"
-
-
-class _FakeSlotScheduler:
- def __init__(self, lr: float) -> None:
- self.optimizer = SimpleNamespace(param_groups=[{"lr": lr}])
- self.increments: list[int] = []
-
- def step(self, increment: int) -> None:
- self.increments.append(increment)
-
-
-def _adapter_lr_messages(caplog) -> list[str]:
- return [record.getMessage() for record in caplog.records if record.name == _MULTI_LORA_UTILS_LOGGER]
-
-
-class TestStepSteppedAdapterSlots:
- def test_stepped_slots_emit_the_train_tag_with_their_new_learning_rates(self, monkeypatch, caplog):
- """Slots whose adapter batch completes log one train-tagged adapter_lr record per rollout/step."""
- optimizer = SimpleNamespace(miles_slot_schedulers={0: _FakeSlotScheduler(0.5), 1: _FakeSlotScheduler(0.25)})
- monkeypatch.setattr(
- multi_lora_optimizer,
- "step_adapter_slots",
- lambda optimizer, model, step_batch_sizes, clip_grad: {0: 1.5, 1: 0.75},
- )
-
- with caplog.at_level(logging.INFO, logger=_MULTI_LORA_UTILS_LOGGER):
- max_grad_norm = multi_lora_utils.step_stepped_adapter_slots(
- SimpleNamespace(clip_grad=1.0),
- [],
- optimizer,
- {"step_adapter_batch_sizes": {0: 64, 1: 32}},
- rollout_id=4,
- step_id=9,
- )
-
- assert max_grad_norm == 1.5
- assert optimizer.miles_slot_schedulers[0].increments == [64]
- assert optimizer.miles_slot_schedulers[1].increments == [32]
- assert _adapter_lr_messages(caplog) == ["train op=adapter_lr rollout=4 step=9 slot_0=0.5 slot_1=0.25"]
diff --git a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py
index 568718dbce..caa92f15b1 100644
--- a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py
+++ b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py
@@ -169,7 +169,6 @@ def test_save_model_does_not_manage_lifecycle(actor_module, monkeypatch):
reload_groups = Mock()
destroy_groups = Mock()
monkeypatch.setattr(actor_module, "save", save)
- monkeypatch.setattr(actor_module, "is_multi_lora_enabled", lambda _args: False)
monkeypatch.setattr(actor_module, "reload_process_groups", reload_groups)
monkeypatch.setattr(actor_module, "destroy_process_groups", destroy_groups)
@@ -578,7 +577,6 @@ def test_critic_output_roundtrips_into_actor_external_data(actor_module: Any, mo
monkeypatch.setattr(actor_module, "log_rollout_data", lambda *_args, **_kwargs: None)
monkeypatch.setattr(actor_module, "log_train_advantage_computation_event", lambda *_args, **_kwargs: None)
monkeypatch.setattr(actor_module, "log_perf_data", lambda *_args, **_kwargs: None)
- monkeypatch.setattr(actor_module, "is_multi_lora_enabled", lambda _args: False)
monkeypatch.setattr(actor_module.train_dump_utils, "save_debug_train_data", lambda *_args, **_kwargs: None)
monkeypatch.setattr(actor_module.torch.cuda, "current_device", lambda: torch.device("cpu"))
critic_values = [torch.tensor([1.0, 2.0]), torch.tensor([3.0])]
@@ -653,7 +651,6 @@ def __init__(self) -> None:
self.connect_calls: list[dict[str, Any]] = []
self.update_weights_calls: int = 0
self.weight_version: int = 0
- self.multi_lora_adapters: dict[str, Any] = {}
def connect_rollout_engines(
self,
@@ -689,7 +686,6 @@ def _weight_update_worker(actor_module: Any, monkeypatch: pytest.MonkeyPatch) ->
worker._heartbeat = Mock()
worker.weight_updater = _RecordingWeightUpdater()
monkeypatch.setattr(actor_module, "print_memory", Mock())
- monkeypatch.setattr(actor_module, "is_multi_lora_enabled", lambda _args: False)
monkeypatch.setattr(actor_module, "get_gloo_group", lambda: None)
monkeypatch.setattr(actor_module.dist, "barrier", lambda **_kwargs: None)
return worker
@@ -859,27 +855,6 @@ def test_update_weights_reconnects_once_per_rollout_snapshot(
assert not updater.conn_status.needs_reconnect({"cell-0": "hash-b"})
-def test_reconnecting_engines_receive_every_loaded_multi_lora_adapter(
- actor_module: Any, monkeypatch: pytest.MonkeyPatch
-) -> None:
- """Reconnecting engines get all loaded adapters even with none pending; a settled connection gets only pending."""
- worker = _weight_update_worker(actor_module, monkeypatch)
- updater = worker.weight_updater
- monkeypatch.setattr(actor_module, "is_multi_lora_enabled", lambda _args: True)
- worker.loaded_adapters = {"alpha": "alpha-weights", "beta": "beta-weights"}
- worker._multi_lora_pending_push = set()
- worker._is_first_replica_megatron_main_rank = False
- engines = [object()]
-
- worker.update_weights(_updatable_engines(engines, {"cell-0": "hash-a"}, gpu_count=4))
- adapters_on_reconnect = updater.multi_lora_adapters
- worker._multi_lora_pending_push = {"beta"}
- worker.update_weights(_updatable_engines(engines, {"cell-0": "hash-a"}, gpu_count=4))
-
- assert adapters_on_reconnect == {"alpha": "alpha-weights", "beta": "beta-weights"}
- assert updater.multi_lora_adapters == {"beta": "beta-weights"}
-
-
def test_reconfigure_indep_dp_forces_the_next_weight_update_to_reconnect(
actor_module: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
diff --git a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
index 2df14f54b1..0e0f2fcba9 100644
--- a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
+++ b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
@@ -4,7 +4,7 @@
import pytest
import torch
-from miles.backends.megatron_utils.multi_lora_utils import slice_lora_to_rank
+from miles.backends.megatron_utils.lora.slots import slice_lora_to_rank
def _padded(shape, live_rows=None, live_cols=None):
diff --git a/tests/fast/backends/megatron_utils/test_slot_checkpoint_keys.py b/tests/fast/backends/megatron_utils/test_slot_checkpoint_keys.py
new file mode 100644
index 0000000000..7b0cdc3001
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_slot_checkpoint_keys.py
@@ -0,0 +1,13 @@
+"""A slot checkpoint's storage keys are slot-agnostic, so any free slot can reload it."""
+
+from types import SimpleNamespace
+
+from miles.backends.megatron_utils.lora.checkpoint import _canonicalize_slot_keys
+
+
+def test_saved_and_loading_slots_meet_at_the_canonical_key():
+ saved = SimpleNamespace(key="decoder.layers.3.linear_qkv.adapters.2.linear_in.weight")
+ loading = SimpleNamespace(key="decoder.layers.3.linear_qkv.adapters.0.linear_in.weight")
+ _canonicalize_slot_keys({"weights": {"w": saved}}, 2)
+ _canonicalize_slot_keys({"weights": {"w": loading}}, 0)
+ assert saved.key == loading.key == "decoder.layers.3.linear_qkv.adapter.linear_in.weight"
diff --git a/tests/fast/backends/megatron_utils/test_step_adapter_slots.py b/tests/fast/backends/megatron_utils/test_step_adapter_slots.py
new file mode 100644
index 0000000000..87f25a8206
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_step_adapter_slots.py
@@ -0,0 +1,67 @@
+"""Slot steps preserve nonfinite skips and propagate unknown execution failures."""
+
+import pytest
+
+from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer, step_slot_optimizers
+
+
+class _FakeSlotOptimizer(SlotOptimizer):
+ """Duck instance: exercises the choreography without a megatron build."""
+
+ def __init__(self, slot, *, prepare_error=None, step_outcome=None):
+ self.slot = slot
+ self._prepare_error = prepare_error
+ self._step_outcome = step_outcome or {"grad_norm": 1.5}
+ self.adam_params = None
+ self.stepped = False
+ self.zeroed = False
+ self.param_gathers = 0
+
+ def apply_adam_params(self, adam_params):
+ self.adam_params = adam_params
+
+ def prepare_grads(self):
+ if self._prepare_error is not None:
+ raise self._prepare_error
+
+ def clip_and_step(self, clip_grad):
+ if "error" in self._step_outcome:
+ raise RuntimeError(self._step_outcome["error"])
+ self.stepped = "grad_norm" in self._step_outcome
+ return self._step_outcome
+
+ def zero_grads(self):
+ self.zeroed = True
+
+ def allgather_params(self):
+ self.param_gathers += 1
+
+
+def _step(slot_optimizers):
+ adam = {slot: {"learning_rate": 1e-4, "grad_clip_norm": 1.0} for slot in slot_optimizers}
+ return step_slot_optimizers(slot_optimizers, adam)
+
+
+def test_an_unknown_step_failure_escapes_before_success_is_reported():
+ healthy = _FakeSlotOptimizer(0)
+ failing = _FakeSlotOptimizer(1, step_outcome={"error": "boom"})
+ with pytest.raises(RuntimeError, match="boom"):
+ _step({0: healthy, 1: failing})
+ assert healthy.stepped
+ assert healthy.param_gathers == 0
+
+
+def test_a_nonfinite_grad_norm_skips_the_step():
+ """BF16 has no grad scaler, so the all-reduced norm is the only inf/nan gate."""
+ skipped = _FakeSlotOptimizer(0, step_outcome={"skipped_nonfinite": 1.0})
+ outcomes = _step({0: skipped})
+ assert outcomes[0] == {"skipped_nonfinite": 1.0}
+ assert not skipped.stepped and skipped.zeroed and skipped.param_gathers == 0
+
+
+def test_an_unknown_preparation_failure_escapes():
+ broken = _FakeSlotOptimizer(0, prepare_error=RuntimeError("prep died"))
+ healthy = _FakeSlotOptimizer(1)
+ with pytest.raises(RuntimeError, match="prep died"):
+ _step({0: broken, 1: healthy})
+ assert not broken.stepped and not healthy.stepped
diff --git a/tests/fast/backends/sglang_utils/test_sglang_config.py b/tests/fast/backends/sglang_utils/test_sglang_config.py
index eab75cfb21..57352d0509 100644
--- a/tests/fast/backends/sglang_utils/test_sglang_config.py
+++ b/tests/fast/backends/sglang_utils/test_sglang_config.py
@@ -278,13 +278,17 @@ def test_giving_both_group_spellings_is_rejected(self, tmp_path):
class TestPrefillNumServersPath:
- def test_prefill_num_servers_counts_engines_not_gpus(self):
+ @pytest.mark.parametrize("multi_lora", [False, True])
+ def test_prefill_num_servers_counts_engines_not_gpus(self, multi_lora):
"""prefill_num_servers is a server count, so its GPU span scales with the engine width."""
cfg = resolve_sglang_config(
- _make_args(rollout_num_gpus=16, prefill_num_servers=3, rollout_num_gpus_per_engine=2)
+ _make_args(
+ rollout_num_gpus=16, prefill_num_servers=3, rollout_num_gpus_per_engine=2, multi_lora=multi_lora
+ )
)
groups = cfg.models[0].server_groups
assert [(group.worker_type, group.num_gpus) for group in groups] == [("prefill", 6), ("decode", 10)]
+ assert cfg.models[0].update_weights is not multi_lora
def test_prefill_consuming_all_gpus_is_rejected(self):
"""prefill_num_servers leaving no decode gpus fails loudly."""
@@ -384,16 +388,19 @@ def test_the_generated_eval_model_is_numbered_after_every_rollout_engine(self):
class TestNeedsOffload:
- def test_no_offload_flag_means_no_group_needs_offload(self, tmp_path):
+ @pytest.mark.parametrize("multi_lora", [False, True])
+ def test_no_offload_flag_means_no_group_needs_offload(self, tmp_path, multi_lora):
"""With offload_rollout off, no group offloads and no memory-saver override is injected."""
cfg = _resolve_yaml(
tmp_path,
"sglang:\n - name: actor\n server_groups:\n - worker_type: regular\n num_gpus: 8\n",
rollout_num_gpus=8,
+ multi_lora=multi_lora,
)
group = cfg.models[0].server_groups[0]
assert group.needs_offload is False
assert "enable_memory_saver" not in group.overrides
+ assert cfg.models[0].update_weights is not multi_lora
def test_groups_overlapping_megatron_offload_and_the_rest_disable_memory_saver(self, tmp_path):
"""Only groups starting inside the megatron gpu range offload; later ones get enable_memory_saver=False."""
diff --git a/tests/fast/backends/training_utils/loss/test_tinker_explicit_targets.py b/tests/fast/backends/training_utils/loss/test_tinker_explicit_targets.py
new file mode 100644
index 0000000000..959fe9ba14
--- /dev/null
+++ b/tests/fast/backends/training_utils/loss/test_tinker_explicit_targets.py
@@ -0,0 +1,58 @@
+"""target_tokens are explicit gather labels; shifted targets reproduce next-token scoring."""
+
+import torch
+from tests.fast.backends.training_utils.loss.loss_test_utils import make_args, make_parallel_state
+
+from miles.backends.training_utils.loss_hub import math_utils, tinker_losses
+
+VOCAB = 32
+
+
+def _naive_compute_log_probs(logits, tokens, _tp_group, *, sampling_mask=None):
+ return torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
+
+
+def _target_logprobs(batch, logits):
+ args = make_args(true_on_policy_mode=False, rollout_temperature=1.0, log_probs_chunk_size=64)
+ return tinker_losses._target_logprobs(args, batch, logits)
+
+
+def _batch(sequences: list[list[int]], targets: list[list[int]]) -> dict:
+ return {
+ "unconcat_tokens": [torch.tensor(sequence) for sequence in sequences],
+ "target_tokens": targets,
+ "total_lengths": [len(sequence) for sequence in sequences],
+ "response_lengths": [len(target) for target in targets],
+ }
+
+
+def test_targets_replace_the_sequence_as_gather_labels(monkeypatch):
+ monkeypatch.setattr(math_utils, "compute_log_probs", _naive_compute_log_probs)
+ make_parallel_state()
+ generator = torch.Generator().manual_seed(0)
+ sequences = [[1, 2, 3, 4, 5], [6, 7, 8]]
+ targets = [[9, 10, 11], [12, 13]] # rl_loop-style: labels differ from the sequence
+ logits = torch.randn((1, 8, VOCAB), generator=generator)
+
+ log_probs = _target_logprobs(_batch(sequences, targets), logits)
+
+ flat = torch.log_softmax(logits.squeeze(0), dim=-1)
+ offset = 0
+ for sequence, target, log_prob in zip(sequences, targets, log_probs, strict=True):
+ start = offset + len(sequence) - len(target) - 1
+ expected = [flat[start + i, label] for i, label in enumerate(target)]
+ torch.testing.assert_close(log_prob, torch.stack(expected))
+ offset += len(sequence)
+
+
+def test_shifted_targets_degenerate_to_next_token_scoring(monkeypatch):
+ monkeypatch.setattr(math_utils, "compute_log_probs", _naive_compute_log_probs)
+ make_parallel_state()
+ generator = torch.Generator().manual_seed(1)
+ sequence = torch.randint(0, VOCAB, (6,), generator=generator).tolist()
+ logits = torch.randn((1, 6, VOCAB), generator=generator)
+
+ (log_probs,) = _target_logprobs(_batch([sequence], [sequence[1:]]), logits)
+
+ expected = torch.log_softmax(logits.squeeze(0)[:-1], dim=-1).gather(-1, torch.tensor(sequence[1:]).unsqueeze(-1))
+ torch.testing.assert_close(log_probs, expected.squeeze(-1))
diff --git a/tests/fast/backends/training_utils/loss/test_tinker_loss_outputs.py b/tests/fast/backends/training_utils/loss/test_tinker_loss_outputs.py
new file mode 100644
index 0000000000..58d7563c13
--- /dev/null
+++ b/tests/fast/backends/training_utils/loss/test_tinker_loss_outputs.py
@@ -0,0 +1,142 @@
+"""Per-datum outputs belong to their loss pass, including when backward recomputes the loss."""
+
+from argparse import Namespace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from miles.backends.training_utils import loss as loss_module
+from miles.backends.training_utils.loss_hub import tinker_losses
+
+
+@pytest.mark.parametrize("recompute", [False, True], ids=["direct", "recomputed"])
+def test_loss_passes_return_independent_detached_outputs(monkeypatch, recompute):
+ parallel = SimpleNamespace(cp=SimpleNamespace(size=1), intra_dp=SimpleNamespace(size=1))
+ monkeypatch.setattr(loss_module, "get_parallel_state", lambda: parallel)
+ monkeypatch.setattr(loss_module, "get_sum_of_sample_mean", lambda *_args, **_kwargs: None)
+ monkeypatch.setattr(tinker_losses, "_target_logprobs", lambda _args, _batch, logits: [logits])
+ args = Namespace(
+ calculate_per_token_loss=False,
+ qkv_format="thd",
+ recompute_loss_function=recompute,
+ use_dynamic_global_batch_size=True,
+ global_batch_size=1,
+ multi_lora=True,
+ )
+ completed = []
+ for sample_index in (7, 11):
+ logprobs = torch.tensor([-0.5, -0.25], requires_grad=True)
+ batch = {
+ "loss_fn": "cross_entropy",
+ "loss_weights": [[2.0, 3.0]],
+ "loss_masks": [torch.ones(2)],
+ "total_lengths": [3],
+ "response_lengths": [2],
+ "sample_indices": [sample_index],
+ "dynamic_global_batch_size": 1,
+ }
+ loss, _, logging = loss_module.loss_function(args, batch, 1, logprobs)
+ loss.backward()
+ assert logprobs.grad.tolist() == [-2.0, -3.0]
+ assert logging["keys"] == ["loss"]
+ assert logging["values"].tolist() == [1.0, 1.75]
+ completed.append(logging["per_datum"])
+
+ assert [[output["sample_index"] for output in outputs] for outputs in completed] == [[7], [11]]
+ for outputs in completed:
+ assert len(outputs) == 1, "loss recomputation must not append another datum output"
+ assert outputs[0]["loss"].item() == 1.75
+ assert outputs[0]["logprobs"].tolist() == [-0.5, -0.25]
+ assert not outputs[0]["loss"].requires_grad
+ assert not outputs[0]["logprobs"].requires_grad
+
+
+@pytest.mark.parametrize("loss_fn", sorted(tinker_losses.TINKER_LOSS_FUNCTIONS))
+def test_a_zero_loss_mask_removes_the_datum_from_every_objective(monkeypatch, loss_fn):
+ logprobs = torch.tensor([-0.5, -0.25], requires_grad=True)
+ monkeypatch.setattr(tinker_losses, "_target_logprobs", lambda _args, _batch, logits: [logits])
+ batch = {
+ "loss_fn": loss_fn,
+ "loss_weights": [[2.0, 3.0]],
+ "advantages": [[1.0, 1.0]],
+ "rollout_log_probs": [[-1.5, -1.25]], # differs from logprobs so the DRO divergence term is nonzero
+ "loss_masks": [torch.zeros(2)],
+ "total_lengths": [3],
+ "response_lengths": [2],
+ "sample_indices": [0],
+ }
+ loss, _ = tinker_losses.TINKER_LOSS_FUNCTIONS[loss_fn](Namespace(), batch, logprobs, None)
+ loss.backward()
+ assert loss.item() == 0.0
+ assert logprobs.grad.abs().sum().item() == 0.0, "a DP-padding datum must contribute no gradient"
+
+
+@pytest.mark.parametrize("recompute", [False, True], ids=["direct", "recomputed"])
+@pytest.mark.parametrize(
+ "loss_fn, config, token_losses, gradients",
+ [
+ ("importance_sampling", {}, [-1, 1.5, -2, 3, -10, 15], [-1, 1.5, -2, 3, -10, 15]),
+ ("ppo", {}, [-1, 2.4, -2, 3, -2.4, 15], [-1, 0, -2, 3, 0, 15]),
+ (
+ "ppo",
+ {"clip_low_threshold": 0.4, "clip_high_threshold": 6},
+ [-1, 1.5, -2, 3, -10, 15],
+ [-1, 1.5, -2, 3, -10, 15],
+ ),
+ ("cispo", {}, [1, -1.5, 2, -3, 8, -12], [-1, 1.5, -2, 3, -8, 12]),
+ (
+ "cispo",
+ {"clip_low_threshold": 0.75, "clip_high_threshold": 2},
+ [1.5, -2.25, 2, -3, 4, -6],
+ [-1.5, 2.25, -2, 3, -4, 6],
+ ),
+ ("dro", {}, None, None),
+ ("dro", {"beta": 0.5}, None, None),
+ ],
+ ids=["is", "ppo-default", "ppo-override", "cispo-default", "cispo-override", "dro-default", "dro-override"],
+)
+def test_nonzero_objectives_and_gradients(monkeypatch, recompute, loss_fn, config, token_losses, gradients):
+ parallel = SimpleNamespace(cp=SimpleNamespace(size=1), intra_dp=SimpleNamespace(size=1))
+ monkeypatch.setattr(loss_module, "get_parallel_state", lambda: parallel)
+ monkeypatch.setattr(loss_module, "get_sum_of_sample_mean", lambda *_args, **_kwargs: None)
+ monkeypatch.setattr(tinker_losses, "_target_logprobs", lambda _args, _batch, logits: [logits[:3], logits[3:]])
+ args = Namespace(
+ calculate_per_token_loss=False,
+ qkv_format="thd",
+ recompute_loss_function=recompute,
+ use_dynamic_global_batch_size=True,
+ global_batch_size=4,
+ multi_lora=True,
+ )
+ # Both advantage signs cross each clipping boundary; the last two tokens are masked.
+ ratios = torch.tensor([0.5, 0.5, 1, 1, 5, 5, 0.5, 5], dtype=torch.float64)
+ logprobs = torch.full((8,), -1.0, dtype=torch.float64, requires_grad=True)
+ sampling_logprobs = -1 - ratios.log()
+ batch = {
+ "loss_fn": loss_fn,
+ "loss_fn_config": config,
+ "advantages": [[2, -3, 2], [-3, 2, -3, 11, -13]],
+ "rollout_log_probs": [sampling_logprobs[:3], sampling_logprobs[3:]],
+ "loss_masks": [torch.ones(3), torch.tensor([1, 1, 1, 0, 0])],
+ "total_lengths": [4, 6],
+ "response_lengths": [3, 5],
+ "sample_indices": [7, 11],
+ "dynamic_global_batch_size": 4,
+ }
+ if loss_fn == "dro":
+ beta = 0.05 if not config else 0.5
+ advantages = torch.tensor([2, -3, 2, -3, 2, -3], dtype=torch.float64)
+ divergence = ratios[:6].log()
+ token_losses = (advantages + beta / 2 * divergence.square()).tolist()
+ gradients = (-advantages + beta * divergence).tolist()
+ expected_losses = torch.tensor([sum(token_losses[:3]), sum(token_losses[3:])], dtype=torch.float64)
+ loss, _, logging = loss_module.loss_function(args, batch, 1, logprobs)
+ loss.backward()
+ torch.testing.assert_close(loss, expected_losses.sum())
+ torch.testing.assert_close(logprobs.grad, torch.tensor(gradients + [0, 0], dtype=torch.float64))
+ outputs = logging["per_datum"]
+ assert [output["sample_index"] for output in outputs] == [7, 11]
+ torch.testing.assert_close(torch.stack([output["loss"] for output in outputs]), expected_losses)
+ torch.testing.assert_close(torch.cat([output["logprobs"] for output in outputs]), logprobs.detach())
+ assert all(not output["loss"].requires_grad and not output["logprobs"].requires_grad for output in outputs)
diff --git a/tests/fast/backends/training_utils/loss/test_training_logprob_reuse.py b/tests/fast/backends/training_utils/loss/test_training_logprob_reuse.py
index d403aeacdc..093580ee15 100644
--- a/tests/fast/backends/training_utils/loss/test_training_logprob_reuse.py
+++ b/tests/fast/backends/training_utils/loss/test_training_logprob_reuse.py
@@ -276,7 +276,7 @@ def test_loss_dispatcher_uses_args_without_extra_policy_kwargs(monkeypatch, skip
batch = make_batch(inputs, "policy_loss")
logits = deep_clone(inputs["policy_logits"])
policy_loss = Mock(return_value=(logits.sum(), {"loss": logits.new_zeros(())}))
- monkeypatch.setattr(loss_module, "get_loss_function", lambda _args: policy_loss)
+ monkeypatch.setattr(loss_module, "get_loss_function", lambda _args, _loss_fn=None: policy_loss)
loss_module.loss_function(
args,
@@ -296,7 +296,7 @@ def test_loss_dispatcher_does_not_apply_actor_skip_to_value_loss(monkeypatch):
batch = make_batch(inputs, "value_loss")
logits = deep_clone(inputs["value_logits"])
value_loss = Mock(return_value=(logits.sum(), {"loss": logits.new_zeros(())}))
- monkeypatch.setattr(loss_module, "get_loss_function", lambda _args: value_loss)
+ monkeypatch.setattr(loss_module, "get_loss_function", lambda _args, _loss_fn=None: value_loss)
loss_module.loss_function(args, batch, 1, logits)
diff --git a/tests/fast/backends/training_utils/test_checkpoint_io.py b/tests/fast/backends/training_utils/test_checkpoint_io.py
new file mode 100644
index 0000000000..aef0b319c1
--- /dev/null
+++ b/tests/fast/backends/training_utils/test_checkpoint_io.py
@@ -0,0 +1,50 @@
+"""Checkpoint filesystem errors propagate to the trainer cell."""
+
+import multiprocessing
+import os
+from pathlib import Path
+
+import pytest
+
+from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+
+
+@pytest.mark.parametrize("error", [OSError("disk full"), RuntimeError("directory creation failed")])
+def test_directory_errors_propagate(error, tmp_path, monkeypatch):
+ def make_dir(*args, **kwargs):
+ raise error
+
+ monkeypatch.setattr(Path, "mkdir", make_dir)
+ with pytest.raises(type(error), match=str(error)) as caught:
+ write_checkpoint_dir(tmp_path / "checkpoint", lambda _: None)
+ assert caught.value is error
+
+
+@pytest.mark.parametrize("crash_before_publish", [True, False])
+def test_crashed_overwrite_keeps_a_complete_checkpoint(tmp_path, crash_before_publish):
+ checkpoint = tmp_path / "checkpoint"
+ write_checkpoint_dir(checkpoint, lambda directory: (directory / "value").write_text("old"))
+ old_version = checkpoint.resolve()
+
+ def overwrite_and_crash():
+ replace = os.replace
+
+ def crash_at_publish(source, destination):
+ if Path(destination) == checkpoint and crash_before_publish:
+ os._exit(73)
+ replace(source, destination)
+ if Path(source) == checkpoint or Path(destination) == checkpoint:
+ os._exit(73)
+
+ os.replace = crash_at_publish
+ write_checkpoint_dir(checkpoint, lambda directory: (directory / "value").write_text("new"))
+
+ child = multiprocessing.get_context("fork").Process(target=overwrite_and_crash)
+ child.start()
+ child.join(timeout=10)
+ assert child.exitcode == 73
+ assert (checkpoint / "value").read_text() == ("old" if crash_before_publish else "new")
+ assert (old_version / "value").read_text() == "old"
+
+ write_checkpoint_dir(checkpoint, lambda directory: (directory / "value").write_text("retry"))
+ assert (checkpoint / "value").read_text() == "retry"
diff --git a/tests/fast/backends/training_utils/weight_update/test_hf_weight_iterator.py b/tests/fast/backends/training_utils/weight_update/test_hf_weight_iterator.py
index f6d3c0ce50..254aa70d05 100644
--- a/tests/fast/backends/training_utils/weight_update/test_hf_weight_iterator.py
+++ b/tests/fast/backends/training_utils/weight_update/test_hf_weight_iterator.py
@@ -62,10 +62,10 @@ def _iter_hf_param_units(self, weights, *, materialize):
for pair in self._base:
yield [pair]
- def _iter_hf_adapter_units(self, lora_name, adapter, *, materialize):
+ def _iter_hf_adapter_units(self, adapter, *, materialize):
self.export_calls.append(adapter)
for name, tensor in self._exported:
- yield [(f"{lora_name}:{name}", tensor)]
+ yield [(name, tensor)]
class TestIterHfWeightsTemplate:
@@ -89,6 +89,15 @@ def test_include_base_false_streams_adapters_only(self):
names = self._names(iterator.iter_hf_weights(None, include_base=False, adapters=[("miles_lora", None)]))
assert names == [f"miles_lora:{n}" for n, _ in SAMPLE_LORA_WEIGHTS]
+ def test_each_adapter_keeps_its_own_prefix(self):
+ """A genexp in the chaining loop late-binds the name and stamps every adapter with the last one."""
+ iterator = _StubIterator(SAMPLE_LORA_WEIGHTS)
+ names = self._names(
+ iterator.iter_hf_weights(None, include_base=False, adapters=[("run_a@1", None), ("run_b@2", None)])
+ )
+ expected = [f"run_a@1:{n}" for n, _ in SAMPLE_LORA_WEIGHTS] + [f"run_b@2:{n}" for n, _ in SAMPLE_LORA_WEIGHTS]
+ assert names == expected
+
def test_no_adapters_matches_base_stream(self):
iterator = _StubIterator([], base=SAMPLE_BASE_ONLY_WEIGHTS)
names = self._names(iterator.iter_hf_weights(None))
diff --git a/tests/fast/fixtures/driver_fakes.py b/tests/fast/fixtures/driver_fakes.py
index 44281843f0..2cc946f23d 100644
--- a/tests/fast/fixtures/driver_fakes.py
+++ b/tests/fast/fixtures/driver_fakes.py
@@ -95,9 +95,6 @@ async def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
async def export_hf(self, rollout_id: int, path: str) -> None:
self.events.append(f"{self.role}_export_hf:{rollout_id}")
- async def reconcile_adapters(self) -> None:
- self.events.append(f"{self.role}_reconcile_adapters")
-
async def offload(self) -> None:
self.events.append(f"{self.role}_offload")
diff --git a/tests/fast/launch_scripts/py_harness.py b/tests/fast/launch_scripts/py_harness.py
index e7e50356b5..5449544fb2 100644
--- a/tests/fast/launch_scripts/py_harness.py
+++ b/tests/fast/launch_scripts/py_harness.py
@@ -71,7 +71,9 @@ def rel(self) -> str:
def iter_py_launch_scripts() -> list[PyLaunchScript]:
- paths = sorted((REPO_ROOT / "scripts").rglob("run_*.py"))
+ paths = sorted(
+ [*(REPO_ROOT / "scripts").rglob("run_*.py"), REPO_ROOT / "examples/multi_lora/serve_qwen3_30b_a3b_tinker.py"]
+ )
return [PyLaunchScript(path=path, entrypoints=tuple(_entrypoint_names(path))) for path in paths]
diff --git a/tests/fast/ray/multi_lora/test_controller.py b/tests/fast/ray/multi_lora/test_controller.py
deleted file mode 100644
index aa8e6b059f..0000000000
--- a/tests/fast/ray/multi_lora/test_controller.py
+++ /dev/null
@@ -1,68 +0,0 @@
-from types import SimpleNamespace
-
-from tests.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=60, suite="stage-a-cpu")
-
-import pytest
-
-import miles.ray.multi_lora.controller as controller_mod
-from miles.ray.multi_lora.backend import MultiLoRABackend
-from miles.ray.multi_lora.controller import CONTROLLER_NAME, CONTROLLER_NAMESPACE, MultiLoRAController
-
-
-def make_args() -> SimpleNamespace:
- return SimpleNamespace(multi_lora_n_adapters=4)
-
-
-class TestConstructorContract:
- def test_keyword_arguments_build_the_backend_and_the_http_server(self):
- """Keyword construction is the supported form and wires up both owned components."""
- args = make_args()
-
- controller = MultiLoRAController.__ray_actor_class__(args=args, router_url="http://router:1", host="10.0.0.9")
-
- assert isinstance(controller.backend, MultiLoRABackend)
- assert controller.backend.router_url == "http://router:1"
- assert controller.server.host == "10.0.0.9"
-
- def test_positional_arguments_are_rejected(self):
- """The actor is launched by keyword only, so positional construction must fail loudly."""
- with pytest.raises(TypeError):
- MultiLoRAController.__ray_actor_class__(make_args(), "http://router:1")
-
-
-class _RecordingActorClass:
- def __init__(self):
- self.options_kwargs = None
- self.positional_args = None
- self.keyword_args = None
-
- def options(self, **kwargs):
- self.options_kwargs = kwargs
- return self
-
- def remote(self, *args, **kwargs):
- self.positional_args = args
- self.keyword_args = kwargs
- return "controller-handle"
-
-
-class TestCreateMultiLoRAController:
- def test_launches_a_named_head_pinned_actor_with_keyword_arguments(self, monkeypatch):
- """The factory names the actor, pins it to the head node and passes every argument by keyword."""
- recorder = _RecordingActorClass()
- monkeypatch.setattr(controller_mod, "MultiLoRAController", recorder)
- monkeypatch.setattr(controller_mod, "compute_ray_pin_head_options", lambda: {"scheduling_strategy": "head"})
- args = make_args()
-
- handle = controller_mod.create_multilora_controller(args, "http://router:1", "10.0.0.9")
-
- assert handle == "controller-handle"
- assert recorder.options_kwargs == {
- "name": CONTROLLER_NAME,
- "namespace": CONTROLLER_NAMESPACE,
- "scheduling_strategy": "head",
- }
- assert recorder.positional_args == ()
- assert recorder.keyword_args == {"args": args, "router_url": "http://router:1", "host": "10.0.0.9"}
diff --git a/tests/fast/ray/multi_lora/test_controller_backend.py b/tests/fast/ray/multi_lora/test_controller_backend.py
deleted file mode 100644
index fe75f985ef..0000000000
--- a/tests/fast/ray/multi_lora/test_controller_backend.py
+++ /dev/null
@@ -1,392 +0,0 @@
-"""Fast tests for AdapterRegistry + MultiLoRABackend validation
-(no Ray, no HTTP I/O, no SGLang, no torch)."""
-
-from types import SimpleNamespace
-
-from tests.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=60, suite="stage-a-cpu")
-
-import pytest
-
-from miles.ray.multi_lora.backend import MultiLoRABackend
-from miles.ray.multi_lora.registry import AdapterRegistry, AdapterState
-from miles.utils.adapter_config import AdapterRunConfig
-from miles.utils.multi_lora import make_rid, min_groups_per_dp_split, parse_adapter
-
-
-# Registration validates that the data path exists; the test file itself is a
-# convenient always-present stand-in.
-DATA_FILE = __file__
-
-
-def make_args(max_adapters: int = 4, save: str | None = None, dp_size: int = 2) -> SimpleNamespace:
- return SimpleNamespace(
- multi_lora_n_adapters=max_adapters,
- save=save,
- lora_rank=32,
- lora_alpha=32,
- rollout_batch_size=16,
- n_samples_per_prompt=4,
- multi_lora_dp_size=dp_size,
- multi_lora_max_adapter_global_batch_size=256,
- )
-
-
-def make_backend(max_adapters: int = 4, save: str | None = None, dp_size: int = 2) -> MultiLoRABackend:
- return MultiLoRABackend(make_args(max_adapters, save, dp_size), "http://unused")
-
-
-def make_config(save: str | None = None, **overrides) -> AdapterRunConfig:
- kwargs = dict(
- rank=8,
- alpha=16,
- data=DATA_FILE,
- rollout_batch_size=4,
- n_samples_per_prompt=4,
- save=save,
- input_key="text",
- label_key="label",
- rm_type="math",
- )
- kwargs.update(overrides)
- return AdapterRunConfig(**kwargs)
-
-
-def register_and_promote(registry: AdapterRegistry, name: str, config=None) -> None:
- registry.register(name, config)
- registry.record_weight_update([name])
-
-
-def test_rid_roundtrip_preserves_names_with_underscores():
- for name in ["a", "adapter_a", "weird__name", "x_y_z"]:
- assert parse_adapter(make_rid(name)) == name
-
-
-def test_register_starts_pending_and_push_promotes():
- registry = AdapterRegistry(max_adapters=4)
- result = registry.register("A", config={"rm_type": "x"})
- assert result == {"name": "A", "slot": 0}
- assert registry.active_adapters() == {} # pending: not sampleable
-
- registry.record_weight_update(["A"])
- assert registry.active_adapters()["A"].slot == 0
- view = registry.active_adapters()["A"]
- assert view.slot == 0
- assert view.config == {"rm_type": "x"}
- assert view.version == 1
-
-
-def test_snapshot_reports_sets_in_registry_vocabulary():
- registry = AdapterRegistry(max_adapters=4)
- register_and_promote(registry, "A")
- registry.register("B", None)
- snapshot = registry.snapshot()
- assert set(snapshot["active"]) == {"A"}
- assert set(snapshot["pending"]) == {"B"}
- assert snapshot["retiring"] == {}
- assert snapshot["cleanup"] == []
- assert set(registry.active_adapters()) == {"A"} # only active adapters are sampleable
-
-
-def test_slot_version_is_monotonic_across_slot_reuse():
- registry = AdapterRegistry(max_adapters=2)
- register_and_promote(registry, "A") # slot 0, version 1
- registry.record_weight_update(["A"]) # version 2
- registry.deregister("A")
- registry.retire_adapters()
- registry.free_slot("A")
-
- registry.register("A2", None) # reuses slot 0
- assert registry.snapshot()["pending"]["A2"].version == 2 # inherits, not reset
- registry.record_weight_update(["A2"])
- assert registry.active_adapters()["A2"].version == 3
-
-
-def test_record_weight_update_only_touches_reported_names():
- registry = AdapterRegistry(max_adapters=4)
- register_and_promote(registry, "A")
- register_and_promote(registry, "B")
- registry.record_weight_update(["A"])
- assert registry.active_adapters()["A"].version == 2
- assert registry.active_adapters()["B"].version == 1
-
-
-def test_register_name_rejected_until_cleanup_done():
- registry = AdapterRegistry(max_adapters=4)
- register_and_promote(registry, "A")
- registry.deregister("A")
- with pytest.raises(ValueError, match="cleaning up"):
- registry.register("A", None) # retiring
- registry.retire_adapters()
- with pytest.raises(ValueError, match="cleaning up"):
- registry.register("A", None) # cleanup
- registry.free_slot("A")
- assert registry.register("A", None) == {"name": "A", "slot": 0}
-
-
-def test_deregister_retires_but_keeps_serving_until_demoted():
- registry = AdapterRegistry(max_adapters=4)
- register_and_promote(registry, "A")
- registry.deregister("A")
- assert registry.adapter_state("A") == AdapterState.RETIRING
- assert "A" in registry.active_adapters() # still sampleable this iteration
- assert "A" in registry.snapshot()["retiring"]
- assert registry.retire_adapters() == ["A"]
- assert registry.active_adapters() == {}
- assert registry.adapter_state("A") == AdapterState.CLEANUP
- assert registry.retire_adapters() == [] # idempotent
-
-
-# make_config(): rollout_batch_size=4 groups/step, n_samples_per_prompt=4.
-
-
-def test_mark_batch_trained_accumulates_and_steps_on_completion():
- registry = AdapterRegistry(max_adapters=4)
- register_and_promote(registry, "A", make_config())
- register_and_promote(registry, "B", make_config())
-
- # Two partial batches accumulate; the third completes the adapter batch.
- registry.record_batch_adapters(1, {"A": 1, "B": 2}, step_names=[])
- assert registry.mark_batch_trained(1) == []
- assert registry.records["A"].accumulated_groups == 1
- assert registry.records["B"].accumulated_groups == 2
-
- registry.record_batch_adapters(2, {"A": 1}, step_names=[])
- assert registry.mark_batch_trained(2) == []
- assert registry.records["A"].accumulated_groups == 2
-
- registry.record_batch_adapters(3, {"A": 2, "B": 2}, step_names=["A", "B"])
- assert registry.mark_batch_trained(3) == ["A", "B"]
- assert registry.step_count("A") == 1
- assert registry.step_count("B") == 1
- assert registry.records["A"].accumulated_groups == 0
- assert registry.records["B"].accumulated_groups == 0
-
- assert registry.mark_batch_trained(3) == [] # record consumed
-
-
-def test_batch_trained_counts_deregistered_adapter_until_freed():
- registry = AdapterRegistry(max_adapters=4)
- register_and_promote(registry, "A", make_config())
- registry.record_batch_adapters(3, {"A": 4}, step_names=["A"])
- registry.deregister("A") # deregistered while its batch is training
- assert registry.mark_batch_trained(3) == ["A"]
- assert registry.step_count("A") == 1 # final ckpt reads this
- registry.retire_adapters()
- assert registry.step_count("A") == 1 # cleanup record still holds it
- registry.free_slot("A")
- assert registry.step_count("A") == 0
-
-
-def test_set_step_on_resume():
- registry = AdapterRegistry(max_adapters=2)
- registry.register("A", make_config())
- registry.set_step("A", 40)
- registry.record_batch_adapters(1, {"A": 4}, step_names=["A"])
- registry.record_weight_update(["A"])
- registry.mark_batch_trained(1)
- assert registry.step_count("A") == 41
-
-
-def test_num_step_deregisters_on_committed_steps():
- registry = AdapterRegistry(max_adapters=2)
- register_and_promote(registry, "A", make_config(num_step=2))
- registry.record_batch_adapters(1, {"A": 4}, step_names=["A"])
- assert registry.mark_batch_trained(1) == ["A"]
- assert registry.adapter_state("A") == AdapterState.ACTIVE
-
- registry.record_batch_adapters(2, {"A": 4}, step_names=["A"])
- assert registry.mark_batch_trained(2) == ["A"]
- assert registry.step_count("A") == 2
- assert registry.adapter_state("A") == AdapterState.RETIRING
-
-
-def test_num_step_is_relative_to_resume_step():
- registry = AdapterRegistry(max_adapters=2)
- register_and_promote(registry, "A", make_config(num_step=2))
- registry.set_step("A", 40)
-
- registry.record_batch_adapters(1, {"A": 4}, step_names=["A"])
- registry.mark_batch_trained(1)
- assert registry.step_count("A") == 41
- assert registry.adapter_state("A") == AdapterState.ACTIVE
-
- registry.record_batch_adapters(2, {"A": 4}, step_names=["A"])
- registry.mark_batch_trained(2)
- assert registry.step_count("A") == 42
- assert registry.adapter_state("A") == AdapterState.RETIRING
-
-
-def test_min_groups_per_dp_split():
- assert min_groups_per_dp_split(n_samples_per_prompt=4, dp_size=8) == 2 # divisor
- assert min_groups_per_dp_split(n_samples_per_prompt=8, dp_size=8) == 1 # equal
- assert min_groups_per_dp_split(n_samples_per_prompt=16, dp_size=8) == 1 # multiple
- with pytest.raises(ValueError, match="divisor or a multiple"):
- min_groups_per_dp_split(n_samples_per_prompt=6, dp_size=8)
-
-
-@pytest.mark.asyncio
-async def test_register_resolves_batch_shape_defaults(tmp_path):
- backend = make_backend(save=str(tmp_path))
- await backend.register("A", AdapterRunConfig(data=DATA_FILE, rm_type="math"))
- config = backend.registry.records["A"].config
- assert config.rollout_batch_size == 16 # <- args.rollout_batch_size
- assert config.n_samples_per_prompt == 4 # <- args.n_samples_per_prompt
- assert config.rank == 32 and config.alpha == 32
- assert config.adapter_global_batch_size == 64
-
-
-@pytest.mark.asyncio
-async def test_register_rejects_bad_batch_shapes(tmp_path):
- backend = make_backend(save=str(tmp_path), dp_size=8)
- with pytest.raises(ValueError, match="divisor or a multiple"):
- await backend.register("B", make_config(n_samples_per_prompt=6, rollout_batch_size=4))
- with pytest.raises(ValueError, match="min_groups_per_dp_split"):
- # dp=8, n_samples=4 -> multiple of 2 groups; 3 groups is not
- await backend.register("C", make_config(rollout_batch_size=3))
- with pytest.raises(ValueError, match="exceeding"):
- await backend.register("D", make_config(rollout_batch_size=128)) # 512 samples > cap 256
- with pytest.raises(ValueError, match="exceeds the allocated maximum rank"):
- await backend.register("E", make_config(rank=64))
- with pytest.raises(ValueError, match="positive integer"):
- await backend.register("F", make_config(rollout_batch_size=0))
- with pytest.raises(ValueError, match="num_step must be a positive integer"):
- await backend.register("G", make_config(num_step=0))
- with pytest.raises(ValueError, match="num_epoch must be a positive integer"):
- await backend.register("H", make_config(num_epoch=0))
- # A valid shape registers fine.
- await backend.register("OK", make_config(rollout_batch_size=8))
-
-
-def test_deregister_holds_slot_until_free_slot():
- registry = AdapterRegistry(max_adapters=2)
- register_and_promote(registry, "A") # slot 0
- register_and_promote(registry, "B") # slot 1
- registry.deregister("A")
- registry.retire_adapters()
- assert not registry.free_slots # slot 0 held until cleanup
- with pytest.raises(RuntimeError, match="No free adapter slots"):
- registry.register("C", None)
- registry.free_slot("A")
- assert registry.register("C", None) == {"name": "C", "slot": 0}
-
-
-@pytest.mark.asyncio
-async def test_free_slot_reaborts_before_releasing_slot():
- """Requests can survive the single retire-time abort (multi-turn groups
- submitting between turns, engine tokenizer-adapter batch misses); free_slot must
- fire one more abort round before the slot becomes reusable."""
- backend = make_backend()
- aborted: list[str] = []
-
- async def record_abort(name: str) -> None:
- aborted.append(name)
-
- backend.abort_adapter_requests = record_abort
-
- register_and_promote(backend.registry, "A")
- await backend.deregister("A")
- await backend.retire_adapters()
- assert aborted == ["A"]
-
- assert await backend.free_slot("A") == 0
- assert aborted == ["A", "A"]
- assert backend.registry.free_slots == {0, 1, 2, 3}
-
-
-@pytest.mark.asyncio
-async def test_free_slot_skips_abort_when_not_in_cleanup():
- backend = make_backend()
- aborted: list[str] = []
-
- async def record_abort(name: str) -> None:
- aborted.append(name)
-
- backend.abort_adapter_requests = record_abort
-
- register_and_promote(backend.registry, "A") # ACTIVE, not CLEANUP
- assert await backend.free_slot("A") == -1
- assert await backend.free_slot("never-registered") == -1
- assert aborted == []
-
-
-@pytest.mark.asyncio
-async def test_custom_backend_validation_rejects():
- class StrictBackend(MultiLoRABackend):
- async def validate_adapter(self, name, config):
- if not config:
- raise ValueError("adapter config is required")
-
- backend = StrictBackend(make_args(), "http://unused")
- with pytest.raises(ValueError, match="config is required"):
- await backend.register("A", None)
- assert backend.registry.active_adapters() == {}
-
- result = await backend.register("A", {"rm_type": "x"})
- assert result == {"name": "A", "slot": 0}
-
-
-def test_register_rejects_unsafe_names():
- registry = AdapterRegistry(max_adapters=4)
- for bad in ["a/b", "..", "a::b", "a b", ""]:
- with pytest.raises(ValueError, match="invalid"):
- registry.register(bad, None)
- registry.register("ok-name_1.2", None)
-
-
-def test_register_rejects_duplicate_save_dir(tmp_path):
- registry = AdapterRegistry(max_adapters=4)
- registry.register("A", make_config(save=tmp_path / "x"))
- with pytest.raises(ValueError, match="already used by adapter 'A'"):
- registry.register("B", make_config(save=tmp_path / "x"))
- registry.register("C", make_config(save=tmp_path / "y"))
-
-
-@pytest.mark.asyncio
-async def test_save_dir_defaults_under_save_root(tmp_path):
- backend = make_backend(save=str(tmp_path))
- await backend.register("A", make_config())
- saved = backend.registry.records["A"].config.save
- assert saved == tmp_path / "adapters" / "A"
-
-
-@pytest.mark.asyncio
-async def test_explicit_save_dir_wins_over_root(tmp_path):
- backend = make_backend(save=str(tmp_path))
- await backend.register("A", make_config(save=tmp_path / "custom"))
- assert backend.registry.records["A"].config.save == tmp_path / "custom"
-
-
-@pytest.mark.asyncio
-async def test_register_fails_without_any_save_dir():
- backend = make_backend(save=None)
- with pytest.raises(ValueError, match="no save dir"):
- await backend.register("A", make_config())
-
-
-@pytest.mark.asyncio
-async def test_register_rejects_missing_data_path(tmp_path):
- # A nonexistent data path would otherwise kill the shared rollout producer
- # thread at the first get_samples, stalling every adapter.
- backend = make_backend(save=str(tmp_path))
- with pytest.raises(ValueError, match="data path"):
- await backend.register("A", make_config(data=str(tmp_path / "missing.jsonl")))
-
-
-@pytest.mark.asyncio
-async def test_register_rejects_unresolvable_reward_config(tmp_path):
- # No adapter rm_type/custom_rm_path and no process-wide --rm-type: every
- # sample would fail reward computation and be dropped.
- backend = make_backend(save=str(tmp_path))
- with pytest.raises(ValueError, match="reward config"):
- await backend.register("A", make_config(rm_type=None))
-
-
-@pytest.mark.asyncio
-async def test_register_accepts_reward_config_from_global_args(tmp_path):
- args = make_args(save=str(tmp_path))
- args.rm_type = "math"
- backend = MultiLoRABackend(args, "http://unused")
- await backend.register("A", make_config(rm_type=None))
- assert backend.registry.records["A"].config.rm_type is None # resolved at reward time via args
diff --git a/tests/fast/ray/multi_lora/test_controller_http.py b/tests/fast/ray/multi_lora/test_controller_http.py
deleted file mode 100644
index b270017597..0000000000
--- a/tests/fast/ray/multi_lora/test_controller_http.py
+++ /dev/null
@@ -1,239 +0,0 @@
-"""HTTP tests for the MultiLoRAHTTPServer control plane with a mock router
-(no Ray, no SGLang)."""
-
-import json
-from contextlib import asynccontextmanager
-from pathlib import Path
-from types import SimpleNamespace
-
-import aiohttp
-import pytest
-from aiohttp import web
-
-from tests.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=60, suite="stage-a-cpu")
-
-from miles.ray.multi_lora.backend import MultiLoRABackend
-from miles.ray.multi_lora.http_server import MultiLoRAHTTPServer
-from miles.utils.adapter_config import AdapterRunConfig
-from miles.utils.multi_lora import RID_SEPARATOR
-
-
-# Registration validates that the data path exists; the test file itself is a
-# convenient always-present stand-in.
-DATA_FILE = __file__
-
-
-def minimal_config(name: str) -> dict:
- return {"data": DATA_FILE, "rm_type": "math", "save": f"/tmp/adapters/{name}"}
-
-
-class ControllerHarness:
- """Running control plane (backend + API listener) against a mock router
- that serves /list_workers and records /abort_request posts."""
-
- def __init__(self, session: aiohttp.ClientSession, backend: MultiLoRABackend, srv: MultiLoRAHTTPServer):
- self.session = session
- self.backend = backend
- self.srv = srv
- self.aborts: list[dict] = []
-
- @property
- def api_base(self) -> str:
- return f"http://127.0.0.1:{self.srv.actual_api_port}"
-
- async def api_post(self, path: str, payload: dict) -> tuple[int, dict]:
- async with self.session.post(f"{self.api_base}{path}", json=payload) as resp:
- return resp.status, await resp.json()
-
- async def api_get(self, path: str) -> tuple[int, dict, dict]:
- async with self.session.get(f"{self.api_base}{path}") as resp:
- headers = {k.lower(): v for k, v in resp.headers.items()}
- return resp.status, await resp.json(), headers
-
- async def api_delete(self, path: str) -> tuple[int, dict]:
- async with self.session.delete(f"{self.api_base}{path}") as resp:
- return resp.status, await resp.json()
-
- async def register(self, name: str) -> tuple[int, dict]:
- status, body = await self.api_post("/adapter_runs", {"name": name, "config": minimal_config(name)})
- # Registered adapters start pending; a weight push promotes them.
- self.backend.registry.record_weight_update([name])
- return status, body
-
- async def deregister(self, name: str) -> tuple[int, dict]:
- return await self.api_delete(f"/adapter_runs/{name}")
-
- async def active(self) -> dict:
- _, body, _ = await self.api_get("/adapter_runs")
- return {
- s["name"]: {"slot": s["slot"], "version": s["version"], "step": s["step"]}
- for s in body["adapters"]
- if s["state"] == "ACTIVE"
- }
-
-
-@asynccontextmanager
-async def running_controller(server_cls=MultiLoRAHTTPServer):
- router_url = ""
- harness: ControllerHarness | None = None
-
- async def router_handler(request):
- if request.path == "/list_workers":
- return web.json_response({"urls": [router_url]})
- if request.path == "/abort_request":
- harness.aborts.append(json.loads(await request.read()))
- return web.json_response({})
- return web.json_response({}, status=404)
-
- app = web.Application()
- app.router.add_resource("/{tail:.*}").add_route("*", router_handler)
- runner = web.AppRunner(app)
- await runner.setup()
- site = web.TCPSite(runner, "127.0.0.1", 0)
- await site.start()
- router_url = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}"
-
- backend = MultiLoRABackend(
- SimpleNamespace(
- multi_lora_n_adapters=4,
- save=None,
- lora_rank=32,
- lora_alpha=32,
- rollout_batch_size=16,
- n_samples_per_prompt=4,
- multi_lora_dp_size=2,
- multi_lora_max_adapter_global_batch_size=256,
- ),
- router_url,
- )
- srv = server_cls(backend)
- await backend.init()
- await srv.start()
- try:
- async with aiohttp.ClientSession() as session:
- harness = ControllerHarness(session, backend, srv)
- yield harness
- finally:
- await srv.stop()
- await backend.close()
- await runner.cleanup()
-
-
-@pytest.mark.asyncio
-async def test_register_and_active_view():
- async with running_controller() as ctl:
- status, body = await ctl.register("A")
- assert status == 200
- assert body["slot"] == 0
- assert await ctl.active() == {"A": {"slot": 0, "version": 1, "step": 0}}
-
-
-@pytest.mark.asyncio
-async def test_deregister_marks_and_retire_adapters_aborts():
- """Deregistration only marks; the driver-synced apply performs the
- demotion and fans out one prefix abort per worker."""
- async with running_controller() as ctl:
- await ctl.register("A")
- status, _ = await ctl.deregister("A")
- assert status == 200
- assert ctl.aborts == [] # still serving until the sync point
- assert "A" in ctl.backend.registry.active_adapters()
-
- applied = await ctl.backend.retire_adapters()
- assert applied == ["A"]
- assert ctl.aborts == [{"rid": f"A{RID_SEPARATOR}", "prefix": True}]
- assert ctl.backend.registry.active_adapters() == {}
-
-
-@pytest.mark.asyncio
-async def test_register_json_config_validates_to_adapter_config():
- """FastAPI validates the JSON body straight into AdapterRunConfig (422 on bad
- payloads)."""
- async with running_controller() as ctl:
- config = {
- "rank": 8,
- "data": DATA_FILE,
- "save": "/tmp/adapters/A",
- "rm_type": "math",
- }
- status, _ = await ctl.api_post("/adapter_runs", {"name": "A", "config": config})
- assert status == 200
- record = ctl.backend.registry.find("A")
- assert isinstance(record.config, AdapterRunConfig)
- assert record.config.data == DATA_FILE
- assert Path(record.config.save) == Path("/tmp/adapters/A")
- assert record.config.input_key == "text" # dataclass default
-
- status, _ = await ctl.api_post("/adapter_runs", {"name": "B", "config": {"rank": 8}})
- assert status == 422 # data is required
-
- status, _ = await ctl.api_post("/adapter_runs", {"name": "C"})
- assert status == 400 # exactly one of config/yaml_path
-
-
-@pytest.mark.asyncio
-async def test_state_endpoint_reports_lifecycle_and_completed():
- """States walk PENDING -> ACTIVE -> RETIRING -> CLEANUP -> COMPLETED;
- unknown names report null; COMPLETED is retained after free_slot."""
- async with running_controller() as ctl:
- await ctl.api_post("/adapter_runs", {"name": "A", "config": minimal_config("A")})
-
- async def state_of(name):
- _, body, _ = await ctl.api_get(f"/adapter_runs/state?names={name}")
- return body["states"][name]
-
- assert await state_of("A") == "PENDING"
- ctl.backend.registry.record_weight_update(["A"])
- assert await state_of("A") == "ACTIVE"
-
- await ctl.deregister("A")
- assert await state_of("A") == "RETIRING"
- await ctl.backend.retire_adapters()
- assert await state_of("A") == "CLEANUP"
-
- ctl.backend.registry.free_slot("A")
- assert await state_of("A") == "COMPLETED"
- assert await state_of("nope") is None
-
- # GET by name serves the completed record; DELETE of unknown 404s.
- status, body, _ = await ctl.api_get("/adapter_runs/A")
- assert status == 200 and body["state"] == "COMPLETED"
- status, _ = await ctl.api_delete("/adapter_runs/nope")
- assert status == 404
-
- # Re-registration reclaims the name; the completed record is dropped.
- status, _ = await ctl.api_post(
- "/adapter_runs",
- {"name": "A", "config": {"data": DATA_FILE, "rm_type": "math", "save": "/tmp/adapters/A2"}},
- )
- assert status == 200
- assert await state_of("A") == "PENDING"
-
-
-@pytest.mark.asyncio
-async def test_custom_server_subclass_adds_routes():
- class CustomServer(MultiLoRAHTTPServer):
- def create_app(self):
- app = super().create_app()
-
- @app.middleware("http")
- async def tag_response(request, call_next):
- response = await call_next(request)
- response.headers["X-Custom-Server"] = "1"
- return response
-
- return app
-
- def add_routes(self, app):
- super().add_routes(app)
- app.get("/custom_status")(self.custom_status)
-
- async def custom_status(self):
- return {"custom": True, "active": sorted(self.backend.registry.active_adapters())}
-
- async with running_controller(server_cls=CustomServer) as ctl:
- _, body, headers = await ctl.api_get("/custom_status")
- assert headers.get("x-custom-server") == "1"
- assert body == {"custom": True, "active": []}
diff --git a/tests/fast/ray/rollout/test_metrics.py b/tests/fast/ray/rollout/test_metrics.py
index 046fd505c6..df548e9f38 100644
--- a/tests/fast/ray/rollout/test_metrics.py
+++ b/tests/fast/ray/rollout/test_metrics.py
@@ -13,7 +13,7 @@
log_rollout_data,
)
from miles.rollout.session.v2.metrics import SESSION_ROLLOUT_METRICS_KEY
-from miles.utils.types import AdapterRef, Sample, WeightVersionSpan, WeightVersionsPerCall
+from miles.utils.types import Sample, WeightVersionSpan, WeightVersionsPerCall
class TestEpisodeResponseLengthMetrics:
@@ -62,18 +62,6 @@ def test_total_length_counts_masked_and_unmasked_tokens_in_every_sample(self):
assert out["episode_response_length/mean"] == pytest.approx(4.5)
assert out["episode_total_response_length/mean"] == pytest.approx(8.0)
- def test_multi_lora_samples_emit_no_episode_length_metrics(self):
- samples = [
- make_sample(index=0, rollout_id=10, adapter=AdapterRef(name="adapter-a", slot=0)),
- make_sample(index=0, rollout_id=10, adapter=AdapterRef(name="adapter-b", slot=1)),
- ]
-
- assert _compute_episode_response_length_metrics(samples) == {}
- out = _compute_metrics_from_samples(make_args(advantage_estimator="ppo"), samples)
- assert not any(key.startswith("episode_response_length/") for key in out)
- assert "episode_total_response_length/mean" not in out
- assert out["response_len/mean"] == pytest.approx(4.0)
-
def test_removed_sample_has_zero_effective_length_but_keeps_total_length(self):
sample = make_sample(
index=0,
@@ -128,21 +116,6 @@ def test_rollout_ids_are_scoped_by_prompt_group(self):
assert out["episode_raw_reward"] == pytest.approx(0.5)
- def test_rollout_ids_are_scoped_by_adapter(self):
- args = make_args(reward_key=None)
- adapter_a = AdapterRef(name="adapter-a", slot=0)
- adapter_b = AdapterRef(name="adapter-b", slot=1)
- samples = [
- make_sample(group_index=0, rollout_id=10, adapter=adapter_a, reward=1.0),
- make_sample(group_index=0, rollout_id=10, adapter=adapter_a, reward=1.0),
- make_sample(group_index=0, rollout_id=10, adapter=adapter_a, reward=1.0),
- make_sample(group_index=0, rollout_id=10, adapter=adapter_b, reward=0.0),
- ]
-
- out = _compute_training_sample_metrics(args, samples)
-
- assert out["episode_raw_reward"] == pytest.approx(0.5)
-
def test_metadata_raw_reward_and_fallback_identities(self):
args = make_args(reward_key=None)
samples = [
diff --git a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py
deleted file mode 100644
index 2dc88bd5f6..0000000000
--- a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py
+++ /dev/null
@@ -1,313 +0,0 @@
-"""Unit tests for multi-LoRA batch collection (get_groups + collect_batch):
-group-multiple math, adapter batch capping, step stamping, coalesce timeout,
-round-robin fairness, retirement, and staleness filtering. No Ray, no engines:
-the worker is built bare."""
-
-import asyncio
-import threading
-import time
-from collections import defaultdict, deque
-from types import SimpleNamespace
-
-import pytest
-
-from miles.rollout.multi_lora.async_rollout import (
- AsyncMultiLoRAWorker,
- GroupBuffer,
- MultiLoRAWorkerMetrics,
- collect_batch,
- group_adapter_name,
-)
-from miles.utils.adapter_config import AdapterRun, AdapterRunConfig
-from miles.utils.types import AdapterRef, Sample
-
-
-def make_args(**overrides) -> SimpleNamespace:
- args = SimpleNamespace(
- global_batch_size=16,
- multi_lora_dp_size=4,
- multi_lora_max_coalesce_wait_s=0.05,
- max_weight_staleness=None,
- reward_key=None,
- )
- for key, value in overrides.items():
- setattr(args, key, value)
- return args
-
-
-def make_worker(args=None) -> AsyncMultiLoRAWorker:
- worker = AsyncMultiLoRAWorker.__new__(AsyncMultiLoRAWorker)
- worker.args = args or make_args()
- worker.buffer_lock = threading.Lock()
- worker.buffers = defaultdict(GroupBuffer)
- worker.rotation = deque()
- worker.dynamic_filter = None
- worker.metrics = MultiLoRAWorkerMetrics()
- worker.registrations = {}
- worker.failure = None
- worker.state = SimpleNamespace(sampling_params={})
- worker.generate_fn = None
- worker.data_source = None
- return worker
-
-
-def adapter_run(
- name: str,
- slot: int,
- rollout_batch_size: int = 4,
- n_samples_per_prompt: int = 4,
- accumulated_groups: int = 0,
- version: int = 1,
- registration_id: str = "",
-) -> AdapterRun:
- config = AdapterRunConfig(
- data="/d",
- rank=8,
- alpha=16,
- rollout_batch_size=rollout_batch_size,
- n_samples_per_prompt=n_samples_per_prompt,
- )
- return AdapterRun(
- name=name,
- config=config,
- slot=slot,
- version=version,
- step=0,
- accumulated_groups=accumulated_groups,
- registration_id=registration_id,
- )
-
-
-def make_group(
- adapter: AdapterRun, slot_version: int | None = None, registration_id: str | None = None
-) -> list[Sample]:
- samples = []
- for _ in range(adapter.config.n_samples_per_prompt):
- sample = Sample(prompt="p", adapter=AdapterRef(adapter.name, adapter.slot))
- if slot_version is not None:
- sample.metadata["slot_version"] = slot_version
- if registration_id is not None:
- sample.metadata["registration_id"] = registration_id
- samples.append(sample)
- return samples
-
-
-def buffer_groups(
- worker, adapter: AdapterRun, count: int, slot_version: int | None = None, registration_id: str | None = None
-):
- for _ in range(count):
- worker.buffers[adapter.name].put(make_group(adapter, slot_version, registration_id))
-
-
-def snapshot_of(*adapters: AdapterRun, retiring: tuple[AdapterRun, ...] = ()) -> dict:
- return {
- "active": {a.name: a for a in adapters},
- "retiring": {a.name: a for a in retiring},
- "cleanup": [],
- }
-
-
-def collect(worker, snapshot):
- return asyncio.run(collect_batch(worker.args, worker, snapshot))
-
-
-@pytest.mark.asyncio
-async def test_process_and_enqueue_drops_missing_reward(monkeypatch):
- worker = make_worker()
- adapter = adapter_run("A", 0)
- group = make_group(adapter)
-
- async def fake_process_group(*args, **kwargs):
- return group
-
- monkeypatch.setattr("miles.rollout.multi_lora.async_rollout.process_group", fake_process_group)
-
- await worker.process_and_enqueue(group)
-
- assert len(worker.buffers["A"]) == 0
- assert worker.metrics.pop_metrics() == {
- "rollout/dynamic_filter/drop_group_has_missing_reward": 1,
- }
-
-
-def test_no_pop_until_a_whole_group_multiple_is_buffered():
- # dp=8 with n_samples=4 -> multiple = 2 groups; one buffered group is below the multiple.
- worker = make_worker(make_args(multi_lora_dp_size=8))
- a = adapter_run("A", 0, rollout_batch_size=4, n_samples_per_prompt=4)
- buffer_groups(worker, a, count=1)
- groups, counts = worker.get_groups(snapshot_of(a), 16, {})
- assert (groups, counts) == ([], {})
-
- buffer_groups(worker, a, count=1)
- groups, counts = worker.get_groups(snapshot_of(a), 16, {})
- assert len(groups) == 2
- assert counts == {"A": 2}
-
-
-def test_reaching_target_stops_collecting():
- worker = make_worker()
- a = adapter_run("A", 0, rollout_batch_size=8) # adapter batch: 8 groups
- buffer_groups(worker, a, count=5) # 20 samples > 16 target
- start = time.monotonic()
- batch = collect(worker, snapshot_of(a))
- assert time.monotonic() - start < worker.args.multi_lora_max_coalesce_wait_s # no timeout waited
- assert batch.group_counts == {"A": 4} # stops once 16 samples are reached
- assert batch.step_names == [] # adapter batch (8 groups) not complete
- assert len(worker.buffers["A"]) == 1
-
-
-def test_below_target_ships_after_no_progress_timeout():
- worker = make_worker()
- a = adapter_run("A", 0, rollout_batch_size=8)
- buffer_groups(worker, a, count=1) # 4 samples < 16 target
- start = time.monotonic()
- batch = collect(worker, snapshot_of(a))
- assert time.monotonic() - start >= worker.args.multi_lora_max_coalesce_wait_s
- assert batch.group_counts == {"A": 1}
-
-
-def test_collection_capped_at_remaining_groups_and_step_stamped():
- worker = make_worker()
- # Adapter batch = 4 groups; 3 already banked -> 1 remaining, despite 4 buffered.
- a = adapter_run("A", 0, rollout_batch_size=4, accumulated_groups=3)
- buffer_groups(worker, a, count=4)
- batch = collect(worker, snapshot_of(a))
- assert batch.group_counts == {"A": 1}
- assert batch.step_names == ["A"]
- assert batch.step_slots == [0]
- assert len(worker.buffers["A"]) == 3 # surplus stays buffered
-
-
-def test_batch_never_overshoots_adapter_batch_across_fetches():
- """Groups arriving after an adapter's remaining groups are already in the
- batch must not be popped into the same batch."""
- worker = make_worker()
- a = adapter_run("A", 0, rollout_batch_size=2)
- buffer_groups(worker, a, count=2)
- groups, counts = worker.get_groups(snapshot_of(a), 16, {})
- assert len(groups) == 2 # whole remaining batch
-
- buffer_groups(worker, a, count=2) # fresh arrivals mid-collection
- groups, counts = worker.get_groups(snapshot_of(a), 16, counts)
- assert groups == []
-
- groups, _counts = worker.get_groups(snapshot_of(a), 16, {}) # next batch may pop them
- assert len(groups) == 2
-
-
-def test_pops_interleave_adapters_round_robin():
- worker = make_worker()
- a = adapter_run("A", 0, rollout_batch_size=16)
- b = adapter_run("B", 1, rollout_batch_size=16)
- buffer_groups(worker, a, count=2)
- buffer_groups(worker, b, count=2)
- groups, counts = worker.get_groups(snapshot_of(a, b), 16, {})
- assert [group_adapter_name(g) for g in groups] == ["A", "B", "A", "B"]
- assert counts == {"A": 2, "B": 2}
- groups, counts = worker.get_groups(snapshot_of(a, b), 16, counts)
- assert groups == [] # buffers drained
-
-
-def test_cursor_persists_across_batches():
- worker = make_worker(make_args(global_batch_size=8))
- a = adapter_run("A", 0, rollout_batch_size=16)
- b = adapter_run("B", 1, rollout_batch_size=16)
- buffer_groups(worker, a, count=4)
- buffer_groups(worker, b, count=4)
-
- # 8-sample target = 2 groups per batch; collection interleaves A and B.
- batch = collect(worker, snapshot_of(a, b))
- assert batch.group_counts == {"A": 1, "B": 1}
-
- # The next batch continues from the cursor, not from A again.
- batch = collect(worker, snapshot_of(a, b))
- assert batch.group_counts == {"A": 1, "B": 1}
- assert len(worker.buffers["A"]) == 2
- assert len(worker.buffers["B"]) == 2
-
-
-def test_retiring_adapter_remains_selectable_until_retired():
- """RETIRING adapters keep serving until the reconcile sync point (base
- deregistration semantics): buffered groups stay poppable."""
- worker = make_worker()
- a = adapter_run("A", 0, rollout_batch_size=4)
- buffer_groups(worker, a, count=4)
- batch = collect(worker, snapshot_of(retiring=(a,)))
- assert batch.group_counts == {"A": 4}
- assert batch.step_names == ["A"]
-
-
-def test_retired_adapter_buffers_are_discarded():
- """Once an adapter leaves the snapshot (retired at reconcile), its buffered
- tail is dropped."""
- worker = make_worker()
- a = adapter_run("A", 0, rollout_batch_size=4)
- b = adapter_run("B", 1, rollout_batch_size=4)
- buffer_groups(worker, a, count=3)
- groups, _counts = worker.get_groups(snapshot_of(b), 16, {}) # A gone from snapshot
- assert groups == []
- assert "A" not in worker.buffers # tail discarded with the adapter
-
-
-def test_stale_buffered_groups_are_dropped():
- worker = make_worker(make_args(max_weight_staleness=1))
- a = adapter_run("A", 0, rollout_batch_size=4, version=5)
- buffer_groups(worker, a, count=2, slot_version=3) # staleness 2 > 1
- buffer_groups(worker, a, count=1, slot_version=5) # fresh
- batch = collect(worker, snapshot_of(a))
- assert batch.group_counts == {"A": 1} # only the fresh group ships
-
-
-def test_empty_collection_times_out_instead_of_spinning_forever():
- worker = make_worker(make_args(multi_lora_max_empty_wait_s=0.02))
- a = adapter_run("A", 0, rollout_batch_size=4)
- with pytest.raises(RuntimeError, match="No poppable groups collected before empty timeout"):
- collect(worker, snapshot_of(a))
-
-
-def test_re_registered_name_drops_previous_tenant_buffer_and_metrics():
- # A retires while its buffer still holds groups; the driver idles (no
- # generate), then the operator re-registers the same name. The new
- # tenant's first get_groups must not ship the old tenant's groups nor
- # inherit its partial step statistics.
- worker = make_worker()
- old = adapter_run("A", 0, registration_id="reg-old")
- buffer_groups(worker, old, count=2, registration_id="reg-old")
- worker.get_groups(snapshot_of(old), 0, {}) # worker has seen the old tenant
- worker.metrics.step_rewards["A"].append(1.0) # old tenant's partial step stats
-
- new = adapter_run("A", 0, registration_id="reg-new")
- groups, counts = worker.get_groups(snapshot_of(new), 16, {})
-
- assert (groups, counts) == ([], {})
- assert len(worker.buffers["A"]) == 0
- assert "A" not in worker.metrics.step_rewards
-
-
-def test_straggler_group_of_previous_registration_is_dropped():
- # An in-flight generation of the old tenant lands in the buffer after the
- # re-registration sweep already reset it; only the new tenant's groups ship.
- worker = make_worker()
- new = adapter_run("A", 0, registration_id="reg-new")
- worker.get_groups(snapshot_of(new), 0, {}) # sweep records the new registration
- buffer_groups(worker, new, count=1, registration_id="reg-old") # straggler
- buffer_groups(worker, new, count=1, registration_id="reg-new")
-
- groups, counts = worker.get_groups(snapshot_of(new), 16, {})
-
- assert counts == {"A": 1}
- assert [s.metadata["registration_id"] for g in groups for s in g] == ["reg-new"] * 4
-
-
-def test_dead_producer_surfaces_its_cause_instead_of_timing_out():
- # A producer-thread failure (e.g. an adapter whose dataset vanished) stops
- # generation for every adapter; collect_batch must raise the recorded cause
- # immediately, not wait out the empty-batch timeout.
- worker = make_worker(make_args(multi_lora_max_empty_wait_s=30.0))
- worker.failure = RuntimeError("dataset gone")
- a = adapter_run("A", 0)
- start = time.monotonic()
- with pytest.raises(RuntimeError, match="producer thread died") as excinfo:
- collect(worker, snapshot_of(a))
- assert time.monotonic() - start < 1.0 # no timeout wait
- assert "dataset gone" in repr(excinfo.value.__cause__)
diff --git a/tests/fast/ray/rollout/test_multi_lora_process_group.py b/tests/fast/ray/rollout/test_multi_lora_process_group.py
deleted file mode 100644
index b97ff2238a..0000000000
--- a/tests/fast/ray/rollout/test_multi_lora_process_group.py
+++ /dev/null
@@ -1,56 +0,0 @@
-"""Pins process_group's submission-time slot-version stamping: the staleness
-filter compares against the version live when the group was submitted, not
-when it completed."""
-
-import pytest
-
-import miles.rollout.multi_lora.async_rollout as mod
-from miles.rollout.multi_lora.async_rollout import process_group
-from miles.utils.types import AdapterRef, Sample
-
-
-class FakeDataSource:
- def __init__(self) -> None:
- self.added: list = []
-
- def add_samples(self, groups) -> None:
- self.added.extend(groups)
-
-
-class FakeAdapterView:
- def __init__(self, version: int, registration_id: str = "reg-1") -> None:
- self.version = version
- self.registration_id = registration_id
-
-
-class FakeAdaptersCache:
- def __init__(self, versions: dict[str, int]) -> None:
- self.versions = versions
-
- def bump(self, name: str, to: int) -> None:
- self.versions[name] = to
-
- async def get(self, adapter_name: str) -> FakeAdapterView | None:
- version = self.versions.get(adapter_name)
- return FakeAdapterView(version) if version is not None else None
-
-
-@pytest.mark.asyncio
-async def test_process_group_stamps_submission_version(monkeypatch):
- """The stamp is the version live at submission (5), not completion (7)."""
- cache = FakeAdaptersCache({"A": 5})
-
- async def gen(args, group, sampling_params):
- cache.bump("A", 7) # update lands mid-generation
- for s in group:
- s.status = Sample.Status.COMPLETED
- return group
-
- monkeypatch.setattr(mod, "AdaptersCache", lambda: cache)
-
- g = [Sample(prompt="p", adapter=AdapterRef("A", 0))]
- result = await process_group(None, g, {}, gen, FakeDataSource())
-
- assert result is g
- assert g[0].metadata["slot_version"] == 5
- assert g[0].metadata["registration_id"] == "reg-1"
diff --git a/tests/fast/ray/rollout/test_multi_lora_train_data.py b/tests/fast/ray/rollout/test_multi_lora_train_data.py
deleted file mode 100644
index 86f259b3b8..0000000000
--- a/tests/fast/ray/rollout/test_multi_lora_train_data.py
+++ /dev/null
@@ -1,106 +0,0 @@
-"""Multi-LoRA train-data pipeline: batch metadata extraction, exact dynamic
-batch size, per-adapter batch loss scales, step stamping, and per-group reward
-normalization with heterogeneous group sizes."""
-
-import pytest
-
-from tests.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=60, suite="stage-a-cpu")
-
-from tests.fast.ray.rollout.conftest import make_args, make_sample
-
-from miles.ray.rollout.rollout_data_conversion import postprocess_rollout_data
-from miles.ray.rollout.train_data_conversion import convert_samples_to_train_data
-from miles.utils.types import AdapterRef
-
-
-def multi_lora_args(**overrides):
- defaults = dict(
- multi_lora=True,
- use_dynamic_global_batch_size=True,
- grpo_std_normalization=True,
- )
- defaults.update(overrides)
- return make_args(**defaults)
-
-
-def adapter_group(
- name: str,
- slot: int,
- n_samples: int,
- adapter_global_batch_size: int,
- rewards: list[float],
- start_index: int = 0,
-):
- assert len(rewards) == n_samples
- group = []
- for k in range(n_samples):
- sample = make_sample(index=start_index + k, reward=rewards[k])
- sample.adapter = AdapterRef(name, slot)
- sample.metadata = {"adapter_global_batch_size": adapter_global_batch_size}
- group.append(sample)
- return group
-
-
-def make_batch():
- """Two adapters, heterogeneous group sizes: A steps this batch, B doesn't."""
- groups = [
- adapter_group("A", 0, 4, 16, [1.0, 0.0, 1.0, 0.0], start_index=0),
- adapter_group("A", 0, 4, 16, [1.0, 1.0, 1.0, 1.0], start_index=4),
- adapter_group("B", 1, 2, 32, [3.0, 1.0], start_index=8),
- ]
- groups[0][0].metadata["step_slots"] = [0]
- groups[0][0].metadata["step_adapter_names"] = ["A"]
- return groups
-
-
-def run_pipeline(dp_size: int = 2):
- args = multi_lora_args()
- data, metadata = postprocess_rollout_data(args, make_batch(), train_parallel_config={"dp_size": dp_size})
- train_data = convert_samples_to_train_data(
- args,
- data,
- metadata=metadata,
- custom_convert_samples_to_train_data_func=None,
- custom_reward_post_process_func=None,
- )
- return data, metadata, train_data
-
-
-def test_postprocess_extracts_batch_metadata_and_exact_batch_size():
- data, metadata, _ = run_pipeline()
- assert metadata["prompt_group_sizes"] == [4, 4, 2]
- assert metadata["step_slots"] == [0]
- assert metadata["step_adapter_names"] == ["A"]
- assert metadata["dynamic_global_batch_size"] == 10 # exact batch size, no trim
- assert len(data) == 10 # flattened
- assert "step_slots" not in data[0].metadata # lifted out
-
-
-def test_multi_lora_rejects_dp_indivisible_batch():
- args = multi_lora_args()
- with pytest.raises(ValueError, match="not divisible by dp_size"):
- postprocess_rollout_data(args, make_batch(), train_parallel_config={"dp_size": 4})
-
-
-def test_step_fields():
- _, _, train_data = run_pipeline()
- assert train_data["adapter_slots"] == [0] * 8 + [1] * 2
- assert train_data["step_slots"] == [0]
- assert train_data["step_adapter_names"] == ["A"]
- # Only A steps: the trainer scales slot 0's accumulated gradient by 1/16.
- assert train_data["step_adapter_batch_sizes"] == {0: 16}
- assert train_data["prompt_group_sizes"] == [4, 4, 2]
-
-
-def test_rewards_normalize_within_heterogeneous_groups():
- _, _, train_data = run_pipeline()
- rewards = train_data["rewards"]
- # Group boundaries: [0:4], [4:8], [8:10] — each zero-mean.
- for start, end in [(0, 4), (4, 8), (8, 10)]:
- assert sum(rewards[start:end]) == pytest.approx(0.0, abs=1e-6)
- # Constant group (all 1.0) normalizes to zeros, not NaN.
- assert rewards[4:8] == pytest.approx([0.0] * 4)
- # Singleton-free std normalization applied to group 1 (n=4, mixed).
- assert max(abs(r) for r in rewards[0:4]) > 0.5
diff --git a/tests/fast/ray/rollout/test_train_data_conversion.py b/tests/fast/ray/rollout/test_train_data_conversion.py
index 1e839dc086..ff143ac1bc 100644
--- a/tests/fast/ray/rollout/test_train_data_conversion.py
+++ b/tests/fast/ray/rollout/test_train_data_conversion.py
@@ -960,3 +960,21 @@ def test_dynamic_gbs_multi_step(self):
assert shards[0]["num_microbatches"] == [2, 2]
assert shards[0]["dynamic_global_batch_size"] == 8
assert shards[0]["num_rollouts"] == [8, 8]
+
+
+def test_delayed_dp_split_preserves_the_tinker_loss_vectors():
+ """The tinker losses zip loss_weights / advantages / rollout_log_probs per
+ datum; a key missing from the shard whitelist disappears silently and only
+ fails inside the trainer."""
+ from miles.ray.rollout.train_data_conversion import _package_shards
+
+ data = {
+ "tokens": [[1], [2], [3], [4]],
+ "loss_weights": [[1.0]] * 4,
+ "advantages": [[0.5]] * 4,
+ "rollout_log_probs": [[-0.1]] * 4,
+ }
+ shards = _package_shards(None, data, [[0, 2], [1, 3]])
+ for shard in shards:
+ assert shard["loss_weights"] and shard["advantages"] and shard["rollout_log_probs"]
+ assert shards[0]["advantages"] == [[0.5], [0.5]]
diff --git a/tests/fast/ray/specs/test_train.py b/tests/fast/ray/specs/test_train.py
index c6b6a12752..706a8e07c6 100644
--- a/tests/fast/ray/specs/test_train.py
+++ b/tests/fast/ray/specs/test_train.py
@@ -21,6 +21,7 @@ def _make_args(**overrides) -> SimpleNamespace:
use_critic=False,
indep_dp=False,
train_backend="megatron",
+ multi_lora=False,
use_fault_tolerance=False,
kl_coef=0,
use_kl_loss=False,
@@ -175,13 +176,22 @@ def test_independent_dp_cells_share_one_rendezvous_store(self, monkeypatch):
addrs = [spec.ctor_kwargs(_make_context(cell_index=i))["indep_dp_store_addr"] for i in range(2)]
assert addrs == ["10.0.0.1:1234", "10.0.0.1:1234"]
- def test_the_backend_selects_the_worker_class(self):
- """A run must not start Megatron workers for an fsdp job."""
- (megatron_spec,) = specs_trainer(_make_args(train_backend="megatron"))
- (fsdp_spec,) = specs_trainer(_make_args(train_backend="fsdp"))
+ @pytest.mark.parametrize(
+ "backend,multi_lora,actor_class",
+ [
+ ("megatron", False, "miles.backends.megatron_utils.actor.MegatronTrainRayActor"),
+ ("megatron", True, "miles.backends.megatron_utils.lora.actor.MultiLoRATrainRayActor"),
+ ("fsdp", False, "miles.backends.fsdp_utils.actor.FSDPTrainRayActor"),
+ ],
+ ids=["megatron", "multi-lora", "fsdp"],
+ )
+ def test_the_backend_selects_the_worker_class(self, backend, multi_lora, actor_class):
+ actor_spec, critic_spec = specs_trainer(
+ _make_args(train_backend=backend, multi_lora=multi_lora, use_critic=True)
+ )
- assert megatron_spec.worker_class.endswith("MegatronTrainRayActor")
- assert fsdp_spec.worker_class.endswith("FSDPTrainRayActor")
+ assert actor_spec.worker_class == actor_class
+ assert critic_spec.worker_class == train_specs._TRAINER_ACTOR_CLASSES[backend]
class TestConcurrencyGroups:
diff --git a/tests/fast/ray/train/dummy_actor.py b/tests/fast/ray/train/dummy_actor.py
index ca3e794534..98e2313191 100644
--- a/tests/fast/ray/train/dummy_actor.py
+++ b/tests/fast/ray/train/dummy_actor.py
@@ -60,9 +60,6 @@ def train(self, *args: Any, **kwargs: Any) -> Any:
return self._train_return_values_per_attempt.pop(0)
return self._train_return_value
- def reconcile_adapters(self) -> None:
- self._record("reconcile_adapters", (), {})
-
def set_rollout_executor(self, *args: Any, **kwargs: Any) -> None:
self._record("set_rollout_executor", args, kwargs)
diff --git a/tests/fast/ray/train/test_group_reconcile_adapters.py b/tests/fast/ray/train/test_group_reconcile_adapters.py
deleted file mode 100644
index 5fddf132dd..0000000000
--- a/tests/fast/ray/train/test_group_reconcile_adapters.py
+++ /dev/null
@@ -1,44 +0,0 @@
-from types import SimpleNamespace
-
-import pytest
-import ray
-from tests.fast.ray.train.conftest import get_raw_actor_handles, make_alive_cell, make_cell
-
-from miles.ray.train.group import TrainerController
-
-pytestmark = pytest.mark.asyncio
-
-
-def _make_controller(cells: list) -> TrainerController:
- group = object.__new__(TrainerController)
- group._cells_by_id = {cell.cell_id: cell for cell in cells}
- group.args = SimpleNamespace()
- return group
-
-
-def _reconcile_calls_of(cell) -> list:
- return [
- [call for call in ray.get(handle.get_calls.remote()) if call[0] == "reconcile_adapters"]
- for handle in get_raw_actor_handles(cell)
- ]
-
-
-class TestReconcileAdapters:
- async def test_every_worker_of_every_cell_is_asked_to_reconcile(self):
- """Each rank of every cell owns its adapter registry, so all of them must be reached."""
- cells = [make_alive_cell(index, alive_cell_indices=[0, 1]) for index in range(2)]
- group = _make_controller(cells)
-
- await group.reconcile_adapters()
-
- for cell in cells:
- assert _reconcile_calls_of(cell) == [[("reconcile_adapters", (), {})]] * 2
-
- async def test_a_failing_worker_in_a_later_cell_propagates_instead_of_being_swallowed(self):
- """A stale adapter set would corrupt routing, so a failure in any cell must reach the caller."""
- cells = [make_cell(index) for index in range(2)]
- ray.get(get_raw_actor_handles(cells[1])[0].set_fail_methods.remote(["reconcile_adapters"]))
- group = _make_controller(cells)
-
- with pytest.raises(Exception, match="Injected failure"):
- await group.reconcile_adapters()
diff --git a/tests/fast/rollout/generate_utils/test_prefill_logprobs.py b/tests/fast/rollout/generate_utils/test_prefill_logprobs.py
index dcab51e32c..aa736a8194 100644
--- a/tests/fast/rollout/generate_utils/test_prefill_logprobs.py
+++ b/tests/fast/rollout/generate_utils/test_prefill_logprobs.py
@@ -3,7 +3,7 @@
import pytest
from miles.rollout.generate_utils import prefill_logprobs
-from miles.utils.types import AdapterRef, Sample
+from miles.utils.types import Sample
@pytest.mark.asyncio
@@ -113,105 +113,6 @@ async def fake_post(url, payload, action="post", headers=None):
assert calls[1][1]["logprob_start_len"] == 1
-@pytest.mark.asyncio
-async def test_recompute_uses_per_sample_adapter_lora_path(monkeypatch):
- """Multi-LoRA: the scoring request must go to the sample's own slot adapter,
- not the single-adapter name (which is never registered on those engines)."""
- sample = Sample(
- tokens=[10, 11, 20],
- response_length=1,
- status=Sample.Status.COMPLETED,
- adapter=AdapterRef(name="run-a", slot=3),
- )
- # Multi-LoRA forces lora_rank > 0, so is_lora_enabled(args) is always true.
- args = SimpleNamespace(recompute_logprobs_via_prefill=True, lora_rank=8)
- seen = {}
-
- async def fake_post(url, payload, headers=None):
- seen["payload"] = payload
- return {"meta_info": {"input_token_logprobs": [(None, 11), (-0.5, 20)]}}
-
- monkeypatch.setattr(prefill_logprobs, "post", fake_post)
-
- await prefill_logprobs.recompute_rollout_logprobs_via_prefill(
- args,
- sample,
- url="http://localhost/generate",
- sampling_params={},
- )
-
- assert seen["payload"]["lora_path"] == "__miles_slot_3"
- assert sample.rollout_log_probs == [-0.5]
-
-
-@pytest.mark.asyncio
-async def test_recompute_samples_batches_group_by_adapter(monkeypatch):
- """Samples from different adapters must not share a batch payload: each
- batch carries a single lora_path."""
- samples = [
- Sample(
- tokens=[10, 11, 20],
- response_length=1,
- status=Sample.Status.COMPLETED,
- adapter=AdapterRef(name="run-a", slot=0),
- ),
- Sample(
- tokens=[10, 11, 21],
- response_length=1,
- status=Sample.Status.COMPLETED,
- adapter=AdapterRef(name="run-b", slot=1),
- ),
- Sample(
- tokens=[10, 11, 22],
- response_length=1,
- status=Sample.Status.COMPLETED,
- adapter=AdapterRef(name="run-a", slot=0),
- ),
- ]
- args = SimpleNamespace(
- recompute_logprobs_via_prefill=True,
- lora_rank=8,
- sglang_router_policy="round_robin",
- )
- generate_payloads = []
-
- async def fake_post(url, payload, action="post", headers=None):
- if url.endswith("/flush_cache"):
- return {}
- generate_payloads.append(payload)
- return [
- {"meta_info": {"input_token_logprobs": [(None, 11), (-float(tokens[-1]), tokens[-1])]}}
- for tokens in payload["input_ids"]
- ]
-
- monkeypatch.setattr(prefill_logprobs, "post", fake_post)
-
- await prefill_logprobs.recompute_samples_rollout_logprobs_via_prefill(
- args,
- samples,
- url="http://localhost/generate",
- sampling_params={"max_new_tokens": 32},
- )
-
- assert [sample.rollout_log_probs for sample in samples] == [[-20.0], [-21.0], [-22.0]]
- by_lora_path = {payload["lora_path"]: payload["input_ids"] for payload in generate_payloads}
- assert by_lora_path == {
- "__miles_slot_0": [[10, 11, 20], [10, 11, 22]],
- "__miles_slot_1": [[10, 11, 21]],
- }
-
-
-def test_batch_payload_rejects_mixed_lora_paths():
- samples = [
- Sample(tokens=[10, 11, 20], response_length=1, adapter=AdapterRef(name="run-a", slot=0)),
- Sample(tokens=[10, 11, 21], response_length=1, adapter=AdapterRef(name="run-b", slot=1)),
- ]
- args = SimpleNamespace(lora_rank=8)
-
- with pytest.raises(ValueError, match="shared lora_path"):
- prefill_logprobs._build_batch_prefill_scoring_payload(args, samples, {})
-
-
@pytest.mark.asyncio
async def test_recompute_samples_batches_by_logprob_start_len(monkeypatch):
samples = [
diff --git a/tests/fast/test_train_multi_lora_async.py b/tests/fast/test_train_multi_lora_async.py
deleted file mode 100644
index 6574c9fe18..0000000000
--- a/tests/fast/test_train_multi_lora_async.py
+++ /dev/null
@@ -1,189 +0,0 @@
-from types import SimpleNamespace
-from typing import Any
-
-import pytest
-import ray
-import train_multi_lora_async as multi_lora_driver
-from tests.fast.fixtures.driver_fakes import (
- FakeInferenceController,
- FakeRemoteMethod,
- FakeRolloutExecutor,
- FakeTrainingModel,
-)
-
-from miles.utils.multi_lora import EmptyBatchTimeoutError
-
-_ACTIVE_SNAPSHOT = {"pending": [], "active": ["alpha"], "retiring": [], "cleanup": []}
-_EMPTY_SNAPSHOT = {"pending": [], "active": [], "retiring": [], "cleanup": []}
-
-
-class FakeMultiLoRAController:
- def __init__(self, events: list[str], snapshots: list[dict[str, list[str]]]) -> None:
- self.events = events
- self.snapshots = snapshots
- self.router_address: str | None = None
- self.registered_adapters: list[tuple[str, Any]] = []
- self.start = FakeRemoteMethod(self._start)
- self.stop = FakeRemoteMethod(self._stop)
- self.http_host = FakeRemoteMethod(self._http_host)
- self.api_port = FakeRemoteMethod(self._api_port)
- self.snapshot = FakeRemoteMethod(self._snapshot)
- self.register_adapter = FakeRemoteMethod(self._register_adapter)
-
- async def _start(self) -> None:
- self.events.append("controller_start")
-
- async def _stop(self) -> None:
- self.events.append("controller_stop")
-
- async def _http_host(self) -> str:
- return "127.0.0.1"
-
- async def _api_port(self) -> int:
- return 7000
-
- async def _snapshot(self) -> dict[str, list[str]]:
- self.events.append("snapshot")
- return self.snapshots.pop(0)
-
- async def _register_adapter(self, name: str, config: Any) -> None:
- self.registered_adapters.append((name, config))
-
-
-def _make_args(**overrides: Any) -> SimpleNamespace:
- args = SimpleNamespace(
- colocate=False,
- multi_lora_adapters=[],
- multi_lora_idle_poll_s=0.01,
- multi_lora_service_mode=False,
- sglang_router_ip="10.0.0.9",
- sglang_router_port=9000,
- )
- for key, value in overrides.items():
- setattr(args, key, value)
- return args
-
-
-def _install_driver_fakes(
- monkeypatch: pytest.MonkeyPatch,
- events: list[str],
- snapshots: list[dict[str, list[str]]],
-) -> SimpleNamespace:
- components = SimpleNamespace(
- inference_controller=FakeInferenceController(events),
- rollout_executor=FakeRolloutExecutor(events),
- actor_model=FakeTrainingModel(events, "actor"),
- controller=FakeMultiLoRAController(events, snapshots),
- )
-
- def create_multilora_controller(_args: SimpleNamespace, router_address: str) -> FakeMultiLoRAController:
- events.append("controller_created")
- components.controller.router_address = router_address
- return components.controller
-
- async def create_rollout_components(_args: SimpleNamespace) -> tuple[Any, Any, int]:
- return components.inference_controller, components.rollout_executor, 4
-
- async def create_training_models(_args: SimpleNamespace, _controller: Any, _executor: Any) -> tuple[Any, Any]:
- return components.actor_model, None
-
- async def update_weights(_model: Any, _executor: Any, rollout_id: int | None = None) -> None:
- events.append(f"update_weights:{rollout_id}")
-
- monkeypatch.setattr(multi_lora_driver, "configure_logger", lambda *_args, **_kwargs: None)
- monkeypatch.setattr(multi_lora_driver, "launch_worker_manager", lambda _args: None)
- monkeypatch.setattr(multi_lora_driver.object_store, "init_instance", lambda *_args, **_kwargs: None)
- monkeypatch.setattr(multi_lora_driver, "init_tracking", lambda _args: None)
- monkeypatch.setattr(multi_lora_driver, "create_rollout_components", create_rollout_components)
- monkeypatch.setattr(multi_lora_driver, "create_multilora_controller", create_multilora_controller)
- monkeypatch.setattr(multi_lora_driver, "get_multi_lora_controller", lambda: components.controller)
- monkeypatch.setattr(multi_lora_driver, "create_training_models", create_training_models)
- monkeypatch.setattr(multi_lora_driver, "define_new_adapter_metrics", lambda _snapshot: None)
- monkeypatch.setattr(multi_lora_driver, "update_weights", update_weights)
- monkeypatch.setattr(multi_lora_driver, "remove_rollout_data_refs", lambda *_args, **_kwargs: None)
- return components
-
-
-def _task_error(cause: Exception) -> ray.exceptions.RayTaskError:
- return ray.exceptions.RayTaskError(
- function_name="rollout_executor.get",
- traceback_str=f"{type(cause).__name__}: {cause}",
- cause=cause,
- proctitle="ray::RolloutExecutor.get",
- pid=1234,
- ip="127.0.0.1",
- )
-
-
-class TestAdapterLifecycle:
- async def test_one_active_adapter_completes_through_the_refactored_components(
- self, monkeypatch: pytest.MonkeyPatch
- ):
- """One active adapter runs push, generate, train and save once, then the driver tears everything down."""
- events: list[str] = []
- args = _make_args()
- components = _install_driver_fakes(
- monkeypatch, events, snapshots=[_ACTIVE_SNAPSHOT, _ACTIVE_SNAPSHOT, _EMPTY_SNAPSHOT]
- )
-
- await multi_lora_driver.main(args)
-
- assert components.controller.router_address == "http://10.0.0.9:9000"
- assert events == [
- "controller_created",
- "controller_start",
- "snapshot",
- "actor_reconcile_adapters",
- "update_weights:None",
- "snapshot",
- "prepare_rollout:0",
- "generate_start:0",
- "generate_done:0",
- "actor_train:0",
- "actor_save:0",
- "snapshot",
- "executor_dispose",
- "inference_dispose",
- "actor_dispose",
- "controller_stop",
- ]
-
-
-class TestEmptyBatchTimeout:
- async def test_empty_batch_timeout_retries_the_same_rollout_without_training(
- self, monkeypatch: pytest.MonkeyPatch
- ):
- """A generation that timed out with no trainable groups is retried under the same rollout id."""
- events: list[str] = []
- args = _make_args()
- components = _install_driver_fakes(
- monkeypatch,
- events,
- snapshots=[_ACTIVE_SNAPSHOT, _ACTIVE_SNAPSHOT, _ACTIVE_SNAPSHOT, _ACTIVE_SNAPSHOT, _EMPTY_SNAPSHOT],
- )
- components.rollout_executor.generation_errors = [_task_error(EmptyBatchTimeoutError("no trainable groups"))]
-
- await multi_lora_driver.main(args)
-
- assert [event for event in events if event.startswith(("generate_", "actor_train", "actor_save"))] == [
- "generate_start:0",
- "generate_failed:0",
- "generate_start:0",
- "generate_done:0",
- "actor_train:0",
- "actor_save:0",
- ]
- assert components.actor_model.trained == [0]
- assert components.actor_model.saved == [0]
-
- async def test_an_unrelated_task_error_stops_the_run(self, monkeypatch: pytest.MonkeyPatch):
- """Retrying every failed generation task would hide a real rollout crash behind an endless loop."""
- events: list[str] = []
- args = _make_args()
- components = _install_driver_fakes(monkeypatch, events, snapshots=[_ACTIVE_SNAPSHOT, _ACTIVE_SNAPSHOT])
- components.rollout_executor.generation_errors = [_task_error(ValueError("rollout worker died"))]
-
- with pytest.raises(ray.exceptions.RayTaskError, match="rollout worker died"):
- await multi_lora_driver.main(args)
-
- assert components.actor_model.trained == []
diff --git a/tests/fast/tinker/__init__.py b/tests/fast/tinker/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/fast/tinker/conftest.py b/tests/fast/tinker/conftest.py
new file mode 100644
index 0000000000..14be45b30b
--- /dev/null
+++ b/tests/fast/tinker/conftest.py
@@ -0,0 +1,18 @@
+import asyncio
+from contextlib import suppress
+
+import pytest
+
+from tests.fast.tinker.harness import make_service
+
+
+@pytest.fixture
+async def service(tmp_path):
+ gateway = make_service(tmp_path)
+ run_task = asyncio.create_task(gateway.run())
+ try:
+ yield gateway
+ finally:
+ run_task.cancel()
+ with suppress(asyncio.CancelledError):
+ await run_task
diff --git a/tests/fast/tinker/core/__init__.py b/tests/fast/tinker/core/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/fast/tinker/core/test_future.py b/tests/fast/tinker/core/test_future.py
new file mode 100644
index 0000000000..d7d602184c
--- /dev/null
+++ b/tests/fast/tinker/core/test_future.py
@@ -0,0 +1,51 @@
+"""RequestFuture reads enforce tenant ownership and expire finished results."""
+
+import pytest
+
+from miles.tinker.core import future as future_module
+from miles.tinker.core.future import DONE, FAILED, PENDING, RequestFutureStore
+from miles.tinker.core.types import OwnershipError
+
+
+def test_resolve_and_fail_settle_the_state():
+ store = RequestFutureStore()
+ done = store.create("model", "tenant")
+ failed = store.create("model", "tenant")
+ assert done.state == PENDING
+
+ store.resolve(done.request_id, {"op": "optim_step"})
+ store.fail(failed.request_id, "boom", "user")
+
+ assert store.get(done.request_id, "tenant").state == DONE
+ fetched = store.get(failed.request_id, "tenant")
+ assert (fetched.state, fetched.error, fetched.error_category) == (FAILED, "boom", "user")
+
+
+def test_cross_tenant_get_raises_ownership():
+ store = RequestFutureStore()
+ future = store.create("model", "tenant-a")
+ with pytest.raises(OwnershipError):
+ store.get(future.request_id, "tenant-b")
+
+
+def test_an_expired_future_returns_none(monkeypatch):
+ store = RequestFutureStore()
+ future = store.create("model", "tenant")
+ store.resolve(future.request_id, {"op": "optim_step"})
+
+ monkeypatch.setattr(future_module, "_FINISHED_TTL_S", -1.0)
+ assert store.get(future.request_id, "tenant") is None, "finished past TTL must read as unknown (410)"
+
+
+def test_terminal_states_do_not_flip():
+ store = RequestFutureStore()
+ failed = store.create("model", "tenant")
+ store.fail(failed.request_id, "lease expired", "user")
+ store.resolve(failed.request_id, {"op": "forward_backward"})
+ fetched = store.get(failed.request_id, "tenant")
+ assert (fetched.state, fetched.result) == (FAILED, None), "a unit that raced with lease expiry must not revive it"
+
+ done = store.create("model", "tenant")
+ store.resolve(done.request_id, {"op": "optim_step"})
+ store.fail(done.request_id, "late failure", "server")
+ assert store.get(done.request_id, "tenant").state == DONE
diff --git a/tests/fast/tinker/core/test_model_queue.py b/tests/fast/tinker/core/test_model_queue.py
new file mode 100644
index 0000000000..6655a5d55e
--- /dev/null
+++ b/tests/fast/tinker/core/test_model_queue.py
@@ -0,0 +1,47 @@
+"""Model queues order commands and hold barriers until all preceding datums finish."""
+
+from tests.fast.tinker.harness import command, datum, fb_payload
+
+from miles.tinker.core.model_queue import ModelRequestQueue
+
+
+def _model_queue() -> ModelRequestQueue:
+ return ModelRequestQueue("model", "tenant", slot=0)
+
+
+def test_out_of_order_arrival_queues_in_seq_order():
+ model_queue = _model_queue()
+ model_queue.submit(command("model", 2, "forward_backward", fb_payload("model", 2, [datum()]), arrival=1))
+ assert not model_queue.queue, "seq 2 must wait for seq 1"
+ model_queue.submit(command("model", 1, "forward_backward", fb_payload("model", 1, [datum()]), arrival=2))
+ assert [pending.command.seq_id for pending in model_queue.queue] == [1, 2]
+
+
+def test_a_gap_holds_later_commands():
+ model_queue = _model_queue()
+ model_queue.submit(command("model", 1, "forward_backward", fb_payload("model", 1, [datum()]), arrival=1))
+ model_queue.submit(command("model", 3, "forward_backward", fb_payload("model", 3, [datum()]), arrival=2))
+ assert [pending.command.seq_id for pending in model_queue.queue] == [1]
+ model_queue.submit(command("model", 2, "forward_backward", fb_payload("model", 2, [datum()]), arrival=3))
+ assert [pending.command.seq_id for pending in model_queue.queue] == [1, 2, 3]
+
+
+def test_batch_runs_alternate_with_barriers():
+ model_queue = _model_queue()
+ model_queue.submit(command("model", 1, "forward_backward", fb_payload("model", 1, [datum()]), arrival=1))
+ model_queue.submit(command("model", 2, "forward_backward", fb_payload("model", 2, [datum()]), arrival=2))
+ model_queue.submit(
+ command("model", 3, "optim_step", {"model_id": "model", "seq_id": 3, "adam_params": {}}, arrival=3)
+ )
+ model_queue.submit(command("model", 4, "forward_backward", fb_payload("model", 4, [datum()]), arrival=4))
+
+ assert [pending.command.seq_id for pending in model_queue.open_batch_run()] == [1, 2]
+ assert model_queue.ready_barrier() is None, "the barrier must wait for its batch run"
+
+ for pending in list(model_queue.open_batch_run()):
+ model_queue.finish(pending)
+ assert model_queue.ready_barrier().command.seq_id == 3
+ assert not model_queue.open_batch_run(), "the next batch run opens only after the barrier"
+
+ model_queue.finish(model_queue.ready_barrier())
+ assert [pending.command.seq_id for pending in model_queue.open_batch_run()] == [4]
diff --git a/tests/fast/tinker/core/test_scheduler.py b/tests/fast/tinker/core/test_scheduler.py
new file mode 100644
index 0000000000..e6cc545d54
--- /dev/null
+++ b/tests/fast/tinker/core/test_scheduler.py
@@ -0,0 +1,119 @@
+"""The scheduler selects the oldest ready seed and packs compatible work within its budget."""
+
+from tests.fast.tinker.harness import ADAM, command, datum, fb_payload
+
+from miles.tinker.core.model_queue import ModelRequestQueue
+from miles.tinker.core.scheduler import BarrierUnit, BatchUnit, RequestScheduler
+
+
+def _scheduler_with_model_queues(n: int, budget: int = 1_000_000) -> tuple[RequestScheduler, list[ModelRequestQueue]]:
+ scheduler = RequestScheduler(budget)
+ model_queues = [ModelRequestQueue(f"model-{i}", f"tenant-{i}", slot=i) for i in range(n)]
+ for model_queue in model_queues:
+ scheduler.add_model_queue(model_queue)
+ return scheduler, model_queues
+
+
+def _submit_fb(model_queue: ModelRequestQueue, seq_id: int, arrival: int, datums: list[dict], loss_fn="cross_entropy"):
+ model_queue.submit(
+ command(
+ model_queue.model_id,
+ seq_id,
+ "forward_backward",
+ fb_payload(model_queue.model_id, seq_id, datums, loss_fn),
+ arrival,
+ )
+ )
+
+
+def _submit_optim(model_queue: ModelRequestQueue, seq_id: int, arrival: int):
+ payload = {"model_id": model_queue.model_id, "seq_id": seq_id, "adam_params": dict(ADAM)}
+ model_queue.submit(command(model_queue.model_id, seq_id, "optim_step", payload, arrival))
+
+
+def test_same_pack_key_datums_pack_across_model_queues():
+ scheduler, (a, b) = _scheduler_with_model_queues(2)
+ _submit_fb(a, 1, arrival=1, datums=[datum(), datum()])
+ _submit_fb(b, 1, arrival=2, datums=[datum()])
+
+ unit = scheduler.schedule_next()
+ assert isinstance(unit, BatchUnit)
+ assert len(unit.datums) == 3
+ assert {ref.model_queue.model_id for ref in unit.datums} == {"model-0", "model-1"}
+
+
+def test_different_loss_functions_do_not_pack():
+ scheduler, (a, b) = _scheduler_with_model_queues(2)
+ _submit_fb(a, 1, arrival=1, datums=[datum()], loss_fn="cross_entropy")
+ _submit_fb(b, 1, arrival=2, datums=[datum()], loss_fn="ppo")
+
+ unit = scheduler.schedule_next()
+ assert [ref.model_queue.model_id for ref in unit.datums] == ["model-0"]
+
+
+def test_the_token_budget_splits_a_large_request():
+ scheduler, (a,) = _scheduler_with_model_queues(1, budget=10)
+ _submit_fb(a, 1, arrival=1, datums=[datum(4), datum(4), datum(4)]) # 5 tokens each with the appended target
+
+ first, second = scheduler.schedule_next(), scheduler.schedule_next()
+ assert [len(first.datums), len(second.datums)] == [2, 1]
+
+
+def test_an_oversized_datum_still_ships_alone():
+ scheduler, (a,) = _scheduler_with_model_queues(1, budget=2)
+ _submit_fb(a, 1, arrival=1, datums=[datum(10)])
+
+ assert len(scheduler.schedule_next().datums) == 1
+
+
+def test_the_oldest_arrival_wins():
+ scheduler, (a, b) = _scheduler_with_model_queues(2)
+ _submit_optim(a, 1, arrival=1)
+ _submit_fb(b, 1, arrival=2, datums=[datum()])
+
+ assert isinstance(scheduler.schedule_next(), BarrierUnit)
+
+
+def test_optim_barriers_merge_and_save_does_not():
+ scheduler, (a, b, c) = _scheduler_with_model_queues(3)
+ _submit_optim(a, 1, arrival=1)
+ _submit_optim(b, 1, arrival=2)
+ c.submit(command("model-2", 1, "save_state", {"model_id": "model-2", "seq_id": 1, "name": "x"}, arrival=3))
+
+ merged = scheduler.schedule_next()
+ assert merged.op == "optim_step"
+ assert {model_queue.model_id for model_queue, _ in merged.entries} == {"model-0", "model-1"}
+
+ for model_queue, pending in merged.entries:
+ model_queue.finish(pending)
+ save = scheduler.schedule_next()
+ assert save.op == "save_state"
+ assert len(save.entries) == 1
+
+
+def test_issued_datums_are_not_reissued():
+ scheduler, (a,) = _scheduler_with_model_queues(1)
+ _submit_fb(a, 1, arrival=1, datums=[datum()])
+
+ assert len(scheduler.schedule_next().datums) == 1
+ assert scheduler.schedule_next() is None
+
+
+def test_forward_only_packs_only_within_one_loss():
+ scheduler = RequestScheduler(batch_token_budget=1 << 20)
+ for index, (model, loss_fn) in enumerate([("m1", "cross_entropy"), ("m2", "dro")], start=1):
+ model_queue = ModelRequestQueue(model, "tenant", slot=index)
+ scheduler.add_model_queue(model_queue)
+ model_queue.submit(
+ command(
+ model,
+ 1,
+ "forward_only",
+ {"datums": [{"tokens": [1, 2], "target_len": 1}], "loss_fn": loss_fn, "loss_fn_config": {}},
+ arrival=index,
+ )
+ )
+ first = scheduler.schedule_next()
+ assert [ref.request.command.payload["loss_fn"] for ref in first.datums] == [
+ "cross_entropy"
+ ], "a DRO request must not execute under another request's loss"
diff --git a/tests/fast/tinker/core/test_service.py b/tests/fast/tinker/core/test_service.py
new file mode 100644
index 0000000000..c92baca047
--- /dev/null
+++ b/tests/fast/tinker/core/test_service.py
@@ -0,0 +1,850 @@
+"""The service preserves request ordering, tenant isolation, and failure isolation."""
+
+import asyncio
+from contextlib import suppress
+from pathlib import Path
+
+import pytest
+from tests.fast.tinker.harness import (
+ ADAM,
+ await_settled,
+ created_model,
+ datum,
+ fb_payload,
+ make_service,
+ model_payload,
+)
+
+from miles.tinker.core.future import DONE, FAILED
+from miles.tinker.core.types import OwnershipError, UserInputError
+from miles.tinker.core.utils import resolve_checkpoint_dir, resolve_sampler_checkpoint
+
+
+def _optim_payload(model_id: str, seq_id: int) -> dict:
+ return {"model_id": model_id, "seq_id": seq_id, "adam_params": dict(ADAM)}
+
+
+async def test_resubmitted_seq_id_reuses_the_future_and_runs_once(service):
+ model_id = await created_model(service)
+ payload = fb_payload(model_id, 1, [datum()])
+
+ first = service.submit("tenant", "forward_backward", payload)
+ second = service.submit("tenant", "forward_backward", payload)
+ assert first == second
+
+ future = await await_settled(service, "tenant", first)
+ assert future.state == DONE
+ assert service.submit("tenant", "forward_backward", payload) == first, "a resubmit after completion also dedups"
+ assert len(service.backend.named("forward_backward")) == 1, "the resubmit must not run the trainer again"
+
+
+async def test_create_model_is_two_phase(service):
+ request_id, model_id = service.create_model("tenant", model_payload(service, lora_config={"rank": 8}))
+
+ future = await await_settled(service, "tenant", request_id)
+ assert future.result == {"op": "create_model", "model_id": model_id}
+ assert service.backend.named("load_slot")[0]["rank"] == 8
+ assert model_id in service.models
+
+
+async def test_failed_slot_init_returns_the_slot(service):
+ free_before = set(service.free_slots)
+ service.backend.fail_next = {"error": "init blew up"}
+ request_id, model_id = service.create_model("tenant", model_payload(service))
+
+ future = await await_settled(service, "tenant", request_id)
+ assert (future.state, future.error_category) == (FAILED, "server")
+ assert model_id not in service.models
+ assert service.free_slots == free_before
+
+
+async def test_no_free_slots_is_a_user_error(service):
+ for _ in range(service.config.n_slots):
+ await created_model(service)
+ with pytest.raises(UserInputError, match="no free adapter slots"):
+ service.create_model("tenant", model_payload(service))
+
+
+async def test_the_wrong_base_model_is_rejected(service):
+ with pytest.raises(UserInputError, match="serves"):
+ service.create_model("tenant", model_payload(service, base_model="other"))
+
+
+async def test_out_of_order_chunks_complete_and_the_barrier_waits(service):
+ model_id = await created_model(service)
+
+ optim = service.submit("tenant", "optim_step", _optim_payload(model_id, 3))
+ late = service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [datum()]))
+ await asyncio.sleep(0.05)
+ assert not service.backend.named("optim_step"), "the barrier must wait for every batch op ahead"
+
+ early = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ for request_id in (early, late, optim):
+ assert (await await_settled(service, "tenant", request_id)).state == DONE
+
+ calls = [name for name, _ in service.backend.calls if name != "load_slot"]
+ assert calls.index("optim_step") > max(i for i, name in enumerate(calls) if name == "forward_backward")
+
+
+async def test_admission_failure_terminates_the_training_model_queue(service):
+ model_id = await created_model(service)
+
+ oversized = service.submit(
+ "tenant", "forward_backward", fb_payload(model_id, 1, [datum(service.config.max_tokens_per_datum + 1)])
+ )
+ assert (await await_settled(service, "tenant", oversized)).state == FAILED
+
+ with pytest.raises(UserInputError, match="create a new model"):
+ service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [datum()]))
+
+
+async def test_forward_backward_outputs_align_to_datums(service):
+ model_id = await created_model(service)
+ request_id = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum(2), datum(5)]))
+
+ future = await await_settled(service, "tenant", request_id)
+ assert [len(output["logprobs"]) for output in future.result["outputs"]] == [2, 5]
+
+
+async def test_optim_step_returns_the_slot_grad_norm(service):
+ model_id = await created_model(service)
+ slot = service.models[model_id].slot
+ service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ request_id = service.submit("tenant", "optim_step", _optim_payload(model_id, 2))
+
+ future = await await_settled(service, "tenant", request_id)
+ assert future.result["metrics"] == {"grad_norm": 0.5 + slot}
+
+
+async def test_save_then_load_roundtrip_paths(service):
+ model_id = await created_model(service)
+ save = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "ckpt", "overwrite": False}
+ )
+ saved_path = (await await_settled(service, "tenant", save)).result["path"]
+ assert saved_path == f"tinker://{model_id}/weights/ckpt"
+
+ load = service.submit(
+ "tenant", "load_state", {"model_id": model_id, "seq_id": 2, "path": saved_path, "optimizer": False}
+ )
+ assert (await await_settled(service, "tenant", load)).state == DONE
+ weights_only = service.backend.named("load_slot")[-1]
+ assert weights_only["load_optimizer"] is False
+ checkpoint = Path(service.config.checkpoint_root) / model_id / "weights" / "ckpt"
+ assert weights_only["ckpt_path"] == str(checkpoint.resolve())
+
+
+async def test_sampler_save_publishes_successive_versions(service):
+ model_id = await created_model(service)
+ for seq_id in (1, 2):
+ request_id = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": seq_id})
+ future = await await_settled(service, "tenant", request_id)
+ assert future.result["path"] == f"tinker://{model_id}/sampler_weights/{seq_id}"
+ assert [export["path"] for export in service.backend.named("export_slot")] == [
+ resolve_checkpoint_dir(service.config.checkpoint_root, model_id, "sampler_weights", str(version))
+ for version in (1, 2)
+ ]
+
+
+async def test_sampling_resolves_against_the_saved_version(service):
+ model_id = await created_model(service)
+ save = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ sampler_path = (await await_settled(service, "tenant", save)).result["path"]
+
+ request_id, sequence_ids = service.submit_sample(
+ "tenant",
+ {
+ "model_path": sampler_path,
+ "num_samples": 2,
+ "prompt_tokens": [1, 2],
+ "sampling_params": {"max_tokens": 4},
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ },
+ )
+ assert len(sequence_ids) == 2
+ future = await await_settled(service, "tenant", request_id)
+ assert len(future.result["sequences"]) == 2
+ assert service.backend.named("sample")[0]["lora_name"] == f"{model_id}@1"
+
+
+async def test_sampler_requests_carry_the_published_checkpoint_path(service):
+ model_id = await created_model(service)
+ request_id = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ path = (await await_settled(service, "tenant", request_id)).result["path"]
+
+ disk_dir = resolve_checkpoint_dir(service.config.checkpoint_root, model_id, "sampler_weights", "1")
+ assert service.backend.named("export_slot")[0]["path"] == disk_dir
+
+ sample_id, _ = service.submit_sample(
+ "tenant",
+ {
+ "model_path": path,
+ "num_samples": 1,
+ "prompt_tokens": [1],
+ "sampling_params": {"max_tokens": 2},
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ },
+ )
+ await await_settled(service, "tenant", sample_id)
+ assert service.backend.named("sample")[0]["lora_path"] == disk_dir
+
+
+async def test_sample_failure_preserves_the_snapshot_and_training_model_queue(service):
+ model_id = await created_model(service)
+ request_id = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ future = await await_settled(service, "tenant", request_id)
+ assert future.result["path"] == f"tinker://{model_id}/sampler_weights/1"
+
+ service.backend.fail_on["sample"] = {"error": "engine down"}
+ sample_id, _ = service.submit_sample(
+ "tenant",
+ {
+ "model_path": future.result["path"],
+ "num_samples": 1,
+ "prompt_tokens": [1],
+ "sampling_params": {"max_tokens": 2},
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ },
+ )
+ future = await await_settled(service, "tenant", sample_id)
+ assert future.state == FAILED
+ assert future.error == "engine down"
+ assert (
+ resolve_sampler_checkpoint(
+ service.config.checkpoint_root,
+ "tenant",
+ f"tinker://{model_id}/sampler_weights/1",
+ service.config.base_model,
+ )[0]
+ == f"{model_id}@1"
+ )
+ fb = service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [datum()]))
+ assert (await await_settled(service, "tenant", fb)).state == DONE
+
+
+async def test_lease_expiry_reclaims_the_tenant(service):
+ session_id = service.create_session("tenant")
+ model_id = await created_model(service, session_id=session_id)
+ slot = service.models[model_id].slot
+ queued = service.submit("tenant", "optim_step", _optim_payload(model_id, 1))
+ await await_settled(service, "tenant", queued)
+ stale = service.submit("tenant", "forward_backward", fb_payload(model_id, 3, [datum()])) # gapped: stays queued
+
+ service.sessions[session_id].last_heartbeat -= service.config.lease_timeout_s + 1
+ await service._expire_sessions()
+
+ assert model_id not in service.models
+ assert slot in service.free_slots
+ assert service.backend.named("unload_slot") == [{"slot": slot}]
+ assert service.retrieve_future("tenant", stale).state == FAILED
+
+
+async def test_a_fresh_heartbeat_keeps_the_model(service):
+ session_id = service.create_session("tenant")
+ model_id = await created_model(service, session_id=session_id)
+
+ service.heartbeat("tenant", session_id)
+ await service._expire_sessions()
+
+ assert model_id in service.models
+
+
+@pytest.mark.parametrize("name", ["../evil", "a/b", "..", ".hidden"], ids=["parent", "slash", "dotdot", "hidden"])
+async def test_traversal_checkpoint_names_are_rejected(service, name):
+ model_id = await created_model(service)
+ request_id = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": name, "overwrite": False}
+ )
+ future = await await_settled(service, "tenant", request_id)
+ assert (future.state, future.error_category) == (FAILED, "user")
+ assert "invalid checkpoint path segment" in future.error
+ assert not service.backend.named("save_slot"), "nothing may touch the disk for a rejected name"
+
+
+async def test_traversal_sampler_paths_are_rejected(service):
+ model_id = await created_model(service)
+ with pytest.raises(UserInputError):
+ service.submit_sample(
+ "tenant",
+ {
+ "model_path": f"tinker://{model_id}/weights/../../etc",
+ "num_samples": 1,
+ "prompt_tokens": [1],
+ "sampling_params": {"max_tokens": 2},
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ },
+ )
+
+
+async def test_expired_result_on_resubmit_fails_instead_of_410(service, monkeypatch):
+ from miles.tinker.core import future as future_module
+
+ model_id = await created_model(service)
+ first = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ await await_settled(service, "tenant", first)
+ monkeypatch.setattr(future_module, "_FINISHED_TTL_S", -1.0)
+ assert service.retrieve_future("tenant", first) is None, "the result must have aged out"
+
+ resubmitted = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ assert resubmitted != first
+ monkeypatch.setattr(future_module, "_FINISHED_TTL_S", 3600.0) # only the first result aged out
+ future = service.retrieve_future("tenant", resubmitted)
+ assert (future.state, future.error_category) == (FAILED, "user")
+ assert "expired" in future.error
+ assert len(service.backend.named("forward_backward")) == 1, "an executed command must never re-run"
+
+
+async def test_save_state_refuses_to_overwrite_unless_asked(service):
+ model_id = await created_model(service)
+ payload = {"model_id": model_id, "name": "ckpt", "overwrite": False}
+ first = service.submit("tenant", "save_state", payload | {"seq_id": 1})
+ assert (await await_settled(service, "tenant", first)).state == DONE
+
+ clobber = service.submit("tenant", "save_state", payload | {"seq_id": 2})
+ future = await await_settled(service, "tenant", clobber)
+ assert (future.state, future.error_category) == (FAILED, "user")
+ assert "overwrite" in future.error
+
+ replace = service.submit("tenant", "save_state", payload | {"seq_id": 3, "overwrite": True})
+ assert (await await_settled(service, "tenant", replace)).state == DONE
+
+
+async def test_checkpoints_outlive_the_lease(service):
+ session_id = service.create_session("tenant")
+ model_id = await created_model(service, session_id=session_id)
+ save = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "kept", "overwrite": False}
+ )
+ path = (await await_settled(service, "tenant", save)).result["path"]
+
+ service.sessions[session_id].last_heartbeat = -1e9
+ await service._expire_sessions()
+ assert model_id not in service.models, "the lease sweep must reclaim the model"
+
+ service.create_session("tenant")
+ fresh_id = await created_model(service)
+ load = service.submit("tenant", "load_state", {"model_id": fresh_id, "seq_id": 1, "path": path, "optimizer": True})
+ assert (await await_settled(service, "tenant", load)).state == DONE
+
+
+async def test_a_foreign_tenants_checkpoint_does_not_load(service):
+ model_id = await created_model(service)
+ save = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "mine", "overwrite": False}
+ )
+ path = (await await_settled(service, "tenant", save)).result["path"]
+
+ thief_model = await created_model(service, tenant="thief")
+ load = service.submit(
+ "thief", "load_state", {"model_id": thief_model, "seq_id": 1, "path": path, "optimizer": True}
+ )
+ future = await await_settled(service, "thief", load)
+ assert (future.state, future.error_category) == (FAILED, "user")
+ assert "belong" in future.error
+
+
+async def test_a_failed_batch_terminates_the_model_queue(service):
+ model_id = await created_model(service)
+ first = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum(), datum()]))
+ second = service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [datum()]))
+ service.backend.fail_next = {"error": "batch input rejected"}
+ assert (await await_settled(service, "tenant", first)).state == FAILED
+ assert (
+ await await_settled(service, "tenant", second)
+ ).state == FAILED, "the failed accumulation terminates model training"
+ assert model_id not in service.models
+ assert service.backend.named("unload_slot")
+
+
+async def test_merged_optim_settles_each_slot_on_its_own(service):
+ model_a = await created_model(service)
+ model_b = await created_model(service)
+ slot_b = service.models[model_b].slot
+ service.backend.optim_outcomes[slot_b] = {"error": "boom"}
+
+ ok = service.submit("tenant", "optim_step", _optim_payload(model_a, 1))
+ bad = service.submit("tenant", "optim_step", _optim_payload(model_b, 1))
+ resolved = await await_settled(service, "tenant", ok)
+ failed = await await_settled(service, "tenant", bad)
+ assert resolved.state == DONE and "grad_norm" in resolved.result["metrics"]
+ assert (failed.state, failed.error_category) == (FAILED, "server") and "boom" in failed.error
+
+
+async def test_a_nonfinite_step_reports_the_skip(service):
+ model_id = await created_model(service)
+ slot = service.models[model_id].slot
+ service.backend.optim_outcomes[slot] = {"skipped_nonfinite": 1.0}
+ request_id = service.submit("tenant", "optim_step", _optim_payload(model_id, 1))
+ future = await await_settled(service, "tenant", request_id)
+ assert future.state == DONE
+ assert future.result["metrics"] == {"skipped_nonfinite": 1.0}
+
+
+async def test_num_samples_is_capped(service):
+ model_id = await created_model(service)
+ version = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ path = (await await_settled(service, "tenant", version)).result["path"]
+ with pytest.raises(UserInputError):
+ service.submit_sample(
+ "tenant",
+ {
+ "model_path": path,
+ "num_samples": service.config.max_samples_per_request + 1,
+ "prompt_tokens": [1],
+ "sampling_params": {"max_tokens": 2},
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ },
+ )
+
+
+async def test_malformed_loss_inputs_are_rejected_at_admission(service):
+ cases = {
+ "unknown loss_fn": {"loss_fn": "made_up", "datums": [datum()]},
+ "missing input": {"loss_fn": "importance_sampling", "datums": [datum()]},
+ "wrong length": {
+ "loss_fn": "importance_sampling",
+ "datums": [datum() | {"sampling_logprobs": [0.0], "advantages": [1.0, 2.0, 3.0, 4.0]}],
+ },
+ # datum() carries weights, which importance_sampling never reads
+ "unread input": {
+ "loss_fn": "importance_sampling",
+ "datums": [datum() | {"sampling_logprobs": [0.0] * 3, "advantages": [1.0] * 3}],
+ },
+ }
+ for payload in cases.values():
+ model_id = await created_model(service)
+ seq_id = 1
+ request_id = service.submit(
+ "tenant",
+ "forward_backward",
+ {"model_id": model_id, "seq_id": seq_id, "loss_fn_config": {}, **payload},
+ )
+ future = await await_settled(service, "tenant", request_id)
+ assert (future.state, future.error_category) == (FAILED, "user")
+ assert not service.backend.named("forward_backward"), "rejected datums must never reach the trainer"
+
+
+async def test_a_failed_window_stops_already_queued_training(service):
+ model_id = await created_model(service)
+ healthy = await created_model(service)
+ early = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ assert (await await_settled(service, "tenant", early)).state == DONE
+
+ service.backend.fail_next = {"error": "batch rejected"}
+ queued = [service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [datum()]))]
+ queued.append(service.submit("tenant", "optim_step", _optim_payload(model_id, 3)))
+ queued.append(service.submit("tenant", "forward_backward", fb_payload(model_id, 4, [datum()])))
+ queued.append(service.submit("tenant", "optim_step", _optim_payload(model_id, 5)))
+ for request_id in queued:
+ assert (await await_settled(service, "tenant", request_id)).state == FAILED
+ assert service.retrieve_future("tenant", early).state == DONE
+ assert not service.backend.named("optim_step")
+ assert len(service.backend.named("forward_backward")) == 2
+ with pytest.raises(UserInputError, match="create a new model"):
+ service.submit("tenant", "optim_step", _optim_payload(model_id, 6))
+ step = service.submit("tenant", "optim_step", _optim_payload(healthy, 1))
+ assert (await await_settled(service, "tenant", step)).state == DONE
+
+
+async def test_unsupported_lora_configs_are_rejected(service):
+ for lora_config in ({"rank": 8, "seed": 7}, {"rank": 8, "train_unembed": True}, {"rank": 8, "train_mlp": False}):
+ with pytest.raises(UserInputError):
+ service.create_model("tenant", model_payload(service, lora_config=lora_config))
+ _, model_id = service.create_model(
+ "tenant",
+ model_payload(service, lora_config={"rank": 8, "train_attn": True, "train_mlp": True, "train_unembed": False}),
+ )
+ assert model_id in service.models
+
+
+@pytest.mark.parametrize("expire_lease", [False, True])
+async def test_sampler_paths_resolve_independently_of_the_lease(service, expire_lease):
+ session_id = service.create_session("tenant")
+ model_id = await created_model(service, session_id=session_id)
+ save = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ path = (await await_settled(service, "tenant", save)).result["path"]
+
+ if expire_lease:
+ service.sessions[session_id].last_heartbeat = -1e9
+ await service._expire_sessions()
+ assert model_id not in service.models
+ service.create_session("tenant")
+ lora_name, lora_path = resolve_sampler_checkpoint(
+ service.config.checkpoint_root, "tenant", path, service.config.base_model
+ )
+ assert lora_name == f"{model_id}@1" and lora_path.endswith("/sampler_weights/1")
+ with pytest.raises((UserInputError, OwnershipError)):
+ resolve_sampler_checkpoint(service.config.checkpoint_root, "thief", path, service.config.base_model)
+ service.config.base_model = "other"
+ with pytest.raises(UserInputError, match="base_model"):
+ resolve_sampler_checkpoint(service.config.checkpoint_root, "tenant", path, service.config.base_model)
+
+
+async def test_an_unnamed_sampler_save_returns_a_sampling_session(service):
+ model_id = await created_model(service)
+ unnamed = service.submit(
+ "tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1, "sampler_path": None}
+ )
+ result = (await await_settled(service, "tenant", unnamed)).result
+ session = service.sampling_sessions[result["sampling_session_id"]]
+ assert (session.tenant, session.model_path) == ("tenant", result["path"])
+
+ named = service.submit(
+ "tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 2, "sampler_path": "ckpt"}
+ )
+ assert "sampling_session_id" not in (await await_settled(service, "tenant", named)).result
+
+
+async def test_weights_info_reads_the_checkpoint_not_the_lease(service):
+ model_id = await created_model(service)
+ save = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "ck", "overwrite": False}
+ )
+ path = (await await_settled(service, "tenant", save)).result["path"]
+ del service.models[model_id]
+
+ info = service.weights_info("tenant", path)
+ assert info == {
+ "base_model": "base",
+ "is_lora": True,
+ "lora_rank": 8,
+ "train_attn": True,
+ "train_mlp": True,
+ "train_unembed": False,
+ }
+ with pytest.raises(OwnershipError):
+ service.weights_info("thief", path)
+
+
+async def test_a_failed_optim_step_retires_the_model(service):
+ model_id = await created_model(service)
+ slot = service.models[model_id].slot
+ service.backend.optim_outcomes[slot] = {"error": "slot update rejected"}
+
+ step = service.submit("tenant", "optim_step", _optim_payload(model_id, 1))
+ future = await await_settled(service, "tenant", step)
+ assert future.state == FAILED and "slot update rejected" in future.error
+ assert model_id not in service.models, "a half-applied step may have diverged the slot across ranks"
+ assert slot in service.free_slots
+ assert service.backend.named("unload_slot") == [{"slot": slot}]
+
+ with pytest.raises(UserInputError, match="restore from a checkpoint"):
+ service.submit("tenant", "optim_step", _optim_payload(model_id, 2))
+
+
+async def test_an_out_of_order_validation_failure_blocks_later_training(service):
+ model_id = await created_model(service)
+ step = service.submit("tenant", "optim_step", _optim_payload(model_id, 3))
+ rejected = service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [], loss_fn="unknown"))
+ early = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ assert (await await_settled(service, "tenant", early)).state == DONE
+ assert (await await_settled(service, "tenant", rejected)).state == FAILED
+ assert (await await_settled(service, "tenant", step)).state == FAILED
+ assert len(service.backend.named("forward_backward")) == 1
+ assert not service.backend.named("optim_step")
+
+
+async def test_a_backend_level_optim_failure_retires_every_model_in_the_barrier(service):
+ model_a = await created_model(service)
+ model_b = await created_model(service)
+ service.backend.fail_on["optim_step"] = {"error": "slot updates rejected"}
+
+ first = service.submit("tenant", "optim_step", _optim_payload(model_a, 1))
+ second = service.submit("tenant", "optim_step", _optim_payload(model_b, 1))
+ for request_id in (first, second):
+ future = await await_settled(service, "tenant", request_id)
+ assert future.state == FAILED and "restore from a checkpoint" in future.error
+ assert model_a not in service.models and model_b not in service.models
+ assert len(service.backend.named("unload_slot")) == 2
+
+ fresh = await created_model(service) # the dispatch loop survived the failure
+ assert fresh in service.models
+
+
+async def test_a_failed_unload_keeps_the_slot_out_of_the_free_pool(service):
+ session_id = service.create_session("tenant")
+ model_id = await created_model(service, session_id=session_id)
+ slot = service.models[model_id].slot
+ service.backend.fail_on["unload_slot"] = {"error": "engine gone"}
+
+ service.sessions[session_id].last_heartbeat -= service.config.lease_timeout_s + 1
+ await service._expire_sessions()
+
+ assert model_id not in service.models
+ assert slot not in service.free_slots, "a slot whose unload failed is dirty and must not be reused"
+ await service._expire_sessions() # the sweep itself survived
+
+
+async def test_a_named_sampler_save_uses_the_name_and_rejects_reuse(service):
+ model_id = await created_model(service)
+ save = service.submit(
+ "tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1, "sampler_path": "v1"}
+ )
+ result = (await await_settled(service, "tenant", save)).result
+ assert result["path"] == f"tinker://{model_id}/sampler_weights/v1"
+ assert "sampling_session_id" not in result
+
+ request_id, _ = service.submit_sample(
+ "tenant",
+ {
+ "model_path": result["path"],
+ "num_samples": 1,
+ "prompt_tokens": [1],
+ "sampling_params": {"max_tokens": 2},
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ },
+ )
+ await await_settled(service, "tenant", request_id)
+ assert service.backend.named("sample")[0]["lora_name"] == f"{model_id}@v1"
+
+ reuse = service.submit(
+ "tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 2, "sampler_path": "v1"}
+ )
+ future = await await_settled(service, "tenant", reuse)
+ assert future.state == FAILED and "already exist" in future.error, "saved versions are immutable"
+
+
+async def test_checkpoint_meta_stores_a_digest_not_the_credential(service):
+ import json
+ import os
+
+ model_id = await created_model(service)
+ saved = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "ck", "overwrite": False}
+ )
+ await await_settled(service, "tenant", saved)
+ meta_path = os.path.join(
+ resolve_checkpoint_dir(service.config.checkpoint_root, model_id, "weights", "ck"), "META.json"
+ )
+ meta = json.loads(open(meta_path).read())
+ assert "tenant" not in meta and meta["tenant_digest"] != "tenant", "the bearer credential must not be persisted"
+ info = service.weights_info("tenant", f"tinker://{model_id}/weights/ck")
+ assert (info["train_attn"], info["train_mlp"], info["train_unembed"]) == (True, True, False)
+
+
+async def test_a_checkpoint_saved_under_other_settings_does_not_load(service):
+ import json
+ import os
+
+ model_id = await created_model(service)
+ saved = service.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "ck", "overwrite": False}
+ )
+ path = (await await_settled(service, "tenant", saved)).result["path"]
+ meta_path = os.path.join(
+ resolve_checkpoint_dir(service.config.checkpoint_root, model_id, "weights", "ck"), "META.json"
+ )
+ meta = json.loads(open(meta_path).read())
+ meta["lora_alpha"] = meta["lora_alpha"] + 1 # the same tensors would be scaled differently
+ open(meta_path, "w").write(json.dumps(meta))
+
+ loaded = service.submit(
+ "tenant", "load_state", {"model_id": model_id, "seq_id": 2, "path": path, "optimizer": True}
+ )
+ future = await await_settled(service, "tenant", loaded)
+ assert (future.state, future.error_category) == (FAILED, "user") and "lora_alpha" in future.error
+ assert not service.backend.named("load_slot")[1:], "nothing may touch the slot on a mismatch"
+
+
+async def test_a_recycled_slot_belongs_to_a_fresh_model_queue(service):
+ session_id = service.create_session("tenant")
+ model_id = await created_model(service, session_id=session_id)
+ slot = service.models[model_id].slot
+ service.backend.fail_next = {"error": "batch input rejected"}
+ failed = service.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ assert (await await_settled(service, "tenant", failed)).state == FAILED
+
+ service.sessions[session_id].last_heartbeat -= service.config.lease_timeout_s + 1
+ await service._expire_sessions()
+ assert slot in service.free_slots
+
+ assert session_id not in service.sessions, "the sweep takes the expired session with it"
+ fresh = await created_model(service)
+ assert service.models[fresh].slot == slot
+ step_after_fb = service.submit("tenant", "forward_backward", fb_payload(fresh, 1, [datum()]))
+ assert (await await_settled(service, "tenant", step_after_fb)).state == DONE
+ step = service.submit("tenant", "optim_step", _optim_payload(fresh, 2))
+ future = await await_settled(service, "tenant", step)
+ assert future.state == DONE, "a fresh model may use the successfully recycled slot"
+
+
+@pytest.mark.parametrize("source", ["handler", "worker", "create", "sweep", "export", "load"])
+async def test_an_unknown_failure_stops_the_dispatcher(tmp_path, source):
+ from tests.fast.tinker.harness import make_service
+
+ gateway = make_service(tmp_path)
+ run_task = asyncio.create_task(gateway.run())
+ error = (
+ OSError("checkpoint IO failed") if source in ("export", "load") else RuntimeError("fatal execution failure")
+ )
+ try:
+ model_id = await created_model(gateway)
+ if source == "export":
+ gateway.backend.fail_on["export_slot"] = error
+ gateway.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ elif source == "load":
+ saved = gateway.submit(
+ "tenant", "save_state", {"model_id": model_id, "seq_id": 1, "name": "ck", "overwrite": False}
+ )
+ path = (await await_settled(gateway, "tenant", saved)).result["path"]
+ gateway.backend.fail_on["load_slot"] = error
+ gateway.submit(
+ "tenant", "load_state", {"model_id": model_id, "seq_id": 2, "path": path, "optimizer": True}
+ )
+ elif source == "create":
+ gateway.backend.fail_on["load_slot"] = RuntimeError("fatal execution failure")
+ gateway.create_model("tenant", model_payload(gateway))
+ elif source == "sweep":
+ gateway.backend.fail_on["unload_slot"] = RuntimeError("fatal execution failure")
+ for session in gateway.sessions.values():
+ session.last_heartbeat -= gateway.config.lease_timeout_s + 1
+ task = asyncio.create_task(gateway._expire_sessions())
+ task.add_done_callback(gateway._observe_background_task)
+ else:
+ if source == "handler":
+
+ async def broken_handler(unit):
+ raise RuntimeError("fatal execution failure")
+
+ gateway._run_barrier = broken_handler
+ else:
+ gateway.backend.fail_on["optim_step"] = RuntimeError("fatal execution failure")
+ gateway.submit("tenant", "optim_step", _optim_payload(model_id, 1))
+ with pytest.raises(type(error), match=str(error)):
+ await asyncio.wait_for(run_task, timeout=2)
+ finally:
+ if not run_task.done():
+ run_task.cancel()
+ with suppress(asyncio.CancelledError):
+ await run_task
+
+
+async def test_a_dead_trainer_escapes_the_dispatch_loop(tmp_path, monkeypatch):
+ gateway = make_service(tmp_path)
+ run_task = asyncio.create_task(gateway.run())
+ try:
+ model_id = await created_model(gateway)
+ monkeypatch.setattr(gateway.backend, "trainer_dead", lambda: True)
+ gateway.submit("tenant", "forward_backward", fb_payload(model_id, 1, [datum()]))
+ with pytest.raises(RuntimeError, match="trainer workers died"):
+ await asyncio.wait_for(run_task, timeout=2)
+ finally:
+ if not run_task.done():
+ run_task.cancel()
+ with suppress(asyncio.CancelledError):
+ await run_task
+
+
+async def test_a_retried_create_model_returns_the_first_allocation(service):
+ """A network retry must not occupy a second slot."""
+ payload = model_payload(service)
+ first = service.create_model("tenant", payload)
+ await await_settled(service, "tenant", first[0])
+ assert service.create_model("tenant", payload) == first
+ assert len(service.backend.named("load_slot")) == 1
+
+
+async def test_an_unknown_session_is_rejected(service):
+ with pytest.raises(UserInputError, match="unknown session"):
+ service.create_model("tenant", model_payload(service, session_id="session-gone"))
+
+
+async def test_a_retried_sampling_session_returns_the_same_id(service):
+ payload = {"session_id": service.create_session("tenant"), "sampling_session_seq_id": 1}
+ assert service.create_sampling_session("tenant", payload) == service.create_sampling_session("tenant", payload)
+
+
+async def test_a_retried_sample_does_not_sample_again(service):
+ model_id = await created_model(service)
+ save = service.submit("tenant", "save_weights_for_sampler", {"model_id": model_id, "seq_id": 1})
+ path = (await await_settled(service, "tenant", save)).result["path"]
+ sampling_session_id = service.create_sampling_session(
+ "tenant", {"session_id": service.create_session("tenant"), "sampling_session_seq_id": 1, "model_path": path}
+ )
+ payload = {
+ "sampling_session_id": sampling_session_id,
+ "seq_id": 1,
+ "num_samples": 1,
+ "prompt_tokens": [1],
+ "sampling_params": {"max_tokens": 2},
+ }
+ first = service.submit_sample("tenant", payload)
+ assert service.submit_sample("tenant", payload) == first, "the retry must reuse the request and sequence ids"
+ await await_settled(service, "tenant", first[0])
+ assert len(service.backend.named("sample")) == 1
+
+
+async def test_an_illegal_seq_id_is_rejected(service):
+ """seq_id 0 would never enter the ordered model queue; every future would pend."""
+ model_id = await created_model(service)
+ with pytest.raises(UserInputError, match="seq_id"):
+ service.submit("tenant", "optim_step", {"model_id": model_id, "seq_id": 0, "adam_params": dict(ADAM)})
+
+
+async def test_an_expired_lease_sweeps_its_sessions(service):
+ session_id = service.create_session("tenant")
+ service.sessions[session_id].last_heartbeat = -1e9
+ await service._expire_sessions()
+ with pytest.raises(UserInputError, match="unknown session"):
+ service.create_model("tenant", model_payload(service, session_id=session_id))
+
+
+async def test_a_sampling_session_for_another_base_model_is_rejected(service):
+ payload = {"session_id": service.create_session("tenant"), "sampling_session_seq_id": 1, "base_model": "other"}
+ with pytest.raises(UserInputError, match="serves"):
+ service.create_sampling_session("tenant", payload)
+
+
+async def test_a_read_only_batch_failure_preserves_queued_training(service):
+ model_id = await created_model(service)
+ service.backend.fail_on["forward_only"] = {"error": "request rejected"}
+ readonly = service.submit("tenant", "forward_only", fb_payload(model_id, 1, [datum()]))
+ backward = service.submit("tenant", "forward_backward", fb_payload(model_id, 2, [datum()]))
+ step = service.submit("tenant", "optim_step", _optim_payload(model_id, 3))
+ assert (await await_settled(service, "tenant", readonly)).state == FAILED
+ assert (await await_settled(service, "tenant", backward)).state == DONE
+ assert (await await_settled(service, "tenant", step)).state == DONE
+
+
+@pytest.mark.parametrize("shutdown", ["cancel", "failure"])
+async def test_dispatcher_shutdown_stops_model_creation_and_sampling(tmp_path, monkeypatch, shutdown):
+ gateway = make_service(tmp_path)
+ started = asyncio.Queue()
+ fail_sweep = asyncio.Event()
+
+ async def blocked_backend_call(*args):
+ started.put_nowait(asyncio.current_task())
+ await asyncio.Event().wait()
+
+ async def _run_lease_sweeper():
+ await fail_sweep.wait()
+ raise RuntimeError("lease sweeper failed")
+
+ monkeypatch.setattr(gateway.backend, "load_slot", blocked_backend_call)
+ monkeypatch.setattr(gateway.backend, "sample", blocked_backend_call)
+ monkeypatch.setattr(gateway, "_run_lease_sweeper", _run_lease_sweeper)
+ run_task = asyncio.create_task(gateway.run())
+ gateway.create_model("tenant", model_payload(gateway))
+ gateway.submit_sample("tenant", {"num_samples": 1, "prompt_tokens": [1]})
+ tasks = [run_task, *gateway._create_tasks, *(task for task, _ in gateway._sample_tasks.values())]
+ try:
+ backend_tasks = [await asyncio.wait_for(started.get(), timeout=2) for _ in range(2)]
+ if shutdown == "cancel":
+ run_task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await asyncio.wait_for(run_task, timeout=2)
+ else:
+ fail_sweep.set()
+ with pytest.raises(RuntimeError, match="lease sweeper failed"):
+ await asyncio.wait_for(run_task, timeout=2)
+ assert all(task.done() for task in backend_tasks), "backend tasks outlived the dispatcher"
+ finally:
+ for task in tasks:
+ task.cancel()
+ await asyncio.gather(*tasks, return_exceptions=True)
diff --git a/tests/fast/tinker/harness.py b/tests/fast/tinker/harness.py
new file mode 100644
index 0000000000..4df663a818
--- /dev/null
+++ b/tests/fast/tinker/harness.py
@@ -0,0 +1,174 @@
+"""Shared fakes for the tinker gateway suite."""
+
+import asyncio
+
+from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir
+from miles.tinker.core.future import DONE, PENDING, RequestFuture
+from miles.tinker.core.service import TinkerService
+from miles.tinker.core.types import Command, CommandOp, GatewayConfig
+
+ADAM = {
+ "learning_rate": 1e-4,
+ "beta1": 0.9,
+ "beta2": 0.95,
+ "eps": 1e-12,
+ "weight_decay": 0.0,
+ "grad_clip_norm": 1.0,
+}
+
+
+class FakeBackend:
+ """Record calls with deterministic outputs and configurable failures."""
+
+ def __init__(self) -> None:
+ self.calls: list[tuple[str, dict]] = []
+ self.fail_next: Exception | dict | None = None
+ self.dead = False
+ self.optim_outcomes: dict[int, dict] = {}
+ self.fail_on: dict[str, Exception | dict] = {}
+
+ def _record(self, name: str, **kwargs) -> dict | None:
+ self.calls.append((name, kwargs))
+ failure, self.fail_next = self.fail_next, None
+ if failure is None:
+ failure = self.fail_on.pop(name, None)
+ if isinstance(failure, Exception):
+ self.dead = True
+ raise failure
+ return failure
+
+ def named(self, name: str) -> list[dict]:
+ return [kwargs for called, kwargs in self.calls if called == name]
+
+ def trainer_dead(self):
+ return self.dead
+
+ async def load_slot(self, slot, rank, alpha, ckpt_path=None, load_optimizer=True):
+ return self._record(
+ "load_slot", slot=slot, rank=rank, alpha=alpha, ckpt_path=ckpt_path, load_optimizer=load_optimizer
+ )
+
+ async def unload_slot(self, slot):
+ return self._record("unload_slot", slot=slot)
+
+ async def forward_backward(self, batch_id, slot_datums, loss_fn, loss_fn_config):
+ failure = self._record(
+ "forward_backward",
+ batch_id=batch_id,
+ slot_datums=slot_datums,
+ loss_fn=loss_fn,
+ loss_fn_config=loss_fn_config,
+ )
+ if failure is not None:
+ return failure
+ return [{"loss": 1.0, "logprobs": [0.0] * datum["target_len"]} for _, datum in slot_datums]
+
+ async def forward_only(self, batch_id, slot_datums, loss_fn, loss_fn_config):
+ failure = self._record(
+ "forward_only", batch_id=batch_id, slot_datums=slot_datums, loss_fn=loss_fn, loss_fn_config=loss_fn_config
+ )
+ if failure is not None:
+ return failure
+ return [{"loss": 0.0, "logprobs": [0.0] * datum["target_len"]} for _, datum in slot_datums]
+
+ async def optim_step(self, adam_params_by_slot):
+ failure = self._record("optim_step", adam_params_by_slot=adam_params_by_slot)
+ if failure is not None:
+ return {slot: failure for slot in adam_params_by_slot}
+ return {slot: self.optim_outcomes.get(slot, {"grad_norm": 0.5 + slot}) for slot in adam_params_by_slot}
+
+ async def save_slot(self, slot, path, metadata=None):
+ return self._write_checkpoint("save_slot", path, metadata, slot=slot)
+
+ async def export_slot(self, slot, rank, alpha, path, metadata=None):
+ return self._write_checkpoint("export_slot", path, metadata, slot=slot, rank=rank, alpha=alpha)
+
+ def _write_checkpoint(self, name, path, metadata, **kwargs):
+ failure = self._record(name, path=path, metadata=metadata, **kwargs)
+ if failure is None and metadata is not None:
+ write_checkpoint_dir(path, lambda _: None, metadata=metadata, overwrite=name == "save_slot")
+ return failure
+
+ async def sample(self, payload, lora_name, lora_path=None):
+ failure = self._record("sample", payload=payload, lora_name=lora_name, lora_path=lora_path)
+ if failure is not None:
+ return failure
+ return {
+ "sequences": [
+ {"tokens": [1, 2], "logprobs": [0.0, 0.0], "stop_reason": "stop"}
+ for _ in range(payload["num_samples"])
+ ]
+ }
+
+
+def make_config(checkpoint_root, **overrides) -> GatewayConfig:
+ defaults = dict(base_model="base", n_slots=2, checkpoint_root=str(checkpoint_root), vocab_size=128000)
+ return GatewayConfig(**{**defaults, **overrides})
+
+
+def make_service(checkpoint_root, **config_overrides) -> TinkerService:
+ return TinkerService(FakeBackend(), make_config(checkpoint_root, **config_overrides))
+
+
+def datum(tokens: int = 3) -> dict:
+ return {
+ "tokens": list(range(tokens + 1)),
+ "target_tokens": list(range(1, tokens + 1)),
+ "target_len": tokens,
+ "weights": [1.0] * tokens,
+ }
+
+
+def rl_datum(tokens: int = 3) -> dict:
+ """RL losses read logprobs+advantages and reject the cross-entropy weights."""
+ return {
+ "tokens": list(range(tokens + 1)),
+ "target_tokens": list(range(1, tokens + 1)),
+ "target_len": tokens,
+ "sampling_logprobs": [0.0] * tokens,
+ "advantages": [1.0] * tokens,
+ }
+
+
+def fb_payload(model_id: str, seq_id: int, datums: list[dict], loss_fn: str = "cross_entropy") -> dict:
+ return {"model_id": model_id, "seq_id": seq_id, "datums": datums, "loss_fn": loss_fn, "loss_fn_config": {}}
+
+
+def command(model_id: str, seq_id: int, op: str, payload: dict, arrival: int) -> Command:
+ return Command(
+ model_id=model_id,
+ seq_id=seq_id,
+ op=CommandOp(op),
+ payload=payload,
+ request_id=f"req-{seq_id}",
+ arrival=arrival,
+ )
+
+
+def model_payload(service: TinkerService, tenant: str = "tenant", session_id: str | None = None, **overrides) -> dict:
+ return {
+ "base_model": service.config.base_model,
+ "session_id": session_id or service.create_session(tenant),
+ "model_seq_id": 1,
+ **overrides,
+ }
+
+
+async def created_model(service: TinkerService, tenant: str = "tenant", session_id: str | None = None) -> str:
+ request_id, model_id = service.create_model(
+ tenant, model_payload(service, tenant, session_id, lora_config={"rank": 8})
+ )
+ future = await await_settled(service, tenant, request_id)
+ assert future.state == DONE, future.error
+ return model_id
+
+
+async def await_settled(service: TinkerService, tenant: str, request_id: str, timeout: float = 2.0) -> RequestFuture:
+ deadline = asyncio.get_running_loop().time() + timeout
+ while asyncio.get_running_loop().time() < deadline:
+ future = service.retrieve_future(tenant, request_id)
+ assert future is not None, f"future {request_id} expired"
+ if future.state != PENDING:
+ return future
+ await asyncio.sleep(0.005)
+ raise AssertionError(f"future {request_id} still pending after {timeout}s")
diff --git a/tests/fast/tinker/server/__init__.py b/tests/fast/tinker/server/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/fast/tinker/server/test_app.py b/tests/fast/tinker/server/test_app.py
new file mode 100644
index 0000000000..7c11ab686e
--- /dev/null
+++ b/tests/fast/tinker/server/test_app.py
@@ -0,0 +1,183 @@
+"""HTTP responses preserve the SDK conversation and retry/error contract."""
+
+import asyncio
+
+import httpx
+import pytest
+from tests.fast.tinker.harness import ADAM
+
+from miles.tinker.server.app import build_app
+
+
+@pytest.fixture
+async def client(service):
+ transport = httpx.ASGITransport(app=build_app(service))
+ async with httpx.AsyncClient(transport=transport, base_url="http://gateway") as http:
+ http.service = service
+ yield http
+
+
+def _headers(tenant: str = "tenant-a") -> dict:
+ return {"X-API-Key": tenant}
+
+
+async def _poll(client, request_id: str, tenant: str = "tenant-a") -> dict:
+ for _ in range(400):
+ response = await client.post(
+ "/api/v1/retrieve_future", json={"request_id": request_id}, headers=_headers(tenant)
+ )
+ body = response.json()
+ if body.get("type") != "try_again":
+ return body
+ await asyncio.sleep(0.005)
+ raise AssertionError(f"{request_id} never settled")
+
+
+async def _model_body(client, tenant: str = "tenant-a", **extra) -> dict:
+ session = (await client.post("/api/v1/create_session", json={}, headers=_headers(tenant))).json()
+ return {"base_model": "base", "session_id": session["session_id"], "model_seq_id": 1, **extra}
+
+
+def _fb_body(model_id: str, seq_id: int) -> dict:
+ return {
+ "model_id": model_id,
+ "seq_id": seq_id,
+ "forward_backward_input": {
+ "data": [
+ {
+ "model_input": {"chunks": [{"type": "encoded_text", "tokens": [1, 2, 3]}]},
+ "loss_fn_inputs": {"target_tokens": [2, 3, 4], "weights": [1.0, 1.0, 1.0]},
+ }
+ ],
+ "loss_fn": "cross_entropy",
+ },
+ }
+
+
+async def test_the_training_conversation(client):
+ session = await client.post("/api/v1/create_session", json={}, headers=_headers())
+ assert session.json()["session_id"].startswith("session-")
+
+ body = {"base_model": "base", "session_id": session.json()["session_id"], "model_seq_id": 1}
+ created = (await client.post("/api/v1/create_model", json=body, headers=_headers())).json()
+ assert (await _poll(client, created["request_id"]))["type"] == "create_model"
+ model_id = created["model_id"]
+
+ info = (await client.post("/api/v1/get_info", json={"model_id": model_id}, headers=_headers())).json()
+ assert info["model_data"]["model_name"] == "base"
+
+ fb = (await client.post("/api/v1/forward_backward", json=_fb_body(model_id, 1), headers=_headers())).json()
+ fb_result = await _poll(client, fb["request_id"])
+ assert len(fb_result["loss_fn_outputs"][0]["logprobs"]["data"]) == 3
+
+ optim = (
+ await client.post(
+ "/api/v1/optim_step",
+ json={"model_id": model_id, "seq_id": 2, "adam_params": dict(ADAM)},
+ headers=_headers(),
+ )
+ ).json()
+ optim_result = await _poll(client, optim["request_id"])
+ assert "grad_norm" in optim_result["metrics"]
+
+ sampler = (
+ await client.post(
+ "/api/v1/save_weights_for_sampler", json={"model_id": model_id, "seq_id": 3}, headers=_headers()
+ )
+ ).json()
+ sampler_path = (await _poll(client, sampler["request_id"]))["path"]
+
+ sample = (
+ await client.post(
+ "/api/v1/asample",
+ json={
+ "model_path": sampler_path,
+ "num_samples": 2,
+ "prompt": {"chunks": [{"type": "encoded_text", "tokens": [1]}]},
+ "sampling_params": {"max_tokens": 4},
+ },
+ headers=_headers(),
+ )
+ ).json()
+ assert len(sample["sample_sequence_ids"]) == 2, "asample must answer an UntypedAPIFuture"
+ assert len((await _poll(client, sample["request_id"]))["sequences"]) == 2
+
+
+async def test_bad_input_answers_400(client):
+ body = _fb_body("model-missing", 1)
+ body["forward_backward_input"]["data"][0]["loss_fn_inputs"]["target_tokens"] = [9, 9]
+ response = await client.post("/api/v1/forward_backward", json=body, headers=_headers())
+ assert response.status_code == 400
+
+
+async def test_a_foreign_future_answers_403(client):
+ created = (
+ await client.post("/api/v1/create_model", json=await _model_body(client), headers=_headers("tenant-a"))
+ ).json()
+ response = await client.post(
+ "/api/v1/retrieve_future", json={"request_id": created["request_id"]}, headers=_headers("tenant-b")
+ )
+ assert response.status_code == 403
+
+
+async def test_an_unknown_future_answers_410(client):
+ response = await client.post("/api/v1/retrieve_future", json={"request_id": "req-gone"}, headers=_headers())
+ assert response.status_code == 410
+
+
+async def test_a_failed_future_reports_the_category(client):
+ body = await _model_body(client)
+ client.service.backend.fail_next = {"error": "boom"}
+ created = (await client.post("/api/v1/create_model", json=body, headers=_headers())).json()
+ body = await _poll(client, created["request_id"])
+ assert (body["category"], "boom" in body["error"]) == ("server", True)
+
+
+async def test_capabilities_and_telemetry_shapes(client):
+ capabilities = (await client.get("/api/v1/get_server_capabilities")).json()
+ assert capabilities["supported_models"][0]["model_name"] == "base"
+ assert (await client.post("/api/v1/telemetry", json={})).json() == {"status": "accepted"}
+
+
+async def test_a_request_without_an_api_key_is_rejected(client):
+ response = await client.post("/api/v1/create_session", json={})
+ assert response.status_code == 400
+ assert "X-API-Key" in response.json()["error"]
+
+
+async def test_a_bearer_authorization_still_authenticates(client):
+ response = await client.post("/api/v1/create_session", json={}, headers={"Authorization": "Bearer tenant-b"})
+ assert response.status_code == 200
+
+
+async def test_weights_info_answers_the_sdk_resume_probe(client):
+ created = (
+ await client.post(
+ "/api/v1/create_model",
+ json=await _model_body(client, lora_config={"rank": 8}),
+ headers=_headers(),
+ )
+ ).json()
+ await _poll(client, created["request_id"])
+ saved = (
+ await client.post(
+ "/api/v1/save_weights",
+ json={"model_id": created["model_id"], "seq_id": 1, "path": "ck", "overwrite": False},
+ headers=_headers(),
+ )
+ ).json()
+ path = (await _poll(client, saved["request_id"]))["path"]
+ info = (await client.post("/api/v1/weights_info", json={"tinker_path": path}, headers=_headers())).json()
+ assert (info["base_model"], info["is_lora"], info["lora_rank"]) == ("base", True, 8)
+
+
+async def test_retrieve_long_polls_until_settlement(client):
+ future = client.service.futures.create("model", "tenant-a")
+ poll = asyncio.create_task(
+ client.post("/api/v1/retrieve_future", json={"request_id": future.request_id}, headers=_headers())
+ )
+ await asyncio.sleep(0.05)
+ assert not poll.done(), "a pending future must hold the poll open instead of answering try_again"
+ client.service.futures.resolve(future.request_id, {"op": "optim_step", "metrics": {"grad_norm": 1.0}})
+ response = await asyncio.wait_for(poll, timeout=2)
+ assert response.json() == {"type": "optim_step", "metrics": {"grad_norm": 1.0}}
diff --git a/tests/fast/tinker/server/test_encoding.py b/tests/fast/tinker/server/test_encoding.py
new file mode 100644
index 0000000000..098689bb79
--- /dev/null
+++ b/tests/fast/tinker/server/test_encoding.py
@@ -0,0 +1,104 @@
+"""JSON decoding validates datum inputs and renders SDK result shapes."""
+
+from miles.tinker.server.encoding import (
+ ADAM_PARAM_DEFAULTS,
+ build_datum,
+ decode_command,
+ decode_sample_request,
+ render_result,
+ tensor_data_to_list,
+)
+
+
+def _datum(tokens: list[int], **extra_inputs) -> dict:
+ return {
+ "model_input": {"chunks": [{"type": "encoded_text", "tokens": tokens[:-1]}]},
+ "loss_fn_inputs": {"target_tokens": tokens[1:], **extra_inputs},
+ }
+
+
+class TestDecodeForwardBackward:
+ def _decode(self, datum, forward_only=False):
+ payload = {
+ "model_id": "model",
+ "seq_id": 1,
+ "forward_only": forward_only,
+ "forward_backward_input": {"data": [datum], "loss_fn": "cross_entropy"},
+ }
+ return decode_command("forward_backward", payload)
+
+ def test_a_shifted_datum_becomes_one_datum(self):
+ op, decoded = self._decode(_datum([1, 2, 3, 4], weights=[1.0, 1.0, 1.0]))
+ assert op == "forward_backward"
+ assert decoded["datums"] == [
+ {"tokens": [1, 2, 3, 4], "target_len": 3, "target_tokens": [2, 3, 4], "weights": [1.0, 1.0, 1.0]}
+ ]
+
+ def test_forward_only_selects_the_op(self):
+ op, _ = self._decode(_datum([1, 2, 3]), forward_only=True)
+ assert op == "forward_only"
+
+ def test_a_length_mismatch_is_rejected(self):
+ datum = _datum([1, 2, 3])
+ datum["loss_fn_inputs"]["target_tokens"] = [2]
+ _, decoded = self._decode(datum)
+ assert "length" in decoded["validation_error"]
+
+ def test_explicit_targets_pass_through(self):
+ """rl_loop pads the prompt region with dummy zero targets; labels need not be the shifted sequence."""
+ datum = _datum([1, 2, 3])
+ datum["loss_fn_inputs"]["target_tokens"] = [0, 9]
+ _, decoded = self._decode(datum)
+ assert decoded["datums"][0]["tokens"] == [1, 2, 9]
+ assert decoded["datums"][0]["target_tokens"] == [0, 9]
+
+
+def test_build_datum_maps_the_wire_input_names():
+ datum = build_datum([1, 2], {"target_tokens": [2, 3], "logprobs": [0.5, 0.5], "advantages": [1.0, -1.0]}, 0)
+ assert datum["sampling_logprobs"] == [0.5, 0.5]
+ assert datum["advantages"] == [1.0, -1.0]
+
+
+class TestAdamParams:
+ def test_defaults_fill_missing_keys(self):
+ _, decoded = decode_command(
+ "optim_step", {"model_id": "m", "seq_id": 1, "adam_params": {"learning_rate": 3e-4}}
+ )
+ assert decoded["adam_params"]["learning_rate"] == 3e-4
+ assert set(decoded["adam_params"]) == set(ADAM_PARAM_DEFAULTS)
+
+ def test_unknown_keys_are_rejected(self):
+ """The SDK's AdamParams model is the contract; a key it lacks must not be silently dropped."""
+ _, decoded = decode_command("optim_step", {"model_id": "m", "seq_id": 1, "adam_params": {"momentum": 0.9}})
+ assert "momentum" in decoded["validation_error"]
+
+
+def test_an_off_model_key_is_rejected():
+ _, decoded = decode_command("save_state", {"model_id": "m", "seq_id": 1, "bogus": 1})
+ assert "bogus" in decoded["validation_error"]
+
+
+class TestTensorData:
+ def test_dense_tensor_data(self):
+ assert tensor_data_to_list({"data": [1.0, 2.0], "shape": [2]}) == [1.0, 2.0]
+
+ def test_csr_expands_to_dense(self):
+ sparse = {"shape": [4], "sparse_crow_indices": [0, 2], "sparse_col_indices": [1, 3], "data": [5.0, 7.0]}
+ assert tensor_data_to_list(sparse) == [0, 5.0, 0, 7.0]
+
+
+def test_decode_sample_request_defaults():
+ decoded = decode_sample_request(
+ {"prompt": {"chunks": [{"type": "encoded_text", "tokens": [1, 2]}]}, "sampling_params": {"max_tokens": 1}}
+ )
+ assert decoded["num_samples"] == 1
+ assert decoded["prompt_tokens"] == [1, 2]
+ assert decoded["topk_prompt_logprobs"] == 0
+
+
+class TestRenderResult:
+ def test_forward_backward_renders_per_datum_records(self):
+ rendered = render_result({"op": "forward_backward", "outputs": [{"loss": 2.0, "logprobs": [0.1, 0.2]}]})
+ assert rendered["loss_fn_output_type"] == "ArrayRecord"
+ assert rendered["metrics"] == {"loss:sum": 2.0}
+ assert rendered["loss_fn_outputs"][0]["logprobs"] == {"dtype": "float32", "shape": [2], "data": [0.1, 0.2]}
diff --git a/tests/fast/tinker/server/test_proto_codec.py b/tests/fast/tinker/server/test_proto_codec.py
new file mode 100644
index 0000000000..a78da233bb
--- /dev/null
+++ b/tests/fast/tinker/server/test_proto_codec.py
@@ -0,0 +1,130 @@
+"""Protobuf and JSON preserve the same datums and SDK result shapes."""
+
+import pytest
+
+tinker = pytest.importorskip("tinker")
+
+from miles.tinker.server.encoding import decode_command # noqa: E402
+from miles.tinker.server.proto_codec import ( # noqa: E402
+ decode_forward_backward_request,
+ encode_forward_backward_output,
+ encode_sample_response,
+ maybe_decompress,
+)
+from tinker import types # noqa: E402
+from tinker.proto.request_conv import forward_backward_request_to_proto # noqa: E402
+from tinker.proto.response_conv import deserialize_forward_backward_output, deserialize_sample_response # noqa: E402
+
+TOKENS = [3, 1, 4, 1, 5]
+
+
+def _sdk_request(loss_fn_inputs: dict) -> bytes:
+ datum = types.Datum(
+ model_input=types.ModelInput.from_ints(TOKENS[:-1]),
+ loss_fn_inputs={"target_tokens": TOKENS[1:], **loss_fn_inputs},
+ )
+ request = types.ForwardBackwardRequest(
+ model_id="model-x",
+ seq_id=7,
+ forward_backward_input=types.ForwardBackwardInput(
+ data=[datum], loss_fn="cross_entropy", loss_fn_config={"beta": 0.5}
+ ),
+ )
+ return forward_backward_request_to_proto(request).SerializeToString()
+
+
+def test_an_sdk_proto_request_decodes_to_internal_datums():
+ op, decoded = decode_forward_backward_request(_sdk_request({"weights": [1.0, 0.0, 1.0, 1.0]}))
+
+ assert (op, decoded["model_id"], decoded["seq_id"]) == ("forward_backward", "model-x", 7)
+ assert decoded["loss_fn_config"] == {"beta": 0.5}
+ assert decoded["datums"] == [
+ {"tokens": TOKENS, "target_len": 4, "target_tokens": TOKENS[1:], "weights": [1.0, 0.0, 1.0, 1.0]}
+ ]
+
+
+def test_both_wire_halves_produce_the_same_datums():
+ proto_datums = decode_forward_backward_request(_sdk_request({"weights": [1.0, 1.0, 1.0, 1.0]}))[1]["datums"]
+ _, json_decoded = decode_command(
+ "forward_backward",
+ {
+ "model_id": "model-x",
+ "seq_id": 7,
+ "forward_backward_input": {
+ "data": [
+ {
+ "model_input": {"chunks": [{"type": "encoded_text", "tokens": TOKENS[:-1]}]},
+ "loss_fn_inputs": {"target_tokens": TOKENS[1:], "weights": [1.0, 1.0, 1.0, 1.0]},
+ }
+ ],
+ "loss_fn": "cross_entropy",
+ },
+ },
+ )
+ assert proto_datums == json_decoded["datums"]
+
+
+def test_a_sparse_tensor_decodes_dense():
+ sparse = types.TensorData(
+ dtype="float32", shape=[4], data=[2.0], sparse_crow_indices=[0, 1], sparse_col_indices=[2]
+ )
+ _, decoded = decode_forward_backward_request(_sdk_request({"weights": sparse}))
+ assert decoded["datums"][0]["weights"] == [0.0, 0.0, 2.0, 0.0]
+
+
+def test_a_zstd_body_decompresses():
+ zstandard = pytest.importorskip("zstandard")
+ body = _sdk_request({"weights": [1.0, 1.0, 1.0, 1.0]})
+ compressed = zstandard.ZstdCompressor().compress(body)
+
+ assert maybe_decompress(compressed, "zstd") == body
+ assert maybe_decompress(body, None) == body
+
+
+def test_the_sdk_parses_our_forward_backward_output():
+ encoded = encode_forward_backward_output(
+ {
+ "op": "forward_backward",
+ "outputs": [{"loss": 2.5, "logprobs": [-0.1, -0.2]}, {"loss": 1.0, "logprobs": [-0.3]}],
+ }
+ )
+ parsed = deserialize_forward_backward_output(encoded)
+
+ assert parsed.metrics["loss:sum"] == pytest.approx(3.5)
+ assert [list(record["loss:sum"].data) for record in parsed.loss_fn_outputs] == [[2.5], [1.0]]
+ assert list(parsed.loss_fn_outputs[0]["logprobs"].data) == pytest.approx([-0.1, -0.2])
+
+
+def test_the_sdk_parses_our_sample_response():
+ encoded = encode_sample_response(
+ {
+ "op": "sample",
+ "sequences": [{"sequence_id": "s", "tokens": [5, 6], "logprobs": [-0.5, -0.6], "stop_reason": "length"}],
+ "prompt_logprobs": [-1.0, -2.0],
+ "topk_prompt_logprobs": {"token_ids": [[1, 2]], "logprobs": [[-0.1, -0.2]]},
+ }
+ )
+ parsed = deserialize_sample_response(encoded)
+
+ sequence = parsed.sequences[0]
+ assert (list(sequence.tokens), sequence.stop_reason) == ([5, 6], "length")
+ assert list(sequence.logprobs) == pytest.approx([-0.5, -0.6])
+ assert list(parsed.prompt_logprobs) == pytest.approx([-1.0, -2.0])
+
+
+@pytest.mark.parametrize("forward_only", [False, True])
+@pytest.mark.parametrize("invalid", ["missing_targets", "bad_shape"])
+def test_content_errors_preserve_the_model_queue_envelope(forward_only, invalid):
+ from tinker.proto import tinker_public_pb2 as public_pb
+
+ request = public_pb.ForwardBackwardRequest()
+ request.ParseFromString(_sdk_request({"weights": [1.0] * 4}))
+ request.forward_only = forward_only
+ if invalid == "missing_targets":
+ del request.data[0].loss_fn_inputs["target_tokens"]
+ else:
+ request.data[0].loss_fn_inputs["weights"].shape.append(2)
+ op, decoded = decode_forward_backward_request(request.SerializeToString())
+ assert op == ("forward_only" if forward_only else "forward_backward")
+ assert decoded["model_id"] == "model-x" and decoded["seq_id"] == 7
+ assert decoded["validation_error"]
diff --git a/tests/fast/tinker/test_layering.py b/tests/fast/tinker/test_layering.py
new file mode 100644
index 0000000000..f3c467dab8
--- /dev/null
+++ b/tests/fast/tinker/test_layering.py
@@ -0,0 +1,45 @@
+"""Machine enforcement of the gateway layering law: core speaks only the
+internal language, server never touches miles infrastructure, and runtime.py
+is the single miles translator inside the package."""
+
+import ast
+import sys
+from pathlib import Path
+
+TINKER_DIR = Path(__file__).resolve().parents[3] / "miles" / "tinker"
+STDLIB = set(sys.stdlib_module_names)
+
+
+def _import_roots(path: Path) -> set[str]:
+ roots = set()
+ for node in ast.walk(ast.parse(path.read_text())):
+ if isinstance(node, ast.Import):
+ roots.update(alias.name for alias in node.names)
+ elif isinstance(node, ast.ImportFrom) and node.level == 0:
+ roots.add(node.module)
+ return roots
+
+
+def test_core_imports_only_stdlib_and_itself():
+ for path in sorted((TINKER_DIR / "core").glob("*.py")):
+ for module in _import_roots(path):
+ allowed = module.split(".")[0] in STDLIB or module.startswith("miles.tinker.core")
+ assert allowed, f"{path.name} imports {module}; core must stay stdlib-only"
+
+
+def test_server_never_imports_miles_infrastructure():
+ for path in sorted((TINKER_DIR / "server").glob("*.py")):
+ for module in _import_roots(path):
+ assert not module.startswith(
+ ("miles.ray", "miles.backends", "miles.utils")
+ ), f"{path.name} imports {module}; server may only speak SDK wire and core"
+
+
+def test_runtime_is_the_only_miles_speaker():
+ for path in sorted(TINKER_DIR.rglob("*.py")):
+ if path.name == "runtime.py":
+ continue
+ for module in _import_roots(path):
+ assert not module.startswith(
+ ("miles.ray", "miles.backends")
+ ), f"{path.relative_to(TINKER_DIR)} imports {module}; only runtime.py translates miles"
diff --git a/tests/fast/tinker/test_loss_tables.py b/tests/fast/tinker/test_loss_tables.py
new file mode 100644
index 0000000000..63d82be2ad
--- /dev/null
+++ b/tests/fast/tinker/test_loss_tables.py
@@ -0,0 +1,13 @@
+"""The admission tables in core must agree with the trainer's loss registry."""
+
+from miles.backends.training_utils.loss_hub.tinker_losses import TINKER_LOSS_FUNCTIONS
+from miles.tinker.core.types import LOSS_FN_INPUTS, LOSS_INPUT_KEYS
+from miles.tinker.runtime import DATUM_TO_BATCH_KEYS
+
+
+def test_every_registered_loss_has_an_admission_entry():
+ assert set(LOSS_FN_INPUTS) == set(TINKER_LOSS_FUNCTIONS)
+
+
+def test_wire_keys_map_onto_trainer_batch_keys():
+ assert set(LOSS_INPUT_KEYS.values()) == set(DATUM_TO_BATCH_KEYS)
diff --git a/tests/fast/tinker/test_runtime.py b/tests/fast/tinker/test_runtime.py
new file mode 100644
index 0000000000..4c9b9d3ac5
--- /dev/null
+++ b/tests/fast/tinker/test_runtime.py
@@ -0,0 +1,201 @@
+"""Runtime translation preserves datum order, sampling parameters, and token logprobs."""
+
+import pytest
+import torch
+
+from miles.tinker.core.types import UserInputError
+from miles.tinker.runtime import (
+ MilesBackend,
+ _build_train_data,
+ _pad_to_dp_multiple,
+ _prompt_logprobs,
+ _slot_failure,
+ _to_sequence,
+ _topk_prompt_logprobs,
+)
+from miles.tinker.server.proto_codec import encode_sample_response
+from tinker.proto.response_conv import deserialize_sample_response
+
+
+def _datum(tokens: list[int], **extra) -> dict:
+ return {"tokens": tokens, "target_len": len(tokens) - 1, "target_tokens": tokens[1:], **extra}
+
+
+class TestBuildTrainData:
+ def test_datums_become_rollout_batch_keys(self):
+ train_data = _build_train_data([(3, _datum([1, 2, 3])), (5, _datum([4, 5]))])
+ assert train_data["tokens"] == [[1, 2, 3], [4, 5]]
+ assert train_data["target_tokens"] == [[2, 3], [5]]
+ assert train_data["response_lengths"] == [2, 1]
+ assert train_data["total_lengths"] == [3, 2]
+ assert train_data["loss_masks"] == [[1, 1], [1]]
+ assert train_data["adapter_slots"] == [3, 5]
+ assert train_data["sample_indices"] == [0, 1]
+ assert train_data["dynamic_global_batch_size"] == 2
+
+ def test_optional_datum_keys_map_to_batch_names(self):
+ train_data = _build_train_data(
+ [(0, _datum([1, 2], weights=[1.0], advantages=[2.0], sampling_logprobs=[-0.5]))]
+ )
+ assert train_data["loss_weights"] == [[1.0]]
+ assert train_data["advantages"] == [[2.0]]
+ assert train_data["rollout_log_probs"] == [[-0.5]]
+
+
+async def test_forward_backward_merges_worker_replicas():
+ backend = MilesBackend(trainer=None, router_url="http://router")
+ per_datum = [
+ {"sample_index": 1, "loss": 2.0, "logprobs": torch.tensor([-0.2])},
+ {"sample_index": 0, "loss": 1.0, "logprobs": torch.tensor([-0.1])},
+ ]
+
+ async def fake_call_trainer(method, batch_id, train_data):
+ return [{"per_datum": per_datum}, {"per_datum": per_datum}] # two ranks report the same datums
+
+ backend._call_trainer = fake_call_trainer
+ outputs = await backend.forward_backward(1, [(0, _datum([1, 2])), (0, _datum([3, 4]))], "cross_entropy", {})
+ assert outputs == [
+ {"loss": 1.0, "logprobs": [pytest.approx(-0.1)]},
+ {"loss": 2.0, "logprobs": [pytest.approx(-0.2)]},
+ ]
+
+
+class TestGenerateRequest:
+ def _request(self, params: dict, **payload_extra) -> dict:
+ payload = {
+ "prompt_tokens": [1, 2],
+ "sampling_params": params,
+ "num_samples": 1,
+ "prompt_logprobs": False,
+ "topk_prompt_logprobs": 0,
+ **payload_extra,
+ }
+ return MilesBackend(None, "http://router")._generate_request(payload, lora_name="m@1")
+
+ def test_max_tokens_is_required(self):
+ with pytest.raises(UserInputError, match="max_tokens"):
+ self._request({})
+
+ def test_params_map_to_sglang_names(self):
+ request = self._request({"max_tokens": 8, "temperature": 0.0, "seed": 7})
+ assert request["sampling_params"]["max_new_tokens"] == 8
+ assert request["sampling_params"]["sampling_seed"] == 7
+ assert request["lora_path"] == "m@1"
+ assert request["return_logprob"] is True
+
+ def test_stop_splits_token_ids_from_strings(self):
+ by_ids = self._request({"max_tokens": 1, "stop": [7, 8]})
+ by_text = self._request({"max_tokens": 1, "stop": ["\n"]})
+ assert by_ids["sampling_params"]["stop_token_ids"] == [7, 8]
+ assert by_text["sampling_params"]["stop"] == ["\n"]
+
+ def test_an_empty_stop_list_disables_eos(self):
+ request = self._request({"max_tokens": 1, "stop": []})
+ assert request["sampling_params"]["ignore_eos"] is True
+ assert "stop" not in request["sampling_params"] and "stop_token_ids" not in request["sampling_params"]
+ default = self._request({"max_tokens": 1})
+ assert "ignore_eos" not in default["sampling_params"], "stop=None keeps the default EOS behavior"
+
+ def test_topk_requests_prompt_logprobs(self):
+ request = self._request({"max_tokens": 1}, topk_prompt_logprobs=3)
+ assert request["logprob_start_len"] == 0
+ assert request["top_logprobs_num"] == 3
+
+
+class TestEngineResponseParsing:
+ def test_to_sequence_reads_tokens_and_stop_reason(self):
+ response = {
+ "meta_info": {
+ "output_token_logprobs": [(-0.1, 11), (-0.2, 12)],
+ "finish_reason": {"type": "length"},
+ }
+ }
+ sequence = _to_sequence(response)
+ assert (sequence["tokens"], sequence["logprobs"], sequence["stop_reason"]) == (
+ [11, 12],
+ [-0.1, -0.2],
+ "length",
+ )
+
+ def test_prompt_logprobs_pad_the_unscored_first_token(self):
+ response = {"meta_info": {"input_token_logprobs": [(None, 1), (-0.5, 2)]}}
+ first, second = _prompt_logprobs(response)
+ assert first != first, "the unscored position must be NaN"
+ assert second == -0.5
+
+ def test_topk_pads_ragged_positions(self):
+ response = {"meta_info": {"input_top_logprobs": [None, [(-0.1, 5)]]}}
+ topk = _topk_prompt_logprobs(response, k=2)
+ assert topk["token_ids"] == [[0, 0], [5, 0]]
+ assert topk["logprobs"][1][0] == -0.1
+ result = {"sequences": [], "topk_prompt_logprobs": topk}
+ decoded = deserialize_sample_response(encode_sample_response(result)).topk_prompt_logprobs
+ assert decoded[0] is None
+ assert len(decoded[1]) == 1
+ token_id, logprob = decoded[1][0]
+ assert token_id == 5
+ assert logprob == pytest.approx(-0.1)
+
+
+async def test_forward_only_runs_the_requested_loss():
+ backend = MilesBackend(trainer=None, router_url="http://router")
+ captured = {}
+
+ async def fake_call_trainer(method, batch_id, train_data):
+ captured["method"] = method
+ captured["loss_fn"] = train_data["loss_fn"]
+ return [{"per_datum": [{"sample_index": 0, "loss": 3.0, "logprobs": torch.tensor([-0.3])}]}]
+
+ backend._call_trainer = fake_call_trainer
+ outputs = await backend.forward_only(1, [(0, _datum([1, 2]))], "importance_sampling", {})
+ assert captured == {"method": "forward_only", "loss_fn": "importance_sampling"}
+ assert outputs == [{"loss": 3.0, "logprobs": [pytest.approx(-0.3)]}]
+
+
+def test_a_pinned_seed_still_gets_one_model_queue_per_sample():
+ from miles.tinker.runtime import _with_sample_seed
+
+ request = {"sampling_params": {"sampling_seed": 7, "temperature": 0.0}}
+ assert [_with_sample_seed(request, i)["sampling_params"]["sampling_seed"] for i in range(3)] == [7, 8, 9]
+ assert request["sampling_params"]["sampling_seed"] == 7
+ assert "sampling_seed" not in _with_sample_seed({"sampling_params": {}}, 2)["sampling_params"]
+
+
+class TestPadToDpMultiple:
+ def test_padding_replicates_the_last_datum_with_a_padding_marker(self):
+ slot_datums = [(0, _datum([1, 2, 3], weights=[1.0, 1.0], advantages=[2.0, 2.0]))]
+ padded = _pad_to_dp_multiple(slot_datums, 4)
+ assert len(padded) == 4
+ for slot, filler in padded[1:]:
+ assert slot == 0 and filler["tokens"] == [1, 2, 3] and filler["padding"]
+ assert "padding" not in slot_datums[0][1], "the original datum must not be mutated"
+
+
+async def test_forward_backward_pads_the_batch_and_drops_padding_outputs():
+ backend = MilesBackend(trainer=None, router_url="http://router", dp_size=2)
+ seen = {}
+
+ async def fake_call_trainer(method, batch_id, train_data):
+ seen.update(train_data)
+ per_datum = [
+ {"sample_index": index, "loss": float(index), "logprobs": torch.tensor([-0.1])}
+ for index in range(len(train_data["tokens"]))
+ ]
+ return [{"per_datum": per_datum}]
+
+ backend._call_trainer = fake_call_trainer
+ outputs = await backend.forward_backward(1, [(0, _datum([1, 2], weights=[1.0]))], "cross_entropy", {})
+ assert seen["dynamic_global_batch_size"] == 2, "a singleton batch must be padded to the DP size"
+ assert seen["loss_masks"] == [[1], [0]], "the zero mask removes padding from every loss term"
+ assert len(outputs) == 1, "padding outputs are dropped"
+
+
+def test_an_actor_error_verdict_survives_runtime_translation():
+ failure = {"error": "bad shard"}
+ assert _slot_failure([None, failure]) is failure
+
+
+def test_an_aborted_sample_fails_instead_of_passing_as_a_stop():
+ """A truncated sequence fed to RL as a completed sample corrupts training data silently."""
+ response = {"meta_info": {"output_token_logprobs": [(-0.1, 11)], "finish_reason": {"type": "abort"}}}
+ assert "abort" in _to_sequence(response)["error"]
diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py
index 6a750db3e9..ef5bed5c14 100644
--- a/tests/fast/utils/test_arguments.py
+++ b/tests/fast/utils/test_arguments.py
@@ -886,29 +886,6 @@ def test_accepts_default_single_tokenizer_worker(self):
assert args.multi_lora is True
- def test_defaults_rollout_fn_and_data_source_to_multi_lora(self):
- args = self._parse([])
-
- miles_validate_args(args)
-
- assert args.rollout_function_path == "miles.rollout.multi_lora.async_rollout.generate_rollout_multi_lora"
- assert args.data_source_path == "miles.rollout.multi_lora.data_source.MultiLoRAAsyncDataSource"
- assert args.rollout_global_dataset is True
-
- def test_keeps_user_supplied_rollout_fn_and_data_source(self):
- args = self._parse(
- ["--rollout-function-path", "my.custom.rollout_fn", "--data-source-path", "my.custom.DataSource"]
- )
-
- miles_validate_args(args)
-
- assert args.rollout_function_path == "my.custom.rollout_fn"
- assert args.data_source_path == "my.custom.DataSource"
-
- def test_empty_wait_is_a_registered_argument(self):
- assert self._parse([]).multi_lora_max_empty_wait_s == 30.0
- assert self._parse(["--multi-lora-max-empty-wait-s", "5"]).multi_lora_max_empty_wait_s == 5.0
-
def test_rejects_non_adam_optimizer(self):
# Per-slot optimizer isolation (state init, retirement cleanup, step
# clocks) only implements Adam semantics. Muon has its own dedicated
diff --git a/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/prepare.txt b/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/prepare.txt
new file mode 100644
index 0000000000..9ee56f5f0b
--- /dev/null
+++ b/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/prepare.txt
@@ -0,0 +1,6 @@
+### 0
+mkdir -p /root/models
+
+### 1
+hf download Qwen/Qwen3-30B-A3B
+ --local-dir /root/models/Qwen3-30B-A3B
diff --git a/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt b/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt
new file mode 100644
index 0000000000..5922bbcadc
--- /dev/null
+++ b/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt
@@ -0,0 +1,80 @@
+### 0
+pkill -9 sglang; sleep 3; ray stop
+ --force; pkill -9 ray; pkill -9 miles; sleep 3; pkill -9 ray; pkill -9 miles; pkill -9 redis; true;
+
+### 1
+export PYTHONUNBUFFERED=1 && ray start
+ --head
+ --node-ip-address 127.0.0.1
+ --num-gpus 8
+ --disable-usage-stats
+
+### 2
+nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l
+
+### 3
+export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit
+ --address="http://127.0.0.1:8265"
+ --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "0", "no_proxy": "127.0.0.1,127.0.0.1", "MASTER_ADDR": "127.0.0.1", "PYTHONPATH": ":/root/Megatron-LM:/frozen/pythonpath"}}'
+ -- python3 /serve_tinker.py
+ --disable-bias-linear
+ --qk-layernorm
+ --group-query-attention
+ --num-attention-heads 32
+ --num-query-groups 4
+ --kv-channels 128
+ --num-layers 48
+ --hidden-size 2048
+ --ffn-hidden-size 6144
+ --normalization RMSNorm
+ --position-embedding-type rope
+ --norm-epsilon 1e-6
+ --rotary-percent 1.0
+ --swiglu
+ --untie-embeddings-and-output-weights
+ --vocab-size 151936
+ --rotary-base 1000000
+ --moe-ffn-hidden-size 768
+ --moe-router-score-function softmax
+ --moe-token-dispatcher-type alltoall
+ --moe-router-topk 8
+ --moe-layer-freq '[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]'
+ --num-experts 128
+ --moe-grouped-gemm
+ --moe-token-drop-policy probs
+ --moe-router-dtype fp32
+ --moe-permute-fusion
+ --moe-aux-loss-coeff 0
+ --hf-checkpoint /root/models/Qwen3-30B-A3B
+ --megatron-to-hf-mode bridge
+ --lora-rank 32
+ --lora-alpha 64
+ --lora-dropout 0.0
+ --no-gradient-accumulation-fusion
+ --multi-lora-n-adapters 4
+ --tinker-server-port 10613
+ --tinker-checkpoint-root /root/shared_data/checkpoints/260101-000000-000
+ --tinker-train-attn
+ --tinker-train-mlp
+ --tinker-train-unembed
+ --optimizer adam
+ --lr 1e-4
+ --tensor-model-parallel-size 2
+ --sequence-parallel
+ --pipeline-model-parallel-size 1
+ --context-parallel-size 1
+ --expert-model-parallel-size 4
+ --expert-tensor-parallel-size 1
+ --use-dynamic-batch-size
+ --max-tokens-per-gpu 8192
+ --rollout-num-gpus-per-engine 2
+ --sglang-mem-fraction-static 0.7
+ --sglang-lora-backend triton
+ --actor-num-nodes 1
+ --actor-num-gpus-per-node 4
+ --rollout-num-gpus 4
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --accumulate-allreduce-grads-in-fp32
+ --attention-softmax-in-fp32
+ --attention-backend flash
diff --git a/train_multi_lora_async.py b/train_multi_lora_async.py
deleted file mode 100644
index 02863bfc62..0000000000
--- a/train_multi_lora_async.py
+++ /dev/null
@@ -1,109 +0,0 @@
-"""Fully-async multi-LoRA trainer driver."""
-
-import asyncio
-import logging
-from pathlib import Path
-
-import ray
-
-from miles.ray.multi_lora.controller import create_multilora_controller, get_multi_lora_controller
-from miles.ray.placement_group import create_rollout_components, create_training_models, update_weights
-from miles.ray.wiring import launch_worker_manager
-from miles.utils import object_store
-from miles.utils.adapter_config import parse_adapter_run_yaml
-from miles.utils.arguments import parse_args
-from miles.utils.audit_utils.process_identity import MainProcessIdentity
-from miles.utils.data import remove_rollout_data_refs
-from miles.utils.logging_utils import configure_logger
-from miles.utils.multi_lora import EmptyBatchTimeoutError, define_new_adapter_metrics
-from miles.utils.tracking_utils.tracking import init_tracking
-
-logger = logging.getLogger(__name__)
-
-
-def _is_empty_batch_timeout(task_error: ray.exceptions.RayTaskError) -> bool:
- cause = getattr(task_error, "cause", None)
- if isinstance(cause, EmptyBatchTimeoutError):
- return True
- return isinstance(task_error.as_instanceof_cause(), EmptyBatchTimeoutError)
-
-
-async def main(args):
- assert (
- not args.colocate
- ), "Colocation is not supported for fully-async training (generation needs continuous GPU; colocate time-shares)."
- configure_logger(args, source=MainProcessIdentity())
-
- # The multi-LoRA rollout fn / data source / global dataset flags are
- # defaulted by miles_validate_args when --multi-lora-n-adapters > 0.
- _worker_manager = launch_worker_manager(args)
- object_store.init_instance(args, contribute_segment=False)
- init_tracking(args)
- inference_controller, rollout_executor, _num_rollout_per_epoch = await create_rollout_components(args)
-
- # Create a controller nclusing MultiLoRAController and MultiLoRAHTTPServer to manage lora
- controller = create_multilora_controller(args, f"http://{args.sglang_router_ip}:{args.sglang_router_port}")
- await controller.start.remote()
- host = await controller.http_host.remote()
- api_port = await controller.api_port.remote()
- logger.info(f"Multi-LoRA control API listening on http://{host}:{api_port} (head node)")
-
- actor_model, _ = await create_training_models(args, inference_controller, rollout_executor)
-
- # CLI-registered adapters are loaded and pushed by the loop's first
- # reconcile + update_weights.
- for name, path in args.multi_lora_adapters:
- config = parse_adapter_run_yaml(Path(path))
- await controller.register_adapter.remote(name, config)
-
- rollout_id = 0
- while True:
- snapshot = await get_multi_lora_controller().snapshot.remote()
-
- # handle dynamic metrics in tracking backend
- define_new_adapter_metrics(snapshot)
- if not (snapshot["pending"] or snapshot["active"] or snapshot["retiring"] or snapshot["cleanup"]):
- if not args.multi_lora_service_mode:
- logger.info("No adapters; exiting.")
- break
- logger.info(f"No adapters; sleeping for {args.multi_lora_idle_poll_s}s...")
- await asyncio.sleep(args.multi_lora_idle_poll_s)
- continue
-
- # Reconcile + push before generate: the push promotes pending adapters,
- # and only then does the data source sample them. The actor pushes only
- # stale adapter weights (newly loaded, or stepped by the last batch).
- await actor_model.reconcile_adapters()
- await update_weights(actor_model, rollout_executor)
-
- # With nothing active, generate would wait forever.
- post_update = await get_multi_lora_controller().snapshot.remote()
- if not (post_update["active"] or post_update["retiring"]):
- continue
-
- try:
- await inference_controller.prepare_rollout(rollout_id)
- rollout_data = await rollout_executor.get.remote(rollout_id)
- except ray.exceptions.RayTaskError as e:
- if _is_empty_batch_timeout(e):
- logger.warning(f"Generate timed out with no trainable groups; retrying reconcile/update. {e}")
- continue
- raise
- await actor_model.train(rollout_id, rollout_data)
- remove_rollout_data_refs(args, rollout_data)
-
- # Per-adapter save cadence decided inside save_model.
- await actor_model.save_model(rollout_id)
- # TODO: support rollout_executor.save
-
- rollout_id += 1
-
- await rollout_executor.dispose.remote()
- await inference_controller.dispose()
- await actor_model.dispose()
- await controller.stop.remote()
-
-
-if __name__ == "__main__":
- args = parse_args()
- asyncio.run(main(args))