diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md new file mode 100644 index 0000000000..7340d830e8 --- /dev/null +++ b/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`). diff --git a/examples/multi_lora/adapters/dapo_math.yaml b/examples/multi_lora/adapters/dapo_math.yaml new file mode 100644 index 0000000000..5003d82a14 --- /dev/null +++ b/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 diff --git a/examples/multi_lora/adapters/gsm8k.yaml b/examples/multi_lora/adapters/gsm8k.yaml new file mode 100644 index 0000000000..ff2798e214 --- /dev/null +++ b/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 diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py new file mode 100644 index 0000000000..1e240e76f6 --- /dev/null +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -0,0 +1,598 @@ +"""Fully-async multi-LoRA rollout with per-adapter gradient accumulation. + +A background producer generates continuously into per-adapter buffers. Each +train batch is collected by popping groups from the buffers round-robin, in +multiples of the adapter's ``min_groups_per_dp_split`` capped at its remaining +batch, so: + +- any batch splits evenly across data-parallel ranks; +- an adapter's batch (``rollout_batch_size`` prompt groups, i.e. + ``adapter_global_batch_size`` samples) is never overshot; +- adapters whose batch completes here are stamped as stepping. +""" + +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.base_types import call_dynamic_filter +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.metric_utils import compute_statistics, dict_add_prefix +from miles.utils.misc import load_function +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 completed prompt groups: a FIFO queue you can also + len(), and sweep for staleness. Bounded; the oldest group is dropped + when a put exceeds the cap.""" + + 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_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 + if adapter_name is not None: + adapter = await AdaptersCache().get(adapter_name) + submission_version = adapter.version if adapter is not None else None + + if submission_version is not None: + for s in group: + s.metadata["slot_version"] = submission_version + + 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 + + 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: + """The worker's cross-batch metric state, kept out of its buffer + machinery. Two cadences: dynamic-filter drops and staleness drops drain + every batch, per-adapter sample stats flush when the adapter's optimizer + step completes. Has its own lock — the producer thread records drops + 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) + # Shipped-sample values, accumulated per adapter across train batches + # and flushed as step statistics when the adapter steps. Bounded by + # adapter_global_batch_size values per adapter. + 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: total groups shipped, and the + # common reward of each uniform-reward (zero advantage) group. + 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 the shipped batch's rewards and response lengths per + adapter; for adapters stepping with this batch, flush statistics over + their whole adapter batch (accumulated across shipped batches, so the + stats cover all ``adapter_global_batch_size`` samples of the step, not + just this batch's slice). Returns {adapter name: flushed metrics}. + + Counted at ship time, not train commit: a failed train call aborts the + run anyway, so the distinction has no practical effect. + """ + 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" + ) + # Keys are single-segment ("raw_reward_mean", not + # "rollout/raw_reward/mean") so that, prefixed with + # "{name}/", they sit one level under the "{name}/*" glob — + # the same shape as "train/loss" under "train/*". Deeper keys + # never get their axis: glob expansion only matches one + # segment on the server. + 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 completed groups while + # get_groups (trainer side) pops them. + self.buffer_lock = threading.Lock() + self.buffers: dict[str, GroupBuffer] = defaultdict(GroupBuffer) + # Fairness cursor: the adapter whose buffer get_groups visits first. + # Advances past every visited adapter, persisting across calls and + # batches, so adapters are served round-robin. + self.rotation: deque[str] = deque() + self.metrics = MultiLoRAWorkerMetrics() + + @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) + 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_result = call_dynamic_filter(self.dynamic_filter, self.args, result) + if not filter_result.keep: + if filter_result.reason: + self.metrics.record_dynamic_filter_drop(filter_result.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 for the batch being collected. Returns the popped groups + ([] when nothing is poppable right now) and an updated copy of + ``group_counts`` (adapter name -> groups in the batch); passing the + counts back on each fetch is what keeps a batch from overshooting an + adapter's remaining groups. + + Pops round-robin from the cursor, one ``min_groups_per_dp_split`` at a + time, until ``num_samples`` is covered (the final multiple may overshoot + it) or no adapter can contribute. An adapter can't contribute when its + buffer holds less than a whole multiple, or the batch already holds all + its remaining groups. + """ + 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: + # Adapters retired at the last reconcile sync point: their buffered + # tail is discarded (base deregistration semantics), along with any + # partially accumulated reward stats. + for name in list(self.buffers): + if name not in adapters: + self.buffers.pop(name) + self.metrics.discard_adapter(name) + + # 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) + 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: + """Collect one train batch from the worker's buffers (same loop shape as + fully_async's generate_rollout_async): keep popping group multiples until + the batch reaches ``--global-batch-size`` samples, or it is non-empty and + made no progress for ``--multi-lora-max-coalesce-wait-s`` (the target can + be permanently unreachable when the live adapters' remaining batches are + smaller than the target, so ship what there is). + + The remaining-groups math relies on the sequential trainer loop: the + previous batch's ``mark_batch_trained`` has landed before this generate + call, so the snapshot's ``accumulated_groups`` is current. + """ + 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: + 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: generate is only called with live adapters, and the + # sequential loop retires adapters and commits accumulated_groups only + # between generate calls — so one snapshot serves the whole 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 live on the adapter's own optimizer-step axis + # ({name}/step), not rollout/step: one point per completed step, means + # over exactly the samples that step trained on. adapters[name].step is + # the committed count at snapshot time; 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), + # Per-adapter cycle metrics: {name}/perf/* rides rollout/step. Two + # segments under {name}/ keeps these off the step axis (glob expansion + # only reaches one segment); the {name}/perf/* glob catches them. + **{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/examples/multi_lora/multi_lora_data_source_async.py b/examples/multi_lora/multi_lora_data_source_async.py new file mode 100644 index 0000000000..a9b366437d --- /dev/null +++ b/examples/multi_lora/multi_lora_data_source_async.py @@ -0,0 +1,137 @@ +"""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 examples.multi_lora.multi_lora_async_rollout import AsyncMultiLoRAWorker + + AsyncMultiLoRAWorker.stop_global() diff --git a/examples/multi_lora/provision.sh b/examples/multi_lora/provision.sh new file mode 100755 index 0000000000..b8736cb36e --- /dev/null +++ b/examples/multi_lora/provision.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +# Download model weights (Qwen3-4B) +hf download Qwen/Qwen3-4B --local-dir /root/Qwen3-4B + +# Download training dataset (dapo-math-17k) +hf download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /root/dapo-math-17k + +# Download training dataset (gsm8k) +hf download --repo-type dataset zhuzilin/gsm8k \ + --local-dir /root/gsm8k diff --git a/examples/multi_lora/run_job.sh b/examples/multi_lora/run_job.sh new file mode 100755 index 0000000000..c111016811 --- /dev/null +++ b/examples/multi_lora/run_job.sh @@ -0,0 +1,94 @@ +#!/bin/bash +set -ex + +export GPUS_PER_NODE=8 + +pkill sglang || true +ray stop --force || true +sleep 3 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source scripts/models/qwen3-4B.sh + +ray start --head --node-ip-address 127.0.0.1 --num-gpus $GPUS_PER_NODE --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM", + "CUDA_DEVICE_MAX_CONNECTIONS": "1" + } + }' \ + -- python3 examples/multi_lora/train_multi_lora_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + --use-miles-router \ + ${MODEL_ARGS[@]} \ + \ + --hf-checkpoint /root/Qwen3-4B/ \ + --megatron-to-hf-mode bridge \ + \ + --lora-rank 32 \ + --lora-alpha 32 \ + --lora-dropout 0.0 \ + --target-modules "all-linear" \ + --multi-lora-n-adapters 4 \ + --multi-lora-idle-poll-s 5 \ + --multi-lora-adapter "dapo_math" "examples/multi_lora/adapters/dapo_math.yaml" \ + --multi-lora-adapter "gsm8k" "examples/multi_lora/adapters/gsm8k.yaml" \ + --multi-lora-disable-service-mode \ + --pause-generation-mode in_place \ + --max-weight-staleness 3 \ + --use-tis \ + \ + --apply-chat-template \ + --rollout-shuffle \ + --num-rollout 50 \ + --rollout-batch-size 32 \ + --n-samples-per-prompt 8 \ + --rollout-max-response-len 4096 \ + --rollout-temperature 1 \ + --global-batch-size 256 \ + \ + --save /tmp/multi_lora \ + --save-interval 5 \ + \ + --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 adam \ + --lr 1e-5 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --adam-beta1 0.9 \ + --adam-beta2 0.98 \ + \ + --tensor-model-parallel-size 2 \ + --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 \ + \ + --rollout-num-gpus-per-engine 1 \ + --sglang-mem-fraction-static 0.8 \ + \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash + # \ + # --use-wandb \ + # --wandb-host https://wandb.ai/ \ + # --wandb-team staging \ + # --wandb-project miles-multi-lora \ + # --wandb-group qwen3-4B + diff --git a/examples/multi_lora/run_service.sh b/examples/multi_lora/run_service.sh new file mode 100755 index 0000000000..eb4c7d56fe --- /dev/null +++ b/examples/multi_lora/run_service.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# Service mode: no adapters preloaded, idles for registrations (API on :8068). +# Pair with: python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 ... +set -ex + +export GPUS_PER_NODE=8 + +pkill sglang || true +ray stop --force || true +sleep 3 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source scripts/models/qwen3-4B.sh + +ray start --head --node-ip-address 127.0.0.1 --num-gpus $GPUS_PER_NODE --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM", + "CUDA_DEVICE_MAX_CONNECTIONS": "1" + } + }' \ + -- python3 examples/multi_lora/train_multi_lora_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + --use-miles-router \ + ${MODEL_ARGS[@]} \ + \ + --hf-checkpoint /root/Qwen3-4B/ \ + --megatron-to-hf-mode bridge \ + \ + --lora-rank 32 \ + --lora-alpha 32 \ + --lora-dropout 0.0 \ + --target-modules "all-linear" \ + --multi-lora-n-adapters 4 \ + --multi-lora-idle-poll-s 5 \ + --multi-lora-api-port 8068 \ + --pause-generation-mode in_place \ + --max-weight-staleness 3 \ + --use-tis \ + \ + --apply-chat-template \ + --rollout-shuffle \ + --num-rollout 50 \ + --rollout-batch-size 32 \ + --n-samples-per-prompt 8 \ + --rollout-max-response-len 4096 \ + --rollout-temperature 1 \ + --global-batch-size 256 \ + \ + --save /tmp/multi_lora_service \ + --save-interval 5 \ + \ + --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 adam \ + --lr 1e-5 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --adam-beta1 0.9 \ + --adam-beta2 0.98 \ + \ + --tensor-model-parallel-size 2 \ + --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 \ + \ + --rollout-num-gpus-per-engine 1 \ + --sglang-mem-fraction-static 0.8 \ + \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash diff --git a/examples/multi_lora/service_smoke.py b/examples/multi_lora/service_smoke.py new file mode 100644 index 0000000000..1f93a9ec3f --- /dev/null +++ b/examples/multi_lora/service_smoke.py @@ -0,0 +1,170 @@ +"""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/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/examples/multi_lora/tests/__init__.py b/examples/multi_lora/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/multi_lora/train_multi_lora_async.py b/examples/multi_lora/train_multi_lora_async.py new file mode 100644 index 0000000000..7aa45a875e --- /dev/null +++ b/examples/multi_lora/train_multi_lora_async.py @@ -0,0 +1,106 @@ +"""Fully-async multi-LoRA trainer driver.""" + +import asyncio +import logging +from pathlib import Path + +import ray + +from miles.ray.multi_lora_controller import create_controller, get_multi_lora_controller +from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models +from miles.utils.adapter_config import parse_adapter_run_yaml +from miles.utils.arguments import parse_args +from miles.utils.multi_lora import EmptyBatchTimeoutError, define_new_adapter_metrics +from miles.utils.audit_utils.process_identity import MainProcessIdentity +from miles.utils.logging_utils import configure_logger +from miles.utils.tracking_utils.tracking import init_tracking + +logger = logging.getLogger(__name__) + +ROLLOUT_FUNCTION_PATH = "examples.multi_lora.multi_lora_async_rollout.generate_rollout_multi_lora" +DATA_SOURCE_PATH = "examples.multi_lora.multi_lora_data_source_async.MultiLoRAAsyncDataSource" + + +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()) + + args.rollout_function_path = ROLLOUT_FUNCTION_PATH + args.data_source_path = DATA_SOURCE_PATH + args.rollout_global_dataset = True + + pgs = create_placement_groups(args) + init_tracking(args) + rollout_manager, _num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + + router_ip, router_port = await rollout_manager.get_router_address.remote() + args.sglang_router_ip, args.sglang_router_port = router_ip, router_port + controller = create_controller(args, f"http://{router_ip}:{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, pgs, rollout_manager) + + # 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 actor_model.update_weights() + + # 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: + rollout_data = await rollout_manager.generate.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) + + # Per-adapter save cadence decided inside save_model. + await actor_model.save_model(rollout_id) + + rollout_id += 1 + + await rollout_manager.dispose.remote() + await controller.stop.remote() + + +if __name__ == "__main__": + args = parse_args() + asyncio.run(main(args)) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b922417598..c85c78f86d 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -3,6 +3,7 @@ import socket from argparse import Namespace from contextlib import nullcontext +from pathlib import Path from typing import TYPE_CHECKING import ray @@ -21,6 +22,7 @@ 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 @@ -235,6 +237,15 @@ def init( is_lora=is_lora_enabled(args), ) + # Adapters currently loaded into Megatron slots on this rank. + self.loaded_adapters: dict[str, object] = {} + # Adapters whose engine-side weights are stale: newly loaded at + # reconcile, or stepped by the last train call. Consumed (pushed) by + # the next update_weights. Derivable identically on every rank because + # the loaded set only mutates at reconcile and the stepped set ships + # with the train data. + self._multi_lora_pending_push: set[str] = set() + # empty cache after initialization clear_memory() @@ -513,12 +524,77 @@ 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): + # Stepped adapters have new weights: schedule their engine push. + # The stepped set ships with the train data, identical on all ranks. + self._multi_lora_pending_push.update(rollout_data.get("step_adapter_names", [])) + + if is_first_replica_megatron_main_rank(): + from miles.ray.multi_lora_controller import get_multi_lora_controller + + ray.get(get_multi_lora_controller().mark_batch_trained.remote(rollout_id)) + log_perf_data(rollout_id, self.args, extra_metrics=self.weight_updater.pop_metrics()) self._heartbeat.bump() return train_step_outcome @with_logs + @timer + def reconcile_adapters(self) -> None: + """Load what the controller wants served, tear down what it retired. + The snapshot is read once on the main rank and broadcast. + + Deregistered (RETIRING) adapters are retired here, at the sync point + before the next generate — their untrained tail (buffered groups and + any partially accumulated gradients) is discarded. + """ + 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) + 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) + 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() @@ -534,7 +610,35 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: maybe_finalize_async_save(blocking=True) - save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler) + if is_multi_lora_enabled(self.args): + from miles.backends.megatron_utils.multi_lora_utils import save_multi_lora_checkpoints + from miles.ray.multi_lora_controller import get_multi_lora_controller + + # Rank 0 picks adapters at a save-interval multiple without a ckpt + # on disk, and broadcasts so the collective export lines up. + due_buffer = [None] + if is_first_replica_megatron_main_rank() and self.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 % self.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: + if self.args.offload_train: + destroy_process_groups() + return + adapter_steps = {name: adapter.step for name, adapter in due_adapters.items()} + save_multi_lora_checkpoints(self.args, self.model, adapter_steps, due_adapters) + else: + 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) @@ -549,6 +653,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: maybe_finalize_async_save(blocking=True) from megatron.training.checkpointing import get_checkpoint_name + from miles.utils.misc import load_function checkpoint_dir = get_checkpoint_name(self.args.save, rollout_id, return_base_dir=True) @@ -598,11 +703,30 @@ def update_weights(self, info: "EnginesAndLock") -> None: destroy_process_groups() return + version_update_names: list[str] = [] + if is_multi_lora_enabled(self.args): + # Push what's stale: newly loaded + stepped since the last push + # (tracked identically on every rank, so per-adapter TP collectives + # line up). New engines need every loaded adapter. + pending = self._multi_lora_pending_push & set(self.loaded_adapters) + push_names = set(self.loaded_adapters) if has_new_engines else pending + self.weight_updater.multi_lora_adapters = { + name: self.loaded_adapters[name] for name in sorted(push_names) + } + version_update_names = sorted(push_names) + 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): + self._multi_lora_pending_push.clear() + if version_update_names and self._is_first_replica_megatron_main_rank: + from miles.ray.multi_lora_controller import get_multi_lora_controller + + ray.get(get_multi_lora_controller().record_weight_update.remote(version_update_names)) + if self.args.ci_test and len(rollout_engines) > 0 and not is_lora_enabled(self.args): engine = random.choice(rollout_engines) engine_version = ray.get(engine.get_weight_version.remote()) diff --git a/miles/backends/megatron_utils/arguments.py b/miles/backends/megatron_utils/arguments.py index 1ed828513d..cab433038d 100644 --- a/miles/backends/megatron_utils/arguments.py +++ b/miles/backends/megatron_utils/arguments.py @@ -13,6 +13,10 @@ def set_default_megatron_args(args): # Muon currently owns its sharding path, and Megatron's distributed optimizer # only supports Adam-family optimizers. args.use_distributed_optimizer = args.optimizer is None or args.optimizer.lower() == "adam" + # Multi-LoRA: per-slot LayerWise optimizers require plain DDP all-reduce + # (whole-parameter sharding + cross-chunk gradient retention idempotency). + if getattr(args, "multi_lora_n_adapters", 0) > 0: + args.use_distributed_optimizer = False # TODO: maybe change this after megatron has good fp8 support args.bf16 = not args.fp16 # placeholders diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py index b0e8d92387..6169d4e111 100644 --- a/miles/backends/megatron_utils/bridge_lora_helpers.py +++ b/miles/backends/megatron_utils/bridge_lora_helpers.py @@ -12,6 +12,7 @@ from megatron.core.utils import get_attr_wrapped_model from miles.utils.hf_config import load_hf_config +from miles.utils.multi_lora import is_multi_lora_enabled from .lora_utils import create_lora_instance, patch_param_grad_buffer_for_colocate_mode_lora @@ -111,7 +112,12 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: provider.dsa_attention_backend = getattr(args, "dsa_attention_backend", "megatron") provider.finalize() - lora = create_lora_instance(args) + if is_multi_lora_enabled(args): + from miles.backends.megatron_utils.multi_lora_utils import create_multi_lora_instance + + lora = create_multi_lora_instance(args) + else: + lora = create_lora_instance(args) def apply_lora_hook(model_chunks): transformed = lora(model_chunks, training=True) @@ -129,6 +135,10 @@ def apply_lora_hook(model_chunks): provider.register_pre_wrap_hook(_make_value_model_hook(hidden_size, provider.sequence_parallel)) use_distributed_optimizer = "muon" not in (args.optimizer or "").lower() + if is_multi_lora_enabled(args): + # Per-slot LayerWise optimizers: plain DDP all-reduce keeps full grads on + # every rank (whole-param sharding + retained-gradient idempotency). + use_distributed_optimizer = False ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=use_distributed_optimizer) ddp_config.finalize() diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 61a49fe395..b9022e2c6c 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -6,6 +6,7 @@ import math from argparse import Namespace from collections.abc import Callable, Sequence +from contextlib import nullcontext from functools import partial from pathlib import Path @@ -31,6 +32,7 @@ from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase from miles.utils.memory_utils import clear_memory +from miles.utils.multi_lora import is_multi_lora_enabled from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor from miles.utils.tracking_utils.structured_log import log_structured @@ -137,7 +139,11 @@ def setup_model_and_optimizer( assert not args.moe_use_upcycling assert args.load is not None or args.pretrained_checkpoint is not None - if is_lora_enabled(args) and role == "actor" and args.megatron_to_hf_mode == "bridge": + # Multi-LoRA and single-LoRA (actor, bridge) both build via the bridge helper, + # which picks the adapter type internally. + if is_multi_lora_enabled(args) or ( + is_lora_enabled(args) and role == "actor" and args.megatron_to_hf_mode == "bridge" + ): model = _setup_lora_model_via_bridge(args) else: model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder) @@ -165,6 +171,10 @@ def setup_model_and_optimizer( use_gloo_process_groups=args.enable_gloo_process_groups, 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 + + optimizer = build_multi_lora_optimizer(args, config, model) else: optimizer = get_megatron_optimizer( config=config, @@ -280,6 +290,7 @@ def forward_step( "response_lengths", "max_seq_lens", "witness_ids", + "adapter_slots", ], args.data_pad_size_multiplier, args.qkv_format, @@ -291,6 +302,11 @@ def forward_step( total_lengths = batch["total_lengths"] response_lengths = batch["response_lengths"] + if "adapter_token_counts" in batch: + from megatron.bridge.peft.multi_lora_layers import set_tokens_per_adapter_slot + + set_tokens_per_adapter_slot(model, batch["adapter_token_counts"]) + output_tensor = model( input_ids=tokens, position_ids=None, @@ -373,6 +389,10 @@ def train_one_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. @@ -390,12 +410,20 @@ def train_one_step( 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 - # Set grad to zero. - for model_chunk in model: - model_chunk.zero_grad_buffer() - if not disable_optimizer: - optimizer.zero_grad() + # 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: + # Set grad to zero. + for model_chunk in model: + model_chunk.zero_grad_buffer() + if not disable_optimizer: + optimizer.zero_grad() if args.custom_megatron_before_train_step_hook_path: from miles.utils.misc import load_function @@ -439,12 +467,18 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "max_seq_lens", "witness_ids", "opd_reverse_kl", + "adapter_slots", ], args.data_pad_size_multiplier, args.qkv_format, allgather_cp=args.allgather_cp, ) + if "adapter_token_counts" in batch: + from megatron.bridge.peft.multi_lora_layers import set_tokens_per_adapter_slot + + set_tokens_per_adapter_slot(model, batch["adapter_token_counts"]) + from miles.utils.replay_base import all_replay_managers old_stages = [m.stage for m in all_replay_managers] @@ -517,7 +551,11 @@ 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 getattr(args, "check_for_nan_in_loss_and_grad", True)): + if ( + (not disable_optimizer) + and (not multi_lora) + 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 @@ -542,18 +580,40 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p dumper_phase_util.finalize(model) if not disable_optimizer and valid_step: - # Update parameters. - update_successful, grad_norm, num_zeros_in_grad = optimizer.step() + if multi_lora: + from miles.backends.megatron_utils.multi_lora_optimizer import step_adapter_slots + + rollout_data = data_iterator[0].rollout_data + # 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, + ) + grad_norm = max(grad_norms_by_slot.values(), default=0.0) + + # Advance the shared LR schedule by the samples actually consumed + # by the optimizer steps that fired (v1: one shared schedule). + stepped_samples = sum(step_batch_sizes.values()) + if stepped_samples: + opt_param_scheduler.step(increment=stepped_samples) + else: + # Update parameters. + update_successful, grad_norm, num_zeros_in_grad = optimizer.step() - # Update learning rate. - assert update_successful - opt_param_scheduler.step(increment=args.global_batch_size) + # Update learning rate. + assert update_successful + opt_param_scheduler.step(increment=args.global_batch_size) - # release grad - for model_chunk in model: - model_chunk.zero_grad_buffer() - if not disable_optimizer: - optimizer.zero_grad() + # release grad (multi-LoRA retains accumulated grads; stepped slots were + # zeroed selectively inside step_adapter_slots) + if not multi_lora: + for model_chunk in model: + model_chunk.zero_grad_buffer() + if not disable_optimizer: + optimizer.zero_grad() log_structured( logger.info, @@ -909,13 +969,25 @@ def initialize_model_and_optimizer( model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role) model[0].role = role clear_memory() - iteration, _ = load_checkpoint( - model, - optimizer, - opt_param_scheduler, - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=False, - ) + + multi_lora = is_multi_lora_enabled(args) + if multi_lora: + # Hide adapter params so the bridge's conversion-task walk doesn't see them + # while loading the base checkpoint. + from megatron.bridge.peft.multi_lora_layers import hide_adapters + + load_ctx = hide_adapters(model) + else: + load_ctx = nullcontext() + + with load_ctx: + iteration, _ = load_checkpoint( + model, + optimizer, + opt_param_scheduler, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=False, + ) check_peak_gpu_memory_after_load(args) clear_memory() diff --git a/miles/backends/megatron_utils/multi_lora_optimizer.py b/miles/backends/megatron_utils/multi_lora_optimizer.py new file mode 100644 index 0000000000..905d3ddd7c --- /dev/null +++ b/miles/backends/megatron_utils/multi_lora_optimizer.py @@ -0,0 +1,248 @@ +"""Per-slot decoupled optimizers for multi-LoRA training. + +One base Adam per adapter slot, assembled under Megatron's +``LayerWiseDistributedOptimizer`` (whole-parameter ZeRO-1): + +- each slot is an independent chained child with its own Adam state, step + count, and gradient clipping; +- optimizer state is sharded across DP ranks at whole-parameter granularity; +- gradients flow through plain DDP all-reduce (``use_distributed_optimizer`` + must be OFF), which is also what makes cross-batch gradient retention + idempotent: the retained, already-reduced portion of the buffer is identical + on every rank, so re-reducing it is a no-op. + +Selective stepping (``step_adapter_slots``) steps only the slots whose +adapter batch completed, zeroes only their gradients, and leaves every +other slot's accumulated gradients and optimizer state untouched. +""" + +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 Adam per adapter slot under a LayerWiseDistributedOptimizer. + + The returned optimizer is a ``ChainedOptimizer`` whose ``chained_optimizers`` + hold one Float16-wrapped Adam per slot, in slot order; each child's param + groups are tagged with ``miles_multi_lora_slot``. Param groups are narrowed + to this rank's whole-parameter shard by LayerWise, so Adam state exists only + for owned params. + """ + 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)" + + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + + # Delay master-weight creation into LayerWise (post-sharding), so fp32 + # masters exist only for owned params. LayerWise unwraps FP32Optimizer + # children and re-wraps with Float16OptimizerWithFloat16Params itself. + 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.enable_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 + ) + + # LayerWise aggregates grad stats globally (params are scattered across DP + # ranks at whole-parameter granularity), so per-child norm/clip reductions + # must span the world too, not just the model-parallel group. + 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: + """Between-batch replacement for ``DistributedDataParallel.zero_grad_buffer``. + + Resets the per-iteration bookkeeping (``grad_added_to_main_grad`` flags and + bucket-group sync metadata) WITHOUT zeroing the grad buffers, so per-adapter + gradient accumulation survives across train batches. Slot gradients are + zeroed selectively at step time instead. + """ + 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 (every rank holds the full buffer under plain DDP) 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 given slots; retain everyone else's gradients. + + ``step_batch_sizes`` maps slot -> adapter_global_batch_size. Samples enter + the buffers with weight 1, so the accumulated gradient is a sum over the + adapter batch; scaling the fresh master-grad copy by 1/batch_size at step + time yields the adapter-batch mean (the constant is known in advance, and + the copy is remade by ``prepare_grads`` each step, so the scale applies + exactly once). + + Per slot: copy accumulated ``main_grad`` into the owned masters' ``grad``, + scale by 1/batch_size, clip per slot, run Adam on the owned shard, copy + masters back to model params, and zero the slot's gradient state. One param + all-gather at the end propagates updated weights (unchanged slots gather + identical bytes). + + Returns the grad norm per stepped slot. Non-finite gradients get the same + treatment as baseline bf16 training: they poison that adapter's weights + (visible as a NaN grad norm in the logs) — co-tenants are unaffected since + slots' params, gradients, and clipping are disjoint. + """ + grad_norms: dict[int, float] = {} + + for slot, batch_size in step_batch_sizes.items(): + children = _slot_children(optimizer, slot) + # Copy model main_grads -> owned masters' grads (bf16 has no loss + # scaler, so prepare_grads performs no found-inf handling here), then + # normalize the accumulated sum into 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 norm over the union of the slot's children, reduced across + # the world (params are scattered whole across DP ranks; every param is + # counted exactly once on its owner rank). + 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_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py new file mode 100644 index 0000000000..ce8f28dc17 --- /dev/null +++ b/miles/backends/megatron_utils/multi_lora_utils.py @@ -0,0 +1,340 @@ +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 + +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 + + 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 all_megatron_checkpoints_exist(step_dir: Path, tp_size, pp_size) -> bool: + return all( + (step_dir / f"adapter_megatron_tp{tp}_pp{pp}.pt").exists() for tp in range(tp_size) for pp in range(pp_size) + ) + + +def find_latest_checkpoint(ckpt_dir: Path) -> tuple[Path | None, int]: + if not ckpt_dir.exists(): + return None, 0 + + parallel_state = get_parallel_state() + tp_size = parallel_state.tp.size + pp_size = parallel_state.pp.size + tp_rank = parallel_state.tp.rank + pp_rank = parallel_state.pp.rank + + 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, tp_size, pp_size): + return step_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt", 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 + 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: + if "lora_A" in hf_name and adapter_rank < tensor.shape[0]: + remainder = tensor[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[:adapter_rank] + if "lora_B" in hf_name and adapter_rank < tensor.shape[1]: + remainder = tensor[:, 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[:, :adapter_rank] + 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}.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 + is_dp_rank_0 = parallel_state.intra_dp.rank == 0 + is_global_writer = is_dp_rank_0 and tp_rank == 0 and pp_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_dp_rank_0: + 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_dp_rank_0: + 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 / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt" + 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, + ): + # Safetensors format can't save aliased tensors, so need clone() + hf_state[hf_name] = weight.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 + + zero_optimizer_state_for_adapter(optimizer, model, slot) + zero_adapter_slot_grads(model, 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 + resume_steps: dict[str, int] = {} + for adapter in adapters: + resume_steps[adapter.name] = _register_adapter(adapter, model) + 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) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 651c99e487..7e782803a4 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py @@ -122,7 +122,14 @@ def _update_weight_implementation( if pbar: pbar.update(1) - def _update_lora_weight_implementation(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None: + def _update_lora_weight_implementation( + self, + named_tensors: list[tuple[str, torch.Tensor]], + *, + lora_name: str = LORA_ADAPTER_NAME, + lora_config: dict | None = None, + upsert: bool = False, + ) -> None: """Send adapter metadata over Ray, then broadcast the tensors (src=0). Reuses the base broadcast group (``self._model_update_groups`` / @@ -130,24 +137,40 @@ def _update_lora_weight_implementation(self, named_tensors: list[tuple[str, torc sharing the NCCL communicator is safe. No CUDA IPC, so it works across nodes: the engine allocates buffers from the metadata and broadcast-receives in order. + + ``lora_name`` / ``lora_config`` default to the single-adapter values; the + multi-LoRA path passes the per-adapter name and config (carrying that + adapter's own ``r`` / ``lora_alpha``). ``upsert`` switches the + engine RPC to an in-place weight overwrite of an already-loaded adapter + (no unload/register); this is the update path for the fixed multi-LoRA + pool, where every adapter is loaded once and then refreshed in place. """ + if lora_config is None: + lora_config = self._lora_config names = [name for name, _ in named_tensors] dtypes = [param.dtype for _, param in named_tensors] shapes = [list(param.shape) for _, param in named_tensors] + extra_kwargs = {"upsert": True} if upsert else {} refs = [ engine.load_lora_adapter_from_distributed.remote( - lora_name=LORA_ADAPTER_NAME, - config_dict=self._lora_config, + lora_name=lora_name, + config_dict=lora_config, names=names, dtypes=dtypes, shapes=shapes, group_name=self._group_name, + **extra_kwargs, ) for engine in self.rollout_engines ] + # NCCL broadcast requires contiguous buffers, but slice_lora_to_rank yields + # a strided (non-contiguous) view for lora_B (column slice). Materialize + # contiguous copies (no-op when already contiguous) and hold the list so + # the buffers stay alive until the async broadcasts complete. + broadcast_tensors = [param.data.contiguous() for _, param in named_tensors] handles = [ - dist.broadcast(param.data, 0, group=self._model_update_groups, async_op=True) for _, param in named_tensors + dist.broadcast(tensor, 0, group=self._model_update_groups, async_op=True) for tensor in broadcast_tensors ] for handle in handles: handle.wait() diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py index f973808fc4..a5f8961e6b 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py @@ -50,10 +50,13 @@ class DistBucketedWeightUpdateMixin: self._update_weight_implementation(converted_named_tensors, pbar) -> None Transfer a bucket of HF-format ``(name, tensor)`` pairs to rollout engines (via NCCL broadcast, p2p write, etc.). - self._update_lora_weight_implementation(named_tensors) -> None - Transfer the full LoRA adapter (HF-format ``(name, tensor)`` pairs) to + self._update_lora_weight_implementation(named_tensors, *, lora_name, lora_config) -> None + Transfer one LoRA adapter (HF-format ``(name, tensor)`` pairs) to rollout engines. Only required when ``is_lora``; the - unload-before-reload is handled by ``_update_lora_weights``. + unload-before-reload is handled by ``_update_lora_weights`` / + ``_send_one_multi_lora_adapter``. ``lora_name`` / ``lora_config`` + default to the single-adapter values and are overridden per adapter + in the multi-LoRA path. """ def _init_lora( @@ -67,22 +70,20 @@ def _init_lora( ) -> None: """Initialize LoRA-specific state. Call from subclass ``__init__``.""" self.is_lora = is_lora + # Set by the actor before each update_weights call (loaded map at reconcile). + self.multi_lora_adapters = None if self.is_lora: - # Distributed LoRA sync requires the bridge iterator 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})." ) - # The bridge exports adapters per local (PP-stage) model, so a single - # source rank holds the complete adapter only at PP=1. With PP>1 each - # stage would broadcast a partial adapter, so reject it explicitly. + # With PP>1 no single rank holds the complete adapter. assert args.pipeline_model_parallel_size == 1, ( "LoRA weight sync over distributed engines requires " f"--pipeline-model-parallel-size 1 (got {args.pipeline_model_parallel_size})." ) self._lora_config = build_lora_sync_config(args) self._lora_loaded = False - self._lora_base_synced = False self._hf_weight_iterator = HfWeightIteratorBase.create( args=args, model=model, @@ -234,8 +235,6 @@ def _update_lora_weights(self) -> None: but only the source rank transmits. """ # All ranks must iterate the bridge for TP collective participation. - # {} weights: bridge exports adapters directly from self.model and ignores - # this dict (bridge-only is enforced in _init_lora). accumulated_named_tensors: list[tuple[str, torch.Tensor]] = [] for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks({}, weight_type="lora"): accumulated_named_tensors.extend(hf_named_tensors) @@ -263,6 +262,55 @@ def _update_lora_weights(self) -> None: self._update_lora_weight_implementation(accumulated_named_tensors) self._lora_loaded = True + def _update_multi_lora_weights(self) -> None: + """Push the selected adapters (upsert, never unload). The push set is + chosen by the actor before each call and is identical on every rank, so + per-adapter TP collectives line up. Version bumps are recorded by the + actor after the push succeeds.""" + adapters = self.multi_lora_adapters + assert adapters is not None, "actor must set multi_lora_adapters before update_weights" + for name in sorted(adapters): + self._send_one_multi_lora_adapter(adapters[name], upsert=True) + + def _send_one_multi_lora_adapter(self, adapter, upsert: bool) -> None: + """All ranks iterate the bridge (TP collectives); only the source + rank transmits.""" + from megatron.bridge.peft.multi_lora_layers import expose_adapter_slot + + from ...multi_lora_utils import slice_lora_to_rank + + config = adapter.config + adapter_rank = config.rank + lora_config = build_lora_sync_config(self.args) + lora_config["r"] = adapter_rank + lora_config["lora_alpha"] = config.alpha + + accumulated_named_tensors: list[tuple[str, torch.Tensor]] = [] + with expose_adapter_slot(self.model, adapter.slot): + for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks({}, weight_type="lora"): + accumulated_named_tensors.extend( + (n, slice_lora_to_rank(n, t, adapter_rank)) for n, t in hf_named_tensors if is_lora_weight_name(n) + ) + + if not self._is_lora_source: + return + + if not accumulated_named_tensors: + raise RuntimeError( + f"Multi-LoRA weight sync failed for adapter {adapter.name!r}: the weight iterator " + "produced no LoRA weights (no lora_A/lora_B names found). This usually means the " + "Megatron-Bridge or SGLang version is incompatible." + ) + + from miles.utils.multi_lora import slot_lora_name + + self._update_lora_weight_implementation( + accumulated_named_tensors, + lora_name=slot_lora_name(adapter.slot), + lora_config=lora_config, + upsert=upsert, + ) + def _pause_and_prepare_engines(self) -> None: """Pause rollout engines, flush cache, and open the weight-update session.""" if dist.get_rank() == 0: @@ -300,8 +348,9 @@ def update_weights(self) -> None: generation. Full: pause → base non-expert (TP) → base expert (EP) → resume. - LoRA: pause → base weights (first iteration only) → LoRA adapter - (every iteration) → resume. + LoRA: pause → LoRA adapter (every iteration) → resume. The frozen base is + never pushed; the remote rollout engines already load it from + ``hf_checkpoint`` at init. """ self.weight_version += 1 @@ -309,10 +358,20 @@ def update_weights(self) -> None: dist.barrier(group=get_gloo_group()) with timer("update_weights_implementation"): - # Base weight sync model: - # full-param RL: base weights change every step -> always sync. - # LoRA RL: base is frozen -> only sync once, on the first iteration. - if not (self.is_lora and self._lora_base_synced): + from miles.utils.multi_lora import is_multi_lora_enabled + + is_lora = getattr(self, "is_lora", False) + is_multi_lora = is_lora and is_multi_lora_enabled(self.args) + + # Base weight sync: + # Full-param RL: base weights change every step -> always sync. + # LoRA RL: the base is frozen and the (remote) rollout engines already + # hold it (they load ``hf_checkpoint`` at init), so there is nothing + # new to push. Skip it -- this mirrors UpdateWeightFromTensor's + # ``skip_base_sync`` for the distributed case, and avoids routing frozen + # base weights (e.g. a VLM vision tower, unsupported by the direct + # convert_to_hf) through the base path. + if not is_lora: pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_source else None self._gather_and_update_non_expert_weights(self._update_weight_implementation, pbar) @@ -320,12 +379,13 @@ def update_weights(self) -> None: self._gather_and_update_expert_weights(self._update_weight_implementation, pbar) dist.barrier(group=get_gloo_group()) - # LoRA adapter weights: every iteration. - if self.is_lora: - self._update_lora_weights() + # Adapter weights: every iteration. + if is_lora: + if is_multi_lora: + self._update_multi_lora_weights() + else: + self._update_lora_weights() dist.barrier(group=get_gloo_group()) - if not self._lora_base_synced: - self._lora_base_synced = True with timer("finalize_and_resume_engines"): self._finalize_and_resume_engines() diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 51e44d5e61..f418d35508 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -18,6 +18,7 @@ from miles.utils.env_report import collect_and_print_node_env_report from miles.utils.http_utils import get_host_info from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled +from miles.utils.multi_lora import is_multi_lora_enabled logger = logging.getLogger(__name__) @@ -341,14 +342,22 @@ def load_lora_adapter_from_tensors( load_format: str | None = None, pinned: bool = False, added_tokens_config: dict | None = None, + upsert: bool = False, ): - """Load a LoRA adapter. ``serialized_named_tensors[tp_rank]`` is bytes for TP rank N.""" + """Load a LoRA adapter. ``serialized_named_tensors[tp_rank]`` is bytes for TP rank N. + + When ``upsert`` is set, the adapter named ``lora_name`` must already be + loaded and is overwritten in place (no unload/register); used by the + multi-LoRA in-place update path. + """ payload = { "lora_name": lora_name, "config_dict": config_dict, "serialized_named_tensors": serialized_named_tensors, "pinned": pinned, } + if upsert: + payload["upsert"] = True if load_format is not None: payload["load_format"] = load_format if added_tokens_config is not None: @@ -369,6 +378,7 @@ def load_lora_adapter_from_distributed( group_name: str, pinned: bool = False, added_tokens_config: dict | None = None, + upsert: bool = False, ): """Load a LoRA adapter whose weights are broadcast over ``group_name``. @@ -376,6 +386,11 @@ def load_lora_adapter_from_distributed( the tensors arrive via NCCL broadcast (src=0), so no CUDA IPC is used and this works across nodes. ``init_weights_update_group`` must have created ``group_name`` already. + + When ``upsert`` is set, the adapter named ``lora_name`` must + already be loaded on the engines; its weights are overwritten in place + (no unload, no register, no wait_for_unload). This is the in-place update + path for the fixed multi-LoRA pool. """ payload = { "lora_name": lora_name, @@ -385,6 +400,7 @@ def load_lora_adapter_from_distributed( "shapes": shapes, "group_name": group_name, "pinned": pinned, + "upsert": upsert, } if added_tokens_config is not None: payload["added_tokens_config"] = added_tokens_config @@ -707,7 +723,12 @@ def _compute_server_args( kwargs["engine_info_bootstrap_port"] = engine_info_bootstrap_port external_engine_need_check_fields = [k for k in kwargs.keys() if k not in _EXTERNAL_ENGINE_SKIP_CHECK_FIELDS] - if is_lora_enabled(args): + if is_multi_lora_enabled(args): + kwargs["enable_lora"] = True + kwargs["max_loras_per_batch"] = args.multi_lora_n_adapters + kwargs["max_lora_rank"] = max(getattr(args, "lora_rank", 0), 1) + kwargs["lora_target_modules"] = convert_target_modules_to_hf(args.target_modules) + elif is_lora_enabled(args): kwargs["enable_lora"] = True kwargs["max_loras_per_batch"] = 1 kwargs["max_lora_rank"] = max(getattr(args, "lora_rank", 0), 1) diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index b5dccb8061..ae78bf6ef4 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -128,7 +128,7 @@ def get_batch( Steps: - Fetch raw fields via iterator. - Save original token tensors under "unconcat_tokens". - - Slice tokens into two chunks for Context Parallelism (CP), concatenate, and pad to a configurable multiple. + - Slice tokens into two batches for Context Parallelism (CP), concatenate, and pad to a configurable multiple. - Build cu_seqlens and `PackedSeqParams` with T-H-D layout (T: sequence length, H: attention heads, D: head dimension). Args: @@ -167,7 +167,9 @@ def get_batch( if qkv_format == "bshd": max_seqlen = batch["max_seq_lens"][0] assert max([t.size(0) for t in tokens]) <= max_seqlen + if allgather_cp: + assert batch.get("adapter_slots") is None, "allgather CP is currently not supported with multi-LoRA: " assert max_seqlen % cp_size == 0, f"max_seqlen {max_seqlen} not divisible by cp_size {cp_size}" local_len = max_seqlen // cp_size start = parallel_state.cp.rank * local_len @@ -176,21 +178,23 @@ def get_batch( ] else: tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] + sample_token_lengths = [t.size(0) for t in tokens] tokens = torch.stack(tokens) elif qkv_format == "thd": cp_rank = parallel_state.cp.rank if allgather_cp: + assert batch.get("adapter_slots") is None, "allgather CP is currently not supported with multi-LoRA: " # DSA mode: concatenate all sequences first, then slice once with CP. - # We also pad the *global* concatenated stream to make per-rank chunks equal. + # We also pad the *global* concatenated stream to make per-rank batches equal. cu_seqlens_list: list[int] = [0] for t in tokens: cu_seqlens_list.append(cu_seqlens_list[-1] + t.size(0)) tokens = torch.cat(tokens, dim=0) - # Pad global stream so (1) divisible by cp_size (equal chunks), + # Pad global stream so (1) divisible by cp_size (equal batches), # (2) divisible by pad_size (reduce fragmentation). global_pad_size = cp_size * pad_size pad = (global_pad_size - tokens.size(0) % global_pad_size) % global_pad_size @@ -202,6 +206,7 @@ def get_batch( tokens = tokens.chunk(cp_size, dim=0)[cp_rank] else: tokens = [slice_with_cp(t, pad_token_id, qkv_format) for t in tokens] + sample_token_lengths = [t.size(0) for t in tokens] cu_seqlens = [0] for t in tokens: @@ -227,6 +232,21 @@ def get_batch( else: raise ValueError(f"Unsupported qkv_format: {qkv_format}") + # Multi-LoRA: compute per-adapter token counts from post-CP per-sample lengths. + # NOTE: allgather CP is currently not supported + adapter_slots = batch.get("adapter_slots") + if adapter_slots is not None: + assert all( + adapter_slots[i] <= adapter_slots[i + 1] for i in range(len(adapter_slots) - 1) + ), f"adapter_slots not sorted in micro-batch: {adapter_slots}" + n_adapters = data_iterator.rollout_data["n_adapters"] + total_tokens = tokens.numel() + counts = torch.zeros(n_adapters, dtype=torch.int32, device=torch.cuda.current_device()) + for slot, length in zip(adapter_slots, sample_token_lengths, strict=True): + counts[slot] += length + counts[adapter_slots[-1]] += total_tokens - counts.sum().item() + batch["adapter_token_counts"] = counts + batch["tokens"] = tokens def _compute_transform_like_token_ids(ids_list: list): @@ -346,7 +366,7 @@ def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]: - If `micro_batch_indices` is provided, selects rows according to the current index list for each requested key. - - Otherwise, slices a contiguous window of size `micro_batch_size` starting + - Otherwise, slices a contiguous adapter batch of size `micro_batch_size` starting at the current offset. Returns a dict mapping each key to a list subset (or None if absent). @@ -426,6 +446,12 @@ def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices= return data_iterator if not args.use_dynamic_batch_size: + if "adapter_slots" in rollout_data and num_local_gbs % args.micro_batch_size != 0: + raise ValueError( + "A multi-LoRA local batch must be divisible by --micro-batch-size; " + f"got local_batch_size={num_local_gbs}, micro_batch_size={args.micro_batch_size}. " + "Use --use-dynamic-batch-size or choose compatible adapter batch shapes." + ) num_microbatches = [num_local_gbs // args.micro_batch_size for _ in range(num_steps_per_rollout)] data_iterator = _generate_data_iterator(rollout_data, args.micro_batch_size) else: @@ -463,6 +489,10 @@ def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices= for j in range(num_mbs): for k in range(len(partitions[j])): partitions[j][k] += start + # Multi-LoRA: microbatches must be contiguous-by-slot for the + # grouped GEMM's per-adapter token-count math. + if "adapter_slots" in rollout_data: + partitions[j].sort(key=lambda index: rollout_data["adapter_slots"][index]) micro_batch_indices.extend(partitions) assert len(set(sum(micro_batch_indices, []))) == num_local_samples diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py index 45b073ce9f..a4aa866ed7 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -137,6 +137,12 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc "witness_ids", "weight_versions", "metadata", + "n_adapters", + "adapter_slots", + "step_slots", + "step_adapter_names", + "step_adapter_batch_sizes", + "prompt_group_sizes", ]: continue # Upload per sample mean for each rollout value diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index 270fcde3e4..cd9828a9e9 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -11,6 +11,7 @@ from miles.backends.training_utils.loss_hub.opd import apply_opd_kl_to_advantages from miles.backends.training_utils.parallel import get_parallel_state from miles.utils.audit_utils.event_logger.logger import get_event_logger, is_event_logger_initialized +from miles.utils.multi_lora import is_multi_lora_enabled from miles.utils.audit_utils.event_logger.models import TrainAdvantageComputationEvent from miles.utils.types import RolloutBatch @@ -154,6 +155,11 @@ def loss_function( # Here we need to divide by cp_size because to cancel the multiply in Megatron. 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 if not args.calculate_per_token_loss: if apply_megatron_loss_scaling: loss_parallel_size = ( diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index ca47ddff6a..741b583113 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -101,6 +101,11 @@ async def update_weights(self, rollout_id: int | None = None): await self._broadcast("update_weights", info=info) + async def reconcile_adapters(self) -> None: + """Multi-LoRA: reconcile loaded adapters with the controller's active set + (load new, cleanup gone). Called by the trainer before generate.""" + await self._broadcast("reconcile_adapters") + async def onload(self): await self._broadcast("wake_up") diff --git a/miles/ray/multi_lora_controller.py b/miles/ray/multi_lora_controller.py new file mode 100644 index 0000000000..b20eb2b75c --- /dev/null +++ b/miles/ray/multi_lora_controller.py @@ -0,0 +1,121 @@ +"""Named Ray actor wrapping the multi-LoRA backend + HTTP server.""" + +import time +from functools import cache +from typing import Any + +import ray + +from miles.utils.adapter_config import AdapterRun +from miles.utils.misc import SingletonMeta, get_current_node_ip, load_function +from miles.utils.multi_lora import MultiLoRABackend, MultiLoRAHTTPServer +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_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, router_url, host) diff --git a/miles/ray/rollout/rollout_data_conversion.py b/miles/ray/rollout/rollout_data_conversion.py index 539730bd43..9b9b06c7ec 100644 --- a/miles/ray/rollout/rollout_data_conversion.py +++ b/miles/ray/rollout/rollout_data_conversion.py @@ -1,6 +1,7 @@ import itertools import logging +from miles.utils.multi_lora import is_multi_lora_enabled logger = logging.getLogger(__name__) @@ -8,6 +9,15 @@ def postprocess_rollout_data(args, data, train_parallel_config): metadata = {} + # 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)) @@ -34,6 +44,16 @@ def postprocess_rollout_data(args, data, train_parallel_config): return data, metadata +def _first_sample(group): + return _first_sample(group[0]) if isinstance(group[0], list) else group[0] + + +def _nested_sample_count(group) -> int: + if not isinstance(group, list): + return 1 + return sum(_nested_sample_count(item) for item in group) + + def _compute_dynamic_global_batch_size(args, train_parallel_config, num_samples: int) -> int: """Calculate dynamic global_batch_size to ensure only one training step. @@ -43,6 +63,17 @@ 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/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index ecaa9f41e5..2907fa4157 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -98,7 +98,12 @@ def __init__(self, args, pg): # -------------------------- lifecycle ----------------------------- # TODO: may have a `async def init` here later + def get_router_address(self) -> tuple[str, int]: + return self.args.sglang_router_ip, self.args.sglang_router_port + def dispose(self): + if (close := getattr(self.data_source, "close", None)) is not None: + close() event_analyzer.run_analysis_from_args(self.args) if self._metric_checker is not None: self._metric_checker.dispose() diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index cf80d7890a..d09673cfeb 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -22,7 +22,10 @@ def convert_samples_to_train_data( return f(args, samples) raw_rewards, rewards = _post_process_rewards( - args, samples, custom_reward_post_process_func=custom_reward_post_process_func + args, + samples, + custom_reward_post_process_func=custom_reward_post_process_func, + prompt_group_sizes=metadata.get("prompt_group_sizes"), ) assert len(raw_rewards) == len(samples) @@ -85,6 +88,27 @@ 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] + # Adapters whose adapter batch completes with this batch (the + # collection loop's step decision, lifted from sample metadata in postprocess), + # with their adapter batch sizes: the trainer scales each stepping slot's + # accumulated gradient by 1/adapter batch and advances the LR schedule by the + # summed adapter batches. + 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 + if samples[0].opd_reverse_kl is not None: train_data["opd_reverse_kl"] = [sample.opd_reverse_kl for sample in samples] @@ -96,7 +120,12 @@ def convert_samples_to_train_data( return train_data -def _post_process_rewards(args, samples: list[Sample] | list[list[Sample]], custom_reward_post_process_func): +def _post_process_rewards( + args, + samples: list[Sample] | list[list[Sample]], + custom_reward_post_process_func, + prompt_group_sizes: list[int] | None = None, +): if (f := custom_reward_post_process_func) is not None: return f(args, samples) @@ -104,6 +133,23 @@ def _post_process_rewards(args, samples: list[Sample] | list[list[Sample]], cust if args.advantage_estimator in ["grpo", "gspo", "reinforce_plus_plus_baseline"] and args.rewards_normalization: # group norm rewards = torch.tensor(raw_rewards, dtype=torch.float) + if prompt_group_sizes is not None: + # Multi-LoRA: groups may have heterogeneous sizes (per-adapter + # n_samples_per_prompt), so normalize within explicit boundaries. + assert sum(prompt_group_sizes) == len( + raw_rewards + ), f"prompt group sizes sum to {sum(prompt_group_sizes)}, but got {len(raw_rewards)} rewards" + normalized_groups = [] + for group_rewards in rewards.split(prompt_group_sizes): + centered = group_rewards - group_rewards.mean() + if ( + args.advantage_estimator in ["grpo", "gspo"] + and args.grpo_std_normalization + and group_rewards.numel() > 1 + ): + centered = centered / (group_rewards.std() + 1e-6) + normalized_groups.append(centered) + return raw_rewards, torch.cat(normalized_groups).tolist() if rewards.shape[-1] == args.n_samples_per_prompt * args.rollout_batch_size: rewards = rewards.reshape(-1, args.n_samples_per_prompt) else: @@ -137,6 +183,12 @@ def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> l else: partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] + # Multi-LoRA: sort partitions by adapter slot so each microbatch is + # contiguous-by-slot (required by the per-adapter token-count math). + adapter_slots = data.get("adapter_slots") + if adapter_slots is not None: + partitions = [sorted(p, key=lambda i: adapter_slots[i]) for p in partitions] + ans = [] for i in range(dp_size): @@ -160,6 +212,7 @@ def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> l "opd_reverse_kl", "seq_witness_ids", "weight_versions", + "adapter_slots", ]: if key not in data: continue @@ -170,9 +223,15 @@ def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> l "raw_reward", "total_lengths", "dynamic_global_batch_size", + "step_slots", + "step_adapter_names", + "step_adapter_batch_sizes", + "prompt_group_sizes", ]: if key not in data: continue rollout_data[key] = data[key] + if "adapter_slots" in rollout_data: + rollout_data["n_adapters"] = args.multi_lora_n_adapters ans.append(rollout_data) return ans diff --git a/miles/rollout/generate_utils/prefill_logprobs.py b/miles/rollout/generate_utils/prefill_logprobs.py index d11752fd1f..333485c9fa 100644 --- a/miles/rollout/generate_utils/prefill_logprobs.py +++ b/miles/rollout/generate_utils/prefill_logprobs.py @@ -6,10 +6,25 @@ from miles.utils.http_utils import post from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_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: + """The adapter name this sample's scoring request must be served with. + + Multi-LoRA samples carry an adapter ref and must score under their own + slot's adapter — the engines only know ``__miles_slot_{N}`` names there. + Single-adapter LoRA uses the fixed name; base models use none. + """ + if sample.adapter is not None: + return slot_lora_name(sample.adapter.slot) + if is_lora_enabled(args): + return LORA_ADAPTER_NAME + return None + + def _build_prefill_scoring_payload( args: Any, sample: Sample, @@ -37,8 +52,8 @@ def _build_prefill_scoring_payload( "logprob_start_len": prompt_len - 1, } - if is_lora_enabled(args): - payload["lora_path"] = LORA_ADAPTER_NAME + if (lora_path := _lora_path_for_sample(args, sample)) is not None: + payload["lora_path"] = lora_path if sample.multimodal_inputs and sample.multimodal_inputs.get("images"): image_data = sample.multimodal_inputs["images"] @@ -63,14 +78,21 @@ def _build_batch_prefill_scoring_payload( if any(payload["logprob_start_len"] != logprob_start_len for payload in payloads): raise ValueError("Batched SGLang prefill scoring requires a shared logprob_start_len") + lora_paths = {payload.get("lora_path") for payload in payloads} + if len(lora_paths) > 1: + # A batch payload carries a single lora_path; scoring a mixed-adapter + # batch under one adapter would silently corrupt the other samples' + # logprobs. Callers must group by adapter first. + raise ValueError("Batched SGLang prefill scoring requires a shared lora_path") + batch_payload: dict[str, Any] = { "input_ids": [payload["input_ids"] for payload in payloads], "sampling_params": payloads[0]["sampling_params"], "return_logprob": True, "logprob_start_len": logprob_start_len, } - if "lora_path" in payloads[0]: - batch_payload["lora_path"] = payloads[0]["lora_path"] + if (lora_path := next(iter(lora_paths))) is not None: + batch_payload["lora_path"] = lora_path return batch_payload @@ -137,12 +159,15 @@ async def recompute_samples_rollout_logprobs_via_prefill( flush_url = url.rsplit("/", 1)[0] + "/flush_cache" if _can_batch_prefill_score(args, samples_to_score): - samples_by_logprob_start_len: dict[int, list[Sample]] = defaultdict(list) + # A batch shares one logprob_start_len and one lora_path, so group by + # both: mixed-adapter batches (multi-LoRA) must not score every sample + # under whichever adapter came first. + samples_by_batch_key: dict[tuple[int, str | None], list[Sample]] = defaultdict(list) for sample in samples_to_score: prompt_len = len(sample.tokens) - sample.response_length - samples_by_logprob_start_len[prompt_len - 1].append(sample) + samples_by_batch_key[(prompt_len - 1, _lora_path_for_sample(args, sample))].append(sample) - for batch_samples in samples_by_logprob_start_len.values(): + for batch_samples in samples_by_batch_key.values(): # SGLang can serve scoring requests from radix/KV cache. Flush before # each scoring group so every group uses the same clean-prefill path. await post(flush_url, {}) diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 1594dac28e..9766963d09 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -142,6 +142,8 @@ 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"), session_id=_merge_equal_value("session_id"), non_generation_time=_merge_equal_value("non_generation_time"), spec_info=_merge_spec_info(a.spec_info, b.spec_info), diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py index 962c421755..4bc1c95b54 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -5,6 +5,7 @@ import aiohttp from miles.utils.misc 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 @@ -28,15 +29,24 @@ async def remote_rm(args, sample: Sample): return await resp.json() +def _resolve_reward_config(args, sample: Sample) -> tuple[str | None, str]: + if sample.reward_spec is not None: + return sample.reward_spec.custom_rm_path, (sample.reward_spec.rm_type or "").strip() + metadata = sample.metadata if isinstance(sample.metadata, dict) else {} + rm_type = (metadata.get("rm_type") or getattr(args, "rm_type", None) or "").strip() + return getattr(args, "custom_rm_path", None), rm_type + + async def async_rm(args, sample: Sample, **kwargs): - if args.custom_rm_path is not None: - rm_function = load_function(args.custom_rm_path) + custom_rm_path, rm_type = _resolve_reward_config(args, sample) + + if custom_rm_path is not None: + rm_function = load_function(custom_rm_path) return await rm_function(args, sample, **kwargs) - metadata = sample.metadata if isinstance(sample.metadata, dict) else {} - rm_type = (metadata.get("rm_type") or args.rm_type or "").strip() response = sample.response label = sample.label + metadata = sample.metadata if isinstance(sample.metadata, dict) else {} if rm_type.startswith("boxed_"): response = extract_boxed_answer(response) or "" rm_type = rm_type[len("boxed_") :] @@ -88,8 +98,7 @@ async def batched_async_rm( sample.reward = reward return None - if args.custom_rm_path is not None: - # Ensure the custom reward function is implemented in batch mode + if args.custom_rm_path is not None and not is_multi_lora_enabled(args): 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 e9bf9aa6a7..cff83be153 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -23,6 +23,7 @@ from miles.utils.http_utils import get, post from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled from miles.utils.misc import SingletonMeta, call_agent_abort_hook, load_function +from miles.utils.multi_lora import make_rid, slot_lora_name from miles.utils.processing_utils import ( call_processor, encode_image_for_rollout_engine, @@ -172,7 +173,24 @@ 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 is_lora_enabled(args): + 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 no longer sampleable (deregistered / cleaned up). Don't + # POST a request the retire-time abort round can no longer see: + # once the slot is reused, such an orphan would keep decoding under + # the next tenant's weights and silently 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 is_lora_enabled(args): payload["lora_path"] = LORA_ADAPTER_NAME if args.use_rollout_routing_replay: diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py new file mode 100644 index 0000000000..203fe82619 --- /dev/null +++ b/miles/utils/adapter_config.py @@ -0,0 +1,90 @@ +"""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 + + +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 858b5c0efd..2614661551 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -537,7 +537,7 @@ def add_rollout_arguments(parser): default=512 * 1024**2, help=( "buffer size for update weight, in bytes. " - "This is used for updating weights by chunk and should be useful for MoE models." + "This is used for updating weights by batch and should be useful for MoE models." ), ) parser.add_argument( @@ -1390,6 +1390,75 @@ def add_lora_arguments(parser): "down lora_B shared across experts, expert_dim=1). Matches SGLang " "PR #21466's experts_shared_outer_loras=True serving contract.", ) + parser.add_argument( + "--multi-lora-n-adapters", + type=int, + 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)." + ), + ) return parser def add_router_arguments(parser): @@ -2502,6 +2571,77 @@ def miles_validate_args(args): "shared-outer" if args.experts_shared_outer_loras else "per-expert", ) + # Multi-LoRA flag — adapter configs are loaded later by the controller + args.multi_lora = getattr(args, "multi_lora_n_adapters", 0) > 0 + if args.multi_lora: + 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" + 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 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" + ) + assert not args.offload_train, ( + "Multi-LoRA retains per-adapter gradient accumulation in GPU buffers between " + "train calls; --offload-train would destroy it. Disable offload for multi-LoRA." + ) + assert not getattr(args, "enable_witness", False), ( + "Multi-LoRA runs without the distributed optimizer (per-slot LayerWise " + "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." + ) + 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" + if args.multi_lora_max_adapter_global_batch_size is None and hasattr(args, "global_batch_size"): + 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" + + # Effective data-parallel size of the trainer; adapter batch shapes are + # validated against it at registration (min_groups_per_dp_split). Guarded for + # harnesses that validate miles args without the megatron arg 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" + ) + args.use_dynamic_global_batch_size = True + args.megatron_to_hf_mode = "bridge" + assert not (args.kl_coef != 0 and args.kl_loss_coef != 0), "Only one of kl_coef and kl_loss_coef can be set" if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]: @@ -2670,7 +2810,9 @@ def miles_validate_args(args): ) args.global_batch_size = global_batch_size - if args.n_samples_per_prompt == 1: + # Multi-LoRA adapters carry their own n_samples_per_prompt; the per-group + # normalization path already skips std for singleton groups. + if args.n_samples_per_prompt == 1 and not args.multi_lora: args.grpo_std_normalization = False logger.info("n_samples_per_prompt is set to 1, grpo_std_normalization will be set to False.") diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py new file mode 100644 index 0000000000..b8bd0bf98c --- /dev/null +++ b/miles/utils/multi_lora.py @@ -0,0 +1,609 @@ +"""Multi-LoRA adapter registry, backend, and control-plane HTTP server.""" + +import asyncio +import logging +import re +import uuid +from dataclasses import asdict, dataclass, replace +from enum import Enum +from pathlib import Path +from typing import Any + +import httpx +import uvicorn +from fastapi import FastAPI, HTTPException, Query, Request +from fastapi.responses import JSONResponse +from pydantic import BaseModel + +from miles.utils.adapter_config import AdapterRun, AdapterRunConfig, parse_adapter_run_yaml + +logger = logging.getLogger(__name__) + +__all__ = [ + "AdapterRegistry", + "AdapterState", + "EmptyBatchTimeoutError", + "MultiLoRABackend", + "MultiLoRAHTTPServer", + "RID_SEPARATOR", + "define_new_adapter_metrics", + "is_multi_lora_enabled", + "make_rid", + "parse_adapter", + "slot_lora_name", +] + + +# Must not appear in adapter names so rid prefix aborts can't cross adapters. +RID_SEPARATOR = "::" + +VALID_ADAPTER_NAME = re.compile(r"^[A-Za-z0-9._-]+$") + + +class EmptyBatchTimeoutError(RuntimeError): + """No trainable groups arrived before empty-wait timeout.""" + + +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 adapters not seen before ({name}/* -> + {name}/step, {name}/perf/* -> rollout/step); already-declared adapters + are skipped internally, so calling this every snapshot is free. + + Glob expansion only reaches one path segment, so {name}/perf/* keys are + out of {name}/*'s reach despite the shared prefix — each key group must + stay exactly one segment under its glob. + + Must run in the the primary tracking writer, whose wandb + definitions are the only ones that reliably persist — and before the + adapter's first metrics, which is guaranteed at snapshot time: an adapter + can't ship step metrics until it has been promoted and trained a full + adapter batch. + """ + # 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") + + +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" + ) + + +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 + + +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]: + """A train call over this batch succeeded: bank each adapter's groups, fire steps. + + This is the only place accumulation/step state advances, so a failed or + retried train call leaves the registry untouched. Returns the adapters + that stepped. + """ + 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, + ) + + 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)), + } + + +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") + + 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. + + The abort in ``retire_adapters`` fires once at the RETIRING->CLEANUP + flip, but requests can survive it: a multi-turn group between turns + submits its next turn only after that round, and a request still inside + the engine's tokenizer adapter batch can be missed by the scheduler-side + matching. Aborting again here — right before the slot becomes reusable — + closes those escapes, so a later tenant of the slot cannot serve a + retired adapter's orphaned requests. + """ + 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 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") + + +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/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index f1f67dc7dc..a3b3651fb0 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -32,11 +32,21 @@ def log(self, metrics: dict[str, Any], step: int | None = None, **kwargs) -> Non @abstractmethod def finish(self) -> None: ... + def define_step_key_metric_group(self, prefix: str, step_key: str) -> None: + """Declare a runtime metric group (``{prefix}/*``) plotted against its + own step key. Only meaningful for backends where the chart axis is + configuration rather than per-call data; the default is a no-op — + tensorboard/mlflow take the step numerically on every log call.""" + return + # Thin adapters for backwards compatibility to keep wandb_utils and tensorboard_utils untouched. class WandbBackend(TrackingBackend): # Delegates to the existing ``wandb_utils`` helpers. + def __init__(self) -> None: + self._defined_step_key_groups: set[tuple[str, str]] = set() + def init(self, args, *, primary: bool = True, **kwargs) -> None: from . import wandb_utils @@ -50,6 +60,18 @@ def log(self, metrics: dict[str, Any], step: int | None = None, **kwargs) -> Non wandb.log(metrics) + def define_step_key_metric_group(self, prefix: str, step_key: str) -> None: + # Runtime analog of the static groups in _init_wandb_common. Call from + # the primary tracking process: definitions asserted by secondary + # shared-mode writers are lost nondeterministically. + if (prefix, step_key) in self._defined_step_key_groups: + return + import wandb + + wandb.define_metric(step_key) + wandb.define_metric(f"{prefix}/*", step_metric=step_key) + self._defined_step_key_groups.add((prefix, step_key)) + def finish(self) -> None: import wandb @@ -148,6 +170,10 @@ def log(self, metrics: dict[str, Any], step: int | None = None, step_key: str | for backend in self._backends: backend.log(metrics, step=step, step_key=step_key) + def define_step_key_metric_group(self, prefix: str, step_key: str) -> None: + for backend in self._backends: + backend.define_step_key_metric_group(prefix, step_key) + def finish(self) -> None: for backend in self._backends: try: diff --git a/miles/utils/tracking_utils/tracking.py b/miles/utils/tracking_utils/tracking.py index 1b9b7740cf..f7a7b310f9 100644 --- a/miles/utils/tracking_utils/tracking.py +++ b/miles/utils/tracking_utils/tracking.py @@ -15,6 +15,18 @@ def init_tracking(args, primary: bool = True, **kwargs): _manager.init(args, primary=primary, **kwargs) +def define_step_key_metric_group(prefix: str, step_key: str) -> None: + """Declare a runtime metric group plotted against its own step key + (e.g. a multi-LoRA adapter's ``{name}/*`` against ``{name}/step``). + + Fans out to the active backends; only wandb does anything (chart axes are + configuration there), and it deduplicates internally. Must be called from + the primary tracking process (the driver): definitions asserted by + secondary shared-mode writers are lost nondeterministically. + """ + _manager.define_step_key_metric_group(prefix, step_key) + + def log(args, metrics, step_key: str): step = metrics.get(step_key) _manager.log(metrics, step=step, step_key=step_key) diff --git a/miles/utils/types.py b/miles/utils/types.py index cd7637d4c4..87b949254f 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -6,6 +6,30 @@ import torch +@dataclass(frozen=True) +class AdapterRef: + """Per-sample handle identifying which LoRA adapter this sample is bound to. + + Set by the multi-LoRA data source and consumed at training (slot routing) + and inference (per-request lora_path) sites. ``None`` means no adapter. + """ + + name: str + slot: int + + +@dataclass(frozen=True) +class RewardSpec: + """Per-sample handle describing how this sample's response is scored. + + Decoupled from ``AdapterRef`` because reward dispatch is a separate concern + from adapter routing — single-adapter or non-LoRA flows can use this too. + """ + + rm_type: str | None = None + custom_rm_path: str | None = None + + @dataclass class Sample: """The sample generated""" @@ -52,6 +76,11 @@ 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) + reward_spec: RewardSpec | None = None + # Session ID for consistent hashing routing (used when router policy is consistent_hashing) # TODO: Its definition needs to merge with the session server's session id in the new rollout function. session_id: str | None = None diff --git a/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py b/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py new file mode 100644 index 0000000000..8e9afd8a87 --- /dev/null +++ b/tests/fast/backends/training_utils/test_get_batch_multi_lora_cp.py @@ -0,0 +1,154 @@ +"""CP=2 correctness for get_batch's multi-LoRA per-adapter token counts. + +get_batch's CP paths previously had zero coverage, which let two silent +regressions ship (a ``chunk``->``batch`` method typo on the allgather path and +a double ``slice_with_cp`` on the bshd path). These tests pin the normal-CP +(zigzag) contract: + +- per-rank ``adapter_token_counts`` equals the post-slice per-sample lengths + summed per slot, with the rank-local stream padding attributed to the last + slot, and sums to the rank-local padded token count; +- zigzag slicing gives every CP rank the same per-sample lengths, so counts + must be identical across ranks; +- unsorted ``adapter_slots`` in a micro-batch is rejected; +- (strict xfail) bshd tokens must match a single zigzag slice — data.py + currently slices bshd tokens twice under CP>1. + +CUDA is stubbed out so the tests run on CPU-only workers. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import miles.backends.training_utils.cp_utils as cp_utils_mod +import miles.backends.training_utils.data as data_mod +from miles.backends.training_utils.cp_utils import slice_with_cp +from miles.backends.training_utils.data import get_batch + + +def _parallel_state(cp_rank: int, cp_size: int, tp_size: int = 1) -> SimpleNamespace: + return SimpleNamespace( + cp=SimpleNamespace(rank=cp_rank, size=cp_size), + tp=SimpleNamespace(size=tp_size), + ) + + +class _FakeIterator: + def __init__(self, batch: dict, n_adapters: int): + self._batch = batch + self.rollout_data = {"n_adapters": n_adapters} + + def get_next(self, keys): + return {key: self._batch[key] for key in keys} + + +KEYS = ["tokens", "loss_masks", "total_lengths", "response_lengths", "adapter_slots"] + +# 5 samples over 3 of 4 slots (sorted), ragged/odd lengths to exercise the +# per-sample zigzag padding (2 * cp_size chunking). +LENGTHS = [7, 13, 5, 9, 4] +SLOTS = [0, 0, 1, 2, 2] +N_ADAPTERS = 4 + + +def _make_batch(max_seqlen: int | None = None) -> dict: + # Token values start at 1 so zigzag pad (value 0) is distinguishable. + tokens = [torch.arange(1, length + 1, dtype=torch.long) for length in LENGTHS] + response_lengths = [max(1, length // 2) for length in LENGTHS] + batch = { + "tokens": tokens, + "loss_masks": [torch.ones(r, dtype=torch.int) for r in response_lengths], + "total_lengths": list(LENGTHS), + "response_lengths": response_lengths, + "adapter_slots": list(SLOTS), + } + if max_seqlen is not None: + batch["max_seq_lens"] = [max_seqlen] * len(LENGTHS) + return batch + + +@pytest.fixture(autouse=True) +def _stub_cuda(monkeypatch): + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, *args, **kwargs: self, raising=False) + monkeypatch.setattr(torch.cuda, "current_device", lambda: "cpu", raising=False) + + +def _patch_state(monkeypatch, state: SimpleNamespace) -> None: + # data.get_batch and cp_utils.slice_with_cp each resolve the state themselves. + monkeypatch.setattr(data_mod, "get_parallel_state", lambda: state) + monkeypatch.setattr(cp_utils_mod, "get_parallel_state", lambda: state) + + +def _expected_thd_counts(state: SimpleNamespace, local_total: int) -> torch.Tensor: + sliced_lengths = [ + slice_with_cp(t, 0, "thd", parallel_state=state).numel() + for t in (torch.arange(1, length + 1, dtype=torch.long) for length in LENGTHS) + ] + expected = torch.zeros(N_ADAPTERS, dtype=torch.int32) + for slot, sliced_length in zip(SLOTS, sliced_lengths, strict=True): + expected[slot] += sliced_length + stream_pad = local_total - int(sum(sliced_lengths)) + assert stream_pad >= 0, "rank-local stream shorter than the sliced samples" + expected[SLOTS[-1]] += stream_pad + return expected + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_thd_cp2_adapter_token_counts(monkeypatch, cp_rank): + state = _parallel_state(cp_rank, cp_size=2) + _patch_state(monkeypatch, state) + + out = get_batch(_FakeIterator(_make_batch(), N_ADAPTERS), KEYS, pad_multiplier=8, qkv_format="thd") + + counts = out["adapter_token_counts"] + local_total = out["tokens"].numel() + assert counts.dtype == torch.int32 + assert counts.tolist() == _expected_thd_counts(state, local_total).tolist() + assert int(counts.sum()) == local_total, "counts must cover every rank-local token incl. padding" + + +def test_thd_cp2_counts_identical_across_ranks(monkeypatch): + # Zigzag gives each rank the same padded share of every sample, so the + # grouped-GEMM routing counts must not depend on the CP rank. + per_rank = [] + for cp_rank in (0, 1): + state = _parallel_state(cp_rank, cp_size=2) + _patch_state(monkeypatch, state) + out = get_batch(_FakeIterator(_make_batch(), N_ADAPTERS), KEYS, pad_multiplier=8, qkv_format="thd") + per_rank.append(out["adapter_token_counts"].tolist()) + assert per_rank[0] == per_rank[1] + + +def test_unsorted_adapter_slots_rejected(monkeypatch): + _patch_state(monkeypatch, _parallel_state(cp_rank=0, cp_size=2)) + batch = _make_batch() + batch["adapter_slots"] = [2, 0, 1, 0, 2] + with pytest.raises(AssertionError, match="not sorted"): + get_batch(_FakeIterator(batch, N_ADAPTERS), KEYS, pad_multiplier=8, qkv_format="thd") + + +def test_bshd_cp2_tokens_are_single_sliced(monkeypatch): + # Regression test: the bshd path used to slice tokens twice under CP>1 + # (pre-slice at the top of the branch, then again in the non-allgather + # arm), silently replacing the back half of each rank's tokens with pad + # zeros while keeping the shape unchanged. + max_seqlen = 16 # divisible by 2 * cp_size + state = _parallel_state(cp_rank=0, cp_size=2) + _patch_state(monkeypatch, state) + + out = get_batch( + _FakeIterator(_make_batch(max_seqlen=max_seqlen), N_ADAPTERS), + KEYS + ["max_seq_lens"], + pad_multiplier=8, + qkv_format="bshd", + ) + + expected = torch.stack( + [ + slice_with_cp(torch.arange(1, length + 1, dtype=torch.long), 0, "bshd", max_seqlen, parallel_state=state) + for length in LENGTHS + ] + ) + assert torch.equal(out["tokens"], expected) diff --git a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py new file mode 100644 index 0000000000..7cd66e2a60 --- /dev/null +++ b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py @@ -0,0 +1,228 @@ +"""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 examples.multi_lora.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, + ) + 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() + 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, +) -> 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 + ) + + +def make_group(adapter: AdapterRun, slot_version: int | 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 + samples.append(sample) + return samples + + +def buffer_groups(worker, adapter: AdapterRun, count: int, slot_version: int | None = None): + for _ in range(count): + worker.buffers[adapter.name].put(make_group(adapter, slot_version)) + + +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)) + + +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)) diff --git a/tests/fast/ray/rollout/test_multi_lora_process_group.py b/tests/fast/ray/rollout/test_multi_lora_process_group.py new file mode 100644 index 0000000000..ff83c7e3a4 --- /dev/null +++ b/tests/fast/ray/rollout/test_multi_lora_process_group.py @@ -0,0 +1,54 @@ +"""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 examples.multi_lora.multi_lora_async_rollout as mod +import pytest +from examples.multi_lora.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) -> None: + self.version = version + + +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 diff --git a/tests/fast/ray/rollout/test_multi_lora_train_data.py b/tests/fast/ray/rollout/test_multi_lora_train_data.py new file mode 100644 index 0000000000..86f259b3b8 --- /dev/null +++ b/tests/fast/ray/rollout/test_multi_lora_train_data.py @@ -0,0 +1,106 @@ +"""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/rollout/generate_utils/test_prefill_logprobs.py b/tests/fast/rollout/generate_utils/test_prefill_logprobs.py index aa736a8194..dcab51e32c 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 Sample +from miles.utils.types import AdapterRef, Sample @pytest.mark.asyncio @@ -113,6 +113,105 @@ 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/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 20a7bfb758..8eef68eb4d 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -172,6 +172,41 @@ def test_custom_megatron_post_save_hook_path_requires_save(): miles_validate_args(args) +class TestMultiLoRAValidation: + def _parse(self, extra): + parser = argparse.ArgumentParser() + get_miles_extra_args_provider()(parser) + return parser.parse_args( + [ + "--multi-lora-n-adapters", + "2", + "--lora-rank", + "8", + "--target-modules", + "linear_qkv", + "--num-rollout", + "1", + ] + + extra + + REQUIRED_ARGS + ) + + def test_rejects_multiple_tokenizer_workers(self): + # Each sglang tokenizer worker holds its own LoRA registry, so per-step + # upserts fail non-deterministically; fail at launch, not first push. + args = self._parse(["--sglang-tokenizer-worker-num", "2"]) + + with pytest.raises(AssertionError, match="sglang-tokenizer-worker-num 1"): + miles_validate_args(args) + + def test_accepts_default_single_tokenizer_worker(self): + args = self._parse([]) + + miles_validate_args(args) + + assert args.multi_lora is True + + class TestResolveFtComponents: def test_disabled_with_no_components_returns_empty_without_warning(self, caplog) -> None: """use_fault_tolerance off and no ft_components yields an empty list and no warning.""" diff --git a/tests/fast/utils/test_controller_backend.py b/tests/fast/utils/test_controller_backend.py new file mode 100644 index 0000000000..f1a9c0f38c --- /dev/null +++ b/tests/fast/utils/test_controller_backend.py @@ -0,0 +1,365 @@ +"""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.utils.adapter_config import AdapterRunConfig +from miles.utils.multi_lora import ( + AdapterRegistry, + AdapterState, + MultiLoRABackend, + min_groups_per_dp_split, + make_rid, + parse_adapter, +) + + +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="/d", + 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="/d")) + 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()) diff --git a/tests/fast/utils/test_controller_http.py b/tests/fast/utils/test_controller_http.py new file mode 100644 index 0000000000..2340efd60e --- /dev/null +++ b/tests/fast/utils/test_controller_http.py @@ -0,0 +1,231 @@ +"""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.utils.adapter_config import AdapterRunConfig +from miles.utils.multi_lora import RID_SEPARATOR, MultiLoRABackend, MultiLoRAHTTPServer + + +def minimal_config(name: str) -> dict: + return {"data": f"/data/{name}.parquet", "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/train.parquet", + "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/train.parquet" + 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/A2.parquet", "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": []}