feat: multi lora async

examples/multi_lora/README.md

@@ -0,0 +1,134 @@
+# Multi-LoRA Training Example (fully-async)
+
+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).
+
+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`
+
+## Layout
+
+```
+provision.sh                         # one-time: download model + datasets
+run_job.sh                           # entrypoint: bounded run, exits when done
+run_service.sh                       # service mode: idles for registrations (port 8068)
+service_smoke.py                     # register/deregister smoke test against the API
+train_multi_lora_async.py            # trainer (entry point)
+multi_lora_async_rollout.py          # fully-async rollout function
+multi_lora_data_source_async.py      # data source (reads controller)
+adapters/
+  gsm8k.yaml
+  dapo_math.yaml
+```
+
+Controller code lives in the library: `miles/utils/multi_lora.py` (registry +
+backend + HTTP API, torch-free) and `miles/ray/multi_lora_controller.py` (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)
+
+```bash
+bash examples/multi_lora/provision.sh
+```
+
+Downloads `Qwen/Qwen3-4B`, `zhuzilin/dapo-math-17k`, and `zhuzilin/gsm8k`.
+
+## Run
+
+```bash
+bash examples/multi_lora/run_job.sh
+```
+
+Registers the two adapters from CLI flags and trains until each hits its `num_step`
+(or `--num-rollout`), then exits.
+
+## 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/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, ...
+```
+
+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`).

examples/multi_lora/adapters/dapo_math.yaml

@@ -0,0 +1,9 @@
+rank: 32
+alpha: 32
+rollout_batch_size: 8       # prompt groups per optimizer step
+n_samples_per_prompt: 8     # -> 64 samples per step
+data: /root/dapo-math-17k/dapo-math-17k.jsonl
+input_key: prompt
+label_key: label
+rm_type: deepscaler
+num_step: 500

examples/multi_lora/adapters/gsm8k.yaml

@@ -0,0 +1,9 @@
+rank: 16
+alpha: 16
+rollout_batch_size: 32      # prompt groups per optimizer step
+n_samples_per_prompt: 4     # -> 128 samples per step
+data: /root/gsm8k/train.parquet
+input_key: messages
+label_key: label
+rm_type: math
+num_step: 400