From def38414ac6514a236cdb7ab25863541207416f3 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 19:50:51 -0700 Subject: [PATCH 01/31] [feat] multi lora fully async --- examples/multi_lora/README.md | 85 ++++ examples/multi_lora/adapters/dapo_math.yaml | 7 + examples/multi_lora/adapters/gsm8k.yaml | 7 + .../multi_lora/multi_lora_async_rollout.py | 254 +++++++++++ .../multi_lora_data_source_async.py | 145 ++++++ examples/multi_lora/provision.sh | 12 + examples/multi_lora/run_job.sh | 95 ++++ examples/multi_lora/run_service.sh | 88 ++++ examples/multi_lora/service_smoke.py | 142 ++++++ examples/multi_lora/tests/__init__.py | 0 .../multi_lora/tests/test_controller_http.py | 216 +++++++++ .../multi_lora/tests/test_controller_logic.py | 209 +++++++++ .../tests/test_multi_lora_async_rollout.py | 106 +++++ examples/multi_lora/train_multi_lora_async.py | 85 ++++ .../megatron_utils/multi_lora_utils.py | 329 ++++++++++++++ miles/ray/multi_lora_controller.py | 121 +++++ miles/utils/adapter_config.py | 71 +++ miles/utils/multi_lora.py | 414 ++++++++++++++++++ 18 files changed, 2386 insertions(+) create mode 100644 examples/multi_lora/README.md create mode 100644 examples/multi_lora/adapters/dapo_math.yaml create mode 100644 examples/multi_lora/adapters/gsm8k.yaml create mode 100644 examples/multi_lora/multi_lora_async_rollout.py create mode 100644 examples/multi_lora/multi_lora_data_source_async.py create mode 100755 examples/multi_lora/provision.sh create mode 100755 examples/multi_lora/run_job.sh create mode 100755 examples/multi_lora/run_service.sh create mode 100644 examples/multi_lora/service_smoke.py create mode 100644 examples/multi_lora/tests/__init__.py create mode 100644 examples/multi_lora/tests/test_controller_http.py create mode 100644 examples/multi_lora/tests/test_controller_logic.py create mode 100644 examples/multi_lora/tests/test_multi_lora_async_rollout.py create mode 100644 examples/multi_lora/train_multi_lora_async.py create mode 100644 miles/backends/megatron_utils/multi_lora_utils.py create mode 100644 miles/ray/multi_lora_controller.py create mode 100644 miles/utils/adapter_config.py create mode 100644 miles/utils/multi_lora.py diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md new file mode 100644 index 0000000000..78145629ba --- /dev/null +++ b/examples/multi_lora/README.md @@ -0,0 +1,85 @@ +# 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, deregisters at num_row) +tests/ # controller logic + HTTP tests (no torch) +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 (no drain, no state machine) + +- **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`). +- **No drain / no rollout-id / no train_steps / no PENDING-DRAINING-DRAINED states.** + The data source deregisters an adapter at `num_row`; the trainer's + `reconcile_adapters` (before each generate) cleans up gone adapters (save ckpt + + clear Megatron slot) and loads new ones. `update_weights` upserts active adapters' + weights in place (SGLang page table, `upsert=True`). +- **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_row` +(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 +data: /root/gsm8k/train.parquet +input_key: messages +label_key: label +rm_type: math +num_row: 400 # stop adapter after N rows +# optional: save, num_epoch, custom_rm_path, ... +``` diff --git a/examples/multi_lora/adapters/dapo_math.yaml b/examples/multi_lora/adapters/dapo_math.yaml new file mode 100644 index 0000000000..dfcabeeb2b --- /dev/null +++ b/examples/multi_lora/adapters/dapo_math.yaml @@ -0,0 +1,7 @@ +rank: 32 +alpha: 32 +data: /root/dapo-math-17k/dapo-math-17k.jsonl +input_key: prompt +label_key: label +rm_type: deepscaler +num_row: 500 diff --git a/examples/multi_lora/adapters/gsm8k.yaml b/examples/multi_lora/adapters/gsm8k.yaml new file mode 100644 index 0000000000..2a734590bf --- /dev/null +++ b/examples/multi_lora/adapters/gsm8k.yaml @@ -0,0 +1,7 @@ +rank: 16 +alpha: 16 +data: /root/gsm8k/train.parquet +input_key: messages +label_key: label +rm_type: math +num_row: 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..18b2a5abb2 --- /dev/null +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -0,0 +1,254 @@ +"""Fully-async multi-LoRA rollout: continuous background producer + collect-a-batch.""" + +import asyncio +import itertools +import logging +import queue +import threading +import time +from collections.abc import Callable +from typing import Any + +from miles.rollout.base_types import RolloutFnTrainOutput +from miles.rollout.filter_hub.base_types import MetricGatherer, 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.misc import load_function + +from miles.ray.multi_lora_controller import AdaptersCache, get_multi_lora_controller + +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] + + +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 AsyncMultiLoRAWorker: + """Background producer: continuously generate groups into a thread-safe queue.""" + + 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.output_queue: queue.Queue = queue.Queue(maxsize=1000) + self.worker_thread: threading.Thread | None = None + self.state = GenerateState(args) + + @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) + + 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 + group = samples[0] + active.add(asyncio.create_task(self.process_and_enqueue(group))) + + await asyncio.sleep(0) + finally: + if active: + await asyncio.wait(active) + + 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 not None: + self.output_queue.put(result) + + def queue_size(self) -> int: + return self.output_queue.qsize() + + +async def generate_rollout_multi_lora_async( + args, rollout_id: int, data_source, generate_fn: GenerateFn = generate_and_rm_group +) -> tuple[RolloutFnTrainOutput, list[list[Sample]]]: + """Fully-async multi-LoRA rollout. Collect a batch from the background worker, + then run the same postprocess as ``generate_rollout_async``.""" + assert args.rollout_global_dataset + + state = GenerateState(args) + + dynamic_filter = ( + load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path else None + ) + metric_gatherer = MetricGatherer() + target_data_size = args.rollout_batch_size + + worker = AsyncMultiLoRAWorker.get_or_create(args, data_source, generate_fn) + + # Groups whose submission-time slot version fell too far behind are dropped. + max_staleness = getattr(args, "max_weight_staleness", None) + + data: list[Group] = [] + stale_dropped = 0 + staleness_values: list[int] = [] + start_time = time.time() + last_progress = start_time + queue_length = worker.queue_size() + while len(data) < target_data_size: + made_progress = False + current_adapters = await AdaptersCache().get_all() + # Pop one at a time so surplus groups stay queued for the next batch. + while len(data) < target_data_size: + try: + group = worker.output_queue.get_nowait() + except queue.Empty: + break + head = first_sample(group) if group else None + adapter_name = head.adapter.name if head is not None and head.adapter else None + if adapter_name not in current_adapters: + continue # adapter deregistered; drop + if max_staleness is not None: + stamped = head.metadata.get("slot_version") + if stamped is not None: + staleness = current_adapters[adapter_name].version - stamped + if staleness > max_staleness: + for s in iter_group_samples(group): + s.reset_for_retry() + stale_dropped += 1 + staleness_values.append(staleness) + logger.info( + f"Dropped stale group (adapter={adapter_name}, " + f"stamped={stamped}, current={current_adapters[adapter_name].version}, " + f"staleness={staleness} > max={max_staleness})" + ) + continue + f = call_dynamic_filter(dynamic_filter, args, group) + if not f.keep: + metric_gatherer.on_dynamic_filter_drop(reason=f.reason) + continue + data.append(group) + made_progress = True + + if made_progress: + last_progress = time.time() + elif time.time() - last_progress > 30: + logger.warning(f"No progress for 30s. queue={worker.queue_size()} collected={len(data)}/{target_data_size}") + last_progress = time.time() + + if len(data) < target_data_size: + await asyncio.sleep(0.01) + + if stale_dropped: + logger.info( + f"Staleness stats: dropped={stale_dropped}, " + f"avg_staleness={sum(staleness_values) / len(staleness_values):.1f}, " + f"max_staleness={max(staleness_values)}" + ) + + data = sorted(data, key=lambda g: first_sample(g).index) + + batch_adapters = sorted( + {first_sample(g).adapter.name for g in data if g and first_sample(g).adapter} + ) + if batch_adapters: + await get_multi_lora_controller().record_batch_adapters.remote(rollout_id, batch_adapters) + + 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, + ) + + metrics = { + **metric_gatherer.collect(), + "perf/fully_async/queue_length": queue_length, + "perf/fully_async/batch_wait_time": time.time() - start_time, + "perf/fully_async/stale_dropped": stale_dropped, + } + if staleness_values: + metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(staleness_values) / len(staleness_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..2be7b78c78 --- /dev/null +++ b/examples/multi_lora/multi_lora_data_source_async.py @@ -0,0 +1,145 @@ +"""Round-robin per-adapter data source; deregisters adapters at num_row.""" + +import copy +import logging +from argparse import Namespace +from collections import deque +from concurrent.futures import ThreadPoolExecutor + +import ray + +from miles.rollout.data_source import DataSource, RolloutDataSource +from miles.utils.adapter_config import AdapterRun +from miles.utils.types import AdapterRef, RewardSpec, Sample + +from miles.ray.multi_lora_controller import get_multi_lora_controller + +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}'") + 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.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) -> list[list[Sample]]: + snapshot = fetch_snapshot() + adapters = sampleable(snapshot) + self.reconcile(adapters) + if not self.sources: + return [] + self.update_queue(set(self.sources)) + + refs = {name: AdapterRef(name=name, slot=adapters[name].slot) for name in self.sources} + reward_specs = { + name: RewardSpec( + rm_type=adapters[name].config.rm_type, + custom_rm_path=adapters[name].config.custom_rm_path, + ) + for name in self.sources + } + + samples_per_adapter, remainder = divmod(num_samples, len(self.source_queue)) + all_samples: list[list[Sample]] = [] + to_deregister: list[str] = [] + + for i in range(len(self.source_queue)): + extra = int(i < remainder) + samples_needed = samples_per_adapter + extra + if samples_needed == 0: + break + name = self.source_queue.popleft() + config = adapters[name].config + self.source_queue.append(name) + source = self.sources[name] + adapter_samples = source.get_samples(samples_needed) + ref = refs[name] + reward_spec = reward_specs[name] + for group in adapter_samples: + for sample in group: + sample.adapter = ref + sample.reward_spec = reward_spec + sample.metadata = {**config.metadata, **sample.metadata} + all_samples.extend(adapter_samples) + + default_num_row = (getattr(config, "num_epoch", 1) or 1) * len(source.dataset) + num_row = config.num_row or default_num_row + if source.sample_group_index >= num_row and name not in snapshot["retiring"]: + logger.info(f"Adapter '{name}' reached num_row={num_row}, deregistering") + to_deregister.append(name) + + for name in to_deregister: + ray.get(get_multi_lora_controller().deregister_adapter.remote(name)) + + return all_samples + + 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) 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..e65a7efc59 --- /dev/null +++ b/examples/multi_lora/run_job.sh @@ -0,0 +1,95 @@ +#!/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 \ + --calculate-per-token-loss \ + --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-multilora \ + # --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..e9a9695e21 --- /dev/null +++ b/examples/multi_lora/run_service.sh @@ -0,0 +1,88 @@ +#!/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 \ + --calculate-per-token-loss \ + --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..bfc9c3f07f --- /dev/null +++ b/examples/multi_lora/service_smoke.py @@ -0,0 +1,142 @@ +"""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 active_adapters(self) -> dict: + response = self.http.get(f"{self.api_url}/adapter_runs") + response.raise_for_status() + return { + status["name"]: {"slot": status["slot"], "version": status["version"], "step": status["step"]} + for status in response.json()["adapters"] + if status["state"] == "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: + self.wait_for( + f"'{name}' to reach step {min_step}", + lambda adapters: 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("--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_a; expect promotion + training progress") + client.register_when_allowed("smoke_a", config("smoke_a")) + client.wait_for_step("smoke_a", args.steps) + + print("phase 3: 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 4: 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 5: 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 6: 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/tests/test_controller_http.py b/examples/multi_lora/tests/test_controller_http.py new file mode 100644 index 0000000000..1e828758fa --- /dev/null +++ b/examples/multi_lora/tests/test_controller_http.py @@ -0,0 +1,216 @@ +"""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 + +import aiohttp +import pytest +from aiohttp import web + +from types import SimpleNamespace + +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), 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": []} diff --git a/examples/multi_lora/tests/test_controller_logic.py b/examples/multi_lora/tests/test_controller_logic.py new file mode 100644 index 0000000000..fecf6234f9 --- /dev/null +++ b/examples/multi_lora/tests/test_controller_logic.py @@ -0,0 +1,209 @@ +"""Fast tests for AdapterRegistry + MultiLoRABackend validation +(no Ray, no HTTP I/O, no SGLang, no torch).""" + +from types import SimpleNamespace + +import pytest + +from miles.utils.adapter_config import AdapterRunConfig +from miles.utils.multi_lora import AdapterRegistry, AdapterState, MultiLoRABackend, make_rid, parse_adapter + + +def make_args(max_adapters: int = 4, save: str | None = None) -> SimpleNamespace: + return SimpleNamespace(multi_lora_n_adapters=max_adapters, save=save) + + +def make_backend(max_adapters: int = 4, save: str | None = None) -> MultiLoRABackend: + return MultiLoRABackend(make_args(max_adapters, save), "http://unused") + + +def make_config(save: str | None = None) -> AdapterRunConfig: + return AdapterRunConfig(rank=8, alpha=16, data="/d", save=save, input_key="text", label_key="label", rm_type="math") + + +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 + + +def test_batch_record_counts_steps_on_confirmation(): + registry = AdapterRegistry(max_adapters=4) + register_and_promote(registry, "A") + register_and_promote(registry, "B") + + registry.record_batch_adapters(7, ["A"]) + assert registry.step_count("A") == 0 # recorded, not yet trained + assert registry.mark_batch_trained(7) == ["A"] + assert registry.step_count("A") == 1 + assert registry.step_count("B") == 0 + assert registry.mark_batch_trained(7) == [] # record consumed + + +def test_batch_trained_counts_deregistered_adapter_until_freed(): + registry = AdapterRegistry(max_adapters=4) + register_and_promote(registry, "A") + registry.record_batch_adapters(3, ["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", None) + registry.set_step("A", 40) + registry.record_batch_adapters(1, ["A"]) + registry.record_weight_update(["A"]) + registry.mark_batch_trained(1) + assert registry.step_count("A") == 41 + + +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_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/examples/multi_lora/tests/test_multi_lora_async_rollout.py b/examples/multi_lora/tests/test_multi_lora_async_rollout.py new file mode 100644 index 0000000000..16d3fd31aa --- /dev/null +++ b/examples/multi_lora/tests/test_multi_lora_async_rollout.py @@ -0,0 +1,106 @@ +"""Tests for the testable core of the multi-LoRA async rollout (process_group): +keep-vs-recycle plus submission-time slot-version stamping.""" + +import pytest + +from miles.utils.types import AdapterRef, Sample + +import examples.multi_lora.multi_lora_async_rollout as mod +from examples.multi_lora.multi_lora_async_rollout import process_group + + +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_all(self) -> dict[str, FakeAdapterView]: + return {name: FakeAdapterView(version) for name, version in self.versions.items()} + + async def get(self, adapter_name: str) -> FakeAdapterView | None: + return (await self.get_all()).get(adapter_name) + + +def group(adapter: str = "A", slot: int = 0) -> list[Sample]: + return [Sample(prompt="p", adapter=AdapterRef(adapter, slot))] + + +async def gen_completed(args, group, sampling_params): + for s in group: + s.status = Sample.Status.COMPLETED + return group + + +@pytest.mark.asyncio +async def test_process_group_keeps_completed(): + ds = FakeDataSource() + g = group("A") + result = await process_group(None, g, {}, gen_completed, ds) + + assert result is g + assert ds.added == [] + + +@pytest.mark.asyncio +async def test_process_group_recycles_aborted(): + async def gen(args, group, sampling_params): + for s in group: + s.status = Sample.Status.ABORTED + return group + + ds = FakeDataSource() + g = group("A") + result = await process_group(None, g, {}, gen, ds) + + assert result is None + assert len(ds.added) == 1 + + +@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 + return await gen_completed(args, group, sampling_params) + + monkeypatch.setattr(mod, "AdaptersCache", lambda: cache) + + ds = FakeDataSource() + g = group("A") + result = await process_group(None, g, {}, gen, ds) + + assert result is g + assert g[0].metadata["slot_version"] == 5 + + +@pytest.mark.asyncio +async def test_process_group_no_adapter_skips_stamp(monkeypatch): + class FailingCache: + async def get(self, adapter_name): + raise AssertionError("adapters cache should not be queried for adapter-less group") + + monkeypatch.setattr(mod, "AdaptersCache", FailingCache) + + ds = FakeDataSource() + g = [Sample(prompt="p", adapter=None)] + result = await process_group(None, g, {}, gen_completed, ds) + + assert result is g + assert "slot_version" not in g[0].metadata 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..3c67ba6ad2 --- /dev/null +++ b/examples/multi_lora/train_multi_lora_async.py @@ -0,0 +1,85 @@ +"""Fully-async multi-LoRA trainer driver.""" + +import asyncio +import logging +from pathlib import Path + +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.audit_utils.process_identity import MainProcessIdentity +from miles.utils.logging_utils import configure_logger +from miles.utils.tracking_utils.tracking import init_tracking + +from miles.ray.multi_lora_controller import create_controller, get_multi_lora_controller + +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" + + +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() + 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. + 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 + + rollout_data = await rollout_manager.generate.remote(rollout_id) + 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/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py new file mode 100644 index 0000000000..9470efe6db --- /dev/null +++ b/miles/backends/megatron_utils/multi_lora_utils.py @@ -0,0 +1,329 @@ +import json +import logging +import os +from argparse import Namespace +from collections.abc import Mapping +from pathlib import Path + +import ray +import torch +import torch.distributed as dist + +from miles.backends.training_utils.parallel import get_parallel_state +from miles.ray.multi_lora_controller import get_multi_lora_controller +from miles.utils.adapter_config import AdapterRun +from miles.utils.multi_lora import is_multi_lora_enabled as is_multi_lora_enabled + +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) + 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_() + + +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 chunk in model + for name, param in chunk.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. + zero_optimizer_state_for_adapter(optimizer, model, slot) + optimizer.reload_model_params() + logger.info(f"{log_prefix} cleared optimizer state 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/ray/multi_lora_controller.py b/miles/ray/multi_lora_controller.py new file mode 100644 index 0000000000..df10b1d338 --- /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() + + def free_slot(self, name: str) -> int: + return self.backend.registry.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, names: list[str]) -> None: + self.backend.registry.record_batch_adapters(rollout_id, names) + + def mark_batch_trained(self, rollout_id: int) -> list[str]: + return self.backend.registry.mark_batch_trained(rollout_id) + + 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/utils/adapter_config.py b/miles/utils/adapter_config.py new file mode 100644 index 0000000000..78e3fb2eda --- /dev/null +++ b/miles/utils/adapter_config.py @@ -0,0 +1,71 @@ +"""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 + + 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 + + num_epoch: int | None = None + num_row: int | None = None + + metadata: dict[str, Any] = field(default_factory=dict) + +@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 + + +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"], + 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_epoch=raw.get("num_epoch"), + num_row=raw.get("num_row"), + metadata=raw.get("metadata") or {}, + ) diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py new file mode 100644 index 0000000000..61ad9c26fe --- /dev/null +++ b/miles/utils/multi_lora.py @@ -0,0 +1,414 @@ +"""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 AdapterRunConfig, AdapterRun, parse_adapter_run_yaml + +logger = logging.getLogger(__name__) + +__all__ = [ + "AdapterRegistry", + "AdapterState", + "MultiLoRABackend", + "MultiLoRAHTTPServer", + "RID_SEPARATOR", + "is_multi_lora_enabled", + "make_rid", + "parse_adapter", +] + + +# 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._-]+$") + + +def is_multi_lora_enabled(args: Any) -> bool: + return getattr(args, "multi_lora", False) + + +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] + + +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 + 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_adapters: dict[int, list[str]] = {} + + 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, names: list[str]) -> None: + self.batch_adapters[rollout_id] = list(names) + while len(self.batch_adapters) > MAX_BATCH_RECORDS: + self.batch_adapters.pop(next(iter(self.batch_adapters))) + + def mark_batch_trained(self, rollout_id: int) -> list[str]: + trained = [] + for name in self.batch_adapters.pop(rollout_id, []): + record = self.records.get(name) + if record is not None and record.state in ( + AdapterState.ACTIVE, + AdapterState.RETIRING, + AdapterState.CLEANUP, + ): + record.step += 1 + trained.append(name) + return trained + + def set_step(self, name: str, step: int) -> None: + if (record := self.find(name)) is not None: + record.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, + ) + + 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_save_dir(self, name: str, config: Any) -> Any: + if config is None or not hasattr(config, "save"): + return config + if config.save is not None: + return config + 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" + ) + return replace(config, save=Path(self.args.save) / "adapters" / name) + + async def register(self, name: str, config: Any) -> dict: + await self.validate_adapter(name, config) + config = self.resolve_save_dir(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 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 + + +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] = Query(default_factory=list)) -> 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} From 2f4b8519f936327edac14e604a12c47666a384bb Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 19:51:00 -0700 Subject: [PATCH 02/31] [feat] multi lora fully async --- miles/backends/megatron_utils/actor.py | 87 +++++++++++++- .../megatron_utils/bridge_lora_helpers.py | 8 +- miles/backends/megatron_utils/model.py | 46 +++++-- .../broadcast.py | 30 ++++- .../update_weight_from_distributed/mixin.py | 113 ++++++++++++++---- .../update_weight_from_tensor.py | 5 +- miles/backends/sglang_utils/sglang_engine.py | 30 ++++- miles/backends/training_utils/data.py | 21 ++++ miles/backends/training_utils/log_utils.py | 2 + miles/ray/actor_group.py | 5 + miles/ray/rollout/metrics.py | 18 +++ miles/ray/rollout/rollout_manager.py | 3 + miles/ray/rollout/train_data_conversion.py | 13 ++ miles/rollout/generate_utils/sample_utils.py | 2 + miles/rollout/rm_hub/__init__.py | 21 +++- miles/rollout/sglang_rollout.py | 10 +- miles/utils/arguments.py | 61 ++++++++++ miles/utils/types.py | 29 +++++ 18 files changed, 452 insertions(+), 52 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 6d52cf7b53..635a7dee12 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -2,6 +2,7 @@ import random import socket from argparse import Namespace +from pathlib import Path from contextlib import nullcontext from typing import TYPE_CHECKING @@ -50,6 +51,7 @@ from .ft.indep_dp import reconfigure_indep_dp_group from .initialize import init, is_first_replica_megatron_main_rank from .lora_utils import is_lora_enabled +from .multi_lora_utils import is_multi_lora_enabled from .model import TrainStepOutcome, forward_only, initialize_model_and_optimizer, save, train from .parallel import verify_megatron_parallel_state from .replay_utils import register_replay_list_moe @@ -235,6 +237,9 @@ def init( is_lora=is_lora_enabled(args), ) + # Adapters currently loaded into Megatron slots on this rank. + self.loaded_adapters: dict[str, object] = {} + # empty cache after initialization clear_memory() @@ -513,12 +518,61 @@ def train_actor( logger.info(f"Updating ref model at rollout_id {rollout_id}") self.weights_backuper.backup("ref") + if is_multi_lora_enabled(self.args) and 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.""" + if not is_multi_lora_enabled(self.args): + return + from miles.backends.megatron_utils.multi_lora_utils import ( + cleanup_adapters as _cleanup_adapters, + 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) + adapters_to_load = [adapter for name, adapter in should_be_loaded.items() if name not in loaded_names] + adapters_to_clean_up = [ + self.loaded_adapters[n] for n in loaded_names if n in cleanup_names or n not in should_be_loaded + ] + 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.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.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 +588,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) @@ -582,6 +664,9 @@ def update_weights(self, info: "EnginesAndLock") -> None: destroy_process_groups() return + if is_multi_lora_enabled(self.args): + self.weight_updater.multi_lora_adapters = dict(self.loaded_adapters) + with torch_memory_saver.disable() if self.args.offload_train else nullcontext(): print_memory("before update_weights") self.weight_updater.update_weights() diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py index b0e8d92387..b20a56833e 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) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 9910df915f..352eb9c226 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 @@ -33,6 +34,7 @@ from miles.utils.memory_utils import clear_memory from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor from miles.utils.tracking_utils.structured_log import log_structured +from miles.utils.multi_lora import is_multi_lora_enabled from ...utils.misc import filter_keys from ..training_utils.ci_utils import check_grad_norm, check_kl @@ -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) @@ -280,6 +286,7 @@ def forward_step( "response_lengths", "max_seq_lens", "witness_ids", + "adapter_slots", ], args.data_pad_size_multiplier, args.qkv_format, @@ -291,6 +298,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, @@ -439,12 +451,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] @@ -909,13 +927,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/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 651c99e487..93d3b189fc 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,39 @@ 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] 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, + upsert=upsert, ) 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..bbc647394f 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) @@ -256,13 +255,63 @@ def _update_lora_weights(self) -> None: "(no lora_A/lora_B names found). Check weight iterator." ) - if self._lora_loaded: - ray.get( - [engine.unload_lora_adapter.remote(lora_name=LORA_ADAPTER_NAME) for engine in self.rollout_engines] - ) - self._update_lora_weight_implementation(accumulated_named_tensors) + self._update_lora_weight_implementation( + accumulated_named_tensors, + upsert=self._lora_loaded, + ) self._lora_loaded = True + def _update_multi_lora_weights(self) -> None: + """Push every loaded adapter (upsert, never unload), then report the + set to the controller. The push set is the reconcile-time loaded map, + identical on every rank, so per-adapter TP collectives line up.""" + from miles.ray.multi_lora_controller import get_multi_lora_controller + + 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) + + if self._is_lora_source and adapters: + ray.get(get_multi_lora_controller().record_weight_update.remote(sorted(adapters))) + + 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." + ) + + self._update_lora_weight_implementation( + accumulated_named_tensors, + lora_name=f"__miles_slot_{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 +349,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 +359,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 ...multi_lora_utils 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 +380,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/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 304f6cd0e0..9173c388fa 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -58,6 +58,7 @@ def __init__( self.quantization_config = quantization_config self.weight_version = 0 self.is_lora = is_lora + self._hf_weight_iterator = HfWeightIteratorBase.create( args=args, model=model, @@ -212,7 +213,8 @@ def update_weights(self) -> None: if rank == 0: mode = self.args.pause_generation_mode ray.get([engine.pause_generation.remote(mode=mode) for engine in self.rollout_engines]) - ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + if mode not in ("in_place"): + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) if not skip_base_sync: begin_weight_update(self.rollout_engines) dist.barrier(group=get_gloo_group()) @@ -303,7 +305,6 @@ def _send_lora_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: self._lora_loaded = True return refs or [], long_lived_tensors - def _send_to_colocated_engine( hf_named_tensors: list[tuple[str, torch.Tensor]], *, diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 51e44d5e61..2cca141dfa 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -13,7 +13,11 @@ from sglang.srt.utils import kill_process_tree from urllib3.exceptions import NewConnectionError -from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_hf, lora_base_cpu_backup_enabled +from miles.backends.megatron_utils.lora_utils import ( + convert_target_modules_to_hf, + lora_base_cpu_backup_enabled, +) +from miles.backends.megatron_utils.multi_lora_utils import is_multi_lora_enabled from miles.ray.ray_actor import RayActor from miles.utils.env_report import collect_and_print_node_env_report from miles.utils.http_utils import get_host_info @@ -341,14 +345,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 +381,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 +389,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 +403,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 +726,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..a2ff6e693a 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -167,7 +167,10 @@ 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 + tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] + 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,12 +179,14 @@ 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. cu_seqlens_list: list[int] = [0] @@ -202,6 +207,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 +233,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): diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py index 45b073ce9f..a2b2742f92 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -137,6 +137,8 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc "witness_ids", "weight_versions", "metadata", + "n_adapters", + "adapter_slots", ]: continue # Upload per sample mean for each rollout value 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/rollout/metrics.py b/miles/ray/rollout/metrics.py index 9ee7ae41e9..72faa02d00 100644 --- a/miles/ray/rollout/metrics.py +++ b/miles/ray/rollout/metrics.py @@ -14,6 +14,7 @@ from miles.utils.misc import load_function from miles.utils.tracking_utils import tracking from miles.utils.types import Sample +from miles.utils.multi_lora import is_multi_lora_enabled logger = logging.getLogger(__name__) @@ -51,6 +52,20 @@ def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] return log_dict +def _compute_per_adapter_metrics(args, samples: list[Sample]) -> dict: + """Compute reward and response length metrics grouped by adapter name.""" + by_adapter = group_by(samples, lambda s: s.adapter.name if s.adapter else None) + log_dict = {} + for name, adapter_samples in by_adapter.items(): + if name is None: + continue + rewards = [s.get_reward_value(args) for s in adapter_samples] + response_lengths = [s.effective_response_length for s in adapter_samples] + prefix = f"{name}/rollout/" + log_dict |= dict_add_prefix(compute_statistics(rewards), f"{prefix}raw_reward/") + log_dict |= dict_add_prefix(compute_statistics(response_lengths), f"{prefix}response_len/") + return log_dict + def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): if (x := args.custom_rollout_log_function_path) is not None: custom_log_func = load_function(x) @@ -63,6 +78,9 @@ def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_t log_dict = {**(rollout_extra_metrics or {})} log_dict |= dict_add_prefix(_compute_metrics_from_samples(args, samples), "rollout/") log_dict |= dict_add_prefix(_compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/") + + if is_multi_lora_enabled(args): + log_dict |= _compute_per_adapter_metrics(args, samples) logger.info(f"perf {rollout_id}: {log_dict}") step = compute_rollout_step(args, rollout_id) log_dict["rollout/step"] = step diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index ecaa9f41e5..0252a6948e 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -98,6 +98,9 @@ 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): event_analyzer.run_analysis_from_args(self.args) if self._metric_checker is not None: diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index cf80d7890a..b3cabe32ca 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -85,6 +85,9 @@ 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): + train_data["adapter_slots"] = [sample.adapter.slot for sample in samples] + if samples[0].opd_reverse_kl is not None: train_data["opd_reverse_kl"] = [sample.opd_reverse_kl for sample in samples] @@ -137,6 +140,13 @@ 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] + + rollout_data_refs = [] ans = [] for i in range(dp_size): @@ -160,6 +170,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 @@ -174,5 +185,7 @@ def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> l 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/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 68eb57e726..e96d1f07c9 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -127,6 +127,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..ffdaa340d6 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -6,6 +6,7 @@ from miles.utils.misc import load_function from miles.utils.types import Sample +from miles.utils.multi_lora import is_multi_lora_enabled from .deepscaler import get_deepscaler_rule_based_reward, get_gemma_math_reward from .f1 import f1_score @@ -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 64662ed522..8af85a8686 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, load_function +from miles.utils.multi_lora import make_rid from miles.utils.processing_utils import ( call_processor, encode_image_for_rollout_engine, @@ -172,7 +173,14 @@ 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 + + payload["lora_path"] = f"__miles_slot_{sample.adapter.slot}" + payload["rid"] = make_rid(sample.adapter.name) + if (adapter := await AdaptersCache().get(sample.adapter.name)) is not None: + 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/arguments.py b/miles/utils/arguments.py index 32ce45a305..4104e0a96b 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1380,6 +1380,56 @@ 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.", + ) return parser def add_router_arguments(parser): @@ -2486,6 +2536,17 @@ 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 not args.colocate, ( + "Multi-LoRA requires disaggregated rollout engines: weight sync is only " + "implemented for the distributed path, not the colocated tensor path." + ) + 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"]: 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 From c1134f838b857d61f69983af5d36dc04842c6521 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 19:57:32 -0700 Subject: [PATCH 03/31] [misc] revert minor changes --- .../update_weight/update_weight_from_tensor.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 9173c388fa..304f6cd0e0 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -58,7 +58,6 @@ def __init__( self.quantization_config = quantization_config self.weight_version = 0 self.is_lora = is_lora - self._hf_weight_iterator = HfWeightIteratorBase.create( args=args, model=model, @@ -213,8 +212,7 @@ def update_weights(self) -> None: if rank == 0: mode = self.args.pause_generation_mode ray.get([engine.pause_generation.remote(mode=mode) for engine in self.rollout_engines]) - if mode not in ("in_place"): - ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) if not skip_base_sync: begin_weight_update(self.rollout_engines) dist.barrier(group=get_gloo_group()) @@ -305,6 +303,7 @@ def _send_lora_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: self._lora_loaded = True return refs or [], long_lived_tensors + def _send_to_colocated_engine( hf_named_tensors: list[tuple[str, torch.Tensor]], *, From 930a1d2215ffecc0d7a0ccdfea93c04ae860a8a4 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 20:04:05 -0700 Subject: [PATCH 04/31] [fix] deterministic ordering --- miles/backends/megatron_utils/actor.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 635a7dee12..0f48a028ae 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -553,10 +553,16 @@ def reconcile_adapters(self) -> None: cleanup_names = set(snapshot["cleanup"]) loaded_names = set(self.loaded_adapters) - adapters_to_load = [adapter for name, adapter in should_be_loaded.items() if name not in loaded_names] - adapters_to_clean_up = [ - self.loaded_adapters[n] for n in loaded_names if n in cleanup_names or n not in should_be_loaded - ] + # 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: From cf37c3f0d89b45c782c384b3368adc6c8a265105 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 20:24:26 -0700 Subject: [PATCH 05/31] [chore] precommit --- .../multi_lora/multi_lora_async_rollout.py | 26 ++++++------------- .../multi_lora_data_source_async.py | 3 +-- .../multi_lora/tests/test_controller_http.py | 3 +-- .../multi_lora/tests/test_controller_logic.py | 4 ++- .../tests/test_multi_lora_async_rollout.py | 5 ++-- examples/multi_lora/train_multi_lora_async.py | 7 ++--- miles/backends/megatron_utils/actor.py | 10 +++---- miles/backends/megatron_utils/model.py | 2 +- .../megatron_utils/multi_lora_utils.py | 1 - .../update_weight_from_distributed/mixin.py | 1 - miles/backends/sglang_utils/sglang_engine.py | 5 +--- miles/ray/multi_lora_controller.py | 5 +--- miles/ray/rollout/metrics.py | 3 ++- miles/ray/rollout/train_data_conversion.py | 1 - miles/rollout/rm_hub/__init__.py | 2 +- miles/utils/adapter_config.py | 1 + miles/utils/multi_lora.py | 21 ++++++--------- 17 files changed, 38 insertions(+), 62 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index 18b2a5abb2..d930d01b7a 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -9,19 +9,13 @@ from collections.abc import Callable 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 MetricGatherer, 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.rollout.sglang_rollout import GenerateState, generate_and_rm_group, get_model_url from miles.utils.async_utils import run from miles.utils.misc import load_function - -from miles.ray.multi_lora_controller import AdaptersCache, get_multi_lora_controller - from miles.utils.types import Sample logger = logging.getLogger(__name__) @@ -131,9 +125,7 @@ async def run_loop(self) -> None: await asyncio.wait(active) 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 - ) + result = await process_group(self.args, group, self.state.sampling_params, self.generate_fn, self.data_source) if result is not None: self.output_queue.put(result) @@ -150,9 +142,7 @@ async def generate_rollout_multi_lora_async( state = GenerateState(args) - dynamic_filter = ( - load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path else None - ) + dynamic_filter = load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path else None metric_gatherer = MetricGatherer() target_data_size = args.rollout_batch_size @@ -205,7 +195,9 @@ async def generate_rollout_multi_lora_async( if made_progress: last_progress = time.time() elif time.time() - last_progress > 30: - logger.warning(f"No progress for 30s. queue={worker.queue_size()} collected={len(data)}/{target_data_size}") + logger.warning( + f"No progress for 30s. queue={worker.queue_size()} collected={len(data)}/{target_data_size}" + ) last_progress = time.time() if len(data) < target_data_size: @@ -220,9 +212,7 @@ async def generate_rollout_multi_lora_async( data = sorted(data, key=lambda g: first_sample(g).index) - batch_adapters = sorted( - {first_sample(g).adapter.name for g in data if g and first_sample(g).adapter} - ) + batch_adapters = sorted({first_sample(g).adapter.name for g in data if g and first_sample(g).adapter}) if batch_adapters: await get_multi_lora_controller().record_batch_adapters.remote(rollout_id, batch_adapters) diff --git a/examples/multi_lora/multi_lora_data_source_async.py b/examples/multi_lora/multi_lora_data_source_async.py index 2be7b78c78..1e2d57aac5 100644 --- a/examples/multi_lora/multi_lora_data_source_async.py +++ b/examples/multi_lora/multi_lora_data_source_async.py @@ -8,12 +8,11 @@ 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 -from miles.ray.multi_lora_controller import get_multi_lora_controller - logger = logging.getLogger(__name__) MAX_RECONCILE_WORKERS = 16 diff --git a/examples/multi_lora/tests/test_controller_http.py b/examples/multi_lora/tests/test_controller_http.py index 1e828758fa..3c06a8a78f 100644 --- a/examples/multi_lora/tests/test_controller_http.py +++ b/examples/multi_lora/tests/test_controller_http.py @@ -4,13 +4,12 @@ import json from contextlib import asynccontextmanager from pathlib import Path +from types import SimpleNamespace import aiohttp import pytest from aiohttp import web -from types import SimpleNamespace - from miles.utils.adapter_config import AdapterRunConfig from miles.utils.multi_lora import RID_SEPARATOR, MultiLoRABackend, MultiLoRAHTTPServer diff --git a/examples/multi_lora/tests/test_controller_logic.py b/examples/multi_lora/tests/test_controller_logic.py index fecf6234f9..8cb4ce405e 100644 --- a/examples/multi_lora/tests/test_controller_logic.py +++ b/examples/multi_lora/tests/test_controller_logic.py @@ -18,7 +18,9 @@ def make_backend(max_adapters: int = 4, save: str | None = None) -> MultiLoRABac def make_config(save: str | None = None) -> AdapterRunConfig: - return AdapterRunConfig(rank=8, alpha=16, data="/d", save=save, input_key="text", label_key="label", rm_type="math") + return AdapterRunConfig( + rank=8, alpha=16, data="/d", save=save, input_key="text", label_key="label", rm_type="math" + ) def register_and_promote(registry: AdapterRegistry, name: str, config=None) -> None: diff --git a/examples/multi_lora/tests/test_multi_lora_async_rollout.py b/examples/multi_lora/tests/test_multi_lora_async_rollout.py index 16d3fd31aa..06dfa57aac 100644 --- a/examples/multi_lora/tests/test_multi_lora_async_rollout.py +++ b/examples/multi_lora/tests/test_multi_lora_async_rollout.py @@ -1,13 +1,12 @@ """Tests for the testable core of the multi-LoRA async rollout (process_group): keep-vs-recycle plus submission-time slot-version stamping.""" +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 -import examples.multi_lora.multi_lora_async_rollout as mod -from examples.multi_lora.multi_lora_async_rollout import process_group - class FakeDataSource: def __init__(self) -> None: diff --git a/examples/multi_lora/train_multi_lora_async.py b/examples/multi_lora/train_multi_lora_async.py index 3c67ba6ad2..c3d6384ed4 100644 --- a/examples/multi_lora/train_multi_lora_async.py +++ b/examples/multi_lora/train_multi_lora_async.py @@ -4,6 +4,7 @@ import logging from pathlib import Path +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 @@ -11,8 +12,6 @@ from miles.utils.logging_utils import configure_logger from miles.utils.tracking_utils.tracking import init_tracking -from miles.ray.multi_lora_controller import create_controller, get_multi_lora_controller - logger = logging.getLogger(__name__) ROLLOUT_FUNCTION_PATH = "examples.multi_lora.multi_lora_async_rollout.generate_rollout_multi_lora" @@ -20,7 +19,9 @@ async def main(args): - assert not args.colocate, "Colocation is not supported for fully-async training (generation needs continuous GPU; colocate time-shares)." + 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 diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 0f48a028ae..151dfcda2a 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -2,8 +2,8 @@ import random import socket from argparse import Namespace -from pathlib import Path from contextlib import nullcontext +from pathlib import Path from typing import TYPE_CHECKING import ray @@ -51,8 +51,8 @@ from .ft.indep_dp import reconfigure_indep_dp_group from .initialize import init, is_first_replica_megatron_main_rank from .lora_utils import is_lora_enabled -from .multi_lora_utils import is_multi_lora_enabled from .model import TrainStepOutcome, forward_only, initialize_model_and_optimizer, save, train +from .multi_lora_utils import is_multi_lora_enabled from .parallel import verify_megatron_parallel_state from .replay_utils import register_replay_list_moe from .update_weight.common import named_params_and_buffers @@ -535,10 +535,8 @@ def reconcile_adapters(self) -> None: The snapshot is read once on the main rank and broadcast.""" if not is_multi_lora_enabled(self.args): return - from miles.backends.megatron_utils.multi_lora_utils import ( - cleanup_adapters as _cleanup_adapters, - load_adapters as _load_adapters, - ) + 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] diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 352eb9c226..990a39904a 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -32,9 +32,9 @@ 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 -from miles.utils.multi_lora import is_multi_lora_enabled from ...utils.misc import filter_keys from ..training_utils.ci_utils import check_grad_norm, check_kl diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py index 9470efe6db..81ffdcfddc 100644 --- a/miles/backends/megatron_utils/multi_lora_utils.py +++ b/miles/backends/megatron_utils/multi_lora_utils.py @@ -12,7 +12,6 @@ from miles.backends.training_utils.parallel import get_parallel_state from miles.ray.multi_lora_controller import get_multi_lora_controller from miles.utils.adapter_config import AdapterRun -from miles.utils.multi_lora import is_multi_lora_enabled as is_multi_lora_enabled logger = logging.getLogger(__name__) 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 bbc647394f..335e636e31 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 @@ -9,7 +9,6 @@ from miles.backends.training_utils.parallel import get_parallel_state from miles.utils.distributed_utils import get_gloo_group -from miles.utils.lora import LORA_ADAPTER_NAME from miles.utils.timer import timer from ...lora_utils import _is_adapter_param_name, build_lora_sync_config, is_lora_weight_name diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 2cca141dfa..e988e1700d 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -13,10 +13,7 @@ from sglang.srt.utils import kill_process_tree from urllib3.exceptions import NewConnectionError -from miles.backends.megatron_utils.lora_utils import ( - convert_target_modules_to_hf, - lora_base_cpu_backup_enabled, -) +from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_hf, lora_base_cpu_backup_enabled from miles.backends.megatron_utils.multi_lora_utils import is_multi_lora_enabled from miles.ray.ray_actor import RayActor from miles.utils.env_report import collect_and_print_node_env_report diff --git a/miles/ray/multi_lora_controller.py b/miles/ray/multi_lora_controller.py index df10b1d338..2d74867f06 100644 --- a/miles/ray/multi_lora_controller.py +++ b/miles/ray/multi_lora_controller.py @@ -47,7 +47,6 @@ 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 @@ -62,9 +61,7 @@ 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) - ) + 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() diff --git a/miles/ray/rollout/metrics.py b/miles/ray/rollout/metrics.py index 72faa02d00..9f90a2e537 100644 --- a/miles/ray/rollout/metrics.py +++ b/miles/ray/rollout/metrics.py @@ -12,9 +12,9 @@ has_repetition, ) from miles.utils.misc import load_function +from miles.utils.multi_lora import is_multi_lora_enabled from miles.utils.tracking_utils import tracking from miles.utils.types import Sample -from miles.utils.multi_lora import is_multi_lora_enabled logger = logging.getLogger(__name__) @@ -66,6 +66,7 @@ def _compute_per_adapter_metrics(args, samples: list[Sample]) -> dict: log_dict |= dict_add_prefix(compute_statistics(response_lengths), f"{prefix}response_len/") return log_dict + def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): if (x := args.custom_rollout_log_function_path) is not None: custom_log_func = load_function(x) diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index b3cabe32ca..0f81c8799c 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -146,7 +146,6 @@ def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> l if adapter_slots is not None: partitions = [sorted(p, key=lambda i: adapter_slots[i]) for p in partitions] - rollout_data_refs = [] ans = [] for i in range(dp_size): diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py index ffdaa340d6..4bc1c95b54 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -5,8 +5,8 @@ import aiohttp from miles.utils.misc import load_function -from miles.utils.types import Sample 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 from .f1 import f1_score diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py index 78e3fb2eda..ad3826628e 100644 --- a/miles/utils/adapter_config.py +++ b/miles/utils/adapter_config.py @@ -35,6 +35,7 @@ class AdapterRunConfig: metadata: dict[str, Any] = field(default_factory=dict) + @dataclass(frozen=True) class AdapterRun: """Read-only join view of a run's static config and current slot.""" diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index 61ad9c26fe..ecc66e3cfd 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -15,7 +15,7 @@ from fastapi.responses import JSONResponse from pydantic import BaseModel -from miles.utils.adapter_config import AdapterRunConfig, AdapterRun, parse_adapter_run_yaml +from miles.utils.adapter_config import AdapterRun, AdapterRunConfig, parse_adapter_run_yaml logger = logging.getLogger(__name__) @@ -103,9 +103,7 @@ def is_active(self, name: str) -> bool: 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 '-'" - ) + 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") @@ -222,7 +220,6 @@ def views(state: AdapterState) -> dict[str, AdapterRun]: } - class MultiLoRABackend: """Registry + engine-facing aborts, shared by the Ray actor and HTTP server. Subclass via --multi-lora-backend-path.""" @@ -250,9 +247,7 @@ def resolve_save_dir(self, name: str, config: Any) -> Any: if config.save is not None: return config 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" - ) + raise ValueError(f"Adapter '{name}' has no save dir: set 'save' in the adapter config or pass --save") return replace(config, save=Path(self.args.save) / "adapters" / name) async def register(self, name: str, config: Any) -> dict: @@ -294,10 +289,7 @@ async def abort_adapter_requests(self, adapter_name: str) -> None: 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 - ), + *(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): @@ -312,6 +304,9 @@ class RegisterAdapterRequest(BaseModel): 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).""" @@ -388,7 +383,7 @@ def adapter_statuses(self) -> list[dict]: async def list_adapters(self) -> dict: return {"adapters": self.adapter_statuses()} - async def adapter_states(self, names: list[str] = Query(default_factory=list)) -> dict: + 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: From 629dc851af02593dd39edbfd0efe35c828bd15eb Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 21:03:13 -0700 Subject: [PATCH 06/31] [fix] import paths for tests --- miles/backends/megatron_utils/actor.py | 3 ++- .../update_weight/update_weight_from_distributed/mixin.py | 2 +- miles/backends/sglang_utils/sglang_engine.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index dbc199ce93..8ef82cee83 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -22,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 @@ -52,7 +53,6 @@ from .initialize import init, is_first_replica_megatron_main_rank from .lora_utils import is_lora_enabled from .model import TrainStepOutcome, forward_only, initialize_model_and_optimizer, save, train -from .multi_lora_utils import is_multi_lora_enabled from .parallel import verify_megatron_parallel_state from .replay_utils import register_replay_list_moe from .update_weight.common import named_params_and_buffers @@ -635,6 +635,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) 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 335e636e31..2ae8eb857e 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 @@ -358,7 +358,7 @@ def update_weights(self) -> None: dist.barrier(group=get_gloo_group()) with timer("update_weights_implementation"): - from ...multi_lora_utils import is_multi_lora_enabled + 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) diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index e988e1700d..f418d35508 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -14,11 +14,11 @@ from urllib3.exceptions import NewConnectionError from miles.backends.megatron_utils.lora_utils import convert_target_modules_to_hf, lora_base_cpu_backup_enabled -from miles.backends.megatron_utils.multi_lora_utils import is_multi_lora_enabled from miles.ray.ray_actor import RayActor 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__) From 9ca4536f92fdd905647e2f16f8ed3a4a50683f38 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 22:11:33 -0700 Subject: [PATCH 07/31] [test] move tests --- examples/multi_lora/README.md | 1 - .../fast/utils/test_controller_backend.py | 4 ++++ .../tests => tests/fast/utils}/test_controller_http.py | 4 ++++ 3 files changed, 8 insertions(+), 1 deletion(-) rename examples/multi_lora/tests/test_controller_logic.py => tests/fast/utils/test_controller_backend.py (98%) rename {examples/multi_lora/tests => tests/fast/utils}/test_controller_http.py (98%) diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md index 78145629ba..063b5de3c7 100644 --- a/examples/multi_lora/README.md +++ b/examples/multi_lora/README.md @@ -19,7 +19,6 @@ service_smoke.py # register/deregister smoke test against th 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, deregisters at num_row) -tests/ # controller logic + HTTP tests (no torch) adapters/ gsm8k.yaml dapo_math.yaml diff --git a/examples/multi_lora/tests/test_controller_logic.py b/tests/fast/utils/test_controller_backend.py similarity index 98% rename from examples/multi_lora/tests/test_controller_logic.py rename to tests/fast/utils/test_controller_backend.py index 8cb4ce405e..580e932c2a 100644 --- a/examples/multi_lora/tests/test_controller_logic.py +++ b/tests/fast/utils/test_controller_backend.py @@ -3,6 +3,10 @@ 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 diff --git a/examples/multi_lora/tests/test_controller_http.py b/tests/fast/utils/test_controller_http.py similarity index 98% rename from examples/multi_lora/tests/test_controller_http.py rename to tests/fast/utils/test_controller_http.py index 3c06a8a78f..eb401c3e1f 100644 --- a/examples/multi_lora/tests/test_controller_http.py +++ b/tests/fast/utils/test_controller_http.py @@ -10,6 +10,10 @@ 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 From e371b2e097df186f11e4a2ce9fe9a8c287691393 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 11 Jul 2026 22:30:25 -0700 Subject: [PATCH 08/31] [fix] tests + keep original lora behavior for non multi-lora --- .../update_weight_from_distributed/broadcast.py | 3 ++- .../update_weight_from_distributed/mixin.py | 10 ++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) 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 93d3b189fc..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 @@ -151,6 +151,7 @@ def _update_lora_weight_implementation( 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_name, @@ -159,7 +160,7 @@ def _update_lora_weight_implementation( dtypes=dtypes, shapes=shapes, group_name=self._group_name, - upsert=upsert, + **extra_kwargs, ) for engine in self.rollout_engines ] 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 2ae8eb857e..27e2f74141 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 @@ -9,6 +9,7 @@ from miles.backends.training_utils.parallel import get_parallel_state from miles.utils.distributed_utils import get_gloo_group +from miles.utils.lora import LORA_ADAPTER_NAME from miles.utils.timer import timer from ...lora_utils import _is_adapter_param_name, build_lora_sync_config, is_lora_weight_name @@ -254,10 +255,11 @@ def _update_lora_weights(self) -> None: "(no lora_A/lora_B names found). Check weight iterator." ) - self._update_lora_weight_implementation( - accumulated_named_tensors, - upsert=self._lora_loaded, - ) + if self._lora_loaded: + ray.get( + [engine.unload_lora_adapter.remote(lora_name=LORA_ADAPTER_NAME) for engine in self.rollout_engines] + ) + self._update_lora_weight_implementation(accumulated_named_tensors) self._lora_loaded = True def _update_multi_lora_weights(self) -> None: From 8bdd7d49fb57a94407f127d38f35083b82218b31 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Sun, 12 Jul 2026 23:02:46 -0700 Subject: [PATCH 09/31] [fix] recompute-logprobs uses per-sample adapter lora_path in multi-LoRA _build_prefill_scoring_payload keyed LoRA solely on is_lora_enabled(args) (always true in multi-LoRA) and sent the single-adapter name miles_lora, which is never registered on multi-LoRA engines -> every step crashed with 'adapters are not loaded' when --recompute-logprobs-via-prefill is on; the batched path also applied payloads[0]'s lora_path to a mixed-adapter batch. - resolve lora_path per sample: adapter samples score under their own __miles_slot_{N}, single-adapter LoRA keeps miles_lora, base keeps none - batched scoring groups by (logprob_start_len, lora_path) and rejects mixed-adapter batch payloads - centralize the engine-side slot adapter name as slot_lora_name(), shared by rollout, prefill scoring, and the weight-push mixin Signed-off-by: Yusheng Su --- .../update_weight_from_distributed/mixin.py | 4 +- .../generate_utils/prefill_logprobs.py | 39 +++++-- miles/rollout/sglang_rollout.py | 4 +- miles/utils/multi_lora.py | 7 ++ .../generate_utils/test_prefill_logprobs.py | 101 +++++++++++++++++- 5 files changed, 144 insertions(+), 11 deletions(-) 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 27e2f74141..dc13b7b07c 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 @@ -306,9 +306,11 @@ def _send_one_multi_lora_adapter(self, adapter, upsert: bool) -> None: "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=f"__miles_slot_{adapter.slot}", + lora_name=slot_lora_name(adapter.slot), lora_config=lora_config, upsert=upsert, ) 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/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 8af85a8686..647af14baa 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -23,7 +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, load_function -from miles.utils.multi_lora import make_rid +from miles.utils.multi_lora import make_rid, slot_lora_name from miles.utils.processing_utils import ( call_processor, encode_image_for_rollout_engine, @@ -176,7 +176,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A if sample.adapter is not None: from miles.ray.multi_lora_controller import AdaptersCache - payload["lora_path"] = f"__miles_slot_{sample.adapter.slot}" + payload["lora_path"] = slot_lora_name(sample.adapter.slot) payload["rid"] = make_rid(sample.adapter.name) if (adapter := await AdaptersCache().get(sample.adapter.name)) is not None: payload["extra_key"] = f"{sample.adapter.name}:v{adapter.version}" diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index ecc66e3cfd..fe5e36aeb7 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -28,6 +28,7 @@ "is_multi_lora_enabled", "make_rid", "parse_adapter", + "slot_lora_name", ] @@ -49,6 +50,12 @@ 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}" + + class AdapterState(str, Enum): PENDING = "PENDING" ACTIVE = "ACTIVE" 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 = [ From 1e93ea40930be334f136641241baa5d3a75d44fd Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Sun, 12 Jul 2026 23:03:15 -0700 Subject: [PATCH 10/31] [fix] harden retired-adapter teardown against orphaned rollout requests The retire-time prefix abort fires exactly once (RETIRING->CLEANUP); requests can survive it (a multi-turn group POSTs its next turn after the round, or a request sits in the engine's tokenizer window) and, once the slot is reused and the next adapter's weights are upserted into the same __miles_slot_{N}, keep decoding under the wrong weights with no error. Batch-time collection filters already drop such groups from training data; these changes shrink the window and stop wasting decode on them: - MultiLoRABackend.free_slot: re-run the prefix abort right before the slot is released for reuse (second round, after a full step of settling) - sglang_rollout.generate: refuse to POST for an adapter that is no longer sampleable (deregistered/cleaned up) and abort the sample locally instead - the tokenizer-window escape itself is fixed engine-side (sgl-project#30912) Signed-off-by: Yusheng Su --- miles/ray/multi_lora_controller.py | 4 +-- miles/rollout/sglang_rollout.py | 14 ++++++-- miles/utils/multi_lora.py | 16 +++++++++ tests/fast/utils/test_controller_backend.py | 39 +++++++++++++++++++++ 4 files changed, 69 insertions(+), 4 deletions(-) diff --git a/miles/ray/multi_lora_controller.py b/miles/ray/multi_lora_controller.py index 2d74867f06..1f3913ee40 100644 --- a/miles/ray/multi_lora_controller.py +++ b/miles/ray/multi_lora_controller.py @@ -81,8 +81,8 @@ async def deregister_adapter(self, name: str) -> None: async def retire_adapters(self) -> list[str]: return await self.backend.retire_adapters() - def free_slot(self, name: str) -> int: - return self.backend.registry.free_slot(name) + 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) diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 647af14baa..74e0358caa 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -176,10 +176,20 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A 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) - if (adapter := await AdaptersCache().get(sample.adapter.name)) is not None: - payload["extra_key"] = f"{sample.adapter.name}:v{adapter.version}" + payload["extra_key"] = f"{sample.adapter.name}:v{adapter.version}" elif is_lora_enabled(args): payload["lora_path"] = LORA_ADAPTER_NAME diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index fe5e36aeb7..1ce388ddc7 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -275,6 +275,22 @@ async def retire_adapters(self) -> list[str]: 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 window 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 ( diff --git a/tests/fast/utils/test_controller_backend.py b/tests/fast/utils/test_controller_backend.py index 580e932c2a..34aca8a131 100644 --- a/tests/fast/utils/test_controller_backend.py +++ b/tests/fast/utils/test_controller_backend.py @@ -161,6 +161,45 @@ def test_deregister_holds_slot_until_free_slot(): 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-window 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): From 50f6ff14f1339195369fc656dc3e0f6236cf12c8 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Mon, 13 Jul 2026 11:09:02 -0700 Subject: [PATCH 11/31] [test] mark recycle-aborted as xfail until re-queuing lands Re-queuing aborted groups (add_samples) is intended but deferred to a future PR (confirmed by mathewjhan; left unimplemented for lack of testing time). strict=True so wiring it up forces marker removal. Signed-off-by: Yusheng Su --- examples/multi_lora/tests/test_multi_lora_async_rollout.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/multi_lora/tests/test_multi_lora_async_rollout.py b/examples/multi_lora/tests/test_multi_lora_async_rollout.py index 06dfa57aac..ae971bda9a 100644 --- a/examples/multi_lora/tests/test_multi_lora_async_rollout.py +++ b/examples/multi_lora/tests/test_multi_lora_async_rollout.py @@ -56,6 +56,12 @@ async def test_process_group_keeps_completed(): @pytest.mark.asyncio +@pytest.mark.xfail( + reason="Re-queuing aborted groups is not wired up yet (the per-adapter " + "source is read-only); planned for a future PR. This test pins the " + "intended end-state behavior.", + strict=True, +) async def test_process_group_recycles_aborted(): async def gen(args, group, sampling_params): for s in group: From a75a5b4639f52988bf03bf7410067a241d8dfa9b Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Mon, 13 Jul 2026 17:10:42 -0700 Subject: [PATCH 12/31] [fix] reject multi-LoRA with --sglang-tokenizer-worker-num > 1 at launch Each sglang tokenizer worker process holds its own LoRA registry with no cross-worker sync (sgl-project/sglang#31084), so per-step adapter upserts resolve on whichever worker the router picks and fail non-deterministically. sglang rejects the upsert at runtime anyway; failing at launch avoids burning GPU time until the first weight push. Signed-off-by: Yusheng Su --- miles/utils/arguments.py | 7 ++++++ tests/fast/utils/test_arguments.py | 35 ++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 54831662c5..003c5f213e 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2558,6 +2558,13 @@ def miles_validate_args(args): "Multi-LoRA requires disaggregated rollout engines: weight sync is only " "implemented for the distributed path, not the colocated tensor path." ) + 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." + ) 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" 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.""" From c5cc40a098a1be160be3dad589493952e343c830 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Wed, 15 Jul 2026 16:25:13 -0700 Subject: [PATCH 13/31] [feat] optimizer changes initial commit --- examples/multi_lora/README.md | 57 ++- examples/multi_lora/adapters/dapo_math.yaml | 2 + examples/multi_lora/adapters/gsm8k.yaml | 2 + .../multi_lora/multi_lora_async_rollout.py | 389 ++++++++++++++---- .../multi_lora_data_source_async.py | 64 ++- examples/multi_lora/run_job.sh | 1 - examples/multi_lora/run_service.sh | 1 - .../tests/test_multi_lora_batch_assembler.py | 239 +++++++++++ examples/multi_lora/train_multi_lora_async.py | 3 +- miles/backends/megatron_utils/actor.py | 46 ++- miles/backends/megatron_utils/arguments.py | 4 + .../megatron_utils/bridge_lora_helpers.py | 4 + miles/backends/megatron_utils/model.py | 78 +++- .../megatron_utils/multi_lora_optimizer.py | 248 +++++++++++ .../megatron_utils/multi_lora_utils.py | 20 +- .../update_weight_from_distributed/mixin.py | 12 +- miles/backends/training_utils/data.py | 22 +- miles/backends/training_utils/log_utils.py | 4 + miles/backends/training_utils/loss.py | 6 + miles/ray/multi_lora_controller.py | 4 +- miles/ray/rollout/rollout_data_conversion.py | 31 ++ miles/ray/rollout/rollout_manager.py | 2 + miles/ray/rollout/train_data_conversion.py | 51 ++- miles/utils/adapter_config.py | 17 + miles/utils/arguments.py | 80 +++- miles/utils/multi_lora.py | 150 ++++++- .../ray/rollout/test_multi_lora_train_data.py | 106 +++++ tests/fast/utils/test_controller_backend.py | 124 ++++-- tests/fast/utils/test_controller_http.py | 14 +- 29 files changed, 1560 insertions(+), 221 deletions(-) create mode 100644 examples/multi_lora/tests/test_multi_lora_batch_assembler.py create mode 100644 miles/backends/megatron_utils/multi_lora_optimizer.py create mode 100644 tests/fast/ray/rollout/test_multi_lora_train_data.py diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md index 063b5de3c7..1bac3d3900 100644 --- a/examples/multi_lora/README.md +++ b/examples/multi_lora/README.md @@ -28,18 +28,50 @@ 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 (no drain, no state machine) +## 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`). -- **No drain / no rollout-id / no train_steps / no PENDING-DRAINING-DRAINED states.** - The data source deregisters an adapter at `num_row`; the trainer's - `reconcile_adapters` (before each generate) cleans up gone adapters (save ckpt + - clear Megatron slot) and loads new ones. `update_weights` upserts active adapters' - weights in place (SGLang page table, `upsert=True`). +- **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. +- The data source deregisters an adapter at `num_row`; 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. TODO: revisit + num_row semantics (the tail means slightly fewer trained rows than + configured). - **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. @@ -75,6 +107,8 @@ Per-adapter `rank` in `adapter.yaml` must be `<= --lora-rank`. ```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 @@ -82,3 +116,14 @@ rm_type: math num_row: 400 # stop adapter after N rows # 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 index dfcabeeb2b..135bcde61c 100644 --- a/examples/multi_lora/adapters/dapo_math.yaml +++ b/examples/multi_lora/adapters/dapo_math.yaml @@ -1,5 +1,7 @@ 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 diff --git a/examples/multi_lora/adapters/gsm8k.yaml b/examples/multi_lora/adapters/gsm8k.yaml index 2a734590bf..8c1185a88a 100644 --- a/examples/multi_lora/adapters/gsm8k.yaml +++ b/examples/multi_lora/adapters/gsm8k.yaml @@ -1,5 +1,7 @@ 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 diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index d930d01b7a..ff36a72deb 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -1,21 +1,34 @@ -"""Fully-async multi-LoRA rollout: continuous background producer + collect-a-batch.""" +"""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 queue 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 MetricGatherer, call_dynamic_filter +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.misc import load_function +from miles.utils.multi_lora import min_groups_per_dp_split from miles.utils.types import Sample logger = logging.getLogger(__name__) @@ -34,6 +47,79 @@ 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: @@ -64,7 +150,8 @@ async def process_group( class AsyncMultiLoRAWorker: - """Background producer: continuously generate groups into a thread-safe queue.""" + """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() @@ -75,9 +162,22 @@ def __init__(self, args, data_source, generate_fn: GenerateFn, concurrency: int self.generate_fn = generate_fn self.concurrency = concurrency or args.rollout_batch_size self.running = True - self.output_queue: queue.Queue = queue.Queue(maxsize=1000) 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.dynamic_filter_drop_counts: dict[str, int] = defaultdict(int) + self.stale_dropped = 0 + self.staleness_values: list[int] = [] @classmethod def get_or_create(cls, args, data_source, generate_fn: GenerateFn, concurrency: int = None): @@ -96,6 +196,14 @@ def stop(self) -> None: 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()) @@ -116,105 +224,213 @@ async def run_loop(self) -> None: samples = self.data_source.get_samples(1) if not samples: break - group = samples[0] - active.add(asyncio.create_task(self.process_and_enqueue(group))) + active.add(asyncio.create_task(self.process_and_enqueue(samples[0]))) - await asyncio.sleep(0) + await asyncio.sleep(0.01) finally: + for task in active: + task.cancel() if active: - await asyncio.wait(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 not None: - self.output_queue.put(result) + 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: + with self.buffer_lock: + self.dynamic_filter_drop_counts[filter_result.reason] += 1 + 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: - return self.output_queue.qsize() + with self.buffer_lock: + return sum(len(buffer) for buffer in self.buffers.values()) + + 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). + for name in list(self.buffers): + if name not in adapters: + self.buffers.pop(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.stale_dropped += len(dropped) + self.staleness_values += dropped + min_groups_per_pop = min_groups_per_dp_split(adapter.config.n_samples_per_prompt, dp_size) + available_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, available_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 + + def pop_metrics(self) -> dict[str, float]: + with self.buffer_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() + metrics["perf/fully_async/stale_dropped"] = self.stale_dropped + if self.staleness_values: + metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(self.staleness_values) / len( + self.staleness_values + ) + metrics["perf/fully_async/stale_dropped_max_staleness"] = max(self.staleness_values) + self.stale_dropped = 0 + self.staleness_values = [] + return metrics + + +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 RuntimeError( + "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 -) -> tuple[RolloutFnTrainOutput, list[list[Sample]]]: - """Fully-async multi-LoRA rollout. Collect a batch from the background worker, - then run the same postprocess as ``generate_rollout_async``.""" +) -> RolloutFnTrainOutput: + """Collect one train batch and record its contents on the controller.""" assert args.rollout_global_dataset state = GenerateState(args) - - dynamic_filter = load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path else None - metric_gatherer = MetricGatherer() - target_data_size = args.rollout_batch_size - worker = AsyncMultiLoRAWorker.get_or_create(args, data_source, generate_fn) - - # Groups whose submission-time slot version fell too far behind are dropped. - max_staleness = getattr(args, "max_weight_staleness", None) - - data: list[Group] = [] - stale_dropped = 0 - staleness_values: list[int] = [] start_time = time.time() - last_progress = start_time queue_length = worker.queue_size() - while len(data) < target_data_size: - made_progress = False - current_adapters = await AdaptersCache().get_all() - # Pop one at a time so surplus groups stay queued for the next batch. - while len(data) < target_data_size: - try: - group = worker.output_queue.get_nowait() - except queue.Empty: - break - head = first_sample(group) if group else None - adapter_name = head.adapter.name if head is not None and head.adapter else None - if adapter_name not in current_adapters: - continue # adapter deregistered; drop - if max_staleness is not None: - stamped = head.metadata.get("slot_version") - if stamped is not None: - staleness = current_adapters[adapter_name].version - stamped - if staleness > max_staleness: - for s in iter_group_samples(group): - s.reset_for_retry() - stale_dropped += 1 - staleness_values.append(staleness) - logger.info( - f"Dropped stale group (adapter={adapter_name}, " - f"stamped={stamped}, current={current_adapters[adapter_name].version}, " - f"staleness={staleness} > max={max_staleness})" - ) - continue - f = call_dynamic_filter(dynamic_filter, args, group) - if not f.keep: - metric_gatherer.on_dynamic_filter_drop(reason=f.reason) - continue - data.append(group) - made_progress = True - - if made_progress: - last_progress = time.time() - elif time.time() - last_progress > 30: - logger.warning( - f"No progress for 30s. queue={worker.queue_size()} collected={len(data)}/{target_data_size}" - ) - last_progress = time.time() - if len(data) < target_data_size: - await asyncio.sleep(0.01) + # 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" - if stale_dropped: - logger.info( - f"Staleness stats: dropped={stale_dropped}, " - f"avg_staleness={sum(staleness_values) / len(staleness_values):.1f}, " - f"max_staleness={max(staleness_values)}" - ) + batch = await collect_batch(args, worker, snapshot) - data = sorted(data, key=lambda g: first_sample(g).index) + 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, + ), + ) - batch_adapters = sorted({first_sample(g).adapter.name for g in data if g and first_sample(g).adapter}) - if batch_adapters: - await get_multi_lora_controller().record_batch_adapters.remote(rollout_id, batch_adapters) + # 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) @@ -227,13 +443,14 @@ async def generate_rollout_multi_lora_async( ) metrics = { - **metric_gatherer.collect(), + **worker.pop_metrics(), "perf/fully_async/queue_length": queue_length, "perf/fully_async/batch_wait_time": time.time() - start_time, - "perf/fully_async/stale_dropped": stale_dropped, + "perf/fully_async/batch_adapters": len(batch.group_counts), + "perf/fully_async/batch_prompt_groups": len(data), + "perf/fully_async/batch_samples": sum(group_sample_count(group) for group in data), + "perf/fully_async/batch_step_count": len(batch.step_names), } - if staleness_values: - metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(staleness_values) / len(staleness_values) return RolloutFnTrainOutput(samples=data, metrics=metrics) diff --git a/examples/multi_lora/multi_lora_data_source_async.py b/examples/multi_lora/multi_lora_data_source_async.py index 1e2d57aac5..b7ee6fa34d 100644 --- a/examples/multi_lora/multi_lora_data_source_async.py +++ b/examples/multi_lora/multi_lora_data_source_async.py @@ -59,6 +59,7 @@ def create_source(self, adapter: AdapterRun) -> RolloutDataSource: 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) @@ -74,56 +75,44 @@ def update_queue(self, active_names: set[str]) -> None: new_queue.append(name) self.source_queue = new_queue - def get_samples(self, num_samples: int) -> list[list[Sample]]: + 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) - if not self.sources: - return [] self.update_queue(set(self.sources)) - refs = {name: AdapterRef(name=name, slot=adapters[name].slot) for name in self.sources} - reward_specs = { - name: RewardSpec( - rm_type=adapters[name].config.rm_type, - custom_rm_path=adapters[name].config.custom_rm_path, - ) - for name in self.sources - } - - samples_per_adapter, remainder = divmod(num_samples, len(self.source_queue)) - all_samples: list[list[Sample]] = [] - to_deregister: list[str] = [] - - for i in range(len(self.source_queue)): - extra = int(i < remainder) - samples_needed = samples_per_adapter + extra - if samples_needed == 0: - break + for _ in range(len(self.source_queue)): name = self.source_queue.popleft() - config = adapters[name].config self.source_queue.append(name) source = self.sources[name] - adapter_samples = source.get_samples(samples_needed) - ref = refs[name] - reward_spec = reward_specs[name] - for group in adapter_samples: - for sample in group: - sample.adapter = ref - sample.reward_spec = reward_spec - sample.metadata = {**config.metadata, **sample.metadata} - all_samples.extend(adapter_samples) + 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} default_num_row = (getattr(config, "num_epoch", 1) or 1) * len(source.dataset) num_row = config.num_row or default_num_row if source.sample_group_index >= num_row and name not in snapshot["retiring"]: logger.info(f"Adapter '{name}' reached num_row={num_row}, deregistering") - to_deregister.append(name) + ray.get(get_multi_lora_controller().deregister_adapter.remote(name)) - for name in to_deregister: - ray.get(get_multi_lora_controller().deregister_adapter.remote(name)) + return groups - return all_samples + return [] def add_samples(self, samples: list[list[Sample]]) -> None: """Recycle retried/aborted groups; drop groups for deregistered adapters.""" @@ -142,3 +131,8 @@ def save(self, 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/run_job.sh b/examples/multi_lora/run_job.sh index e65a7efc59..80a29da846 100755 --- a/examples/multi_lora/run_job.sh +++ b/examples/multi_lora/run_job.sh @@ -23,7 +23,6 @@ ray job submit --address="http://127.0.0.1:8265" \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 4 \ --rollout-num-gpus 4 \ - --calculate-per-token-loss \ --use-miles-router \ ${MODEL_ARGS[@]} \ \ diff --git a/examples/multi_lora/run_service.sh b/examples/multi_lora/run_service.sh index e9a9695e21..eb4c7d56fe 100755 --- a/examples/multi_lora/run_service.sh +++ b/examples/multi_lora/run_service.sh @@ -25,7 +25,6 @@ ray job submit --address="http://127.0.0.1:8265" \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 4 \ --rollout-num-gpus 4 \ - --calculate-per-token-loss \ --use-miles-router \ ${MODEL_ARGS[@]} \ \ diff --git a/examples/multi_lora/tests/test_multi_lora_batch_assembler.py b/examples/multi_lora/tests/test_multi_lora_batch_assembler.py new file mode 100644 index 0000000000..d3af21ee1f --- /dev/null +++ b/examples/multi_lora/tests/test_multi_lora_batch_assembler.py @@ -0,0 +1,239 @@ +"""Unit tests for 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, + 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.dynamic_filter_drop_counts = defaultdict(int) + worker.stale_dropped = 0 + worker.staleness_values = [] + 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_n_samples_multiple_of_dp_pops_single_groups(): + worker = make_worker(make_args(multi_lora_dp_size=4, global_batch_size=8)) + a = adapter_run("A", 0, rollout_batch_size=2, n_samples_per_prompt=8) # 8 % 4 == 0 + buffer_groups(worker, a, count=1) + batch = collect(worker, snapshot_of(a)) # 8 samples = target + assert batch.group_counts == {"A": 1} + assert sum(1 for g in batch.groups for _ in g) == 8 + + +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 worker.stale_dropped == 2 + assert batch.group_counts == {"A": 1} + + +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/examples/multi_lora/train_multi_lora_async.py b/examples/multi_lora/train_multi_lora_async.py index c3d6384ed4..e5435a2a07 100644 --- a/examples/multi_lora/train_multi_lora_async.py +++ b/examples/multi_lora/train_multi_lora_async.py @@ -60,7 +60,8 @@ async def main(args): continue # Reconcile + push before generate: the push promotes pending adapters, - # and only then does the data source sample them. + # 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() diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 8ef82cee83..69336232b8 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -239,6 +239,12 @@ def init( # 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() @@ -518,10 +524,15 @@ def train_actor( logger.info(f"Updating ref model at rollout_id {rollout_id}") self.weights_backuper.backup("ref") - if is_multi_lora_enabled(self.args) and is_first_replica_megatron_main_rank(): - from miles.ray.multi_lora_controller import get_multi_lora_controller + 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)) + 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()) @@ -532,7 +543,14 @@ def train_actor( @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.""" + 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. TODO: revisit num_row + semantics — the tail means an adapter trains on slightly fewer rows + than configured. + """ if not is_multi_lora_enabled(self.args): return from miles.backends.megatron_utils.multi_lora_utils import cleanup_adapters as _cleanup_adapters @@ -565,11 +583,13 @@ def reconcile_adapters(self) -> None: _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. @@ -685,14 +705,30 @@ def update_weights(self, info: "EnginesAndLock") -> None: destroy_process_groups() return + version_update_names: list[str] = [] if is_multi_lora_enabled(self.args): - self.weight_updater.multi_lora_adapters = dict(self.loaded_adapters) + # 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 b20a56833e..6169d4e111 100644 --- a/miles/backends/megatron_utils/bridge_lora_helpers.py +++ b/miles/backends/megatron_utils/bridge_lora_helpers.py @@ -135,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 990a39904a..140852d583 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -171,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, @@ -385,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. @@ -402,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 @@ -535,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 @@ -560,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() - - # 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() + 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) + + # 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, 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 index 81ffdcfddc..ce8f28dc17 100644 --- a/miles/backends/megatron_utils/multi_lora_utils.py +++ b/miles/backends/megatron_utils/multi_lora_utils.py @@ -90,6 +90,8 @@ def zero_optimizer_state_for_adapter(optimizer, model, idx: int) -> None: 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 @@ -97,6 +99,12 @@ def zero_optimizer_state_for_adapter(optimizer, model, idx: int) -> None: 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: @@ -174,8 +182,8 @@ def save_multi_lora_checkpoints( if is_dp_rank_0: shard: dict[str, torch.Tensor] = { name: param.data.cpu() - for chunk in model - for name, param in chunk.named_parameters() + 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" @@ -281,10 +289,14 @@ def _deregister_adapter(adapter: AdapterRun, args, model, optimizer) -> None: clear_adapter_slot(model, slot) logger.info(f"{log_prefix} cleared adapter slot {slot}") - # Prevent future slot tenants from inheriting optimizer momentum. + # 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 for slot {slot}") + logger.info(f"{log_prefix} cleared optimizer state and retained grads for slot {slot}") def load_adapters(args, model, optimizer, adapters) -> int: 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 dc13b7b07c..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 @@ -263,19 +263,15 @@ def _update_lora_weights(self) -> None: self._lora_loaded = True def _update_multi_lora_weights(self) -> None: - """Push every loaded adapter (upsert, never unload), then report the - set to the controller. The push set is the reconcile-time loaded map, - identical on every rank, so per-adapter TP collectives line up.""" - from miles.ray.multi_lora_controller import get_multi_lora_controller - + """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) - if self._is_lora_source and adapters: - ray.get(get_multi_lora_controller().record_weight_update.remote(sorted(adapters))) - def _send_one_multi_lora_adapter(self, adapter, upsert: bool) -> None: """All ranks iterate the bridge (TP collectives); only the source rank transmits.""" diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index a2ff6e693a..5643b7adda 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: @@ -188,14 +188,14 @@ def get_batch( 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 @@ -204,7 +204,7 @@ def get_batch( cu_seqlens_list.append(cu_seqlens_list[-1] + pad) cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=torch.cuda.current_device()) - tokens = tokens.chunk(cp_size, dim=0)[cp_rank] + tokens = tokens.batch(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] @@ -307,7 +307,7 @@ def _compute_transform_like_token_ids(ids_list: list): loss_masks = torch.cat(loss_masks, dim=0) if pad != 0: loss_masks = F.pad(loss_masks, (0, pad), value=0) - loss_masks = loss_masks.chunk(cp_size, dim=0)[cp_rank].unsqueeze(0) + loss_masks = loss_masks.batch(cp_size, dim=0)[cp_rank].unsqueeze(0) elif qkv_format == "thd": loss_masks = torch.cat(loss_masks) loss_masks = F.pad(loss_masks, (0, pad), value=0).unsqueeze(0) @@ -367,7 +367,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). @@ -447,6 +447,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: @@ -484,6 +490,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 a2b2742f92..a4aa866ed7 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -139,6 +139,10 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc "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/multi_lora_controller.py b/miles/ray/multi_lora_controller.py index 1f3913ee40..eba731dc96 100644 --- a/miles/ray/multi_lora_controller.py +++ b/miles/ray/multi_lora_controller.py @@ -87,8 +87,8 @@ async def free_slot(self, name: str) -> int: def record_weight_update(self, names: list[str]) -> None: self.backend.registry.record_weight_update(names) - def record_batch_adapters(self, rollout_id: int, names: list[str]) -> None: - self.backend.registry.record_batch_adapters(rollout_id, 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) 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 0252a6948e..2907fa4157 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -102,6 +102,8 @@ 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 0f81c8799c..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) @@ -86,7 +89,25 @@ def convert_samples_to_train_data( 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] @@ -99,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) @@ -107,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: @@ -180,6 +223,10 @@ 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 diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py index ad3826628e..89955358cc 100644 --- a/miles/utils/adapter_config.py +++ b/miles/utils/adapter_config.py @@ -21,6 +21,13 @@ class AdapterRunConfig: 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" @@ -35,6 +42,12 @@ class AdapterRunConfig: 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: @@ -45,6 +58,8 @@ class AdapterRun: 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: @@ -60,6 +75,8 @@ def parse_adapter_run_yaml(path: Path) -> 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"), diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 003c5f213e..9f7c10272a 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -273,7 +273,7 @@ def add_train_arguments(parser): help="Whether to enable recompute loss function to save memory during training.", ) parser.add_argument( - "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" + "--log-probs-batch-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) parser.add_argument( "--indep-dp", @@ -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( @@ -1440,6 +1440,25 @@ def add_lora_arguments(parser): 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): @@ -2554,10 +2573,26 @@ def miles_validate_args(args): 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 " @@ -2565,6 +2600,43 @@ def miles_validate_args(args): "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" @@ -2735,7 +2807,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 index 1ce388ddc7..e574d3b2d3 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -56,6 +56,26 @@ def slot_lora_name(slot: int) -> str: 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" @@ -79,6 +99,9 @@ class AdapterRecord: slot: int config: Any 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 @@ -95,7 +118,7 @@ def __init__(self, max_adapters: int) -> None: self.free_slots: set[int] = set(range(max_adapters)) self.slot_versions: list[int] = [0] * max_adapters self.records: dict[str, AdapterRecord] = {} - self.batch_adapters: dict[int, list[str]] = {} + 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} @@ -172,23 +195,51 @@ def record_weight_update(self, names: list[str]) -> None: if record.state is AdapterState.PENDING: record.state = AdapterState.ACTIVE - def record_batch_adapters(self, rollout_id: int, names: list[str]) -> None: - self.batch_adapters[rollout_id] = list(names) - while len(self.batch_adapters) > MAX_BATCH_RECORDS: - self.batch_adapters.pop(next(iter(self.batch_adapters))) + 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]: - trained = [] - for name in self.batch_adapters.pop(rollout_id, []): + """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 = [] + for name, n_groups in record_entry["groups"].items(): record = self.records.get(name) - if record is not None and record.state in ( + 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 - trained.append(name) - return trained + record.accumulated_groups = 0 + stepped.append(name) + return stepped def set_step(self, name: str, step: int) -> None: if (record := self.find(name)) is not None: @@ -205,6 +256,7 @@ def view(self, record: AdapterRecord) -> AdapterRun: slot=record.slot, version=self.slot_versions[record.slot], step=record.step, + accumulated_groups=record.accumulated_groups, ) def active_adapters(self) -> dict[str, AdapterRun]: @@ -248,18 +300,78 @@ async def close(self) -> None: async def validate_adapter(self, name: str, config: Any) -> None: """Override to reject adapter registrations (raise ValueError).""" - def resolve_save_dir(self, name: str, config: Any) -> Any: - if config is None or not hasattr(config, "save"): - return config - if config.save is not None: + 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 - 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") - return replace(config, save=Path(self.args.save) / "adapters" / name) + + 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") + 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) - config = self.resolve_save_dir(name, config) result = self.registry.register(name, config) resolved = getattr(config, "save", None) if resolved is not None: @@ -281,7 +393,7 @@ async def free_slot(self, name: str) -> int: 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 window can be missed by the scheduler-side + 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. 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/utils/test_controller_backend.py b/tests/fast/utils/test_controller_backend.py index 34aca8a131..30ebdf6efe 100644 --- a/tests/fast/utils/test_controller_backend.py +++ b/tests/fast/utils/test_controller_backend.py @@ -10,21 +10,47 @@ import pytest from miles.utils.adapter_config import AdapterRunConfig -from miles.utils.multi_lora import AdapterRegistry, AdapterState, MultiLoRABackend, make_rid, parse_adapter - - -def make_args(max_adapters: int = 4, save: str | None = None) -> SimpleNamespace: - return SimpleNamespace(multi_lora_n_adapters=max_adapters, save=save) +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) -> MultiLoRABackend: - return MultiLoRABackend(make_args(max_adapters, save), "http://unused") +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) -> AdapterRunConfig: - return AdapterRunConfig( - rank=8, alpha=16, data="/d", save=save, input_key="text", label_key="label", rm_type="math" +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: @@ -112,23 +138,38 @@ def test_deregister_retires_but_keeps_serving_until_demoted(): assert registry.retire_adapters() == [] # idempotent -def test_batch_record_counts_steps_on_confirmation(): +# 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") - register_and_promote(registry, "B") + 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(7, ["A"]) - assert registry.step_count("A") == 0 # recorded, not yet trained - assert registry.mark_batch_trained(7) == ["A"] + 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") == 0 - assert registry.mark_batch_trained(7) == [] # record consumed + 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") - registry.record_batch_adapters(3, ["A"]) + 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 @@ -140,14 +181,51 @@ def test_batch_trained_counts_deregistered_adapter_until_freed(): def test_set_step_on_resume(): registry = AdapterRegistry(max_adapters=2) - registry.register("A", None) + registry.register("A", make_config()) registry.set_step("A", 40) - registry.record_batch_adapters(1, ["A"]) + 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_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)) + # 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 @@ -164,7 +242,7 @@ def test_deregister_holds_slot_until_free_slot(): @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-window misses); free_slot must + 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] = [] diff --git a/tests/fast/utils/test_controller_http.py b/tests/fast/utils/test_controller_http.py index eb401c3e1f..2340efd60e 100644 --- a/tests/fast/utils/test_controller_http.py +++ b/tests/fast/utils/test_controller_http.py @@ -88,7 +88,19 @@ async def router_handler(request): 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), router_url) + 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() From 5c4a4e84b89cc9421a47f290f720f1c580e4c7e5 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Wed, 15 Jul 2026 18:43:14 -0700 Subject: [PATCH 14/31] [fix] typo lol --- miles/utils/arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 9f7c10272a..421198731a 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -273,7 +273,7 @@ def add_train_arguments(parser): help="Whether to enable recompute loss function to save memory during training.", ) parser.add_argument( - "--log-probs-batch-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" + "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) parser.add_argument( "--indep-dp", From f115610364da53febc2da7f32a8925f103397ff9 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Wed, 15 Jul 2026 21:12:53 -0700 Subject: [PATCH 15/31] [fix] handle empty --- .../multi_lora/multi_lora_async_rollout.py | 10 +++++++--- examples/multi_lora/train_multi_lora_async.py | 18 +++++++++++++++++- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index ff36a72deb..0aea0be7c2 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -62,6 +62,10 @@ def group_sample_count(group: Group) -> int: EMPTY_BATCH_TIMEOUT_S = 30.0 +class EmptyBatchTimeoutError(RuntimeError): + """No trainable groups arrived before empty-wait timeout.""" + + 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 @@ -300,9 +304,9 @@ def get_groups( self.stale_dropped += len(dropped) self.staleness_values += dropped min_groups_per_pop = min_groups_per_dp_split(adapter.config.n_samples_per_prompt, dp_size) - available_groups = len(buffer) // min_groups_per_pop * min_groups_per_pop + 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, available_groups, remaining_allowed_groups) + 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)) @@ -366,7 +370,7 @@ async def collect_batch(args, worker: AsyncMultiLoRAWorker, snapshot: dict) -> T if collected and stalled_s > wait_s: break if not collected and stalled_s > empty_wait_s: - raise RuntimeError( + 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'])}" diff --git a/examples/multi_lora/train_multi_lora_async.py b/examples/multi_lora/train_multi_lora_async.py index e5435a2a07..350bc877b8 100644 --- a/examples/multi_lora/train_multi_lora_async.py +++ b/examples/multi_lora/train_multi_lora_async.py @@ -4,6 +4,9 @@ import logging from pathlib import Path +import ray + +from examples.multi_lora.multi_lora_async_rollout import EmptyBatchTimeoutError 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 @@ -18,6 +21,13 @@ 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 @@ -70,7 +80,13 @@ async def main(args): if not (post_update["active"] or post_update["retiring"]): continue - rollout_data = await rollout_manager.generate.remote(rollout_id) + 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. From 52b0410797e198d6665683515475c0ef3b33d4c1 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 00:42:52 -0700 Subject: [PATCH 16/31] [feat] use num_step instead of num_row --- examples/multi_lora/README.md | 16 +++---- examples/multi_lora/adapters/dapo_math.yaml | 2 +- examples/multi_lora/adapters/gsm8k.yaml | 2 +- .../multi_lora_data_source_async.py | 13 +++--- examples/multi_lora/service_smoke.py | 23 +++++++--- miles/backends/megatron_utils/actor.py | 4 +- miles/utils/adapter_config.py | 2 + miles/utils/multi_lora.py | 26 ++++++++++++ tests/fast/utils/test_controller_backend.py | 42 +++++++++++++++++++ 9 files changed, 107 insertions(+), 23 deletions(-) diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md index 1bac3d3900..e8fd255059 100644 --- a/examples/multi_lora/README.md +++ b/examples/multi_lora/README.md @@ -18,7 +18,7 @@ run_service.sh # service mode: idles for registrations (po 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, deregisters at num_row) +multi_lora_data_source_async.py # data source (reads controller, legacy num_row fallback) adapters/ gsm8k.yaml dapo_math.yaml @@ -65,13 +65,15 @@ Ray actor, pinned to the head node). - **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. -- The data source deregisters an adapter at `num_row`; the trainer's +- 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. The data source still supports legacy `num_row` + deregistration when configured. 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. TODO: revisit - num_row semantics (the tail means slightly fewer trained rows than - configured). + 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. @@ -90,7 +92,7 @@ Downloads `Qwen/Qwen3-4B`, `zhuzilin/dapo-math-17k`, and `zhuzilin/gsm8k`. bash examples/multi_lora/run_job.sh ``` -Registers the two adapters from CLI flags and trains until each hits its `num_row` +Registers the two adapters from CLI flags and trains until each hits its `num_step` (or `--num-rollout`), then exits. ## Multi-LoRA CLI flags @@ -113,7 +115,7 @@ data: /root/gsm8k/train.parquet input_key: messages label_key: label rm_type: math -num_row: 400 # stop adapter after N rows +num_step: 400 # stop adapter after N optimizer steps # optional: save, num_epoch, custom_rm_path, ... ``` diff --git a/examples/multi_lora/adapters/dapo_math.yaml b/examples/multi_lora/adapters/dapo_math.yaml index 135bcde61c..5003d82a14 100644 --- a/examples/multi_lora/adapters/dapo_math.yaml +++ b/examples/multi_lora/adapters/dapo_math.yaml @@ -6,4 +6,4 @@ data: /root/dapo-math-17k/dapo-math-17k.jsonl input_key: prompt label_key: label rm_type: deepscaler -num_row: 500 +num_step: 500 diff --git a/examples/multi_lora/adapters/gsm8k.yaml b/examples/multi_lora/adapters/gsm8k.yaml index 8c1185a88a..ff2798e214 100644 --- a/examples/multi_lora/adapters/gsm8k.yaml +++ b/examples/multi_lora/adapters/gsm8k.yaml @@ -6,4 +6,4 @@ data: /root/gsm8k/train.parquet input_key: messages label_key: label rm_type: math -num_row: 400 +num_step: 400 diff --git a/examples/multi_lora/multi_lora_data_source_async.py b/examples/multi_lora/multi_lora_data_source_async.py index b7ee6fa34d..b75b7deb34 100644 --- a/examples/multi_lora/multi_lora_data_source_async.py +++ b/examples/multi_lora/multi_lora_data_source_async.py @@ -1,4 +1,4 @@ -"""Round-robin per-adapter data source; deregisters adapters at num_row.""" +"""Round-robin per-adapter data source; legacy num_row-based deregistration.""" import copy import logging @@ -104,11 +104,12 @@ def get_samples(self, num_samples: int = 1) -> list[list[Sample]]: sample.reward_spec = reward_spec sample.metadata = {**config.metadata, **sample.metadata} - default_num_row = (getattr(config, "num_epoch", 1) or 1) * len(source.dataset) - num_row = config.num_row or default_num_row - if source.sample_group_index >= num_row and name not in snapshot["retiring"]: - logger.info(f"Adapter '{name}' reached num_row={num_row}, deregistering") - ray.get(get_multi_lora_controller().deregister_adapter.remote(name)) + if config.num_step is None: + default_num_row = (getattr(config, "num_epoch", 1) or 1) * len(source.dataset) + num_row = config.num_row or default_num_row + if source.sample_group_index >= num_row and name not in snapshot["retiring"]: + logger.info(f"Adapter '{name}' reached num_row={num_row}, deregistering") + ray.get(get_multi_lora_controller().deregister_adapter.remote(name)) return groups diff --git a/examples/multi_lora/service_smoke.py b/examples/multi_lora/service_smoke.py index bfc9c3f07f..9566b16669 100644 --- a/examples/multi_lora/service_smoke.py +++ b/examples/multi_lora/service_smoke.py @@ -87,6 +87,12 @@ def main() -> int: 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() @@ -108,25 +114,32 @@ def config(name: str) -> dict: print("phase 1: api reachable, no active adapters expected") client.wait_for("api reachable", lambda adapters: True) - print("phase 2: register smoke_a; expect promotion + training progress") + 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 3: register smoke_b mid-run; both must train") + 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 4: deregister smoke_a mid-run; smoke_b must keep training") + 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 5: re-register the name smoke_a (waits out cleanup, reuses slot)") + 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 6: deregister everything; service should drain to idle") + 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) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 69336232b8..c85c78f86d 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -547,9 +547,7 @@ def reconcile_adapters(self) -> None: 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. TODO: revisit num_row - semantics — the tail means an adapter trains on slightly fewer rows - than configured. + any partially accumulated gradients) is discarded. """ if not is_multi_lora_enabled(self.args): return diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py index 89955358cc..b57613e352 100644 --- a/miles/utils/adapter_config.py +++ b/miles/utils/adapter_config.py @@ -38,6 +38,7 @@ class AdapterRunConfig: custom_rm_path: str | None = None num_epoch: int | None = None + num_step: int | None = None num_row: int | None = None metadata: dict[str, Any] = field(default_factory=dict) @@ -84,6 +85,7 @@ def parse_adapter_run_yaml(path: Path) -> AdapterRunConfig: rm_type=raw.get("rm_type"), custom_rm_path=raw.get("custom_rm_path"), num_epoch=raw.get("num_epoch"), + num_step=raw.get("num_step"), num_row=raw.get("num_row"), metadata=raw.get("metadata") or {}, ) diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index e574d3b2d3..a886d4c94b 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -99,6 +99,8 @@ class AdapterRecord: 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 @@ -220,6 +222,7 @@ def mark_batch_trained(self, rollout_id: int) -> list[str]: 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 ( @@ -239,11 +242,24 @@ def mark_batch_trained(self, rollout_id: int) -> list[str]: 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 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) @@ -334,6 +350,16 @@ def resolve_adapter_config(self, name: str, config: Any) -> Any: 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_row is not None and (type(config.num_row) is not int or config.num_row <= 0): + raise ValueError(f"Adapter '{name}' num_row must be a positive integer") + if config.num_step is not None and config.num_row is not None: + logger.warning( + f"Adapter '{name}' sets both num_step and num_row; num_step takes precedence and num_row is ignored" + ) + elif config.num_step is None and config.num_row is not None: + logger.warning(f"Adapter '{name}' uses deprecated num_row={config.num_row}; prefer num_step") 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: diff --git a/tests/fast/utils/test_controller_backend.py b/tests/fast/utils/test_controller_backend.py index 30ebdf6efe..caffa52b17 100644 --- a/tests/fast/utils/test_controller_backend.py +++ b/tests/fast/utils/test_controller_backend.py @@ -189,6 +189,35 @@ def test_set_step_on_resume(): 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 @@ -222,10 +251,23 @@ async def test_register_rejects_bad_batch_shapes(tmp_path): 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_row must be a positive integer"): + await backend.register("H", make_config(num_row=0)) # A valid shape registers fine. await backend.register("OK", make_config(rollout_batch_size=8)) +@pytest.mark.asyncio +async def test_num_step_takes_precedence_over_num_row(tmp_path): + backend = make_backend(save=str(tmp_path)) + await backend.register("A", make_config(num_step=10, num_row=1000)) + config = backend.registry.records["A"].config + assert config.num_step == 10 + assert config.num_row == 1000 + + def test_deregister_holds_slot_until_free_slot(): registry = AdapterRegistry(max_adapters=2) register_and_promote(registry, "A") # slot 0 From d81bf523ff1faf5947cdf8a91849ce4d48ee0d69 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 01:33:50 -0700 Subject: [PATCH 17/31] [fix] missing import --- examples/multi_lora/multi_lora_async_rollout.py | 6 +----- examples/multi_lora/train_multi_lora_async.py | 2 +- miles/utils/multi_lora.py | 5 +++++ 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index 0aea0be7c2..19b89cc954 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -28,7 +28,7 @@ from miles.rollout.sglang_rollout import GenerateState, generate_and_rm_group, get_model_url from miles.utils.async_utils import run from miles.utils.misc import load_function -from miles.utils.multi_lora import min_groups_per_dp_split +from miles.utils.multi_lora import EmptyBatchTimeoutError, min_groups_per_dp_split from miles.utils.types import Sample logger = logging.getLogger(__name__) @@ -62,10 +62,6 @@ def group_sample_count(group: Group) -> int: EMPTY_BATCH_TIMEOUT_S = 30.0 -class EmptyBatchTimeoutError(RuntimeError): - """No trainable groups arrived before empty-wait timeout.""" - - 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 diff --git a/examples/multi_lora/train_multi_lora_async.py b/examples/multi_lora/train_multi_lora_async.py index 350bc877b8..d3b641623c 100644 --- a/examples/multi_lora/train_multi_lora_async.py +++ b/examples/multi_lora/train_multi_lora_async.py @@ -6,11 +6,11 @@ import ray -from examples.multi_lora.multi_lora_async_rollout import EmptyBatchTimeoutError 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 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 diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index a886d4c94b..05f2296268 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -22,6 +22,7 @@ __all__ = [ "AdapterRegistry", "AdapterState", + "EmptyBatchTimeoutError", "MultiLoRABackend", "MultiLoRAHTTPServer", "RID_SEPARATOR", @@ -38,6 +39,10 @@ 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) From 5f1e83f5995a07d159840d8432d5bfca4061dd04 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 01:50:16 -0700 Subject: [PATCH 18/31] [fix] smoke --- examples/multi_lora/service_smoke.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/examples/multi_lora/service_smoke.py b/examples/multi_lora/service_smoke.py index 9566b16669..1f93a9ec3f 100644 --- a/examples/multi_lora/service_smoke.py +++ b/examples/multi_lora/service_smoke.py @@ -24,15 +24,24 @@ def __init__(self, api_url: str, timeout_s: float): self.timeout_s = timeout_s self.http = httpx.Client(timeout=30.0) - def active_adapters(self) -> dict: + 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"]} + status["name"]: { + "slot": status["slot"], + "version": status["version"], + "step": status["step"], + "state": status["state"], + } for status in response.json()["adapters"] - if status["state"] == "ACTIVE" + 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}) @@ -57,9 +66,15 @@ def wait_for(self, description: str, predicate) -> dict: 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 adapters: name in adapters and adapters[name]["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: From 7245240581d7364b00f4eebaa2477e788228d8f9 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 02:23:00 -0700 Subject: [PATCH 19/31] [test] move tests to test dir --- .../fast/ray/rollout/test_multi_lora_batch_collection.py | 2 +- .../fast/ray/rollout/test_multi_lora_process_group.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename examples/multi_lora/tests/test_multi_lora_batch_assembler.py => tests/fast/ray/rollout/test_multi_lora_batch_collection.py (99%) rename examples/multi_lora/tests/test_multi_lora_async_rollout.py => tests/fast/ray/rollout/test_multi_lora_process_group.py (97%) diff --git a/examples/multi_lora/tests/test_multi_lora_batch_assembler.py b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py similarity index 99% rename from examples/multi_lora/tests/test_multi_lora_batch_assembler.py rename to tests/fast/ray/rollout/test_multi_lora_batch_collection.py index d3af21ee1f..a57181e978 100644 --- a/examples/multi_lora/tests/test_multi_lora_batch_assembler.py +++ b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py @@ -1,4 +1,4 @@ -"""Unit tests for batch collection (get_groups + collect_batch): +"""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.""" diff --git a/examples/multi_lora/tests/test_multi_lora_async_rollout.py b/tests/fast/ray/rollout/test_multi_lora_process_group.py similarity index 97% rename from examples/multi_lora/tests/test_multi_lora_async_rollout.py rename to tests/fast/ray/rollout/test_multi_lora_process_group.py index ae971bda9a..7118b5b667 100644 --- a/examples/multi_lora/tests/test_multi_lora_async_rollout.py +++ b/tests/fast/ray/rollout/test_multi_lora_process_group.py @@ -1,4 +1,4 @@ -"""Tests for the testable core of the multi-LoRA async rollout (process_group): +"""Tests for the multi-LoRA async rollout process_group core: keep-vs-recycle plus submission-time slot-version stamping.""" import examples.multi_lora.multi_lora_async_rollout as mod From 9aea7566caae617bf87860897dcd46f64b5f2691 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 14:55:12 -0700 Subject: [PATCH 20/31] [chore] improve metrics --- .../multi_lora/multi_lora_async_rollout.py | 109 +++++++++++++----- 1 file changed, 83 insertions(+), 26 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index 19b89cc954..f883fd4325 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -149,6 +149,81 @@ async def process_group( return result +class MultiLoRAWorkerMetrics: + """The worker's cross-batch metric state, kept out of its buffer + machinery: dynamic-filter drop counts, staleness drops, and per-adapter + reward accumulation flushed as a step mean. Has its own lock — the + producer thread records drops while the trainer thread flushes.""" + + def __init__(self) -> None: + self.lock = threading.Lock() + self.dynamic_filter_drop_counts: dict[str, int] = defaultdict(int) + self.stale_dropped = 0 + self.staleness_values: list[int] = [] + # Rewards of shipped samples, accumulated per adapter across train + # batches and flushed as a step-mean metric when the adapter steps. + self.reward_sums: dict[str, float] = defaultdict(float) + self.reward_counts: dict[str, int] = defaultdict(int) + + def record_dynamic_filter_drop(self, reason: str) -> None: + with self.lock: + self.dynamic_filter_drop_counts[reason] += 1 + + def record_stale_drops(self, staleness_values: list[int]) -> None: + with self.lock: + self.stale_dropped += len(staleness_values) + self.staleness_values += staleness_values + + def record_shipped_rewards(self, args, data: list[Group], step_names: list[str]) -> dict[str, float]: + """Accumulate the shipped batch's rewards per adapter; for adapters + stepping with this batch, flush the mean over their whole adapter + batch (accumulated across shipped batches, so it covers all + ``adapter_global_batch_size`` samples of the step, not just this + batch's slice). + + 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 + for sample in iter_group_samples(group): + self.reward_sums[name] += sample.get_reward_value(args) + self.reward_counts[name] += 1 + + metrics: dict[str, float] = {} + for name in step_names: + if (count := self.reward_counts.pop(name, 0)) > 0: + metrics[f"{name}/rollout/raw_reward/step_mean"] = self.reward_sums.pop(name) / count + metrics[f"{name}/rollout/raw_reward/step_n"] = count + return metrics + + def discard_adapter(self, name: str) -> None: + """Drop a retired adapter's partial reward accumulation.""" + with self.lock: + self.reward_sums.pop(name, None) + self.reward_counts.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() + metrics["perf/fully_async/stale_dropped"] = self.stale_dropped + if self.staleness_values: + metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(self.staleness_values) / len( + self.staleness_values + ) + metrics["perf/fully_async/stale_dropped_max_staleness"] = max(self.staleness_values) + self.stale_dropped = 0 + self.staleness_values = [] + return metrics + + class AsyncMultiLoRAWorker: """Background producer filling bounded per-adapter completed-group buffers; the collection loop pops from them via ``get_groups``.""" @@ -175,9 +250,7 @@ def __init__(self, args, data_source, generate_fn: GenerateFn, concurrency: int # Advances past every visited adapter, persisting across calls and # batches, so adapters are served round-robin. self.rotation: deque[str] = deque() - self.dynamic_filter_drop_counts: dict[str, int] = defaultdict(int) - self.stale_dropped = 0 - self.staleness_values: list[int] = [] + self.metrics = MultiLoRAWorkerMetrics() @classmethod def get_or_create(cls, args, data_source, generate_fn: GenerateFn, concurrency: int = None): @@ -241,8 +314,7 @@ async def process_and_enqueue(self, group: list[Sample]) -> None: filter_result = call_dynamic_filter(self.dynamic_filter, self.args, result) if not filter_result.keep: if filter_result.reason: - with self.buffer_lock: - self.dynamic_filter_drop_counts[filter_result.reason] += 1 + self.metrics.record_dynamic_filter_drop(filter_result.reason) return adapter_name = group_adapter_name(result) @@ -279,10 +351,12 @@ def get_groups( with self.buffer_lock: # Adapters retired at the last reconcile sync point: their buffered - # tail is discarded (base deregistration semantics). + # 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) @@ -297,8 +371,7 @@ def get_groups( adapter = adapters[name] buffer = self.buffers[name] if dropped := buffer.drop_stale(adapter.version, max_staleness): - self.stale_dropped += len(dropped) - self.staleness_values += dropped + self.metrics.record_stale_drops(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)) @@ -314,23 +387,6 @@ def get_groups( break # a full pass over rotation yielded nothing return popped, group_counts - def pop_metrics(self) -> dict[str, float]: - with self.buffer_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() - metrics["perf/fully_async/stale_dropped"] = self.stale_dropped - if self.staleness_values: - metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(self.staleness_values) / len( - self.staleness_values - ) - metrics["perf/fully_async/stale_dropped_max_staleness"] = max(self.staleness_values) - self.stale_dropped = 0 - self.staleness_values = [] - return metrics - async def collect_batch(args, worker: AsyncMultiLoRAWorker, snapshot: dict) -> TrainBatch: """Collect one train batch from the worker's buffers (same loop shape as @@ -443,7 +499,8 @@ async def generate_rollout_multi_lora_async( ) metrics = { - **worker.pop_metrics(), + **worker.metrics.pop_metrics(), + **worker.metrics.record_shipped_rewards(args, data, batch.step_names), "perf/fully_async/queue_length": queue_length, "perf/fully_async/batch_wait_time": time.time() - start_time, "perf/fully_async/batch_adapters": len(batch.group_counts), From c4bed1de37ddedcef519d53706362749d52fc095 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 14:55:17 -0700 Subject: [PATCH 21/31] [test] clean up some tests --- .../test_multi_lora_batch_collection.py | 16 +--- .../rollout/test_multi_lora_process_group.py | 77 +++---------------- tests/fast/utils/test_controller_backend.py | 9 --- 3 files changed, 13 insertions(+), 89 deletions(-) diff --git a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py index a57181e978..c88940b7ab 100644 --- a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py +++ b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py @@ -14,6 +14,7 @@ from examples.multi_lora.multi_lora_async_rollout import ( AsyncMultiLoRAWorker, GroupBuffer, + MultiLoRAWorkerMetrics, collect_batch, group_adapter_name, ) @@ -41,9 +42,7 @@ def make_worker(args=None) -> AsyncMultiLoRAWorker: worker.buffers = defaultdict(GroupBuffer) worker.rotation = deque() worker.dynamic_filter = None - worker.dynamic_filter_drop_counts = defaultdict(int) - worker.stale_dropped = 0 - worker.staleness_values = [] + worker.metrics = MultiLoRAWorkerMetrics() return worker @@ -190,15 +189,6 @@ def test_cursor_persists_across_batches(): assert len(worker.buffers["B"]) == 2 -def test_n_samples_multiple_of_dp_pops_single_groups(): - worker = make_worker(make_args(multi_lora_dp_size=4, global_batch_size=8)) - a = adapter_run("A", 0, rollout_batch_size=2, n_samples_per_prompt=8) # 8 % 4 == 0 - buffer_groups(worker, a, count=1) - batch = collect(worker, snapshot_of(a)) # 8 samples = target - assert batch.group_counts == {"A": 1} - assert sum(1 for g in batch.groups for _ in g) == 8 - - def test_retiring_adapter_remains_selectable_until_retired(): """RETIRING adapters keep serving until the reconcile sync point (base deregistration semantics): buffered groups stay poppable.""" @@ -228,7 +218,7 @@ def test_stale_buffered_groups_are_dropped(): 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 worker.stale_dropped == 2 + assert worker.metrics.stale_dropped == 2 assert batch.group_counts == {"A": 1} diff --git a/tests/fast/ray/rollout/test_multi_lora_process_group.py b/tests/fast/ray/rollout/test_multi_lora_process_group.py index 7118b5b667..ff83c7e3a4 100644 --- a/tests/fast/ray/rollout/test_multi_lora_process_group.py +++ b/tests/fast/ray/rollout/test_multi_lora_process_group.py @@ -1,5 +1,6 @@ -"""Tests for the multi-LoRA async rollout process_group core: -keep-vs-recycle plus submission-time slot-version stamping.""" +"""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 @@ -28,52 +29,9 @@ def __init__(self, versions: dict[str, int]) -> None: def bump(self, name: str, to: int) -> None: self.versions[name] = to - async def get_all(self) -> dict[str, FakeAdapterView]: - return {name: FakeAdapterView(version) for name, version in self.versions.items()} - async def get(self, adapter_name: str) -> FakeAdapterView | None: - return (await self.get_all()).get(adapter_name) - - -def group(adapter: str = "A", slot: int = 0) -> list[Sample]: - return [Sample(prompt="p", adapter=AdapterRef(adapter, slot))] - - -async def gen_completed(args, group, sampling_params): - for s in group: - s.status = Sample.Status.COMPLETED - return group - - -@pytest.mark.asyncio -async def test_process_group_keeps_completed(): - ds = FakeDataSource() - g = group("A") - result = await process_group(None, g, {}, gen_completed, ds) - - assert result is g - assert ds.added == [] - - -@pytest.mark.asyncio -@pytest.mark.xfail( - reason="Re-queuing aborted groups is not wired up yet (the per-adapter " - "source is read-only); planned for a future PR. This test pins the " - "intended end-state behavior.", - strict=True, -) -async def test_process_group_recycles_aborted(): - async def gen(args, group, sampling_params): - for s in group: - s.status = Sample.Status.ABORTED - return group - - ds = FakeDataSource() - g = group("A") - result = await process_group(None, g, {}, gen, ds) - - assert result is None - assert len(ds.added) == 1 + version = self.versions.get(adapter_name) + return FakeAdapterView(version) if version is not None else None @pytest.mark.asyncio @@ -83,29 +41,14 @@ async def test_process_group_stamps_submission_version(monkeypatch): async def gen(args, group, sampling_params): cache.bump("A", 7) # update lands mid-generation - return await gen_completed(args, group, sampling_params) + for s in group: + s.status = Sample.Status.COMPLETED + return group monkeypatch.setattr(mod, "AdaptersCache", lambda: cache) - ds = FakeDataSource() - g = group("A") - result = await process_group(None, g, {}, gen, ds) + 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 - - -@pytest.mark.asyncio -async def test_process_group_no_adapter_skips_stamp(monkeypatch): - class FailingCache: - async def get(self, adapter_name): - raise AssertionError("adapters cache should not be queried for adapter-less group") - - monkeypatch.setattr(mod, "AdaptersCache", FailingCache) - - ds = FakeDataSource() - g = [Sample(prompt="p", adapter=None)] - result = await process_group(None, g, {}, gen_completed, ds) - - assert result is g - assert "slot_version" not in g[0].metadata diff --git a/tests/fast/utils/test_controller_backend.py b/tests/fast/utils/test_controller_backend.py index caffa52b17..b3a155dd45 100644 --- a/tests/fast/utils/test_controller_backend.py +++ b/tests/fast/utils/test_controller_backend.py @@ -259,15 +259,6 @@ async def test_register_rejects_bad_batch_shapes(tmp_path): await backend.register("OK", make_config(rollout_batch_size=8)) -@pytest.mark.asyncio -async def test_num_step_takes_precedence_over_num_row(tmp_path): - backend = make_backend(save=str(tmp_path)) - await backend.register("A", make_config(num_step=10, num_row=1000)) - config = backend.registry.records["A"].config - assert config.num_step == 10 - assert config.num_row == 1000 - - def test_deregister_holds_slot_until_free_slot(): registry = AdapterRegistry(max_adapters=2) register_and_promote(registry, "A") # slot 0 From 2100479890809903bbb45a8933ad12f03c84cfcb Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 15:36:32 -0700 Subject: [PATCH 22/31] [feat] improve metrics --- .../multi_lora/multi_lora_async_rollout.py | 94 +++++++++++++------ miles/ray/rollout/metrics.py | 19 ---- miles/utils/tracking_utils/base.py | 14 ++- 3 files changed, 78 insertions(+), 49 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index f883fd4325..61341810aa 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -29,6 +29,7 @@ from miles.utils.async_utils import run 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__) @@ -152,34 +153,41 @@ async def process_group( class MultiLoRAWorkerMetrics: """The worker's cross-batch metric state, kept out of its buffer machinery: dynamic-filter drop counts, staleness drops, and per-adapter - reward accumulation flushed as a step mean. Has its own lock — the - producer thread records drops while the trainer thread flushes.""" + sample stats flushed when the adapter's optimizer step completes. Has its + own lock — the producer thread records drops while the trainer thread + flushes.""" def __init__(self) -> None: self.lock = threading.Lock() self.dynamic_filter_drop_counts: dict[str, int] = defaultdict(int) self.stale_dropped = 0 self.staleness_values: list[int] = [] - # Rewards of shipped samples, accumulated per adapter across train - # batches and flushed as a step-mean metric when the adapter steps. - self.reward_sums: dict[str, float] = defaultdict(float) - self.reward_counts: dict[str, int] = defaultdict(int) + # Shipped-sample stats and stale drops, accumulated per adapter across + # train batches and flushed when the adapter steps. + self.step_reward_sums: dict[str, float] = defaultdict(float) + self.step_response_len_sums: dict[str, float] = defaultdict(float) + self.step_sample_counts: dict[str, int] = defaultdict(int) + self.step_staleness_values: dict[str, list[int]] = 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, staleness_values: list[int]) -> None: + def record_stale_drops(self, name: str, staleness_values: list[int]) -> None: with self.lock: self.stale_dropped += len(staleness_values) self.staleness_values += staleness_values + self.step_staleness_values[name] += staleness_values - def record_shipped_rewards(self, args, data: list[Group], step_names: list[str]) -> dict[str, float]: - """Accumulate the shipped batch's rewards per adapter; for adapters - stepping with this batch, flush the mean over their whole adapter - batch (accumulated across shipped batches, so it covers all - ``adapter_global_batch_size`` samples of the step, not just this - batch's slice). + def record_shipped_samples(self, args, data: list[Group], step_names: list[str]) -> dict[str, dict[str, float]]: + """Accumulate the shipped batch's rewards and response lengths per + adapter; for adapters stepping with this batch, flush means over their + whole adapter batch (accumulated across shipped batches, so each mean + covers all ``adapter_global_batch_size`` samples of the step, not just + this batch's slice). Returns {adapter name: flushed metrics}. + + ``n`` should always equal the adapter's ``adapter_global_batch_size``; + any deviation means batch accounting drifted. Counted at ship time, not train commit: a failed train call aborts the run anyway, so the distinction has no practical effect. @@ -190,21 +198,32 @@ def record_shipped_rewards(self, args, data: list[Group], step_names: list[str]) if name is None: continue for sample in iter_group_samples(group): - self.reward_sums[name] += sample.get_reward_value(args) - self.reward_counts[name] += 1 + self.step_reward_sums[name] += sample.get_reward_value(args) + self.step_response_len_sums[name] += sample.effective_response_length + self.step_sample_counts[name] += 1 - metrics: dict[str, float] = {} + flushed: dict[str, dict[str, float]] = {} for name in step_names: - if (count := self.reward_counts.pop(name, 0)) > 0: - metrics[f"{name}/rollout/raw_reward/step_mean"] = self.reward_sums.pop(name) / count - metrics[f"{name}/rollout/raw_reward/step_n"] = count - return metrics + if (count := self.step_sample_counts.pop(name, 0)) > 0: + flushed[name] = { + "rollout/raw_reward/mean": self.step_reward_sums.pop(name) / count, + "rollout/response_len/mean": self.step_response_len_sums.pop(name) / count, + "rollout/n": count, + } + staleness = self.step_staleness_values.pop(name, []) + flushed[name]["rollout/stale_dropped"] = len(staleness) + if staleness: + flushed[name]["rollout/stale_dropped_avg_staleness"] = sum(staleness) / len(staleness) + flushed[name]["rollout/stale_dropped_max_staleness"] = max(staleness) + return flushed def discard_adapter(self, name: str) -> None: - """Drop a retired adapter's partial reward accumulation.""" + """Drop a retired adapter's partial step accumulation.""" with self.lock: - self.reward_sums.pop(name, None) - self.reward_counts.pop(name, None) + self.step_reward_sums.pop(name, None) + self.step_response_len_sums.pop(name, None) + self.step_sample_counts.pop(name, None) + self.step_staleness_values.pop(name, None) def pop_metrics(self) -> dict[str, float]: with self.lock: @@ -327,6 +346,11 @@ 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]]: @@ -371,7 +395,7 @@ def get_groups( adapter = adapters[name] buffer = self.buffers[name] if dropped := buffer.drop_stale(adapter.version, max_staleness): - self.metrics.record_stale_drops(dropped) + 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)) @@ -498,15 +522,27 @@ async def generate_rollout_multi_lora_async( 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. + if flushed := worker.metrics.record_shipped_samples(args, data, batch.step_names): + queue_sizes = worker.queue_sizes() + for name, step_metrics in flushed.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()} + log_dict[f"{name}/rollout/queue_length"] = queue_sizes.get(name, 0) + tracking.log(args, log_dict, step_key=step_key) + metrics = { **worker.metrics.pop_metrics(), - **worker.metrics.record_shipped_rewards(args, data, batch.step_names), "perf/fully_async/queue_length": queue_length, "perf/fully_async/batch_wait_time": time.time() - start_time, - "perf/fully_async/batch_adapters": len(batch.group_counts), - "perf/fully_async/batch_prompt_groups": len(data), - "perf/fully_async/batch_samples": sum(group_sample_count(group) for group in data), - "perf/fully_async/batch_step_count": len(batch.step_names), + "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), } return RolloutFnTrainOutput(samples=data, metrics=metrics) diff --git a/miles/ray/rollout/metrics.py b/miles/ray/rollout/metrics.py index 9f90a2e537..9ee7ae41e9 100644 --- a/miles/ray/rollout/metrics.py +++ b/miles/ray/rollout/metrics.py @@ -12,7 +12,6 @@ has_repetition, ) from miles.utils.misc import load_function -from miles.utils.multi_lora import is_multi_lora_enabled from miles.utils.tracking_utils import tracking from miles.utils.types import Sample @@ -52,21 +51,6 @@ def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] return log_dict -def _compute_per_adapter_metrics(args, samples: list[Sample]) -> dict: - """Compute reward and response length metrics grouped by adapter name.""" - by_adapter = group_by(samples, lambda s: s.adapter.name if s.adapter else None) - log_dict = {} - for name, adapter_samples in by_adapter.items(): - if name is None: - continue - rewards = [s.get_reward_value(args) for s in adapter_samples] - response_lengths = [s.effective_response_length for s in adapter_samples] - prefix = f"{name}/rollout/" - log_dict |= dict_add_prefix(compute_statistics(rewards), f"{prefix}raw_reward/") - log_dict |= dict_add_prefix(compute_statistics(response_lengths), f"{prefix}response_len/") - return log_dict - - def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): if (x := args.custom_rollout_log_function_path) is not None: custom_log_func = load_function(x) @@ -79,9 +63,6 @@ def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_t log_dict = {**(rollout_extra_metrics or {})} log_dict |= dict_add_prefix(_compute_metrics_from_samples(args, samples), "rollout/") log_dict |= dict_add_prefix(_compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/") - - if is_multi_lora_enabled(args): - log_dict |= _compute_per_adapter_metrics(args, samples) logger.info(f"perf {rollout_id}: {log_dict}") step = compute_rollout_step(args, rollout_id) log_dict["rollout/step"] = step diff --git a/miles/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index f1f67dc7dc..040288388a 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -37,6 +37,9 @@ def finish(self) -> None: ... class WandbBackend(TrackingBackend): # Delegates to the existing ``wandb_utils`` helpers. + def __init__(self) -> None: + self._defined_step_keys: set[str] = set() + def init(self, args, *, primary: bool = True, **kwargs) -> None: from . import wandb_utils @@ -45,9 +48,18 @@ def init(self, args, *, primary: bool = True, **kwargs) -> None: else: wandb_utils.init_wandb_secondary(args, **kwargs) - def log(self, metrics: dict[str, Any], step: int | None = None, **kwargs) -> None: + def log(self, metrics: dict[str, Any], step: int | None = None, *, step_key: str | None = None, **kwargs) -> None: import wandb + if step_key is not None and step_key not in self._defined_step_keys: + # Declare the axis on first sight so the key's family plots + # against it instead of wandb's global step. Static families + # (train/rollout/eval) are already declared at init; this makes + # runtime axes (e.g. a per-adapter "{name}/step") work the same. + wandb.define_metric(step_key) + if "/" in step_key: + wandb.define_metric(f"{step_key.rsplit('/', 1)[0]}/*", step_metric=step_key) + self._defined_step_keys.add(step_key) wandb.log(metrics) def finish(self) -> None: From d972436b98090d7429cf62e1de7f658dfea79e38 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 16:07:34 -0700 Subject: [PATCH 23/31] [fix] metrics for queue don't need step --- .../multi_lora/multi_lora_async_rollout.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index 61341810aa..ea009e7c28 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -478,7 +478,7 @@ async def generate_rollout_multi_lora_async( state = GenerateState(args) worker = AsyncMultiLoRAWorker.get_or_create(args, data_source, generate_fn) start_time = time.time() - queue_length = worker.queue_size() + 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 @@ -526,18 +526,16 @@ async def generate_rollout_multi_lora_async( # ({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. - if flushed := worker.metrics.record_shipped_samples(args, data, batch.step_names): - queue_sizes = worker.queue_sizes() - for name, step_metrics in flushed.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()} - log_dict[f"{name}/rollout/queue_length"] = queue_sizes.get(name, 0) - tracking.log(args, log_dict, step_key=step_key) + for name, step_metrics in worker.metrics.record_shipped_samples(args, data, batch.step_names).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) metrics = { **worker.metrics.pop_metrics(), - "perf/fully_async/queue_length": queue_length, + "perf/fully_async/queue_length": sum(queue_sizes.values()), + **{f"perf/fully_async/queue_length/{name}": size for name, size in queue_sizes.items()}, "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), From ce6c2dd1947630b3210b46004d52e1d706d3b3da Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 16:07:39 -0700 Subject: [PATCH 24/31] [fix] metrics glob expansion --- miles/utils/tracking_utils/base.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/miles/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index 040288388a..a522400ccc 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -38,7 +38,7 @@ class WandbBackend(TrackingBackend): # Delegates to the existing ``wandb_utils`` helpers. def __init__(self) -> None: - self._defined_step_keys: set[str] = set() + self._defined_metrics: set[str] = set() def init(self, args, *, primary: bool = True, **kwargs) -> None: from . import wandb_utils @@ -51,15 +51,21 @@ def init(self, args, *, primary: bool = True, **kwargs) -> None: def log(self, metrics: dict[str, Any], step: int | None = None, *, step_key: str | None = None, **kwargs) -> None: import wandb - if step_key is not None and step_key not in self._defined_step_keys: - # Declare the axis on first sight so the key's family plots - # against it instead of wandb's global step. Static families - # (train/rollout/eval) are already declared at init; this makes - # runtime axes (e.g. a per-adapter "{name}/step") work the same. - wandb.define_metric(step_key) - if "/" in step_key: - wandb.define_metric(f"{step_key.rsplit('/', 1)[0]}/*", step_metric=step_key) - self._defined_step_keys.add(step_key) + if step_key is not None: + # Pin every logged key to its axis with an exact-name definition. + # Glob definitions (like the "rollout/*" ones at init) are not + # expanded client-side anymore; the raw glob is sent to the server + # for expansion, which not every server version supports + # (wandb#11533) — charts then silently fall back to the global + # step axis. Exact names always work, including for axes that + # only exist at runtime (e.g. a per-adapter "{name}/step"). + if step_key not in self._defined_metrics: + wandb.define_metric(step_key) + self._defined_metrics.add(step_key) + for key in metrics: + if key != step_key and key not in self._defined_metrics: + wandb.define_metric(key, step_metric=step_key) + self._defined_metrics.add(key) wandb.log(metrics) def finish(self) -> None: From 54e13f895d59b537cd13f4847921e343c291d319 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 16:30:08 -0700 Subject: [PATCH 25/31] [fix] steps --- miles/utils/tracking_utils/base.py | 22 +++++++--------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/miles/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index a522400ccc..8ed4a0d312 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -38,7 +38,7 @@ class WandbBackend(TrackingBackend): # Delegates to the existing ``wandb_utils`` helpers. def __init__(self) -> None: - self._defined_metrics: set[str] = set() + self._defined_step_keys: set[str] = set() def init(self, args, *, primary: bool = True, **kwargs) -> None: from . import wandb_utils @@ -51,21 +51,13 @@ def init(self, args, *, primary: bool = True, **kwargs) -> None: def log(self, metrics: dict[str, Any], step: int | None = None, *, step_key: str | None = None, **kwargs) -> None: import wandb - if step_key is not None: - # Pin every logged key to its axis with an exact-name definition. - # Glob definitions (like the "rollout/*" ones at init) are not - # expanded client-side anymore; the raw glob is sent to the server - # for expansion, which not every server version supports - # (wandb#11533) — charts then silently fall back to the global - # step axis. Exact names always work, including for axes that + if step_key is not None and step_key not in self._defined_step_keys: + # Same glob pattern as _init_wandb_common, for axes whose names # only exist at runtime (e.g. a per-adapter "{name}/step"). - if step_key not in self._defined_metrics: - wandb.define_metric(step_key) - self._defined_metrics.add(step_key) - for key in metrics: - if key != step_key and key not in self._defined_metrics: - wandb.define_metric(key, step_metric=step_key) - self._defined_metrics.add(key) + wandb.define_metric(step_key) + if "/" in step_key: + wandb.define_metric(f"{step_key.rsplit('/', 1)[0]}/*", step_metric=step_key) + self._defined_step_keys.add(step_key) wandb.log(metrics) def finish(self) -> None: From 07d9e9b7873044f397b9af202a1cc81844d5f2c5 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 16:31:42 -0700 Subject: [PATCH 26/31] [chore] minor naming --- examples/multi_lora/run_job.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/multi_lora/run_job.sh b/examples/multi_lora/run_job.sh index 80a29da846..c111016811 100755 --- a/examples/multi_lora/run_job.sh +++ b/examples/multi_lora/run_job.sh @@ -89,6 +89,6 @@ ray job submit --address="http://127.0.0.1:8265" \ # --use-wandb \ # --wandb-host https://wandb.ai/ \ # --wandb-team staging \ - # --wandb-project miles-multilora \ + # --wandb-project miles-multi-lora \ # --wandb-group qwen3-4B From 031c086eb1e23c11fe58976acbe1ebad09644b8a Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 16:56:52 -0700 Subject: [PATCH 27/31] Revert "[fix] steps" This reverts commit 54e13f895d59b537cd13f4847921e343c291d319. --- miles/utils/tracking_utils/base.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/miles/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index 8ed4a0d312..a522400ccc 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -38,7 +38,7 @@ class WandbBackend(TrackingBackend): # Delegates to the existing ``wandb_utils`` helpers. def __init__(self) -> None: - self._defined_step_keys: set[str] = set() + self._defined_metrics: set[str] = set() def init(self, args, *, primary: bool = True, **kwargs) -> None: from . import wandb_utils @@ -51,13 +51,21 @@ def init(self, args, *, primary: bool = True, **kwargs) -> None: def log(self, metrics: dict[str, Any], step: int | None = None, *, step_key: str | None = None, **kwargs) -> None: import wandb - if step_key is not None and step_key not in self._defined_step_keys: - # Same glob pattern as _init_wandb_common, for axes whose names + if step_key is not None: + # Pin every logged key to its axis with an exact-name definition. + # Glob definitions (like the "rollout/*" ones at init) are not + # expanded client-side anymore; the raw glob is sent to the server + # for expansion, which not every server version supports + # (wandb#11533) — charts then silently fall back to the global + # step axis. Exact names always work, including for axes that # only exist at runtime (e.g. a per-adapter "{name}/step"). - wandb.define_metric(step_key) - if "/" in step_key: - wandb.define_metric(f"{step_key.rsplit('/', 1)[0]}/*", step_metric=step_key) - self._defined_step_keys.add(step_key) + if step_key not in self._defined_metrics: + wandb.define_metric(step_key) + self._defined_metrics.add(step_key) + for key in metrics: + if key != step_key and key not in self._defined_metrics: + wandb.define_metric(key, step_metric=step_key) + self._defined_metrics.add(key) wandb.log(metrics) def finish(self) -> None: From 8fa298a5f3a399e24c3ae017d7d7b6c73dcd2bb9 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 18:37:17 -0700 Subject: [PATCH 28/31] [fix] metric naming --- .../multi_lora/multi_lora_async_rollout.py | 76 +++++++++++-------- .../test_multi_lora_batch_collection.py | 3 +- 2 files changed, 44 insertions(+), 35 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index ea009e7c28..241efe81a2 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -152,22 +152,21 @@ async def process_group( class MultiLoRAWorkerMetrics: """The worker's cross-batch metric state, kept out of its buffer - machinery: dynamic-filter drop counts, staleness drops, and per-adapter - sample stats flushed when the adapter's optimizer step completes. Has its - own lock — the producer thread records drops while the trainer thread - flushes.""" + 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) - self.stale_dropped = 0 - self.staleness_values: list[int] = [] - # Shipped-sample stats and stale drops, accumulated per adapter across - # train batches and flushed when the adapter steps. + # Staleness of dropped groups per adapter, drained every batch. + self.staleness_values: dict[str, list[int]] = defaultdict(list) + # Shipped-sample stats, accumulated per adapter across train batches + # and flushed as step means when the adapter steps. self.step_reward_sums: dict[str, float] = defaultdict(float) self.step_response_len_sums: dict[str, float] = defaultdict(float) self.step_sample_counts: dict[str, int] = defaultdict(int) - self.step_staleness_values: dict[str, list[int]] = defaultdict(list) def record_dynamic_filter_drop(self, reason: str) -> None: with self.lock: @@ -175,20 +174,24 @@ def record_dynamic_filter_drop(self, reason: str) -> None: def record_stale_drops(self, name: str, staleness_values: list[int]) -> None: with self.lock: - self.stale_dropped += len(staleness_values) - self.staleness_values += staleness_values - self.step_staleness_values[name] += staleness_values + self.staleness_values[name] += staleness_values - def record_shipped_samples(self, args, data: list[Group], step_names: list[str]) -> dict[str, dict[str, float]]: + 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 means over their whole adapter batch (accumulated across shipped batches, so each mean covers all ``adapter_global_batch_size`` samples of the step, not just this batch's slice). Returns {adapter name: flushed metrics}. - ``n`` should always equal the adapter's ``adapter_global_batch_size``; - any deviation means batch accounting drifted. - Counted at ship time, not train commit: a failed train call aborts the run anyway, so the distinction has no practical effect. """ @@ -205,16 +208,16 @@ def record_shipped_samples(self, args, data: list[Group], step_names: list[str]) flushed: dict[str, dict[str, float]] = {} for name in step_names: if (count := self.step_sample_counts.pop(name, 0)) > 0: + expected = adapters[name].config.adapter_global_batch_size + if count != expected: + logger.warning( + f"Adapter '{name}' stepped with {count} shipped samples, expected " + f"adapter_global_batch_size={expected}; batch accounting drifted" + ) flushed[name] = { "rollout/raw_reward/mean": self.step_reward_sums.pop(name) / count, "rollout/response_len/mean": self.step_response_len_sums.pop(name) / count, - "rollout/n": count, } - staleness = self.step_staleness_values.pop(name, []) - flushed[name]["rollout/stale_dropped"] = len(staleness) - if staleness: - flushed[name]["rollout/stale_dropped_avg_staleness"] = sum(staleness) / len(staleness) - flushed[name]["rollout/stale_dropped_max_staleness"] = max(staleness) return flushed def discard_adapter(self, name: str) -> None: @@ -223,7 +226,7 @@ def discard_adapter(self, name: str) -> None: self.step_reward_sums.pop(name, None) self.step_response_len_sums.pop(name, None) self.step_sample_counts.pop(name, None) - self.step_staleness_values.pop(name, None) + self.staleness_values.pop(name, None) def pop_metrics(self) -> dict[str, float]: with self.lock: @@ -232,14 +235,6 @@ def pop_metrics(self) -> dict[str, float]: for reason, count in self.dynamic_filter_drop_counts.items() } self.dynamic_filter_drop_counts.clear() - metrics["perf/fully_async/stale_dropped"] = self.stale_dropped - if self.staleness_values: - metrics["perf/fully_async/stale_dropped_avg_staleness"] = sum(self.staleness_values) / len( - self.staleness_values - ) - metrics["perf/fully_async/stale_dropped_max_staleness"] = max(self.staleness_values) - self.stale_dropped = 0 - self.staleness_values = [] return metrics @@ -526,22 +521,37 @@ async def generate_rollout_multi_lora_async( # ({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).items(): + 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()), - **{f"perf/fully_async/queue_length/{name}": size for name, size in queue_sizes.items()}, + "perf/fully_async/stale_dropped": len(all_staleness), + # Adapter-prefixed keys land in the adapter's dashboard section; the + # exact-name define pins them to this call's rollout/step axis (unlike + # the {name}/step-axis metrics, which ship in their own step-keyed + # calls above). + **{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) diff --git a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py index c88940b7ab..7cd66e2a60 100644 --- a/tests/fast/ray/rollout/test_multi_lora_batch_collection.py +++ b/tests/fast/ray/rollout/test_multi_lora_batch_collection.py @@ -218,8 +218,7 @@ def test_stale_buffered_groups_are_dropped(): 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 worker.metrics.stale_dropped == 2 - assert batch.group_counts == {"A": 1} + assert batch.group_counts == {"A": 1} # only the fresh group ships def test_empty_collection_times_out_instead_of_spinning_forever(): From 75d753879db13b9260d0ab811019c0dc166289e5 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Thu, 16 Jul 2026 19:23:49 -0700 Subject: [PATCH 29/31] [feat] improve adapter metrics + adapter metric use real adapter steps --- .../multi_lora/multi_lora_async_rollout.py | 104 ++++++++++++------ examples/multi_lora/train_multi_lora_async.py | 5 +- miles/utils/multi_lora.py | 20 ++++ miles/utils/tracking_utils/base.py | 42 ++++--- miles/utils/tracking_utils/tracking.py | 12 ++ 5 files changed, 131 insertions(+), 52 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index 241efe81a2..8561827f83 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -27,6 +27,7 @@ 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 @@ -162,11 +163,17 @@ def __init__(self) -> None: 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 stats, accumulated per adapter across train batches - # and flushed as step means when the adapter steps. - self.step_reward_sums: dict[str, float] = defaultdict(float) - self.step_response_len_sums: dict[str, float] = defaultdict(float) - self.step_sample_counts: dict[str, int] = defaultdict(int) + # 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: @@ -187,10 +194,10 @@ 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 means over their - whole adapter batch (accumulated across shipped batches, so each mean - covers all ``adapter_global_batch_size`` samples of the step, not just - this batch's slice). Returns {adapter name: flushed metrics}. + 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. @@ -200,32 +207,62 @@ def record_shipped_samples( name = group_adapter_name(group) if name is None: continue + group_rewards = [] for sample in iter_group_samples(group): - self.step_reward_sums[name] += sample.get_reward_value(args) - self.step_response_len_sums[name] += sample.effective_response_length - self.step_sample_counts[name] += 1 + 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: - if (count := self.step_sample_counts.pop(name, 0)) > 0: - expected = adapters[name].config.adapter_global_batch_size - if count != expected: - logger.warning( - f"Adapter '{name}' stepped with {count} shipped samples, expected " - f"adapter_global_batch_size={expected}; batch accounting drifted" - ) - flushed[name] = { - "rollout/raw_reward/mean": self.step_reward_sums.pop(name) / count, - "rollout/response_len/mean": self.step_response_len_sums.pop(name) / count, - } + 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_reward_sums.pop(name, None) - self.step_response_len_sums.pop(name, None) - self.step_sample_counts.pop(name, None) + 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]: @@ -533,12 +570,11 @@ async def generate_rollout_multi_lora_async( **worker.metrics.pop_metrics(), "perf/fully_async/queue_length": sum(queue_sizes.values()), "perf/fully_async/stale_dropped": len(all_staleness), - # Adapter-prefixed keys land in the adapter's dashboard section; the - # exact-name define pins them to this call's rollout/step axis (unlike - # the {name}/step-axis metrics, which ship in their own step-keyed - # calls above). - **{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}, + # Per-adapter cycle metrics live in their own {name}_perf namespace: + # a namespace maps to exactly one x-axis, and these ride rollout/step + # while the {name}/ namespace is on the adapter's step axis. + **{f"{name}_perf/queue_length": size for name, size in queue_sizes.items()}, + **{f"{name}_perf/stale_dropped": len(stale_drops.get(name, [])) for name in adapters}, "perf/fully_async/batch_wait_time": time.time() - start_time, "perf/fully_async/batch_n_adapters": len(batch.group_counts), "perf/fully_async/batch_n_groups": len(data), @@ -550,8 +586,8 @@ async def generate_rollout_multi_lora_async( 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) + 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) diff --git a/examples/multi_lora/train_multi_lora_async.py b/examples/multi_lora/train_multi_lora_async.py index d3b641623c..7aa45a875e 100644 --- a/examples/multi_lora/train_multi_lora_async.py +++ b/examples/multi_lora/train_multi_lora_async.py @@ -10,7 +10,7 @@ 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 +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 @@ -61,6 +61,9 @@ async def main(args): 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.") diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index 05f2296268..65ed99d2d3 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -26,6 +26,7 @@ "MultiLoRABackend", "MultiLoRAHTTPServer", "RID_SEPARATOR", + "define_new_adapter_metrics", "is_multi_lora_enabled", "make_rid", "parse_adapter", @@ -47,6 +48,25 @@ 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. + + 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}" diff --git a/miles/utils/tracking_utils/base.py b/miles/utils/tracking_utils/base.py index a522400ccc..a3b3651fb0 100644 --- a/miles/utils/tracking_utils/base.py +++ b/miles/utils/tracking_utils/base.py @@ -32,13 +32,20 @@ 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_metrics: set[str] = set() + self._defined_step_key_groups: set[tuple[str, str]] = set() def init(self, args, *, primary: bool = True, **kwargs) -> None: from . import wandb_utils @@ -48,26 +55,23 @@ def init(self, args, *, primary: bool = True, **kwargs) -> None: else: wandb_utils.init_wandb_secondary(args, **kwargs) - def log(self, metrics: dict[str, Any], step: int | None = None, *, step_key: str | None = None, **kwargs) -> None: + def log(self, metrics: dict[str, Any], step: int | None = None, **kwargs) -> None: import wandb - if step_key is not None: - # Pin every logged key to its axis with an exact-name definition. - # Glob definitions (like the "rollout/*" ones at init) are not - # expanded client-side anymore; the raw glob is sent to the server - # for expansion, which not every server version supports - # (wandb#11533) — charts then silently fall back to the global - # step axis. Exact names always work, including for axes that - # only exist at runtime (e.g. a per-adapter "{name}/step"). - if step_key not in self._defined_metrics: - wandb.define_metric(step_key) - self._defined_metrics.add(step_key) - for key in metrics: - if key != step_key and key not in self._defined_metrics: - wandb.define_metric(key, step_metric=step_key) - self._defined_metrics.add(key) 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 @@ -166,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) From 793715aa410e31e553a13a4d0b1349e857055b0b Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Fri, 17 Jul 2026 11:35:26 -0700 Subject: [PATCH 30/31] [test] --- examples/multi_lora/multi_lora_async_rollout.py | 14 +++++++------- miles/utils/multi_lora.py | 8 ++++++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/examples/multi_lora/multi_lora_async_rollout.py b/examples/multi_lora/multi_lora_async_rollout.py index 8561827f83..1e240e76f6 100644 --- a/examples/multi_lora/multi_lora_async_rollout.py +++ b/examples/multi_lora/multi_lora_async_rollout.py @@ -570,11 +570,11 @@ async def generate_rollout_multi_lora_async( **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 live in their own {name}_perf namespace: - # a namespace maps to exactly one x-axis, and these ride rollout/step - # while the {name}/ namespace is on the adapter's step axis. - **{f"{name}_perf/queue_length": size for name, size in queue_sizes.items()}, - **{f"{name}_perf/stale_dropped": len(stale_drops.get(name, [])) for name in adapters}, + # 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), @@ -586,8 +586,8 @@ async def generate_rollout_multi_lora_async( 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) + 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) diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index 65ed99d2d3..6ebc0d6c39 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -50,9 +50,13 @@ def is_multi_lora_enabled(args: Any) -> bool: 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 + {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 @@ -64,7 +68,7 @@ def define_new_adapter_metrics(snapshot: dict) -> None: 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") + define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step") def make_rid(adapter_name: str) -> str: From 3c1e9d8ebc00d7e3c66cd13c009613c0ffc94995 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Fri, 17 Jul 2026 13:46:23 -0700 Subject: [PATCH 31/31] [misc] deprecate num_row in favor of num_step --- examples/multi_lora/README.md | 9 ++++--- .../multi_lora_data_source_async.py | 14 +++++----- miles/ray/multi_lora_controller.py | 3 +++ miles/utils/adapter_config.py | 7 +++-- miles/utils/multi_lora.py | 26 +++++++++++++------ tests/fast/utils/test_controller_backend.py | 4 +-- 6 files changed, 38 insertions(+), 25 deletions(-) diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md index e8fd255059..7340d830e8 100644 --- a/examples/multi_lora/README.md +++ b/examples/multi_lora/README.md @@ -18,7 +18,7 @@ run_service.sh # service mode: idles for registrations (po 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, legacy num_row fallback) +multi_lora_data_source_async.py # data source (reads controller) adapters/ gsm8k.yaml dapo_math.yaml @@ -68,8 +68,10 @@ Ray actor, pinned to the head node). - 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. The data source still supports legacy `num_row` - deregistration when configured. The trainer'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 @@ -116,6 +118,7 @@ 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, ... ``` diff --git a/examples/multi_lora/multi_lora_data_source_async.py b/examples/multi_lora/multi_lora_data_source_async.py index b75b7deb34..a9b366437d 100644 --- a/examples/multi_lora/multi_lora_data_source_async.py +++ b/examples/multi_lora/multi_lora_data_source_async.py @@ -1,4 +1,6 @@ -"""Round-robin per-adapter data source; legacy num_row-based deregistration.""" +"""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 @@ -48,6 +50,9 @@ def reconcile(self, adapters: dict[str, AdapterRun]) -> None: 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: @@ -104,13 +109,6 @@ def get_samples(self, num_samples: int = 1) -> list[list[Sample]]: sample.reward_spec = reward_spec sample.metadata = {**config.metadata, **sample.metadata} - if config.num_step is None: - default_num_row = (getattr(config, "num_epoch", 1) or 1) * len(source.dataset) - num_row = config.num_row or default_num_row - if source.sample_group_index >= num_row and name not in snapshot["retiring"]: - logger.info(f"Adapter '{name}' reached num_row={num_row}, deregistering") - ray.get(get_multi_lora_controller().deregister_adapter.remote(name)) - return groups return [] diff --git a/miles/ray/multi_lora_controller.py b/miles/ray/multi_lora_controller.py index eba731dc96..b20eb2b75c 100644 --- a/miles/ray/multi_lora_controller.py +++ b/miles/ray/multi_lora_controller.py @@ -93,6 +93,9 @@ def record_batch_adapters(self, rollout_id: int, groups: dict[str, int], step_na 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) diff --git a/miles/utils/adapter_config.py b/miles/utils/adapter_config.py index b57613e352..203fe82619 100644 --- a/miles/utils/adapter_config.py +++ b/miles/utils/adapter_config.py @@ -37,9 +37,9 @@ class AdapterRunConfig: rm_type: str | None = None custom_rm_path: str | None = None - num_epoch: int | None = None + # Stop after N optimizer steps; derived from num_epoch (default 1) when absent. num_step: int | None = None - num_row: int | None = None + num_epoch: int | None = None metadata: dict[str, Any] = field(default_factory=dict) @@ -84,8 +84,7 @@ def parse_adapter_run_yaml(path: Path) -> AdapterRunConfig: metadata_key=raw.get("metadata_key"), rm_type=raw.get("rm_type"), custom_rm_path=raw.get("custom_rm_path"), - num_epoch=raw.get("num_epoch"), num_step=raw.get("num_step"), - num_row=raw.get("num_row"), + num_epoch=raw.get("num_epoch"), metadata=raw.get("metadata") or {}, ) diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index 6ebc0d6c39..b8bd0bf98c 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -285,6 +285,19 @@ def mark_batch_trained(self, rollout_id: int) -> list[str]: 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 @@ -381,14 +394,11 @@ def resolve_adapter_config(self, name: str, config: Any) -> Any: 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_row is not None and (type(config.num_row) is not int or config.num_row <= 0): - raise ValueError(f"Adapter '{name}' num_row must be a positive integer") - if config.num_step is not None and config.num_row is not None: - logger.warning( - f"Adapter '{name}' sets both num_step and num_row; num_step takes precedence and num_row is ignored" - ) - elif config.num_step is None and config.num_row is not None: - logger.warning(f"Adapter '{name}' uses deprecated num_row={config.num_row}; prefer num_step") + 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: diff --git a/tests/fast/utils/test_controller_backend.py b/tests/fast/utils/test_controller_backend.py index b3a155dd45..f1a9c0f38c 100644 --- a/tests/fast/utils/test_controller_backend.py +++ b/tests/fast/utils/test_controller_backend.py @@ -253,8 +253,8 @@ async def test_register_rejects_bad_batch_shapes(tmp_path): 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_row must be a positive integer"): - await backend.register("H", make_config(num_row=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))