[multi-lora v2 6/7] Tinker gateway runtime and entry point

docs/advanced/lora.md

@@ -289,49 +289,10 @@ stability evidence rather than a released benchmark.
 
 ## Multi-LoRA training
 
-The dataset-driven multi-LoRA v1 backend (fully-async driver, adapter
-controller, per-adapter data sources) has been removed. A Tinker-protocol
-backend replaces it: clients submit explicit forward_backward / optim_step
-operations against adapter slots instead of the server owning datasets and
-schedules. The retained slot mechanics live in
-`miles/backends/megatron_utils/lora/`.
-
-## 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 is incompatible with general Bridge LoRA.
-- **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/lora/bridge.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/backends/megatron_utils/lora/slots.py` and
-  `miles/backends/megatron_utils/lora/optimizer.py` implement the
-  multi-adapter slot mechanics and per-slot optimizers.
+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.

miles/backends/megatron_utils/model.py

@@ -66,7 +66,7 @@ def _has_loadable_ckpt(load_dir: str | None) -> bool:
 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

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]

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",
+    }

miles/tinker/server/encoding.py

@@ -3,11 +3,14 @@
 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 miles.tinker.core.input_validation import validate_save_options
+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()
@@ -179,9 +182,21 @@ def render_result(result: dict) -> dict:
         }
     if op == "sample":
         rendered = {"type": "sample", "sequences": result["sequences"]}
-        for key in ("prompt_logprobs", "topk_prompt_logprobs"):
-            if result.get(key) is not None:
-                rendered[key] = result[key]
+        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"]}