From aa2180b6878145365916de793697601339848a27 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 6 Aug 2026 09:48:20 +0800 Subject: [PATCH 01/20] feat(async-rl): add client-orchestrated multi-tenant workflows --- cookbook/client/async_rl/README.md | 75 + .../async_rl/client_orchestrated_dpo.py | 234 +++ .../async_rl/client_orchestrated_grpo.py | 452 ++++++ .../server/transformer/server_config.yaml | 19 + cookbook/rl/async_multi_lora_dapo_grpo.yaml | 176 +++ .../async_multi_lora_dapo_hparam_sweep.yaml | 204 +++ cookbook/rl/async_multi_lora_grpo.py | 21 + cookbook/rl/async_multi_lora_grpo.yaml | 136 ++ cookbook/rl/async_single_lora_dapo_grpo.yaml | 134 ++ .../rl/async_single_lora_gsm8k_areal.yaml | 191 +++ cookbook/rl/sync_barrier_multi_lora_grpo.py | 608 ++++++++ pyproject.toml | 7 + src/twinkle/loss/base.py | 6 + src/twinkle/loss/chunked_cross_entropy.py | 15 + src/twinkle/loss/cross_entropy.py | 15 + src/twinkle/loss/grpo.py | 32 +- src/twinkle/metric/__init__.py | 3 + src/twinkle/metric/buffer.py | 26 + src/twinkle/metric/reporting.py | 593 ++++++++ src/twinkle/metric/types.py | 42 + src/twinkle/model/__init__.py | 2 + src/twinkle/model/micro_batch.py | 212 +++ .../model/transformers/transformers.py | 166 ++- src/twinkle/reward/__init__.py | 5 +- src/twinkle/reward/boxed_math.py | 54 + src/twinkle/reward/dapo_math.py | 153 ++ src/twinkle/reward/gsm8k.py | 121 ++ .../sampler/vllm_sampler/vllm_engine.py | 14 + .../sampler/vllm_sampler/vllm_sampler.py | 5 + src/twinkle/server/config/__init__.py | 4 +- src/twinkle/server/config/application_spec.py | 15 +- src/twinkle/server/data_plane/__init__.py | 4 + src/twinkle/server/data_plane/app.py | 45 + src/twinkle/server/data_plane/handlers.py | 52 + src/twinkle/server/data_plane/proxy.py | 56 + src/twinkle/server/data_plane/store.py | 133 ++ src/twinkle/server/deployment.py | 6 +- .../server/launcher/builder_registry.py | 7 +- src/twinkle/server/model/app.py | 8 +- src/twinkle/server/model/twinkle_handlers.py | 202 ++- src/twinkle/server/sampler/app.py | 35 +- .../server/sampler/backends/mock_sampler.py | 5 + src/twinkle/server/sampler/tinker_handlers.py | 3 +- .../server/sampler/twinkle_handlers.py | 237 ++- src/twinkle/server/utils/task_queue/mixin.py | 58 +- src/twinkle/server/utils/task_queue/worker.py | 9 +- src/twinkle/tq_utils.py | 43 + src/twinkle_agentic/async_rl/__init__.py | 33 + .../async_rl/context_manager.py | 300 ++++ src/twinkle_agentic/async_rl/data_plane.py | 273 ++++ src/twinkle_agentic/async_rl/metrics.py | 129 ++ src/twinkle_agentic/async_rl/native_tq.py | 187 +++ src/twinkle_agentic/async_rl/pipeline.py | 687 +++++++++ src/twinkle_agentic/async_rl/scheduler.py | 73 + src/twinkle_agentic/async_rl/tq_utils.py | 29 + src/twinkle_agentic/async_rl/types.py | 103 ++ src/twinkle_agentic/async_rl/utils.py | 218 +++ .../async_rl/vllm_sampler_tq.py | 805 ++++++++++ src/twinkle_agentic/async_rl/workers.py | 666 +++++++++ src/twinkle_client/__init__.py | 5 +- src/twinkle_client/async_rl/__init__.py | 4 + src/twinkle_client/async_rl/workers.py | 61 + src/twinkle_client/common/json_utils.py | 33 + src/twinkle_client/common/serialize.py | 1 + src/twinkle_client/data_plane.py | 129 ++ .../model/multi_lora_transformers.py | 95 ++ src/twinkle_client/remote_task.py | 94 ++ src/twinkle_client/rollout/multi_turn.py | 42 +- src/twinkle_client/sampler/vllm_sampler.py | 111 +- src/twinkle_client/types/__init__.py | 15 + src/twinkle_client/types/component.py | 115 ++ tests/loss/test_grpo_gkd.py | 23 + tests/model/test_micro_batch.py | 284 ++++ tests/server/config/test_server_config.py | 26 + .../server/contract/client_api_baseline.json | 114 ++ tests/server/data_plane/test_proxy.py | 55 + tests/server/data_plane/test_store.py | 110 ++ tests/server/gateway/test_future_retrieval.py | 85 ++ .../server/model/test_twinkle_async_inputs.py | 72 + tests/server/sampler/test_mock_sampler.py | 37 +- .../server/sampler/test_twinkle_async_rows.py | 50 + tests/server/state/test_managers.py | 1 - .../test_app_builders_characterization.py | 30 +- tests/server/utils/test_task_queue_mixin.py | 21 + .../twinkle_agentic/test_async_rl_metrics.py | 343 +++++ .../test_async_rl_native_tq.py | 1292 +++++++++++++++++ .../test_vllm_sampler_tq_generation.py | 188 +++ tests/twinkle_client/test_async_components.py | 165 +++ tests/twinkle_client/test_async_rl_workers.py | 69 + .../test_client_multi_turn_rollout.py | 26 + .../test_client_orchestrated_dpo.py | 116 ++ .../test_client_orchestrated_grpo.py | 202 +++ tests/twinkle_client/test_data_plane_async.py | 102 ++ 93 files changed, 12172 insertions(+), 90 deletions(-) create mode 100644 cookbook/client/async_rl/README.md create mode 100644 cookbook/client/async_rl/client_orchestrated_dpo.py create mode 100644 cookbook/client/async_rl/client_orchestrated_grpo.py create mode 100644 cookbook/rl/async_multi_lora_dapo_grpo.yaml create mode 100644 cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml create mode 100644 cookbook/rl/async_multi_lora_grpo.py create mode 100644 cookbook/rl/async_multi_lora_grpo.yaml create mode 100644 cookbook/rl/async_single_lora_dapo_grpo.yaml create mode 100644 cookbook/rl/async_single_lora_gsm8k_areal.yaml create mode 100644 cookbook/rl/sync_barrier_multi_lora_grpo.py create mode 100644 src/twinkle/metric/buffer.py create mode 100644 src/twinkle/metric/reporting.py create mode 100644 src/twinkle/metric/types.py create mode 100644 src/twinkle/model/micro_batch.py create mode 100644 src/twinkle/reward/boxed_math.py create mode 100644 src/twinkle/reward/dapo_math.py create mode 100644 src/twinkle/server/data_plane/__init__.py create mode 100644 src/twinkle/server/data_plane/app.py create mode 100644 src/twinkle/server/data_plane/handlers.py create mode 100644 src/twinkle/server/data_plane/proxy.py create mode 100644 src/twinkle/server/data_plane/store.py create mode 100644 src/twinkle/tq_utils.py create mode 100644 src/twinkle_agentic/async_rl/__init__.py create mode 100644 src/twinkle_agentic/async_rl/context_manager.py create mode 100644 src/twinkle_agentic/async_rl/data_plane.py create mode 100644 src/twinkle_agentic/async_rl/metrics.py create mode 100644 src/twinkle_agentic/async_rl/native_tq.py create mode 100644 src/twinkle_agentic/async_rl/pipeline.py create mode 100644 src/twinkle_agentic/async_rl/scheduler.py create mode 100644 src/twinkle_agentic/async_rl/tq_utils.py create mode 100644 src/twinkle_agentic/async_rl/types.py create mode 100644 src/twinkle_agentic/async_rl/utils.py create mode 100644 src/twinkle_agentic/async_rl/vllm_sampler_tq.py create mode 100644 src/twinkle_agentic/async_rl/workers.py create mode 100644 src/twinkle_client/async_rl/__init__.py create mode 100644 src/twinkle_client/async_rl/workers.py create mode 100644 src/twinkle_client/common/json_utils.py create mode 100644 src/twinkle_client/data_plane.py create mode 100644 src/twinkle_client/remote_task.py create mode 100644 src/twinkle_client/types/component.py create mode 100644 tests/model/test_micro_batch.py create mode 100644 tests/server/data_plane/test_proxy.py create mode 100644 tests/server/data_plane/test_store.py create mode 100644 tests/server/gateway/test_future_retrieval.py create mode 100644 tests/server/model/test_twinkle_async_inputs.py create mode 100644 tests/server/sampler/test_twinkle_async_rows.py create mode 100644 tests/twinkle_agentic/test_async_rl_metrics.py create mode 100644 tests/twinkle_agentic/test_async_rl_native_tq.py create mode 100644 tests/twinkle_agentic/test_vllm_sampler_tq_generation.py create mode 100644 tests/twinkle_client/test_async_components.py create mode 100644 tests/twinkle_client/test_async_rl_workers.py create mode 100644 tests/twinkle_client/test_client_orchestrated_dpo.py create mode 100644 tests/twinkle_client/test_client_orchestrated_grpo.py create mode 100644 tests/twinkle_client/test_data_plane_async.py diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md new file mode 100644 index 000000000..f44653707 --- /dev/null +++ b/cookbook/client/async_rl/README.md @@ -0,0 +1,75 @@ +# Client-orchestrated asynchronous RL + +This directory demonstrates direct orchestration of the server's Model, +Sampler, and TransferQueue DataPlane components. + +The client owns its Dataset, multi-turn rollout, Reward, Advantage, policy +versioning, staleness, and algorithm. There is no central async-RL management, +runtime, tenant submission API, or server-side RL worker involved. + +- client_orchestrated_grpo.py maps each DataLoader batch to a private client-side + rollout partition. Its Rollout, Advantage, and Trainer workers run as + independent asyncio tasks. Prompt groups stream through TQ independently, so + ready groups can train while the remaining groups are still sampling. A + policy is published once, after the whole partition has trained. +- client_orchestrated_dpo.py uses Dataset, Reference, and Trainer workers. It is + a runnable offline DPO loop and shows that Worker is a role lifecycle rather + than a fixed RL stage graph. + +Start the component server with +cookbook/client/server/transformer/server_config.yaml. + +```bash +pip install -e '.[async-rl,client]' +twinkle-server launch -c cookbook/client/server/transformer/server_config.yaml +``` + +Then start one or more independent client orchestrators: + +```bash +python cookbook/client/async_rl/client_orchestrated_grpo.py +``` + +The training loop composes only the low-level component methods: + +- `sampler.submit_sample(...)` / `sampler.asample(...)` +- `model.submit_forward_only(...)` +- `model.submit_forward_backward(...)` +- `model.submit_clip_grad_and_step(...)` +- `model.submit_save(...)` +- `data_plane.put/get/append/release(...)` and + `aput/aget/aget_batch/aappend/arelease(...)` + +There is no additional RL runtime or orchestration protocol. The GRPO example +uses `submit_sample()` so the Sampler's output `DataRef` remains in TQ. Each +generation is one row tagged with its group, generation index, rollout policy, +and status. The Advantage worker reads those rows and appends reward, advantage, +old log-probability, and updated status to the same keys. The Trainer consumes +the resulting `DataRef` and releases it in a local `finally` block. `asample()` +remains available when an algorithm prefers materialized response objects and +does not need to retain the intermediate TQ rows. + +- `ClientMultiTurnRollout.arun()` keeps tool calls and Reward computation in + the client and accepts an explicit `adapter_uri` policy snapshot. + +`_RolloutPartition` is a private client record, not a server resource or SDK +API. The local FIFO limits live DataLoader batches before rollout, ready prompt +groups immediately use the Model primitives above, and the client calls +`submit_save()` once after a whole batch has trained. `WorkerPipeline` only +starts, joins, and fail-fast cancels concrete roles; queues and algorithm state +remain ordinary client Python code. + +Different client processes may run GRPO and DPO against the same component +server. Model adapters are session-scoped. DataRefs are opaque capabilities +whose UUID identifies an independent physical TQ partition; DataPlane storage +does not know about tokens or sessions. The client remains the single writer +and algorithm owner for its adapter. Async Sampler requests share vLLM +continuous batching; this example does not promise strict round-robin fairness +between sampler tenants. + +The original YAML-managed runtime is still separate and can be started with: + +```bash +python cookbook/rl/async_multi_lora_grpo.py \ + --config cookbook/rl/async_multi_lora_grpo.yaml +``` diff --git a/cookbook/client/async_rl/client_orchestrated_dpo.py b/cookbook/client/async_rl/client_orchestrated_dpo.py new file mode 100644 index 000000000..d9c822188 --- /dev/null +++ b/cookbook/client/async_rl/client_orchestrated_dpo.py @@ -0,0 +1,234 @@ +"""Runnable offline DPO composed from the low-level async component clients.""" +from __future__ import annotations + +import asyncio +import inspect +import os +from typing import Any + +from peft import LoraConfig + +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.preprocessor import EmojiDPOProcessor +from twinkle_client import DataPlaneClient, init_twinkle_client +from twinkle_client.async_rl import Worker, WorkerPipeline +from twinkle_client.common.serialize import json_safe +from twinkle_client.model import MultiLoraTransformersModel +from twinkle_client.types import DataRef + +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +MODEL_ID = f'ms://{BASE_MODEL}' +DATASET_ID = os.environ.get('TWINKLE_DPO_DATASET_ID', 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji') +ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-dpo') +MAX_STEPS = int(os.environ.get('TWINKLE_MAX_STEPS', '100')) +BATCH_SIZE = int(os.environ.get('TWINKLE_BATCH_SIZE', '4')) +MAX_LENGTH = int(os.environ.get('TWINKLE_MAX_LENGTH', '2048')) + + +def create_dataset() -> Dataset: + """Load and encode preference pairs in the client process.""" + dataset = Dataset(DatasetMeta(DATASET_ID, data_slice=range(MAX_STEPS * BATCH_SIZE))) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=MAX_LENGTH) + dataset.map(EmojiDPOProcessor, init_args={'system': 'You are a helpful assistant.'}) + dataset.encode() + return dataset + + +def prepare_dpo_batch(batch: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Flatten pairs as ``chosen_0, rejected_0, ...`` for DP-safe slicing.""" + rows: list[dict[str, Any]] = [] + for pair in batch: + common = {key: value for key, value in pair.items() if key not in ('positive', 'negative')} + rows.append({**common, **pair['positive']}) + rows.append({**common, **pair['negative']}) + return json_safe(rows) + + +def _extract_ref_outputs(result: Any, rows: list[dict[str, Any]] | None = None) -> dict[str, Any]: + """Normalize an async forward-only result into the DPOLoss input shape.""" + payload: Any = rows + if payload is None and isinstance(result, dict): + payload = result.get('result', result) + if isinstance(payload, list) and payload and all( + isinstance(row, dict) and row.get('logps') is not None for row in payload): + return {'logps': [row['logps'] for row in payload]} + if isinstance(payload, list) and len(payload) == 1 and isinstance(payload[0], dict): + payload = payload[0].get('result', payload[0]) + if not isinstance(payload, dict) or payload.get('logps') is None: + raise RuntimeError('reference forward did not return per-token logps') + return {'logps': payload['logps']} + + +async def _put_rows(data_plane, rows, *, kind, tags): + try: + return await data_plane.aput(rows, kind=kind, tags=tags) + except TypeError as error: + if 'tags' not in str(error): + raise + return await data_plane.aput(rows, kind=kind) + + +async def _submit(method, *args, **kwargs): + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + task = await asyncio.to_thread(method, *args, **kwargs) + if inspect.isawaitable(task): + return await task + return task + + +class _DatasetWorker(Worker): + + def __init__(self, dataloader, data_plane, output): + super().__init__('dataset') + self.dataloader = dataloader + self.data_plane = data_plane + self.output = output + + async def run(self) -> None: + completed = 0 + for batch in self.dataloader: + if completed >= MAX_STEPS: + break + rows = prepare_dpo_batch(batch) + tags = [] + for index in range(0, len(rows), 2): + source_pair_id = rows[index].get('pair_id', f'pair-{completed}-{index // 2}') + tags.extend(( + { + 'record_type': 'preference', + 'pair_id': str(source_pair_id), + 'pair_role': 'chosen', + 'pair_status': 'DATA_READY', + }, + { + 'record_type': 'preference', + 'pair_id': str(source_pair_id), + 'pair_role': 'rejected', + 'pair_status': 'DATA_READY', + }, + )) + ref = await _put_rows( + self.data_plane, rows, kind='dpo-preference', tags=tags) + await self.output.put(ref) + completed += 1 + await self.output.put(None) + + +class _ReferenceWorker(Worker): + + def __init__(self, model, data_plane, source, output): + super().__init__('reference') + self.model = model + self.data_plane = data_plane + self.source = source + self.output = output + + async def run(self) -> None: + while True: + item = await self.source.get() + if item is None: + await self.output.put(None) + return + ref = item + ref_outputs = await _reference_forward(self.model, self.data_plane, ref) + await self.output.put((ref, ref_outputs)) + + +class _TrainerWorker(Worker): + + def __init__(self, model, data_plane, source): + super().__init__('trainer') + self.model = model + self.data_plane = data_plane + self.source = source + self.completed_steps = 0 + self.saved = None + + async def run(self) -> None: + while True: + item = await self.source.get() + if item is None: + self.saved = await _submit( + self.model.submit_save, + f'dpo-policy-{self.completed_steps}', + save_optimizer=True, + ) + return + ref, ref_outputs = item + try: + await _submit( + self.model.submit_forward_backward, + ref, + ref_outputs=ref_outputs, + ) + await _submit(self.model.submit_clip_grad_and_step, max_grad_norm=1.0) + finally: + await self.data_plane.arelease(ref) + self.completed_steps += 1 + + +async def _reference_forward( + model: MultiLoraTransformersModel, + data_plane: DataPlaneClient, + batch_ref: DataRef, +) -> dict[str, Any]: + """Run the frozen base model and materialize its DataPlane result.""" + result = await _submit(model.submit_forward_only, batch_ref, disable_lora=True) + if not isinstance(result, dict) or not result.get('output_ref'): + return _extract_ref_outputs(result) + + output_ref = DataRef(**result['output_ref']) + try: + rows = await data_plane.aget(output_ref) + return _extract_ref_outputs(result, rows) + finally: + await data_plane.arelease(output_ref) + + +async def run_dpo( + dataloader: DataLoader, + model: MultiLoraTransformersModel, + data_plane: DataPlaneClient, +) -> dict[str, Any]: + """Run client-owned DPO roles over the shared Model and DataPlane services.""" + preference_ready = asyncio.Queue(maxsize=2) + reference_ready = asyncio.Queue(maxsize=2) + trainer = _TrainerWorker(model, data_plane, reference_ready) + await WorkerPipeline(( + _DatasetWorker(dataloader, data_plane, preference_ready), + _ReferenceWorker(model, data_plane, preference_ready, reference_ready), + trainer, + )).run() + return trainer.saved + + +async def train() -> None: + client = init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + model = MultiLoraTransformersModel(MODEL_ID) + data_plane = DataPlaneClient() + + model.add_adapter_to_model( + ADAPTER_NAME, + LoraConfig(target_modules='all-linear', r=8, lora_alpha=32, lora_dropout=0.05), + ) + model.set_template('Qwen3_5Template', model_id=MODEL_ID) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('DPOLoss', beta=0.1, loss_type='sigmoid', reference_free=False) + model.add_metric('DPOMetric', beta=0.1) + model.set_optimizer('AdamW', lr=1e-5) + + try: + dataloader = DataLoader(dataset=create_dataset(), batch_size=BATCH_SIZE, num_workers=0) + saved = await run_dpo(dataloader, model, data_plane) + print(f"saved DPO adapter to {saved['twinkle_path']}") + finally: + client.close() + + +if __name__ == '__main__': + asyncio.run(train()) diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py new file mode 100644 index 000000000..b5892e193 --- /dev/null +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -0,0 +1,452 @@ +"""Client-orchestrated async GRPO built from the low-level component APIs.""" +from __future__ import annotations + +import asyncio +import inspect +import os +from collections import deque +from dataclasses import dataclass, field +from typing import Any + +from peft import LoraConfig + +from twinkle.advantage import GRPOAdvantage +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.preprocessor.llm import GSM8KProcessor +from twinkle.reward import GSM8KAccuracyReward +from twinkle_client import DataPlaneClient, init_twinkle_client +from twinkle_client.async_rl import Worker, WorkerPipeline +from twinkle_client.common.serialize import json_safe +from twinkle_client.model import MultiLoraTransformersModel +from twinkle_client.sampler import vLLMSampler + +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +MODEL_ID = f'ms://{BASE_MODEL}' +ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-grpo') +MAX_PARTITIONS = int(os.environ.get('TWINKLE_MAX_PARTITIONS', '100')) +MAX_STALENESS = int(os.environ.get('TWINKLE_MAX_STALENESS', '2')) +ROLLOUT_CONCURRENCY = int(os.environ.get('TWINKLE_ROLLOUT_CONCURRENCY', '8')) +NUM_GENERATIONS = int(os.environ.get('TWINKLE_NUM_GENERATIONS', '4')) +BATCH_SIZE = int(os.environ.get('TWINKLE_BATCH_SIZE', '8')) +TRAIN_MINI_BATCH_SIZE = int(os.environ.get('TWINKLE_TRAIN_MINI_BATCH_SIZE', '8')) +MICRO_BATCH_SIZE = int(os.environ.get('TWINKLE_MICRO_BATCH_SIZE', '4')) +MAX_TOKENS_PER_MICRO_BATCH = int(os.environ.get('TWINKLE_MAX_TOKENS_PER_MICRO_BATCH', '4096')) + + +@dataclass(frozen=True) +class _Policy: + version: int + adapter_uri: str + + +@dataclass +class _RolloutPartition: + """One DataLoader batch bound to one immutable policy snapshot.""" + + partition_id: int + policy: _Policy + rollouts: list[asyncio.Task[Any]] + ready: asyncio.Queue['_ReadyGroup'] = field(default_factory=asyncio.Queue) + + +@dataclass +class _ReadyGroup: + group_index: int + rows: list[dict[str, Any]] + tags: list[dict[str, Any]] + ref: Any + forward_kwargs: dict[str, Any] + + +@dataclass +class _RolloutResult: + partition: _RolloutPartition + group_index: int + ref: Any + + +class _GRPOState: + + def __init__(self, policy: _Policy): + self.policy = policy + self.live: deque[_RolloutPartition] = deque() + self.input_done = False + self.failure: BaseException | None = None + self.condition = asyncio.Condition() + + async def wait_for_admission(self) -> _Policy: + async with self.condition: + await self.condition.wait_for( + lambda: self.failure is not None or len(self.live) < MAX_STALENESS + 1) + if self.failure is not None: + raise self.failure + return self.policy + + async def add_partition(self, partition: _RolloutPartition) -> None: + async with self.condition: + self.live.append(partition) + self.condition.notify_all() + + async def finish_input(self) -> None: + async with self.condition: + self.input_done = True + self.condition.notify_all() + + async def fail(self, error: BaseException) -> None: + async with self.condition: + if self.failure is None: + self.failure = error + self.condition.notify_all() + + async def oldest_partition(self) -> _RolloutPartition | None: + async with self.condition: + await self.condition.wait_for( + lambda: self.failure is not None or bool(self.live) or self.input_done) + if self.failure is not None: + raise self.failure + return self.live[0] if self.live else None + + async def publish(self, partition: _RolloutPartition, policy: _Policy) -> None: + async with self.condition: + if not self.live or self.live[0] is not partition: + raise RuntimeError(f'partition {partition.partition_id} attempted out-of-order publication') + self.policy = policy + self.live.popleft() + self.condition.notify_all() + + +def create_dataset() -> Dataset: + dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=2048, enable_thinking=False) + dataset.map(GSM8KProcessor(system='Put the final answer within \\boxed{}.')) + dataset.encode(add_generation_prompt=True) + return dataset + + +async def rollout_group( + sampler: vLLMSampler, + prompt: dict[str, Any], + policy: _Policy, + semaphore: asyncio.Semaphore, + group_id: str, +) -> Any: + """Submit one GRPO group and keep its sample-level TQ DataRef alive.""" + async with semaphore: + result = await _submit( + sampler.submit_sample, + [prompt], + adapter_name=ADAPTER_NAME, + adapter_uri=policy.adapter_uri, + policy_version=policy.version, + group_ids=[group_id], + sampling_params={ + 'max_tokens': 1024, + 'temperature': 1.0, + 'top_p': 0.95, + 'logprobs': 1, + }, + num_samples=NUM_GENERATIONS, + ) + if not isinstance(result, dict) or not result.get('output_ref'): + raise RuntimeError('async sampler did not return a DataPlane output_ref') + from twinkle_client.types import DataRef + return DataRef(**result['output_ref']) + + +def start_partition( + partition_id: int, + batch: list[dict[str, Any]], + policy: _Policy, + sampler: vLLMSampler, + semaphore: asyncio.Semaphore, +) -> _RolloutPartition: + """Capture the snapshot before submitting any rollout in this partition.""" + return _RolloutPartition( + partition_id=partition_id, + policy=policy, + rollouts=[ + asyncio.create_task( + rollout_group( + sampler, + prompt, + policy, + semaphore, + f'partition-{partition_id}/group-{group_index}', + )) + for group_index, prompt in enumerate(json_safe(batch)) + ], + ) + + +async def _put_rows(data_plane: DataPlaneClient, rows, *, kind: str, tags): + """Pass native sample tags while remaining friendly to small cookbook fakes.""" + try: + return await data_plane.aput(rows, kind=kind, tags=tags) + except TypeError as error: + if 'tags' not in str(error): + raise + return await data_plane.aput(rows, kind=kind) + + +async def _submit(method, *args, **kwargs): + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + task = await asyncio.to_thread(method, *args, **kwargs) + if inspect.isawaitable(task): + return await task + return task + + +class _RolloutWorker(Worker): + + def __init__(self, dataloader, sampler, state: _GRPOState, output: asyncio.Queue, semaphore): + super().__init__('rollout') + self.dataloader = dataloader + self.sampler = sampler + self.state = state + self.output = output + self.semaphore = semaphore + + async def _collect(self, partition, group_index, task): + try: + ref = await task + await self.output.put(_RolloutResult(partition, group_index, ref)) + except BaseException as error: + await self.state.fail(error) + raise + + async def run(self) -> None: + batches = iter(self.dataloader) + collectors: list[asyncio.Task] = [] + try: + for partition_id in range(MAX_PARTITIONS): + policy = await self.state.wait_for_admission() + batch = next(batches, None) + if batch is None: + break + prompts = batch if isinstance(batch, list) else [batch] + if len(prompts) != BATCH_SIZE: + print(f'dropping incomplete final batch with {len(prompts)} prompts') + break + partition = start_partition( + partition_id, prompts, policy, self.sampler, self.semaphore) + await self.state.add_partition(partition) + collectors.extend( + asyncio.create_task(self._collect(partition, index, task)) + for index, task in enumerate(partition.rollouts) + ) + await self.state.finish_input() + await asyncio.gather(*collectors) + await self.output.put(None) + except BaseException as error: + await self.state.fail(error) + for task in collectors: + if not task.done(): + task.cancel() + await asyncio.gather(*collectors, return_exceptions=True) + raise + + +class _AdvantageWorker(Worker): + + def __init__(self, data_plane, state: _GRPOState, source: asyncio.Queue): + super().__init__('advantage') + self.data_plane = data_plane + self.state = state + self.source = source + + async def run(self) -> None: + try: + while True: + result = await self.source.get() + if result is None: + return + group_id = f'partition-{result.partition.partition_id}/group-{result.group_index}' + try: + batch = await self.data_plane.aget_batch(result.ref) + if len(batch.rows) != NUM_GENERATIONS: + raise RuntimeError( + f'group {group_id} expected {NUM_GENERATIONS} generations, ' + f'got {len(batch.rows)}') + features = [row.get('new_input_feature') for row in batch.rows] + if not all(isinstance(feature, dict) for feature in features): + raise RuntimeError(f'group {group_id} has no trainable new_input_feature') + old_logps = [ + [position[0][1] for position in (row.get('logprobs') or [])] + for row in batch.rows + ] + rewards = await asyncio.to_thread(GSM8KAccuracyReward(), features) + advantages = await asyncio.to_thread( + GRPOAdvantage(), rewards, num_generations=NUM_GENERATIONS) + train_rows = [dict(feature) for feature in features] + if batch.tags and len(batch.tags) != len(train_rows): + raise RuntimeError( + f'group {group_id} returned {len(batch.tags)} tags for ' + f'{len(train_rows)} rows') + source_tags = batch.tags or [{} for _ in train_rows] + tags = [{ + **tag, + 'record_type': 'sample', + 'group_id': group_id, + 'generation_idx': index, + 'rollout_status': 'ROLLOUT_DONE', + 'advantage_status': 'ADVANTAGE_DONE', + 'rollout_policy_version': result.partition.policy.version, + 'rollout_adapter_uri': result.partition.policy.adapter_uri, + } for index, tag in enumerate(source_tags)] + ref = await self.data_plane.aappend(result.ref, train_rows, tags=tags) + # The physical TQ rows retain rollout fields, while this + # reference selects only the trainable InputFeature. + train_ref = ref.model_copy(update={'fields': list(train_rows[0])}) + await result.partition.ready.put( + _ReadyGroup( + result.group_index, + train_rows, + tags, + train_ref, + { + 'old_logps': json_safe(old_logps), + 'advantages': json_safe(advantages), + }, + )) + except BaseException: + await self.data_plane.arelease(result.ref) + raise + except BaseException as error: + await self.state.fail(error) + raise + + +class _TrainerWorker(Worker): + + def __init__(self, model, data_plane, state: _GRPOState): + super().__init__('trainer') + self.model = model + self.data_plane = data_plane + self.state = state + self.optimizer_step = 0 + + async def _train(self, groups: list[_ReadyGroup]) -> None: + rows = [row for group in groups for row in group.rows] + tags = [tag for group in groups for tag in group.tags] + created_batch = len(groups) > 1 + train_ref = ( + await _put_rows(self.data_plane, rows, kind='grpo-train', tags=tags) + if created_batch else groups[0].ref + ) + old_logps = [ + value + for group in groups + for value in group.forward_kwargs['old_logps'] + ] + advantages = [ + value + for group in groups + for value in group.forward_kwargs['advantages'] + ] + try: + await _submit( + self.model.submit_forward_backward, + train_ref, + old_logps=old_logps, + advantages=advantages, + dynamic_batching=True, + micro_batch_size=MICRO_BATCH_SIZE, + max_tokens_per_micro_batch=MAX_TOKENS_PER_MICRO_BATCH, + ) + await _submit(self.model.submit_clip_grad_and_step, max_grad_norm=1.0) + self.optimizer_step += 1 + finally: + if created_batch: + await self.data_plane.arelease(train_ref) + for group in groups: + await self.data_plane.arelease(group.ref) + + async def run(self) -> None: + try: + while True: + partition = await self.state.oldest_partition() + if partition is None: + return + staleness = self.state.policy.version - partition.policy.version + if staleness > MAX_STALENESS: + raise RuntimeError( + f'partition {partition.partition_id} staleness {staleness} exceeds {MAX_STALENESS}') + groups_per_step = TRAIN_MINI_BATCH_SIZE // NUM_GENERATIONS + ready = [] + for _ in range(len(partition.rollouts)): + ready.append(await partition.ready.get()) + if len(ready) == groups_per_step: + await self._train(ready) + ready.clear() + if ready: + raise RuntimeError('partition ended with an incomplete train mini-batch') + publish_version = self.state.policy.version + 1 + saved = await _submit(self.model.submit_save, f'policy-{publish_version}') + policy = _Policy(publish_version, saved['twinkle_path']) + await self.state.publish(partition, policy) + print( + f'partition={partition.partition_id} policy={policy.version} ' + f'optimizer_step={self.optimizer_step} staleness={staleness}') + except BaseException as error: + await self.state.fail(error) + raise + + +async def run_grpo( + dataloader: DataLoader, + model: MultiLoraTransformersModel, + sampler: vLLMSampler, + data_plane: DataPlaneClient, +) -> None: + """Overlap rollout partitions while training and publishing them in FIFO order.""" + if MAX_STALENESS < 0: + raise ValueError('MAX_STALENESS must be non-negative') + if min(ROLLOUT_CONCURRENCY, NUM_GENERATIONS, BATCH_SIZE, TRAIN_MINI_BATCH_SIZE) <= 0: + raise ValueError('rollout concurrency and all batch sizes must be positive') + if TRAIN_MINI_BATCH_SIZE % NUM_GENERATIONS: + raise ValueError('TRAIN_MINI_BATCH_SIZE must be divisible by NUM_GENERATIONS') + groups_per_step = TRAIN_MINI_BATCH_SIZE // NUM_GENERATIONS + if BATCH_SIZE % groups_per_step: + raise ValueError('BATCH_SIZE * NUM_GENERATIONS must be divisible by TRAIN_MINI_BATCH_SIZE') + + initial = await _submit(model.submit_save, 'policy-0') + state = _GRPOState(_Policy(version=0, adapter_uri=initial['twinkle_path'])) + semaphore = asyncio.Semaphore(ROLLOUT_CONCURRENCY) + rollout_results: asyncio.Queue = asyncio.Queue() + await WorkerPipeline(( + _RolloutWorker(dataloader, sampler, state, rollout_results, semaphore), + _AdvantageWorker(data_plane, state, rollout_results), + _TrainerWorker(model, data_plane, state), + )).run() + + +async def train() -> None: + client = init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + try: + model = MultiLoraTransformersModel(MODEL_ID) + sampler = vLLMSampler(MODEL_ID) + data_plane = DataPlaneClient() + + model.add_adapter_to_model( + ADAPTER_NAME, + LoraConfig(target_modules='all-linear', r=8, lora_alpha=32, lora_dropout=0.05), + ) + model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) + model.set_optimizer('AdamW', lr=2e-5) + model.set_processor('InputProcessor', padding_free=True) + model.set_template('Qwen3_5Template', model_id=MODEL_ID) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID) + + dataloader = DataLoader(dataset=create_dataset(), batch_size=BATCH_SIZE, num_workers=0) + await run_grpo(dataloader, model, sampler, data_plane) + finally: + client.close() + + +if __name__ == '__main__': + asyncio.run(train()) diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index d3ddb2adb..b11adfc1c 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -53,6 +53,21 @@ applications: env_vars: TWINKLE_FAIL_FAST: "0" + # TransferQueue-backed data references used by client-orchestrated RL. + - name: data-plane + route_prefix: /api/v1/data-plane + import_path: data_plane + args: + config: + backend: + SimpleStorage: + num_data_storage_units: 2 + deployments: + - name: DataPlaneManagement + num_replicas: 1 + ray_actor_options: + num_cpus: 1 + # 2. Model Service - Hosts the base model for training. - name: models-Qwen3.5-4B route_prefix: /api/v1/model/Qwen/Qwen3.5-4B @@ -60,7 +75,9 @@ applications: args: backend: transformers # Model backend: transformers | megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + max_loras: 8 # Concurrent client-owned training adapters max_length: 10240 + data_plane_url: http://127.0.0.1:8000/api/v1/data-plane nproc_per_node: 1 # Number of GPU processes per node device_group: name: model @@ -94,12 +111,14 @@ applications: import_path: sampler args: model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + data_plane_url: http://127.0.0.1:8000/api/v1/data-plane nproc_per_node: 1 # Number of GPU processes per node sampler_type: vllm # Inference engine: 'vllm' (fast) or 'torch' (TorchSampler) engine_args: # vLLM engine-specific settings max_model_len: 4096 # Maximum sequence length the engine supports gpu_memory_utilization: 0.5 # Fraction of GPU memory to use (0.0-1.0) enable_lora: true # Allow loading LoRA adapters during inference + max_loras: 8 # Published policies cached across tenants logprobs_mode: processed_logprobs # Logprobs mode for sampling results device_group: # Logical device group for the sampler name: sampler diff --git a/cookbook/rl/async_multi_lora_dapo_grpo.yaml b/cookbook/rl/async_multi_lora_dapo_grpo.yaml new file mode 100644 index 000000000..d54641f59 --- /dev/null +++ b/cookbook/rl/async_multi_lora_dapo_grpo.yaml @@ -0,0 +1,176 @@ +runtime: + run_id: async_multi_lora_dapo_grpo + model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3-4B} + mode: ray + model_gpus: 2 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 4 + max_staleness: 2 + max_steps: null + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_multi_lora_dapo_grpo + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/dapo_metrics.jsonl + summary_path: outputs/async_rl/dapo_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: false + +model: + strategy: native_fsdp + fsdp_config: + reshard_after_forward: true + # Two trainer GPUs form one Ulysses SP group. The effective model DP size is 1. + sequence_parallel_size: 2 + padding_free: true + attn_implementation: flash_attention_2 + max_length: 12288 + +sampler: + max_model_len: 12288 + gpu_memory_utilization: 0.8 + max_num_seqs: 32 + max_num_batched_tokens: 16384 + # Keep the long-sequence sampler on eager execution until CUDA-graph stability is verified. + enforce_eager: true + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +evaluation: + enabled: true + interval: 5 + batch_size: 16 + sampling_params: + max_tokens: 8192 + temperature: 0.0 + top_p: 1.0 + +lora: + target_modules: all-linear + r: 16 + alpha: 32 + dropout: 0.0 + # Match verl's stable GRPO-on-DAPO learning rate. No scheduler means constant LR. + learning_rate: 1.0e-6 + +loss: + cls: GRPOLoss + epsilon: 0.2 + normalization: sequence_mean + +lora_contexts: + # Set both environment variables to distinct local parquet splits for a disjoint-tenant experiment. + - tenant_id: tenant_a + training_run_id: dapo_async + adapter_name: tenant_a_dapo_math_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:TENANT_A_DAPO_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 2000 + # Keeps prompt + max rollout tokens within model.max_length. + max_length: 4096 + processor: DAPOMathProcessor + eval_dataset: + name: aime2024 + dataset_id: ${oc.env:AIME2024_DATASET_ID,ms://Maxwell-Jia/AIME_2024} + subset_name: default + split: train + data_num: null + max_length: 4096 + processor: AIME2024Processor + reward: + class_path: twinkle.reward.BoxedMathAccuracyReward + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + # Four prompt groups x eight generations = 32 samples per optimizer step. + mini_batch_size: 32 + micro_batch_size: 1 + dynamic_batching: false + + - tenant_id: tenant_b + training_run_id: dapo_async + adapter_name: tenant_b_dapo_math_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:TENANT_B_DAPO_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 2000 + max_length: 4096 + processor: DAPOMathProcessor + eval_dataset: + name: aime2024 + dataset_id: ${oc.env:AIME2024_DATASET_ID,ms://Maxwell-Jia/AIME_2024} + subset_name: default + split: train + data_num: null + max_length: 4096 + processor: AIME2024Processor + reward: + class_path: twinkle.reward.BoxedMathAccuracyReward + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + mini_batch_size: 32 + micro_batch_size: 1 + dynamic_batching: false + + +# export MODEL_ID=/nas/disk1/Qwen3.5-4B + +# # 可选。不设置时，两个租户默认都读取远程 DAPO-Math-17k。 +# export TENANT_A_DAPO_DATASET_ID=/path/to/tenant_a_dapo.parquet +# export TENANT_B_DAPO_DATASET_ID=/path/to/tenant_b_dapo.parquet + +# python cookbook/rl/async_multi_lora_grpo.py \ +# --config cookbook/rl/async_multi_lora_dapo_grpo.yaml diff --git a/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml b/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml new file mode 100644 index 000000000..02d9168c6 --- /dev/null +++ b/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml @@ -0,0 +1,204 @@ +runtime: + run_id: async_multi_lora_dapo_hparam_sweep + model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3-4B} + mode: ray + # Three-GPU layout: two trainer GPUs in one SP group plus one sampler GPU. + model_gpus: 2 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 4 + max_staleness: 2 + max_steps: null + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_multi_lora_dapo_hparam_sweep + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/dapo_hparam_sweep_metrics.jsonl + summary_path: outputs/async_rl/dapo_hparam_sweep_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: false + +model: + strategy: native_fsdp + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 2 + padding_free: true + attn_implementation: flash_attention_2 + max_length: 12288 + +sampler: + max_model_len: 12288 + gpu_memory_utilization: 0.8 + max_num_seqs: 32 + max_num_batched_tokens: 16384 + enforce_eager: true + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + # Interleave hyperparameter candidates so one adapter cannot monopolize training. + train: {policy: round_robin, max_consecutive_units: 1} + +lora: + target_modules: all-linear + r: 16 + alpha: 32 + dropout: 0.05 + # Used only when a context omits train.learning_rate. + learning_rate: 1.0e-6 + +loss: + cls: GRPOLoss + epsilon: 0.2 + normalization: sequence_mean + +lora_contexts: + # A/B/C isolate learning rate while keeping mini_batch_size fixed at 32 samples. + - tenant_id: tenant_lr1e6_g4 + training_run_id: dapo_hparam_sweep + adapter_name: dapo_lr1e6_g4_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 500 + max_length: 4096 + processor: DAPOMathProcessor + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + learning_rate: 1.0e-6 + mini_batch_size: 32 + micro_batch_size: 1 + dynamic_batching: false + + - tenant_id: tenant_lr5e6_g4 + training_run_id: dapo_hparam_sweep + adapter_name: dapo_lr5e6_g4_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 500 + max_length: 4096 + processor: DAPOMathProcessor + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + learning_rate: 5.0e-6 + mini_batch_size: 32 + micro_batch_size: 1 + dynamic_batching: false + + - tenant_id: tenant_lr1e5_g4 + training_run_id: dapo_hparam_sweep + adapter_name: dapo_lr1e5_g4_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 500 + max_length: 4096 + processor: DAPOMathProcessor + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + learning_rate: 1.0e-5 + mini_batch_size: 32 + micro_batch_size: 1 + dynamic_batching: false + + # B/D isolate train batch size while keeping learning_rate fixed at 5e-6. + - tenant_id: tenant_lr5e6_g16 + training_run_id: dapo_hparam_sweep + adapter_name: dapo_lr5e6_g16_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 500 + max_length: 4096 + processor: DAPOMathProcessor + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + learning_rate: 5.0e-6 + mini_batch_size: 128 + micro_batch_size: 1 + dynamic_batching: false + + +# CUDA_VISIBLE_DEVICES=0,1,2 \ +# MODEL_ID=/nas/disk1/Qwen3-4B \ +# DAPO_HPARAM_DATASET_ID=/path/to/dapo_math_500.parquet \ +# TWINKLE_SEED=42 \ +# python3 cookbook/rl/async_multi_lora_grpo.py \ +# --config cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml diff --git a/cookbook/rl/async_multi_lora_grpo.py b/cookbook/rl/async_multi_lora_grpo.py new file mode 100644 index 000000000..92b7f6598 --- /dev/null +++ b/cookbook/rl/async_multi_lora_grpo.py @@ -0,0 +1,21 @@ +"""Launch native-TQ async multi-LoRA GRPO from one YAML configuration.""" + +from __future__ import annotations + +import argparse + +from omegaconf import OmegaConf + +from twinkle_agentic.async_rl import AsyncMultiLoraGRPOPipeline + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument('--config', default='cookbook/rl/async_multi_lora_grpo.yaml') + args = parser.parse_args() + config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) + print(AsyncMultiLoraGRPOPipeline.from_config(config).run()) + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/async_multi_lora_grpo.yaml b/cookbook/rl/async_multi_lora_grpo.yaml new file mode 100644 index 000000000..4467abb9d --- /dev/null +++ b/cookbook/rl/async_multi_lora_grpo.yaml @@ -0,0 +1,136 @@ +runtime: + run_id: async_multi_lora_grpo + model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3.5-4B} + mode: ray + model_gpus: 2 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 4 + max_staleness: 1 + max_steps: null + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_multi_lora_grpo + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/metrics.jsonl + summary_path: outputs/async_rl/summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Qwen3_5Template + enable_thinking: false + +model: + strategy: native_fsdp + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 1 + padding_free: false + max_length: 8192 + +sampler: + max_model_len: 8192 + gpu_memory_utilization: 0.8 + max_num_seqs: 64 + max_num_batched_tokens: 8192 + enforce_eager: false + +rollout_output: + enabled: true + output_dir: ${runtime.output_dir}/rollouts + include_token_ids: false + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +lora: + target_modules: all-linear + r: 16 + alpha: 32 + dropout: 0.05 + learning_rate: 5.0e-5 + lr_scheduler: + cls: CosineAnnealingLR + # One optimizer step per prompt group with mini_batch_size=4. + T_max: 2000 + eta_min: 0.0 + +loss: + cls: GRPOLoss + epsilon: 0.2 + normalization: sequence_mean + +lora_contexts: + - tenant_id: tenant_a + training_run_id: gsm8k_async + adapter_name: tenant_a_gsm8k_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyBrevityReward + dataset: + dataset_id: ${oc.env:TENANT_A_DATASET_ID} + subset_name: main + split: train + data_num: 128 + max_length: 8192 + processor: GSM8KProcessor + system_prompt: >- + You are a helpful math assistant. Solve the problem with minimal but + correct reasoning and put your final answer within \boxed{}. + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 2048 + temperature: 1.0 + top_p: 0.95 + train: + mini_batch_size: 4 + micro_batch_size: 1 + dynamic_batching: false + + - tenant_id: tenant_b + training_run_id: gsm8k_async + adapter_name: tenant_b_gsm8k_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + dataset: + dataset_id: ${oc.env:TENANT_B_DATASET_ID} + subset_name: main + split: train + data_num: 128 + max_length: 8192 + processor: GSM8KProcessor + system_prompt: >- + You are a helpful math assistant. Solve the problem with minimal but + correct reasoning and put your final answer within \boxed{}. + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 2048 + temperature: 1.0 + top_p: 0.95 + train: + mini_batch_size: 4 + micro_batch_size: 1 + dynamic_batching: false diff --git a/cookbook/rl/async_single_lora_dapo_grpo.yaml b/cookbook/rl/async_single_lora_dapo_grpo.yaml new file mode 100644 index 000000000..004639886 --- /dev/null +++ b/cookbook/rl/async_single_lora_dapo_grpo.yaml @@ -0,0 +1,134 @@ +runtime: + run_id: async_single_lora_dapo_grpo + model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3-4B} + mode: ray + # Same three-GPU layout as the four-context sweep: two trainer GPUs plus one sampler GPU. + model_gpus: 2 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 4 + max_staleness: 0 + max_steps: null + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_single_lora_dapo_grpo + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/single_lora_dapo_metrics.jsonl + summary_path: outputs/async_rl/single_lora_dapo_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: false + +model: + strategy: native_fsdp + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 2 + padding_free: true + attn_implementation: flash_attention_2 + max_length: 12288 + +sampler: + max_model_len: 12288 + gpu_memory_utilization: 0.8 + max_num_seqs: 32 + max_num_batched_tokens: 16384 + enforce_eager: true + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + # Keep scheduler settings identical to the four-context experiment. + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: round_robin, max_consecutive_units: 1} + +evaluation: + enabled: true + interval: 5 + batch_size: 16 + sampling_params: + max_tokens: 8192 + temperature: 0.0 + top_p: 1.0 + +lora: + target_modules: all-linear + r: 16 + alpha: 32 + dropout: 0.05 + learning_rate: 1.0e-6 + +loss: + cls: GRPOLoss + epsilon: 0.2 + normalization: sequence_mean + +lora_contexts: + - tenant_id: tenant_single_dapo + training_run_id: dapo_single_tenant + adapter_name: dapo_single_lr1e6_g4_lora + reward: + class_path: twinkle.reward.DAPOMathReward + kwargs: + max_response_length: ${...rollout.max_tokens} + # Intentionally identical to tenant_lr1e6_g4 in the four-context sweep. + overlong_buffer_length: 4096 + overlong_penalty_factor: 1.0 + score_tail_chars: 300 + dataset: + dataset_id: ${oc.env:DAPO_SINGLE_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} + subset_name: default + split: train + data_num: 500 + max_length: 4096 + processor: DAPOMathProcessor + eval_dataset: + name: aime2024 + dataset_id: ${oc.env:AIME2024_DATASET_ID,ms://Maxwell-Jia/AIME_2024} + subset_name: default + split: train + data_num: null + max_length: 4096 + processor: AIME2024Processor + reward: + class_path: twinkle.reward.BoxedMathAccuracyReward + rollout: + batch_size: 16 + num_generations: 8 + max_tokens: 8192 + temperature: 1.0 + top_p: 0.95 + train: + learning_rate: 1.0e-6 + mini_batch_size: 32 + micro_batch_size: 1 + dynamic_batching: false + + +# CUDA_VISIBLE_DEVICES=0,1,2 \ +# MODEL_ID=/nas/disk1/Qwen3-4B \ +# DAPO_SINGLE_DATASET_ID=/path/to/the_same_dapo_math_500.parquet \ +# TWINKLE_SEED=42 \ +# python3 cookbook/rl/async_multi_lora_grpo.py \ +# --config cookbook/rl/async_single_lora_dapo_grpo.yaml diff --git a/cookbook/rl/async_single_lora_gsm8k_areal.yaml b/cookbook/rl/async_single_lora_gsm8k_areal.yaml new file mode 100644 index 000000000..81ac1141d --- /dev/null +++ b/cookbook/rl/async_single_lora_gsm8k_areal.yaml @@ -0,0 +1,191 @@ +runtime: + run_id: async_single_lora_gsm8k_areal + model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} + mode: ray + model_gpus: 1 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 1 + seed: 1 + max_staleness: 0 + max_steps: null + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_single_lora_gsm8k_areal + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/single_lora_gsm8k_areal_metrics.jsonl + summary_path: outputs/async_rl/single_lora_gsm8k_areal_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: true + +model: + strategy: native_fsdp + attn_implementation: flash_attention_2 + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 1 + padding_free: false + max_length: 2048 + +sampler: + max_model_len: 2048 + gpu_memory_utilization: 0.8 + max_num_seqs: 64 + enforce_eager: false + +rollout_output: + enabled: true + output_dir: ${runtime.output_dir}/rollouts + include_token_ids: false + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +evaluation: + enabled: false + interval: 10 + batch_size: 16 + sampling_params: + max_tokens: 2048 + temperature: 0.6 + top_p: 1.0 + +lora: + target_modules: all-linear + r: 16 + alpha: 16 + dropout: 0.0 + learning_rate: 1.7e-5 + +loss: + cls: GRPOLoss + epsilon: 0.2 + normalization: token_mean + +lora_contexts: + - tenant_id: tenant_single + training_run_id: gsm8k_areal_pk + adapter_name: gsm8k_areal_pk_lora + reward: + class_path: twinkle.reward.MathVerifyAccuracyReward + dataset: + dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID} + subset_name: main + split: train + data_num: 2000 + max_length: 1024 + processor: AReaLGSM8KProcessor + eval_dataset: + name: gsm8k/test + dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID} + subset_name: main + split: test + data_num: null + max_length: 1024 + processor: AReaLGSM8KProcessor + reward: + class_path: twinkle.reward.MathVerifyAccuracyReward + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 1024 + temperature: 1.0 + top_p: 1.0 + train: + mini_batch_size: 64 + micro_batch_size: 64 + dynamic_batching: true + max_tokens_per_micro_batch: 4096 + packing_algorithm: ffd + + +# export VIRTUAL_ENV=/opt/.venv +# export PATH=/opt/.venv/bin:$PATH +# export CUDA_VISIBLE_DEVICES=2,3 +# export AREAL_MODEL_PATH=/nas/disk1/Qwen3-4B +# export AREAL_ADMIN_API_KEY="$( +# python3 -c 'import secrets; print(secrets.token_urlsafe(32))' +# )" +# export AREAL_TRIAL="pk-$(date +%Y%m%d-%H%M%S)" + +# python3 examples/math/gsm8k_rl.py \ +# --config examples/math/gsm8k_grpo_lora.yaml \ +# scheduler.type=local \ +# +rollout.agent.admin_api_key="$AREAL_ADMIN_API_KEY" \ +# seed=1 \ +# cluster.n_nodes=1 \ +# cluster.n_gpus_per_node=2 \ +# actor.path="$AREAL_MODEL_PATH" \ +# +actor.attn_impl=flash_attention_2 \ +# rollout.backend=vllm:d1 \ +# actor.backend=fsdp:d1 \ +# train_dataset.path=/model/ljl/project/data/gsm8k \ +# valid_dataset.path=/model/ljl/project/data/gsm8k \ +# train_dataset.batch_size=16 \ +# train_dataset.shuffle=false \ +# total_train_epochs=1 \ +# +total_train_steps=125 \ +# rollout.consumer_batch_size=16 \ +# rollout.max_concurrent_rollouts=16 \ +# rollout.max_head_offpolicyness=0 \ +# gconfig.n_samples=4 \ +# gconfig.max_new_tokens=1024 \ +# gconfig.max_tokens=2048 \ +# gconfig.temperature=1.0 \ +# +gconfig.top_p=1.0 \ +# +actor.mb_spec.n_mbs=64 \ +# ++actor.mb_spec.max_tokens_per_mb=null \ +# actor.use_lora=true \ +# actor.lora_rank=16 \ +# actor.lora_alpha=16 \ +# actor.optimizer.lr=1.7e-4 \ +# actor.optimizer.weight_decay=0.01 \ +# actor.optimizer.lr_scheduler_type=constant \ +# actor.eps_clip=0.2 \ +# actor.ppo_n_minibatches=1 \ +# actor.reward_scaling=1.0 \ +# actor.reward_bias=0.0 \ +# actor.adv_norm=null \ +# actor.kl_ctl=0.0 \ +# actor.recompute_logprob=false \ +# actor.use_decoupled_loss=false \ +# actor.rejection_sampling=null \ +# ++vllm.max_model_len=2048 \ +# ++vllm.gpu_memory_utilization=0.8 \ +# ++vllm.max_num_seqs=64 \ +# ++vllm.max_loras=1 \ +# ++vllm.enforce_eager=false \ +# evaluator.freq_epochs=null \ +# evaluator.freq_steps=null \ +# evaluator.freq_secs=null \ +# trial_name="$AREAL_TRIAL" \ +# +stats_logger.tensorboard.path="/tmp/areal/tensorboard/gsm8k-grpo/$AREAL_TRIAL" \ +# cluster.fileroot=/nas/disk1/areal-experiments \ +# saver.freq_epochs=null \ +# saver.freq_steps=125 \ +# saver.freq_secs=null diff --git a/cookbook/rl/sync_barrier_multi_lora_grpo.py b/cookbook/rl/sync_barrier_multi_lora_grpo.py new file mode 100644 index 000000000..2914d0ec3 --- /dev/null +++ b/cookbook/rl/sync_barrier_multi_lora_grpo.py @@ -0,0 +1,608 @@ +"""Synchronous barrier baseline for native async multi-LoRA GRPO. + +The model, sampler, datasets, rewards, batch semantics, and checkpoint cadence +match ``async_multi_lora_grpo.py``. The only intentional difference is the +execution schedule: every round finishes rollout for all active contexts +before any context starts training. +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import time +from dataclasses import dataclass +from typing import Any, Iterator, Sequence + +from omegaconf import OmegaConf + +from twinkle.metric import MetricRecord, create_metrics_reporter +from twinkle_agentic.async_rl.metrics import advantage_signal_metrics, rollout_metrics +from twinkle_agentic.async_rl.pipeline import (_prompt_batches, _reward_for_context, _train_batch) +from twinkle_agentic.async_rl.tq_utils import REQUIRED_MODEL_INPUT_FIELDS, columns_to_tq_fields +from twinkle_agentic.async_rl.types import LoraContext, PartitionAdmission +from twinkle_agentic.async_rl.utils import ( + TrainBatchConfig, + build_native_fsdp_model_kwargs, + configure_lora_lr_scheduler, + resolve_context_learning_rate, + resolve_context_lora_target_modules, + resolve_context_loss_config, + resolve_model_attention_implementation, + resolve_sequence_parallel_size, + sample_responses_to_rollout_rows, + sampler_data_parallel_size, + validate_context_batch_config, +) +from twinkle_agentic.async_rl.vllm_sampler_tq import _compute_reward_metrics + + +@dataclass +class SyncContextState: + context: LoraContext + prompt_batches: Iterator[Sequence[dict[str, Any]]] + rollout_batch_size: int + num_generations: int + sampling_params: Any + mini_batch_size: int + reward_fn: Any + adapter_path: str + adapter_history: list[str] + partition_step: int = 0 + optimizer_steps: int = 0 + policy_version: int = 0 + exhausted: bool = False + + +@dataclass +class SyncPartition: + admission: PartitionAdmission + state: SyncContextState + rows: list[dict[str, Any]] + rewards: list[float] + advantages: list[float] | None = None + + +class SyncBarrierMultiLoraGRPO: + + def __init__(self, raw_config: dict[str, Any]): + import twinkle + from peft import LoraConfig + from twinkle import DeviceGroup, DeviceMesh + from twinkle.data_format import SamplingParams + from twinkle.model import MultiLoraTransformersModel + from twinkle.processor import InputProcessor + from twinkle.sampler import vLLMSampler + + raw_config = OmegaConf.to_container(OmegaConf.create(raw_config), resolve=True) + if not isinstance(raw_config, dict): + raise TypeError('sync RL config must resolve to a mapping') + + runtime = raw_config['runtime'] + model_config = raw_config['model'] + lora_data = raw_config['lora'] + loss_data = raw_config.get('loss') + template_data = raw_config.get('template', {}) + template_cls = template_data.get('cls', 'Qwen3_5Template') + enable_thinking = bool(template_data.get('enable_thinking', False)) + model_gpus = int(runtime['model_gpus']) + sampler_gpus = int(runtime['sampler_gpus']) + sampler_tp = int(runtime['sampler_tp']) + sampler_dp = sampler_data_parallel_size(sampler_gpus, sampler_tp) + sequence_parallel_size = resolve_sequence_parallel_size( + model_gpus, + int(model_config['sequence_parallel_size']), + ) + padding_free = bool(model_config['padding_free']) + attn_implementation = resolve_model_attention_implementation( + model_config, + padding_free=padding_free, + sequence_parallel_size=sequence_parallel_size, + ) + model_max_length = int(model_config['max_length']) + sampler_config = raw_config['sampler'] + total_gpus = model_gpus + sampler_gpus + + twinkle.initialize( + mode='ray', + nproc_per_node=total_gpus, + groups=[ + DeviceGroup('model', list(range(model_gpus)), device_type='GPU'), + DeviceGroup( + 'sampler', + list(range(model_gpus, total_gpus)), + device_type='GPU', + gpus_per_worker=sampler_tp, + ), + ], + lazy_collect=False, + ) + model_mesh = DeviceMesh.from_sizes( + world_size=model_gpus, + dp_size=model_gpus, + ulysses_size=sequence_parallel_size, + ) + model_data_parallel_size = model_mesh.data_world_size + self.model_data_parallel_size = model_data_parallel_size + sampler_mesh = DeviceMesh.from_sizes( + world_size=sampler_gpus, + dp_size=sampler_dp, + tp_size=sampler_tp, + ) + model_kwargs = build_native_fsdp_model_kwargs(model_config) + if attn_implementation is not None: + model_kwargs['attn_implementation'] = attn_implementation + self.model = MultiLoraTransformersModel( + model_id=runtime['model_id'], + device_mesh=model_mesh, + remote_group='model', + max_length=model_max_length, + **model_kwargs, + ) + self.train_batch_configs: dict[str, TrainBatchConfig] = {} + self.states: list[SyncContextState] = [] + for item in raw_config['lora_contexts']: + context = LoraContext( + item['tenant_id'], + item['training_run_id'], + runtime['model_id'], + item['adapter_name'], + ) + rollout = item['rollout'] + train = item['train'] + rollout_batch_size = int(rollout['batch_size']) + num_generations = int(rollout['num_generations']) + train_batch_config = TrainBatchConfig( + mini_batch_size=int(train['mini_batch_size']), + micro_batch_size=int(train['micro_batch_size']), + dynamic_batching=bool(train.get('dynamic_batching', False)), + max_tokens_per_micro_batch=( + int(train['max_tokens_per_micro_batch']) + if train.get('max_tokens_per_micro_batch') is not None else None + ), + packing_algorithm=str(train.get('packing_algorithm', 'ffd')), + ) + validate_context_batch_config( + context.key, + rollout_groups=rollout_batch_size, + num_generations=num_generations, + train=train_batch_config, + sampler_dp=sampler_dp, + model_dp=model_data_parallel_size, + ) + adapter_lora_config = LoraConfig( + target_modules=resolve_context_lora_target_modules(item, lora_data), + r=lora_data['r'], + lora_alpha=lora_data['alpha'], + lora_dropout=lora_data['dropout'], + ) + self.model.add_adapter_to_model( + context.adapter_name, + adapter_lora_config, + gradient_accumulation_steps=1, + ) + self.model.set_optimizer( + 'AdamW', + lr=resolve_context_learning_rate(train, lora_data), + adapter_name=context.adapter_name, + ) + configure_lora_lr_scheduler(self.model, context.adapter_name, lora_data) + loss_cls, loss_kwargs = resolve_context_loss_config(item, loss_data) + self.model.set_loss( + loss_cls, + adapter_name=context.adapter_name, + **loss_kwargs, + ) + self.model.set_processor( + InputProcessor, + adapter_name=context.adapter_name, + padding_free=padding_free, + ) + self.model.set_template( + template_cls, + model_id=runtime['model_id'], + adapter_name=context.adapter_name, + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + initial_path = self.model.save( + f'sync-{context.adapter_name}-initial', + output_dir=runtime['output_dir'], + adapter_name=context.adapter_name, + ) + state = SyncContextState( + context=context, + prompt_batches=iter( + _prompt_batches( + item['dataset'], + model_id=runtime['model_id'], + batch_size=rollout_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + )), + rollout_batch_size=rollout_batch_size, + num_generations=num_generations, + sampling_params=SamplingParams( + max_tokens=rollout['max_tokens'], + temperature=rollout['temperature'], + top_p=rollout['top_p'], + repetition_penalty=float(rollout.get('repetition_penalty', 1.0)), + logprobs=1, + num_samples=1, + ), + mini_batch_size=train_batch_config.mini_batch_size, + reward_fn=_reward_for_context( + item.get('reward'), + context_key=context.key, + ), + adapter_path=initial_path, + adapter_history=[initial_path], + ) + self.states.append(state) + self.train_batch_configs[context.key] = train_batch_config + + sampler_engine_args = { + 'tensor_parallel_size': sampler_tp, + 'enable_lora': True, + 'max_loras': int(runtime['sampler_max_loras']), + 'max_lora_rank': lora_data['r'], + 'max_model_len': int(sampler_config['max_model_len']), + 'gpu_memory_utilization': float(sampler_config['gpu_memory_utilization']), + 'max_num_seqs': int(sampler_config['max_num_seqs']), + 'enforce_eager': bool(sampler_config['enforce_eager']), + } + if sampler_config.get('max_num_batched_tokens') is not None: + sampler_engine_args['max_num_batched_tokens'] = int(sampler_config['max_num_batched_tokens']) + self.sampler = vLLMSampler( + model_id=runtime['model_id'], + remote_group='sampler', + device_mesh=sampler_mesh, + engine_args=sampler_engine_args, + ) + self.sampler.set_template( + template_cls, + model_id=runtime['model_id'], + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + self.output_dir = runtime['output_dir'] + self.max_steps = runtime.get('max_steps') + self.max_steps = None if self.max_steps is None else int(self.max_steps) + self.keep_adapter_versions = max(0, int(runtime.get('keep_adapter_versions', 0))) + self.metrics = create_metrics_reporter( + raw_config.get('metrics'), + run_id=str(runtime.get('run_id', 'sync_barrier_multi_lora_grpo')), + ) + self.completed_partitions = 0 + self._creation_order = 0 + + def _record_metric( + self, + stage: str, + *, + admission: PartitionAdmission | None = None, + context: LoraContext | None = None, + values: dict[str, Any] | None = None, + status: str = 'completed', + attributes: dict[str, Any] | None = None, + optimizer_step: int | None = None, + policy_version: int | None = None, + ) -> None: + if self.metrics is None: + return + self.metrics.record(MetricRecord( + stage=stage, + values=dict(values or {}), + context_key=( + admission.context.key if admission is not None + else context.key if context is not None else None + ), + partition_id=admission.partition_id if admission is not None else None, + partition_index=admission.step if admission is not None else None, + optimizer_step=optimizer_step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) + + def run(self) -> dict[str, Any]: + started = time.perf_counter() + try: + while self.max_steps is None or self.completed_partitions < self.max_steps: + partitions = self._rollout_round() + if not partitions: + break + self._advantage_round(partitions) + self._train_round(partitions) + except Exception as exc: + self._record_metric( + 'run', + status='failed', + values={'wall_time_s': time.perf_counter() - started}, + attributes={'error': f'{type(exc).__name__}: {exc}'}, + ) + if self.metrics is not None: + self.metrics.close() + raise + result = { + 'trained_partitions': self.completed_partitions, + 'wall_time_s': time.perf_counter() - started, + 'per_context': { + state.context.key: { + 'optimizer_steps': state.optimizer_steps, + 'policy_version': state.policy_version, + 'adapter_path': state.adapter_path, + } + for state in self.states + }, + } + self._record_metric( + 'run', + values={ + 'trained_partitions': result['trained_partitions'], + 'wall_time_s': result['wall_time_s'], + }, + ) + if self.metrics is not None: + self.metrics.flush() + result['metrics_health'] = self.metrics.health() + self.metrics.close() + return result + + def _rollout_round(self) -> list[SyncPartition]: + partitions = [] + for state in self.states: + if state.exhausted: + continue + if self.max_steps is not None and self.completed_partitions + len(partitions) >= self.max_steps: + break + prompts = next(state.prompt_batches, None) + if prompts is None or len(prompts) != state.rollout_batch_size: + state.exhausted = True + continue + admission = PartitionAdmission( + context=state.context, + partition_id=state.context.partition_id(state.partition_step), + step=state.partition_step, + target_groups=state.rollout_batch_size, + num_generations=state.num_generations, + created_order=self._creation_order, + ) + self._creation_order += 1 + self._record_metric( + 'rollout', + admission=admission, + status='submitted', + policy_version=state.policy_version, + values={ + 'prompt_count': admission.target_groups, + 'sample_count': admission.sample_count, + 'num_generations': admission.num_generations, + }, + attributes={'scope': 'partition'}, + ) + rollout_started = time.perf_counter() + sources = [{ + **dict(prompt), + 'group_id': f'{admission.partition_id}/group_{group_index}', + 'generation_idx': generation_index, + } for group_index, prompt in enumerate(prompts) + for generation_index in range(state.num_generations)] + responses = self.sampler.sample( + [dict(prompt) for prompt in prompts for _ in range(state.num_generations)], + state.sampling_params, + adapter_name=state.context.adapter_name, + adapter_path=state.adapter_path, + ) + rows = sample_responses_to_rollout_rows( + sources, + responses, + policy_version=state.policy_version, + ) + if len(rows) != admission.sample_count: + raise ValueError( + f'{admission.partition_id} expected {admission.sample_count} samples, got {len(rows)}') + for row in rows: + row.update({ + 'rollout_adapter_path': state.adapter_path, + 'rollout_policy_versions': [state.policy_version], + 'initial_policy_version': state.policy_version, + 'final_policy_version': state.policy_version, + 'policy_version_span': 0, + }) + rewards = [float(value) for value in state.reward_fn(rows, context=state.context)] + if len(rewards) != len(rows): + raise ValueError(f'{admission.partition_id} reward count does not match sample count') + rollout_latency_s = time.perf_counter() - rollout_started + self._record_rollout_groups(state, admission, rows, rewards) + self._record_metric( + 'rollout', + admission=admission, + policy_version=state.policy_version, + values=rollout_metrics( + completion_lengths=[int(row['completion_length']) for row in rows], + stop_reasons=[row.get('stop_reason') for row in rows], + rollout_latency_s=rollout_latency_s, + ), + attributes={'scope': 'partition'}, + ) + partitions.append(SyncPartition(admission, state, rows, rewards)) + state.partition_step += 1 + return partitions + + def _record_rollout_groups( + self, + state: SyncContextState, + admission: PartitionAdmission, + rows: list[dict[str, Any]], + rewards: list[float], + ) -> None: + for group_index in range(admission.target_groups): + start = group_index * admission.num_generations + end = start + admission.num_generations + group_rows = rows[start:end] + group_rewards = rewards[start:end] + metrics = { + **_compute_reward_metrics( + {state.context.key: state.reward_fn}, + state.context, + group_rows, + group_rewards, + ), + **rollout_metrics( + rewards={'reward': group_rewards}, + completion_lengths=[int(row['completion_length']) for row in group_rows], + stop_reasons=[row.get('stop_reason') for row in group_rows], + ), + } + self._record_metric( + 'rollout', + admission=admission, + policy_version=state.policy_version, + values=metrics, + attributes={ + 'scope': 'group', + 'group_id': f'{admission.partition_id}/group_{group_index}', + }, + ) + + def _advantage_round(self, partitions: list[SyncPartition]) -> None: + from twinkle.advantage import GRPOAdvantage + + advantage_fn = GRPOAdvantage() + for partition in partitions: + admission = partition.admission + partition.advantages = advantage_fn( + partition.rewards, + num_generations=admission.num_generations, + scale='group', + ).tolist() + samples_per_batch = admission.num_generations + for start in range(0, len(partition.rows), samples_per_batch): + end = min(start + samples_per_batch, len(partition.rows)) + self._record_metric( + 'advantage', + admission=admission, + policy_version=partition.state.policy_version, + values={ + 'sample_count': end - start, + **advantage_signal_metrics( + partition.rewards[start:end], + partition.advantages[start:end], + num_generations=admission.num_generations, + ), + }, + ) + + def _train_round(self, partitions: list[SyncPartition]) -> None: + for partition in partitions: + admission = partition.admission + state = partition.state + assert partition.advantages is not None + samples_per_batch = state.mini_batch_size + for start in range(0, len(partition.rows), samples_per_batch): + end = start + samples_per_batch + batch = self._training_batch( + partition.rows[start:end], + partition.rewards[start:end], + partition.advantages[start:end], + ) + train_started = time.perf_counter() + metrics = _train_batch( + self.model, + self.train_batch_configs, + batch, + admission, + model_data_parallel_size=self.model_data_parallel_size, + ) + state.optimizer_steps += 1 + metrics.update({ + 'sample_count': end - start, + 'train_latency_s': time.perf_counter() - train_started, + 'policy_version_gap_mean': 0.0, + 'policy_version_gap_p95': 0.0, + 'policy_version_gap_max': 0, + 'rollout_policy_span_mean': 0.0, + 'rollout_policy_span_max': 0, + }) + self._record_metric( + 'train', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values=metrics, + ) + finalize_started = time.perf_counter() + next_policy_version = state.policy_version + 1 + save_started = time.perf_counter() + state.adapter_path = self.model.save( + f'sync-{state.context.adapter_name}-v{next_policy_version}', + output_dir=self.output_dir, + adapter_name=state.context.adapter_name, + ) + adapter_save_latency_s = time.perf_counter() - save_started + publish_started = time.perf_counter() + state.policy_version = next_policy_version + policy_publish_latency_s = time.perf_counter() - publish_started + state.adapter_history.append(state.adapter_path) + self._record_metric( + 'policy', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + }, + attributes={'operation': 'publish', 'adapter_path': state.adapter_path}, + ) + prune_started = time.perf_counter() + self._prune_adapter_history(state) + adapter_prune_latency_s = time.perf_counter() - prune_started + self.completed_partitions += 1 + self._record_metric( + 'partition', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + 'adapter_prune_latency_s': adapter_prune_latency_s, + 'partition_finalize_latency_s': time.perf_counter() - finalize_started, + }, + ) + + @staticmethod + def _training_batch(rows: list[dict[str, Any]], rewards: list[float], advantages: list[float]): + fields = { + name: [row[name] for row in rows] + for name in (*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs') + } + fields.update({'rewards': rewards, 'advantages': advantages}) + return columns_to_tq_fields(fields, len(rows)) + + def _prune_adapter_history(self, state: SyncContextState) -> None: + retained_count = max(1, self.keep_adapter_versions) + stale = state.adapter_history[:-retained_count] + state.adapter_history = state.adapter_history[-retained_count:] + for path in stale: + if os.path.isdir(path): + shutil.rmtree(path) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument('--config', default='cookbook/rl/sync_barrier_multi_lora_grpo.yaml') + args = parser.parse_args() + config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) + print(SyncBarrierMultiLoraGRPO(config).run()) + + +if __name__ == '__main__': + main() + +# MODEL_ID=/path/to/model \ +# DATASET_ID=/path/to/gsm8k \ +# python cookbook/rl/async_multi_lora_grpo.py diff --git a/pyproject.toml b/pyproject.toml index 12fb2d929..71216ff04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,13 @@ rl = [ "vllm>=0.11", "ray[serve]" ] +async-rl = [ + "vllm>=0.11", + "ray[serve]", + "TransferQueue>=0.1.9.dev0", + "math-verify==0.8.0", + "swanlab>=0.8", +] client = [ "textual>=1.0.0", "plotext>=5.2.0", diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index 5fd046ae7..34fbfbefe 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -10,3 +10,9 @@ class Loss: def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: ... + + def micro_batch_scale(self, inputs: list[InputFeature], indices: list[int]) -> float: + if len(indices) == len(inputs): + return 1.0 + raise NotImplementedError( + f'{self.__class__.__name__} does not support micro-batching, including dynamic batching') diff --git a/src/twinkle/loss/chunked_cross_entropy.py b/src/twinkle/loss/chunked_cross_entropy.py index 0df7b6472..1e829a692 100644 --- a/src/twinkle/loss/chunked_cross_entropy.py +++ b/src/twinkle/loss/chunked_cross_entropy.py @@ -129,6 +129,21 @@ def __init__(self, chunk_size: int, ignore_index: int = -100, reduction: str = ' self.reduction = reduction self.dft = dft + def micro_batch_scale(self, inputs, indices): + if self.reduction == 'sum': + return 1.0 + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens + def __call__(self, inputs, outputs, **kwargs): labels = inputs['labels'] logps = outputs.get('logps') diff --git a/src/twinkle/loss/cross_entropy.py b/src/twinkle/loss/cross_entropy.py index c1b5225d6..8d3627c45 100644 --- a/src/twinkle/loss/cross_entropy.py +++ b/src/twinkle/loss/cross_entropy.py @@ -12,6 +12,21 @@ def __init__(self, ignore_index: int = -100, reduction='mean', dft: bool = False self.reduction = reduction self.dft = dft + def micro_batch_scale(self, inputs, indices): + if self.reduction == 'sum': + return 1.0 + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens + def __call__(self, inputs, outputs, **kwargs): labels = inputs['labels'] logps = outputs.get('logps') diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 781b22060..2136db2d7 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -32,7 +32,6 @@ def __init__( beta: float = 0.0, entropy_coef: float = 0.0, ignore_index: int = -100, - **kwargs, ): self.epsilon = epsilon self.epsilon_high = epsilon_high if epsilon_high is not None else epsilon @@ -42,6 +41,9 @@ def __init__( self.require_entropy = entropy_coef > 0.0 self.ignore_index = ignore_index + def micro_batch_scale(self, inputs, indices): + return len(indices) / len(inputs) + def _compute_log_importance_weights( self, per_token_logps: 'torch.Tensor', @@ -102,8 +104,7 @@ def _aggregate_loss( """ Aggregate per-token loss to scalar. - Override this method in subclasses for different normalization. - Default: mean over sequences, then mean over batch. + Mean over response tokens within each sequence, then over sequences. Args: per_token_loss: [batch, seq_len] per-token loss values @@ -113,7 +114,6 @@ def _aggregate_loss( Returns: loss: scalar loss value """ - # Per-sequence mean, then batch mean (aligned with Swift/TRL GRPO). # Each sequence contributes equally regardless of length. return ((per_token_loss * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).mean() @@ -366,6 +366,18 @@ class CISPOLoss(GRPOLoss): Clamps the IS weight and uses policy gradient. """ + def micro_batch_scale(self, inputs, indices): + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens def _compute_per_token_loss( self, @@ -397,6 +409,18 @@ class BNPOLoss(GRPOLoss): Normalizes by total completion tokens across batch. """ + def micro_batch_scale(self, inputs, indices): + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens def _aggregate_loss( self, diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index baeb6c1c9..a70a85142 100644 --- a/src/twinkle/metric/__init__.py +++ b/src/twinkle/metric/__init__.py @@ -1,9 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .accuracy import Accuracy from .base import Metric +from .buffer import MetricBuffer from .completion_and_reward import CompletionRewardMetric from .dpo import DPOMetric from .embedding import EmbeddingMetric from .grpo import CISPOMetric, GRPOMetric, GSPOMetric from .loss import LossMetric +from .reporting import MetricsReporter, create_metrics_reporter from .train_metric import TrainMetric +from .types import MetricRecord diff --git a/src/twinkle/metric/buffer.py b/src/twinkle/metric/buffer.py new file mode 100644 index 000000000..74d532a8c --- /dev/null +++ b/src/twinkle/metric/buffer.py @@ -0,0 +1,26 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Worker-local metric buffering.""" + +from __future__ import annotations + +import threading + +from .types import MetricRecord + + +class MetricBuffer: + """Thread-safe, destructive worker-local metric buffer.""" + + def __init__(self): + self._records: list[MetricRecord] = [] + self._lock = threading.Lock() + + def record(self, record: MetricRecord) -> None: + with self._lock: + self._records.append(record) + + def drain(self) -> list[MetricRecord]: + with self._lock: + records = self._records + self._records = [] + return records diff --git a/src/twinkle/metric/reporting.py b/src/twinkle/metric/reporting.py new file mode 100644 index 000000000..ce7d20104 --- /dev/null +++ b/src/twinkle/metric/reporting.py @@ -0,0 +1,593 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Asynchronous JSONL and SwanLab metric reporting.""" + +from __future__ import annotations + +import json +import logging +import math +import os +import re +import threading +import time +from collections import Counter, defaultdict, deque +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +from .types import MetricRecord + +logger = logging.getLogger(__name__) + + +def _finite_number(value: Any) -> float | int | None: + if isinstance(value, bool): + return int(value) + if isinstance(value, int): + return value + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, str): + match = re.fullmatch(r'\s*tensor\(([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)\)\s*', value) + candidate = match.group(1) if match else value + try: + number = float(candidate) + except ValueError: + return None + return number if math.isfinite(number) else None + try: + scalar = value.item() + except (AttributeError, RuntimeError, ValueError): + return None + return _finite_number(scalar) + + +def _safe_name(value: str) -> str: + return re.sub(r'[^A-Za-z0-9_.-]+', '_', value).strip('_') or 'default' + + +@dataclass +class _ScalarSummary: + count: int = 0 + total: float = 0.0 + minimum: float = math.inf + maximum: float = -math.inf + last: float = 0.0 + + def add(self, value: float | int) -> None: + number = float(value) + self.count += 1 + self.total += number + self.minimum = min(self.minimum, number) + self.maximum = max(self.maximum, number) + self.last = number + + def as_dict(self) -> dict[str, float | int]: + return { + 'count': self.count, + 'last': self.last, + 'mean': self.total / self.count, + 'min': self.minimum, + 'max': self.maximum, + } + + +class _SummaryReducer: + + def __init__(self, run_id: str, started_at: float): + self.run_id = run_id + self.started_at = started_at + self.record_counts: Counter[str] = Counter() + self.context_counts: dict[str, Counter[str]] = defaultdict(Counter) + self.context_rollout_groups: Counter[str] = Counter() + self.context_rollout_samples: Counter[str] = Counter() + self.context_trained_samples: Counter[str] = Counter() + self.context_optimizer_steps: dict[str, int] = {} + self.context_policy_versions: dict[str, int] = {} + self.metric_summaries: dict[str, _ScalarSummary] = defaultdict(_ScalarSummary) + self.run_status = 'running' + self.result: dict[str, Any] = {} + + def add(self, record: MetricRecord) -> None: + record_key = f'{record.stage}:{record.status}' + self.record_counts[record_key] += 1 + if record.context_key: + counts = self.context_counts[record.context_key] + counts[record_key] += 1 + if record.optimizer_step is not None: + self.context_optimizer_steps[record.context_key] = record.optimizer_step + if record.policy_version is not None: + previous_version = self.context_policy_versions.get(record.context_key) + self.context_policy_versions[record.context_key] = ( + record.policy_version + if previous_version is None + else max(previous_version, record.policy_version) + ) + sample_count = _finite_number(record.values.get('sample_count')) + if sample_count is not None and record.status == 'completed': + if record.stage == 'rollout' and record.attributes.get('scope', 'group') == 'group': + self.context_rollout_groups[record.context_key] += 1 + self.context_rollout_samples[record.context_key] += int(sample_count) + elif record.stage == 'train': + self.context_trained_samples[record.context_key] += int(sample_count) + summarize_values = ( + record.status == 'completed' + and ( + record.stage != 'rollout' + or record.attributes.get('scope', 'group') == 'group' + ) + ) + if summarize_values: + for name, value in record.values.items(): + number = _finite_number(value) + if number is not None: + self.metric_summaries[f'{record.stage}/{name}'].add(number) + if record.stage == 'run': + self.run_status = record.status + self.result = {**record.values, **record.attributes} + + def as_dict(self, backend_health: Mapping[str, Any]) -> dict[str, Any]: + wall_time = _finite_number(self.result.get('wall_time_s')) + if wall_time is None: + wall_time = time.time() - self.started_at + rollout_groups = sum(self.context_rollout_groups.values()) + train_steps = sum(counts['train:completed'] for counts in self.context_counts.values()) + trained_partitions = sum( + counts['partition:completed'] + for counts in self.context_counts.values() + ) + terminal_partitions = _finite_number(self.result.get('trained_partitions')) + if terminal_partitions is not None: + trained_partitions = int(terminal_partitions) + rollout_samples = sum(self.context_rollout_samples.values()) + trained_samples = sum(self.context_trained_samples.values()) + dropped_records = sum( + int(item.get('dropped_records', 0)) + for item in backend_health.values() + ) + backend_write_latency_s = sum( + float(item.get('write_latency_s', 0.0)) + for item in backend_health.values() + ) + contexts = {} + for context_key, counts in self.context_counts.items(): + contexts[context_key] = { + 'rollout_groups': self.context_rollout_groups[context_key], + 'rollout_samples': self.context_rollout_samples[context_key], + 'train_steps': counts['train:completed'], + 'trained_samples': self.context_trained_samples[context_key], + 'trained_partitions': counts['partition:completed'], + 'optimizer_step': self.context_optimizer_steps.get(context_key), + 'policy_version': self.context_policy_versions.get(context_key), + } + return { + 'run_id': self.run_id, + 'status': self.run_status, + 'wall_time_s': wall_time, + 'record_counts': dict(sorted(self.record_counts.items())), + 'rollout_groups': rollout_groups, + 'rollout_samples': rollout_samples, + 'train_steps': train_steps, + 'trained_samples': trained_samples, + 'trained_partitions': trained_partitions, + 'rollout_groups_per_sec': rollout_groups / wall_time if wall_time > 0 else 0.0, + 'rollout_samples_per_sec': rollout_samples / wall_time if wall_time > 0 else 0.0, + 'train_steps_per_hour': train_steps * 3600 / wall_time if wall_time > 0 else 0.0, + 'trained_samples_per_sec': trained_samples / wall_time if wall_time > 0 else 0.0, + 'train_partitions_per_hour': trained_partitions * 3600 / wall_time if wall_time > 0 else 0.0, + 'dropped_records': dropped_records, + 'backend_write_latency_s': backend_write_latency_s, + 'per_context': contexts, + 'metrics': { + name: summary.as_dict() + for name, summary in sorted(self.metric_summaries.items()) + }, + 'backends': dict(backend_health), + 'result': self.result, + } + + +class _QueuedBackend: + + def __init__( + self, + name: str, + *, + queue_capacity: int, + batch_size: int, + flush_interval_s: float, + ): + if queue_capacity <= 0: + raise ValueError('queue_capacity must be positive') + if batch_size <= 0: + raise ValueError('batch_size must be positive') + if flush_interval_s <= 0: + raise ValueError('flush_interval_s must be positive') + self.name = name + self.queue_capacity = queue_capacity + self.batch_size = batch_size + self.flush_interval_s = flush_interval_s + self._queue: deque[dict[str, Any]] = deque() + self._condition = threading.Condition() + self._closing = False + self._flush_requested = False + self._closed = False + self._disabled = False + self._in_flight = False + self._submitted = 0 + self._written = 0 + self._dropped = 0 + self._write_batches = 0 + self._write_latency_s = 0.0 + self._failures = 0 + self._last_error: str | None = None + self._warning_emitted = False + self._thread = threading.Thread(target=self._run, name=f'twinkle-metrics-{name}', daemon=True) + self._thread.start() + + def submit(self, payload: dict[str, Any]) -> None: + with self._condition: + if self._closing or self._disabled: + self._dropped += 1 + return + if len(self._queue) >= self.queue_capacity: + self._queue.popleft() + self._dropped += 1 + self._queue.append(payload) + self._submitted += 1 + self._condition.notify() + + def flush(self, timeout_s: float | None = None) -> bool: + deadline = None if timeout_s is None else time.monotonic() + timeout_s + with self._condition: + self._flush_requested = True + self._condition.notify_all() + while (self._queue or self._in_flight) and not self._disabled: + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + return False + self._condition.wait(remaining) + self._flush_requested = False + return True + + def close(self, timeout_s: float | None = None) -> bool: + with self._condition: + if self._closed: + return True + self._closing = True + self._condition.notify_all() + self._thread.join(timeout_s) + closed = not self._thread.is_alive() + if closed: + self._closed = True + return closed + + def health(self) -> dict[str, Any]: + with self._condition: + return { + 'enabled': not self._disabled, + 'queue_size': len(self._queue), + 'submitted_records': self._submitted, + 'written_records': self._written, + 'dropped_records': self._dropped, + 'write_batches': self._write_batches, + 'write_latency_s': self._write_latency_s, + 'failure_count': self._failures, + 'last_error': self._last_error, + } + + def _run(self) -> None: + last_write = time.monotonic() + while True: + with self._condition: + while True: + if self._closing and not self._queue: + batch = [] + break + elapsed = time.monotonic() - last_write + should_write = bool(self._queue) and ( + self._closing + or self._flush_requested + or len(self._queue) >= self.batch_size + or elapsed >= self.flush_interval_s + ) + if should_write: + batch = [ + self._queue.popleft() + for _ in range(min(len(self._queue), self.batch_size)) + ] + break + wait_s = ( + max(0.0, self.flush_interval_s - elapsed) + if self._queue else self.flush_interval_s + ) + self._condition.wait(wait_s) + if not batch: + break + self._in_flight = True + started = time.perf_counter() + try: + self._write_batch(batch) + except Exception as exc: + with self._condition: + self._failures += 1 + self._last_error = f'{type(exc).__name__}: {exc}' + self._dropped += len(batch) + len(self._queue) + self._queue.clear() + self._disabled = True + if not self._warning_emitted: + logger.warning('Metrics backend %s failed and was disabled: %s', self.name, exc) + self._warning_emitted = True + else: + elapsed = time.perf_counter() - started + with self._condition: + self._written += len(batch) + self._write_batches += 1 + self._write_latency_s += elapsed + finally: + last_write = time.monotonic() + with self._condition: + self._in_flight = False + self._condition.notify_all() + if self._disabled: + break + try: + self._close_sink() + except Exception as exc: + with self._condition: + self._failures += 1 + self._last_error = f'{type(exc).__name__}: {exc}' + with self._condition: + self._closed = True + self._condition.notify_all() + + def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: + raise NotImplementedError + + def _close_sink(self) -> None: + pass + + +class _JSONLBackend(_QueuedBackend): + + def __init__(self, path: str | Path, **kwargs: Any): + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._stream = self.path.open('w', encoding='utf-8') + super().__init__('jsonl', **kwargs) + + def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: + self._stream.writelines( + json.dumps(payload, ensure_ascii=True, default=str) + '\n' + for payload in batch + ) + self._stream.flush() + + def _close_sink(self) -> None: + self._stream.close() + + +class _SwanLabBackend(_QueuedBackend): + + def __init__( + self, + *, + project: str, + experiment_name: str, + log_dir: str | Path, + mode: str, + **kwargs: Any, + ): + import swanlab + + self._swanlab = swanlab + self._swanlab_run = swanlab.init( + project=project, + experiment_name=experiment_name, + logdir=str(log_dir), + mode=mode, + ) + super().__init__('swanlab', **kwargs) + + def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: + for payload in batch: + prefix = ( + f'context/{_safe_name(payload["context_key"])}' + if payload.get('context_key') else 'global' + ) + stage = _safe_name(payload['stage']) + values = {} + for name, value in payload['values'].items(): + metric_name = str(name) + if metric_name.startswith(f'{stage}/'): + metric_name = metric_name[len(stage) + 1:] + values[f'{prefix}/{stage}/{_safe_name(metric_name)}'] = value + if payload.get('optimizer_step') is not None: + values[f'{prefix}/train/optimizer_step'] = payload['optimizer_step'] + if payload.get('policy_version') is not None: + values[f'{prefix}/policy/version'] = payload['policy_version'] + if payload.get('partition_index') is not None: + values[f'{prefix}/partition/index'] = payload['partition_index'] + if values: + self._swanlab_run.log(values, step=payload['sequence']) + + def _close_sink(self) -> None: + self._swanlab.finish() + + +class MetricsReporter: + """Assign ordering and asynchronously fan metric records out to backends.""" + + def __init__( + self, + *, + run_id: str, + backends: Sequence[_QueuedBackend] = (), + summary_path: str | Path | None = None, + close_timeout_s: float = 10.0, + ): + self.run_id = run_id + self.started_at = time.time() + self.close_timeout_s = close_timeout_s + self.summary_path = Path(summary_path) if summary_path is not None else None + self._backends = tuple(backends) + self._lock = threading.Lock() + self._sequence = 0 + self._closed = False + self._reducer = _SummaryReducer(run_id, self.started_at) + self._initial_errors: dict[str, str] = {} + + def add_backend_error(self, name: str, error: BaseException) -> None: + self._initial_errors[name] = f'{type(error).__name__}: {error}' + + def record(self, record: MetricRecord) -> None: + with self._lock: + if self._closed: + return + self._sequence += 1 + normalized = self._normalize(record, self._sequence) + self._reducer.add(normalized) + payload = self._payload(normalized) + for backend in self._backends: + backend.submit(payload) + + def record_many(self, records: Sequence[MetricRecord]) -> None: + for record in records: + self.record(record) + + def flush(self, timeout_s: float | None = None) -> None: + timeout = self.close_timeout_s if timeout_s is None else timeout_s + deadline = time.monotonic() + timeout + for backend in self._backends: + backend.flush(max(0.0, deadline - time.monotonic())) + + def close(self, timeout_s: float | None = None) -> None: + timeout = self.close_timeout_s if timeout_s is None else timeout_s + with self._lock: + if self._closed: + return + self._closed = True + deadline = time.monotonic() + timeout + for backend in self._backends: + backend.close(max(0.0, deadline - time.monotonic())) + self._write_summary() + + def health(self) -> dict[str, Any]: + backend_health = { + backend.name: backend.health() + for backend in self._backends + } + for name, error in self._initial_errors.items(): + backend_health[name] = { + 'enabled': False, + 'failure_count': 1, + 'last_error': error, + 'dropped_records': self._sequence, + } + return { + 'record_count': self._sequence, + 'dropped_records': sum( + int(item.get('dropped_records', 0)) + for item in backend_health.values() + ), + 'backends': backend_health, + } + + def summary(self) -> dict[str, Any]: + return self._reducer.as_dict(self.health()['backends']) + + def _normalize(self, record: MetricRecord, sequence: int) -> MetricRecord: + values: dict[str, float | int] = {} + non_numeric = {} + for name, value in record.values.items(): + number = _finite_number(value) + if number is None: + non_numeric[name] = value + else: + values[name] = number + attributes = dict(record.attributes) + if non_numeric: + attributes['non_numeric_values'] = non_numeric + return replace(record, sequence=sequence, values=values, attributes=attributes) + + def _payload(self, record: MetricRecord) -> dict[str, Any]: + return { + 'timestamp': record.timestamp, + 'elapsed_s': record.timestamp - self.started_at, + 'sequence': record.sequence, + 'run_id': self.run_id, + 'stage': record.stage, + 'context_key': record.context_key, + 'partition_id': record.partition_id, + 'partition_index': record.partition_index, + 'optimizer_step': record.optimizer_step, + 'policy_version': record.policy_version, + 'status': record.status, + 'values': record.values, + 'attributes': record.attributes, + } + + def _write_summary(self) -> None: + if self.summary_path is None: + return + try: + self.summary_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = self.summary_path.with_suffix(f'{self.summary_path.suffix}.tmp') + with temporary_path.open('w', encoding='utf-8') as stream: + json.dump(self.summary(), stream, ensure_ascii=True, indent=2, default=str) + stream.write('\n') + os.replace(temporary_path, self.summary_path) + except Exception as exc: + logger.warning('Failed to write metrics summary %s: %s', self.summary_path, exc) + + +def create_metrics_reporter(config: Mapping[str, Any] | None, *, run_id: str) -> MetricsReporter | None: + if config is None: + return None + config = dict(config or {}) + if not bool(config.get('enabled', True)): + return None + queue_capacity = int(config.get('queue_capacity', 10000)) + close_timeout_s = float(config.get('close_timeout_s', 10.0)) + jsonl_config = dict(config.get('jsonl') or {}) + swanlab_config = dict(config.get('swanlab') or {}) + backends: list[_QueuedBackend] = [] + backend_errors: dict[str, BaseException] = {} + backend_defaults = { + 'queue_capacity': queue_capacity, + 'batch_size': int(jsonl_config.get('batch_size', 64)), + 'flush_interval_s': float(jsonl_config.get('flush_interval_s', 2.0)), + } + if bool(jsonl_config.get('enabled', True)): + try: + backends.append(_JSONLBackend(jsonl_config['path'], **backend_defaults)) + except Exception as exc: + backend_errors['jsonl'] = exc + logger.warning('JSONL metrics backend could not start: %s', exc) + if bool(swanlab_config.get('enabled', False)) and swanlab_config.get('mode') != 'disabled': + try: + backends.append(_SwanLabBackend( + project=str(swanlab_config.get('project', 'twinkle-rl')), + experiment_name=str(swanlab_config.get('name', run_id)), + log_dir=swanlab_config.get('log_dir', 'outputs/swanlab'), + mode=str(swanlab_config.get('mode', 'local')), + queue_capacity=queue_capacity, + batch_size=int(swanlab_config.get('batch_size', 16)), + flush_interval_s=float(swanlab_config.get('flush_interval_s', 1.0)), + )) + except Exception as exc: + backend_errors['swanlab'] = exc + logger.warning('SwanLab metrics backend could not start: %s', exc) + summary_path = jsonl_config.get('summary_path') + if summary_path is None and jsonl_config.get('path') is not None: + summary_path = Path(jsonl_config['path']).with_name('summary.json') + reporter = MetricsReporter( + run_id=run_id, + backends=backends, + summary_path=summary_path, + close_timeout_s=close_timeout_s, + ) + for name, error in backend_errors.items(): + reporter.add_backend_error(name, error) + return reporter diff --git a/src/twinkle/metric/types.py b/src/twinkle/metric/types.py new file mode 100644 index 000000000..68741dff7 --- /dev/null +++ b/src/twinkle/metric/types.py @@ -0,0 +1,42 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Transport-neutral metric value types.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any + +METRIC_STAGES = frozenset({ + 'rollout', + 'advantage', + 'train', + 'evaluation', + 'partition', + 'policy', + 'run', +}) +METRIC_STATUSES = frozenset({'submitted', 'completed', 'failed'}) + + +@dataclass(frozen=True) +class MetricRecord: + stage: str + values: dict[str, Any] + timestamp: float = field(default_factory=time.time) + sequence: int | None = None + context_key: str | None = None + partition_id: str | None = None + partition_index: int | None = None + optimizer_step: int | None = None + policy_version: int | None = None + status: str = 'completed' + attributes: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.stage not in METRIC_STAGES: + raise ValueError(f'unsupported metric stage: {self.stage!r}') + if self.status not in METRIC_STATUSES: + raise ValueError(f'unsupported metric status: {self.status!r}') + object.__setattr__(self, 'values', dict(self.values)) + object.__setattr__(self, 'attributes', dict(self.attributes)) diff --git a/src/twinkle/model/__init__.py b/src/twinkle/model/__init__.py index 88f544d6c..6930963f6 100644 --- a/src/twinkle/model/__init__.py +++ b/src/twinkle/model/__init__.py @@ -6,11 +6,13 @@ if TYPE_CHECKING: from .base import TwinkleModel from .megatron import MegatronModel, MultiLoraMegatronModel + from .micro_batch import MicroBatchConfig from .transformers import MultiLoraTransformersModel, TransformersModel else: _import_structure = { 'base': ['TwinkleModel'], + 'micro_batch': ['MicroBatchConfig'], 'transformers': ['TransformersModel', 'MultiLoraTransformersModel'], 'megatron': ['MegatronModel', 'MultiLoraMegatronModel'], } diff --git a/src/twinkle/model/micro_batch.py b/src/twinkle/model/micro_batch.py new file mode 100644 index 000000000..561013745 --- /dev/null +++ b/src/twinkle/model/micro_batch.py @@ -0,0 +1,212 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Packing algorithms are adapted from AReaL (Apache-2.0). +from __future__ import annotations + +import heapq +import math +from dataclasses import dataclass +from typing import Any, Literal + + +@dataclass(frozen=True) +class MicroBatchConfig: + micro_batch_size: int + dynamic_batching: bool = False + max_tokens_per_micro_batch: int | None = None + packing_algorithm: Literal['ffd', 'kk'] = 'ffd' + + def __post_init__(self): + if self.micro_batch_size <= 0: + raise ValueError(f'micro_batch_size must be positive, got {self.micro_batch_size}') + if self.packing_algorithm not in ('ffd', 'kk'): + raise ValueError(f'packing_algorithm must be ffd or kk, got {self.packing_algorithm!r}') + if self.dynamic_batching and ( + self.max_tokens_per_micro_batch is None or self.max_tokens_per_micro_batch <= 0): + raise ValueError('max_tokens_per_micro_batch must be positive when dynamic_batching=true') + + @classmethod + def from_kwargs(cls, kwargs: dict[str, Any]) -> 'MicroBatchConfig | None': + option_names = ( + 'micro_batch_size', + 'dynamic_batching', + 'max_tokens_per_micro_batch', + 'packing_algorithm', + ) + if not any(name in kwargs for name in option_names): + return None + if 'micro_batch_size' not in kwargs: + raise ValueError('micro_batch_size is required when configuring micro-batching') + return cls( + micro_batch_size=int(kwargs.pop('micro_batch_size')), + dynamic_batching=bool(kwargs.pop('dynamic_batching', False)), + max_tokens_per_micro_batch=kwargs.pop('max_tokens_per_micro_batch', None), + packing_algorithm=kwargs.pop('packing_algorithm', 'ffd'), + ) + + +def sequence_length(model_input: dict[str, Any]) -> int: + input_ids = model_input['input_ids'] + return int(input_ids.shape[-1]) if hasattr(input_ids, 'shape') else len(input_ids) + + +def _batch_cost(group: list[int], lengths: list[int], padding_free: bool) -> int: + if not group: + return 0 + values = [lengths[index] for index in group] + return sum(values) if padding_free else max(values) * len(values) + + +def _fits(group: list[int], index: int, lengths: list[int], config: MicroBatchConfig, + padding_free: bool) -> bool: + if len(group) >= config.micro_batch_size: + return False + candidate = [*group, index] + return _batch_cost(candidate, lengths, padding_free) <= config.max_tokens_per_micro_batch + + +def _ffd_allocate(lengths: list[int], config: MicroBatchConfig, padding_free: bool, + min_micro_batches: int) -> list[list[int]]: + groups: list[list[int]] = [[] for _ in range(min_micro_batches)] + for index in sorted(range(len(lengths)), key=lengths.__getitem__, reverse=True): + candidates = [ + group_index for group_index, group in enumerate(groups) + if _fits(group, index, lengths, config, padding_free) + ] + if not candidates: + groups.append([index]) + continue + group_index = min( + candidates, + key=lambda candidate: ( + _batch_cost(groups[candidate], lengths, padding_free), + len(groups[candidate]), + ), + ) + groups[group_index].append(index) + return [group for group in groups if group] + + +class _KKSet: + __slots__ = ('total', 'items') + + def __init__(self): + self.total = 0 + self.items: list[int] = [] + + def add(self, index: int, value: int) -> None: + self.items.append(index) + self.total += value + + def merge(self, other: '_KKSet') -> None: + self.items.extend(other.items) + self.total += other.total + + def __lt__(self, other: '_KKSet') -> bool: + return (self.total, len(self.items), self.items) < (other.total, len(other.items), other.items) + + +class _KKState: + __slots__ = ('sets', ) + + def __init__(self, items: list[tuple[int, int]], group_count: int): + self.sets = [_KKSet() for _ in range(group_count)] + for group, (index, value) in zip(self.sets, items): + group.add(index, value) + self.sets.sort(reverse=True) + + @property + def spread(self) -> int: + return self.sets[0].total - self.sets[-1].total + + def merge(self, other: '_KKState') -> None: + for index in range(len(self.sets)): + self.sets[index].merge(other.sets[-1 - index]) + self.sets.sort(reverse=True) + + def __lt__(self, other: '_KKState') -> bool: + return self.spread > other.spread + + +def _kk_partition(lengths: list[int], group_count: int) -> list[list[int]]: + queue = [] + for value, index in sorted((value, index) for index, value in enumerate(lengths)): + heapq.heappush(queue, _KKState([(index, value)], group_count)) + while len(queue) > 1: + first = heapq.heappop(queue) + second = heapq.heappop(queue) + first.merge(second) + heapq.heappush(queue, first) + return [group.items for group in queue[0].sets if group.items] + + +def _kk_allocate(lengths: list[int], config: MicroBatchConfig, padding_free: bool, + min_micro_batches: int) -> list[list[int]]: + capacity = config.max_tokens_per_micro_batch + group_count = max(min_micro_batches, math.ceil(sum(lengths) / capacity)) + while group_count <= len(lengths): + groups = _kk_partition(lengths, group_count) + if all( + len(group) <= config.micro_batch_size + and _batch_cost(group, lengths, padding_free) <= capacity + for group in groups): + return groups + group_count += 1 + raise ValueError('unable to construct a valid KK micro-batch plan') + + +def plan_micro_batches( + inputs: list[dict[str, Any]], + config: MicroBatchConfig, + *, + padding_free: bool, + min_micro_batches: int = 1, +) -> list[list[int]]: + if not inputs: + raise ValueError('cannot plan micro-batches for empty inputs') + if min_micro_batches <= 0 or min_micro_batches > len(inputs): + raise ValueError(f'invalid min_micro_batches={min_micro_batches} for {len(inputs)} inputs') + if not config.dynamic_batching: + group_count = max(min_micro_batches, math.ceil(len(inputs) / config.micro_batch_size)) + base_size, remainder = divmod(len(inputs), group_count) + groups = [] + start = 0 + for group_index in range(group_count): + size = base_size + int(group_index < remainder) + groups.append(list(range(start, start + size))) + start += size + return groups + lengths = [sequence_length(model_input) for model_input in inputs] + capacity = config.max_tokens_per_micro_batch + oversized = [length for length in lengths if length > capacity] + if oversized: + raise ValueError( + f'sequence length {max(oversized)} exceeds max_tokens_per_micro_batch={capacity}') + if config.packing_algorithm == 'ffd': + return _ffd_allocate(lengths, config, padding_free, min_micro_batches) + return _kk_allocate(lengths, config, padding_free, min_micro_batches) + + +def select_batch(value: Any, indices: list[int], batch_size: int) -> Any: + if isinstance(value, list): + return [value[index] for index in indices] if len(value) == batch_size else value + if isinstance(value, tuple): + return tuple(value[index] for index in indices) if len(value) == batch_size else value + if hasattr(value, 'shape') and len(value.shape) > 0 and value.shape[0] == batch_size: + return value[indices] + return value + + +def collect_micro_batch_outputs(outputs: list[dict[str, Any]], device_mesh: Any) -> dict[str, Any]: + from twinkle.infra.collectors import collect_tensor_dict + + result = collect_tensor_dict(outputs, device_mesh) + if len(outputs) <= 1 or 'micro_batch_count' not in outputs[0]: + return result + collected = [output for index, output in enumerate(outputs) if index in device_mesh.get_collect_ranks()] + result['micro_batch_count'] = collected[0]['micro_batch_count'] + result['micro_batch_samples_mean'] = ( + sum(output['micro_batch_samples_mean'] for output in collected) / len(collected)) + result['micro_batch_tokens_mean'] = ( + sum(output['micro_batch_tokens_mean'] for output in collected) / len(collected)) + result['micro_batch_tokens_max'] = max(output['micro_batch_tokens_max'] for output in collected) + return result diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 017a515b3..d7fb1c783 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -33,6 +33,8 @@ from twinkle.loss import CrossEntropyLoss, Loss from twinkle.metric import Accuracy, LossMetric, Metric, TrainMetric from twinkle.model.base import TwinkleModel +from twinkle.model.micro_batch import (MicroBatchConfig, collect_micro_batch_outputs, plan_micro_batches, + select_batch) from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus from twinkle.model.transformers.moe import apply_expert_parallel from twinkle.model.transformers.strategy import AccelerateStrategy, NativeFSDPStrategy @@ -683,8 +685,14 @@ def backward(self, **kwargs): self.set_grad_scaler(adapter_name=adapter_name) scaler = optimizer_config.scaler - optimizer_config.cur_step += 1 - should_sync = optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')) + increment_step = kwargs.pop('_increment_step', True) + if increment_step: + optimizer_config.cur_step += 1 + sync_gradients = kwargs.pop('sync_gradients', None) + should_sync = ( + optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')) + if sync_gradients is None else bool(sync_gradients) + ) import contextlib no_sync_ctx = contextlib.nullcontext() @@ -708,7 +716,143 @@ def backward(self, **kwargs): optimizer_config.train_status.loss_value = None - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) + def _build_micro_batch_plan(self, inputs, config, optimizer_config): + processor = optimizer_config.processor + assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' + optimizer_config._ensure_dp_group() + dp_group = optimizer_config._dp_group + min_micro_batches = 1 + while True: + try: + plan = plan_micro_batches( + inputs, + config, + padding_free=processor.padding_free, + min_micro_batches=min_micro_batches, + ) + planning_error = None + except Exception as exc: + if dp_group is None: + raise + plan = None + planning_error = f'{type(exc).__name__}: {exc}' + + if dp_group is None: + return plan + + local_state = { + 'micro_batch_count': len(plan) if plan is not None else None, + 'input_count': len(inputs), + 'error': planning_error, + } + states = [None] * dist.get_world_size(dp_group) + dist.all_gather_object(states, local_state, group=dp_group) + errors = [ + f'rank {rank}: {state["error"]}' + for rank, state in enumerate(states) + if state['error'] is not None + ] + if errors: + raise RuntimeError( + 'micro-batch planning failed on one or more model DP ranks: ' + + '; '.join(errors)) + + counts = [state['micro_batch_count'] for state in states] + if all(count == len(plan) for count in counts): + return plan + min_micro_batches = max(counts) + if any(min_micro_batches > state['input_count'] for state in states): + raise ValueError( + 'model DP ranks cannot execute the same number of non-empty micro-batches; ' + 'make the input batch divisible by the model data-parallel size') + + def _forward_backward_micro_batch( + self, + *, + inputs, + optimizer_config, + loss_scale, + sync_gradients, + increment_step, + **kwargs, + ): + outputs = self.forward( + inputs=inputs, + router_replay_manual_cleanup=True, + **kwargs, + ) + previous_normalizer = optimizer_config.train_status.num_tokens + loss = self.calculate_loss(**kwargs) + normalizer_delta = optimizer_config.train_status.num_tokens - previous_normalizer + optimizer_config.train_status.loss_value = ( + optimizer_config.train_status.loss_value * loss_scale) + optimizer_config.train_status.num_tokens = ( + previous_normalizer + normalizer_delta * loss_scale) + outputs['loss'] = loss * loss_scale + self.backward( + sync_gradients=sync_gradients, + _increment_step=increment_step, + **kwargs, + ) + return outputs + + def _forward_backward_micro_batches( + self, + *, + inputs, + config, + sync_gradients, + loss_scale, + **kwargs, + ): + adapter_name = kwargs.get('adapter_name') + if adapter_name is None: + adapter_name = self._get_default_group() + optimizer_config = self.optimizer_group[adapter_name] + if isinstance(inputs, dict): + inputs = [inputs] + if self._not_encoded(inputs[0]): + assert optimizer_config.template is not None, \ + 'Use set_template to add a template when trying to input `List[Trajectory]`' + inputs = optimizer_config.template.batch_encode(inputs) + + local_batch_size = len(inputs) + processor: InputProcessor = optimizer_config.processor + plan = self._build_micro_batch_plan(inputs, config, optimizer_config) + outputs = {} + micro_batch_samples = [] + micro_batch_tokens = [] + loss_instance = optimizer_config.loss_instance + for micro_batch_index, indices in enumerate(plan): + micro_kwargs = { + key: select_batch(value, indices, local_batch_size) + for key, value in kwargs.items() + } + micro_loss_scale = loss_scale * loss_instance.micro_batch_scale(inputs, indices) + is_last_micro_batch = micro_batch_index == len(plan) - 1 + outputs = self._forward_backward_micro_batch( + inputs=[inputs[index] for index in indices], + optimizer_config=optimizer_config, + loss_scale=micro_loss_scale, + sync_gradients=sync_gradients if is_last_micro_batch else False, + increment_step=is_last_micro_batch, + **micro_kwargs, + ) + lengths = [ + int(inputs[index]['input_ids'].shape[-1]) + if hasattr(inputs[index]['input_ids'], 'shape') else len(inputs[index]['input_ids']) + for index in indices + ] + micro_batch_samples.append(len(indices)) + micro_batch_tokens.append( + sum(lengths) if processor.padding_free else max(lengths) * len(lengths)) + outputs['micro_batch_count'] = len(plan) + outputs['micro_batch_samples_mean'] = sum(micro_batch_samples) / len(plan) + outputs['micro_batch_tokens_mean'] = sum(micro_batch_tokens) / len(plan) + outputs['micro_batch_tokens_max'] = max(micro_batch_tokens) + return outputs + + @remote_function(dispatch='slice_dp', collect=collect_micro_batch_outputs) def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): """Do forward, calculate loss, and backward. @@ -718,10 +862,26 @@ def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Tr **kwargs: adapter_name: Lora adapter name. gradient_accumulation_steps: Number of gradient accumulation steps. + micro_batch_size: Maximum samples processed per rank in one forward/backward. + dynamic_batching: Pack sequences by token cost instead of fixed sample slices. + max_tokens_per_micro_batch: Per-rank token limit used by dynamic batching. + packing_algorithm: Dynamic packing algorithm. + sync_gradients: Override gradient synchronization on the final micro-batch. + loss_scale: Weight applied to this input batch's loss. Any parameters needed for the specific loss type. Returns: The output of the model forward. """ + micro_batch_config = MicroBatchConfig.from_kwargs(kwargs) + if micro_batch_config is not None: + return self._forward_backward_micro_batches( + inputs=inputs, + config=micro_batch_config, + sync_gradients=kwargs.pop('sync_gradients', None), + loss_scale=float(kwargs.pop('loss_scale', 1.0)), + **kwargs, + ) + outputs = self.forward(inputs=inputs, router_replay_manual_cleanup=True, **kwargs) loss = self.calculate_loss(**kwargs) outputs['loss'] = loss diff --git a/src/twinkle/reward/__init__.py b/src/twinkle/reward/__init__.py index 3ba5babd4..4e8bed086 100644 --- a/src/twinkle/reward/__init__.py +++ b/src/twinkle/reward/__init__.py @@ -1,7 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .base import Reward +from .boxed_math import BoxedMathAccuracyReward +from .dapo_math import DAPOMathAccuracyReward, DAPOMathReward from .format_reward import FormatReward -from .gsm8k import GSM8KAccuracyReward, GSM8KFormatReward +from .gsm8k import (GSM8KAccuracyBrevityReward, GSM8KAccuracyReward, GSM8KBrevityReward, GSM8KFormatReward, + MathVerifyAccuracyReward) from .math_reward import MathReward from .mm_reward import MultiModalAccuracyReward from .olympiad_bench import OlympiadBenchAccuracyReward, OlympiadBenchFormatReward, OlympiadBenchQualityReward diff --git a/src/twinkle/reward/boxed_math.py b/src/twinkle/reward/boxed_math.py new file mode 100644 index 000000000..f9134944f --- /dev/null +++ b/src/twinkle/reward/boxed_math.py @@ -0,0 +1,54 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Accuracy reward for math trajectories with boxed final answers.""" + +from __future__ import annotations + +from decimal import Decimal, InvalidOperation +from typing import Any + +from twinkle.data_format import Trajectory, user_data_get +from twinkle.reward.base import Reward +from twinkle.reward.math_reward import MathReward + + +class BoxedMathAccuracyReward(Reward): + """Compare the final boxed answer with ``user_data.ground_truth``.""" + + def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: + return [self._score(trajectory) for trajectory in trajectories] + + def metric_payload( + self, + trajectories: list[Trajectory], + *, + rewards: list[float], + **kwargs: Any, + ) -> dict[str, float]: + return {'accuracy_reward': sum(rewards) / len(rewards)} + + @classmethod + def _score(cls, trajectory: Trajectory) -> float: + completion = cls._last_assistant_content(trajectory) + if '\\boxed{' not in completion: + return 0.0 + prediction = MathReward.extract_boxed_result(completion).strip() + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')).strip() + if not prediction or not ground_truth: + return 0.0 + if cls._decimal_equal(prediction, ground_truth): + return 1.0 + return float(MathReward.compare_consecutive(prediction, ground_truth)) + + @staticmethod + def _last_assistant_content(trajectory: Trajectory) -> str: + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + return str(message.get('content', '')) + return '' + + @staticmethod + def _decimal_equal(first: str, second: str) -> bool: + try: + return Decimal(first.replace(',', '')) == Decimal(second.replace(',', '')) + except InvalidOperation: + return False diff --git a/src/twinkle/reward/dapo_math.py b/src/twinkle/reward/dapo_math.py new file mode 100644 index 000000000..a9374fd51 --- /dev/null +++ b/src/twinkle/reward/dapo_math.py @@ -0,0 +1,153 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Accuracy reward for the Answer-style format used by DAPO-Math.""" + +from __future__ import annotations + +import re +from decimal import Decimal, InvalidOperation +from typing import Any + +from twinkle.data_format import Trajectory, user_data_get +from twinkle.reward.base import Reward +from twinkle.reward.math_reward import MathReward + + +_ANSWER_LINE = re.compile(r'^\s*Answer\s*:\s*(.+?)\s*$', re.IGNORECASE | re.MULTILINE) + + +class DAPOMathAccuracyReward(Reward): + """Compare the final ``Answer:`` or boxed result with the ground truth.""" + + def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: + return [self._score(trajectory) for trajectory in trajectories] + + def metric_payload( + self, + trajectories: list[Trajectory], + *, + rewards: list[float], + **kwargs: Any, + ) -> dict[str, float]: + return {'accuracy_reward': sum(rewards) / len(rewards)} + + @classmethod + def _score(cls, trajectory: Trajectory) -> float: + completion = cls._last_assistant_content(trajectory) + prediction = cls.extract_answer(completion) + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')).strip() + if not prediction or not ground_truth: + return 0.0 + if cls._decimal_equal(prediction, ground_truth): + return 1.0 + return float(MathReward.compare_consecutive(prediction, ground_truth)) + + @staticmethod + def extract_answer(completion: str) -> str: + if '\\boxed{' in completion: + return MathReward.extract_boxed_result(completion).strip() + matches = _ANSWER_LINE.findall(completion) + if not matches: + return '' + answer = matches[-1].strip().rstrip('.') + if len(answer) >= 2 and answer.startswith('$') and answer.endswith('$'): + answer = answer[1:-1].strip() + return answer + + @staticmethod + def _last_assistant_content(trajectory: Trajectory) -> str: + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + return str(message.get('content', '')) + return '' + + @staticmethod + def _decimal_equal(first: str, second: str) -> bool: + try: + return Decimal(first.replace(',', '')) == Decimal(second.replace(',', '')) + except InvalidOperation: + return False + + +class DAPOMathReward(Reward): + """DAPO math training reward with token-level overlong shaping. + + Accuracy remains a ``0/1`` diagnostic metric, while the optimization score + is ``+1/-1`` with a linear penalty near the response-length limit. + """ + + def __init__( + self, + max_response_length: int, + overlong_buffer_length: int, + overlong_penalty_factor: float = 1.0, + score_tail_chars: int = 300, + ): + if max_response_length <= 0: + raise ValueError('max_response_length must be positive') + if overlong_buffer_length <= 0 or overlong_buffer_length > max_response_length: + raise ValueError('overlong_buffer_length must be in [1, max_response_length]') + if overlong_penalty_factor < 0: + raise ValueError('overlong_penalty_factor must be non-negative') + if score_tail_chars <= 0: + raise ValueError('score_tail_chars must be positive') + self.max_response_length = max_response_length + self.overlong_buffer_length = overlong_buffer_length + self.overlong_penalty_factor = overlong_penalty_factor + self.score_tail_chars = score_tail_chars + + def components(self, trajectories: list[Trajectory]) -> tuple[list[float], list[float]]: + accuracy_rewards = [self._accuracy(trajectory) for trajectory in trajectories] + overlong_rewards = [self._overlong_reward(trajectory) for trajectory in trajectories] + return accuracy_rewards, overlong_rewards + + def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: + accuracy_rewards, overlong_rewards = self.components(trajectories) + return [ + (1.0 if accuracy else -1.0) + overlong + for accuracy, overlong in zip(accuracy_rewards, overlong_rewards) + ] + + def metric_payload( + self, + trajectories: list[Trajectory], + *, + rewards: list[float], + **kwargs: Any, + ) -> dict[str, float]: + accuracy_rewards, overlong_rewards = self.components(trajectories) + size = len(trajectories) + if size == 0: + return { + 'total_reward': 0.0, + 'accuracy_reward': 0.0, + 'overlong_reward': 0.0, + 'overlong_ratio': 0.0, + } + return { + 'total_reward': sum(rewards) / size, + 'accuracy_reward': sum(accuracy_rewards) / size, + 'overlong_reward': sum(overlong_rewards) / size, + 'overlong_ratio': sum(value < 0 for value in overlong_rewards) / size, + } + + def _accuracy(self, trajectory: Trajectory) -> float: + completion = DAPOMathAccuracyReward._last_assistant_content(trajectory) + scored_completion = completion[-self.score_tail_chars:] + prediction = DAPOMathAccuracyReward.extract_answer(scored_completion) + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')).strip() + if not prediction or not ground_truth: + return 0.0 + if DAPOMathAccuracyReward._decimal_equal(prediction, ground_truth): + return 1.0 + return float(MathReward.compare_consecutive(prediction, ground_truth)) + + def _overlong_reward(self, trajectory: Trajectory) -> float: + if 'completion_length' not in trajectory: + raise ValueError('DAPOMathReward requires token-level completion_length on every trajectory') + completion_length = int(trajectory['completion_length']) + expected_length = self.max_response_length - self.overlong_buffer_length + exceed_length = completion_length - expected_length + return min( + -exceed_length / self.overlong_buffer_length * self.overlong_penalty_factor, + 0.0, + ) diff --git a/src/twinkle/reward/gsm8k.py b/src/twinkle/reward/gsm8k.py index 347d49e40..2871e30a7 100644 --- a/src/twinkle/reward/gsm8k.py +++ b/src/twinkle/reward/gsm8k.py @@ -1,6 +1,7 @@ import re from typing import Any, Dict, List +from twinkle.data_format import user_data_get from twinkle.reward.base import Reward @@ -80,6 +81,126 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: return rewards +class MathVerifyAccuracyReward(Reward): + """Use the same math-verify parsing and equivalence check as AReaL.""" + + def __init__(self, *, precision: int = 6, try_extract_without_anchor: bool = True): + self.precision = precision + self.try_extract_without_anchor = try_extract_without_anchor + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + from math_verify.grader import verify + from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig, parse + + extraction_config = ( + ExprExtractionConfig(try_extract_without_anchor=self.try_extract_without_anchor), + LatexExtractionConfig(), + ) + + rewards = [] + for trajectory in trajectories: + completion = '' + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + completion = str(message.get('content', '')) + break + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')) + try: + # Disable signal-based timeouts because rollout rewards run on + # the sampler's background event-loop thread. + gold = parse( + ground_truth, + extraction_config=extraction_config, + parsing_timeout=None, + ) + answer = parse( + completion, + extraction_config=extraction_config, + parsing_timeout=None, + ) + if not gold or not answer: + rewards.append(0.0) + continue + correct = verify( + gold, + answer, + float_rounding=self.precision, + timeout_seconds=None, + ) + rewards.append(1.0 if correct else 0.0) + except Exception: + rewards.append(0.0) + return rewards + + def metric_payload( + self, + trajectories: List[Dict[str, Any]], + *, + rewards: List[float], + **kwargs, + ) -> Dict[str, float]: + return {'accuracy_reward': sum(rewards) / len(rewards)} + + +class GSM8KBrevityReward(Reward): + """Reward concise completions that contain a parseable final answer.""" + + def __init__(self, full_reward_length: int = 300, decay_length: int = 3000): + self.full_reward_length = full_reward_length + self.decay_length = decay_length + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for trajectory in trajectories: + completion = '' + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + completion = message.get('content', '') + break + has_answer = _has_boxed(completion) or bool(re.search(r'####\s*[\-\d,\.]+', completion)) + if not has_answer: + rewards.append(0.0) + continue + excess_length = max(0, len(completion) - self.full_reward_length) + rewards.append(max(0.0, 1.0 - excess_length / self.decay_length)) + return rewards + + +class GSM8KAccuracyBrevityReward(Reward): + """Sum GSM8K answer accuracy and brevity rewards.""" + + def __init__(self, accuracy_weight: float = 1.0, brevity_weight: float = 1.0): + self.accuracy_weight = accuracy_weight + self.brevity_weight = brevity_weight + self.accuracy_reward = GSM8KAccuracyReward() + self.brevity_reward = GSM8KBrevityReward() + + def components(self, trajectories: List[Dict[str, Any]]) -> tuple[list[float], list[float]]: + return self.accuracy_reward(trajectories), self.brevity_reward(trajectories) + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + accuracy_rewards, brevity_rewards = self.components(trajectories) + return [ + self.accuracy_weight * accuracy + self.brevity_weight * brevity + for accuracy, brevity in zip(accuracy_rewards, brevity_rewards) + ] + + def metric_payload( + self, + trajectories: List[Dict[str, Any]], + *, + rewards: List[float], + **kwargs, + ) -> Dict[str, float]: + accuracy_rewards, brevity_rewards = self.components(trajectories) + size = len(trajectories) + return { + 'total_reward': sum(rewards) / size, + 'accuracy_reward': sum(accuracy_rewards) / size, + 'brevity_reward': sum(brevity_rewards) / size, + } + + class GSM8KFormatReward(Reward): """Format reward: checks if output contains \\boxed{} or #### answer format. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index b1e1790de..a38dd6959 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -495,6 +495,20 @@ async def _get_or_load_lora( logger.error(f'Failed to load LoRA from {lora_path}: {e}') return None + async def unload_lora_paths(self, adapter_paths: list[str]) -> None: + """Evict selected LoRA requests without requiring their files to exist.""" + for adapter_path in adapter_paths: + normalized = os.path.abspath(os.path.expanduser(adapter_path)) + request = self._lora_request_cache.pop(normalized, None) + if request is None: + request = self._lora_request_cache.pop(adapter_path, None) + if request is None: + continue + try: + await self.engine.remove_lora(request.lora_int_id) + except Exception as exc: + logger.warning('Failed to unload LoRA %s: %s', adapter_path, exc) + async def sleep(self, level: int = 2) -> None: """ Offload weights and/or KV cache from GPU memory. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 3c7b2f686..1665dddf9 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -486,6 +486,11 @@ async def _receive_and_load(): self._run_in_loop(_receive_and_load()) + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def unload_adapter_paths(self, adapter_paths: list[str]) -> None: + """Unload policy snapshots from vLLM and clear cached requests.""" + self._run_in_loop(self.engine.unload_lora_paths(adapter_paths)) + @remote_function(dispatch='all', collect='first', lazy_collect=False) def shutdown(self): """Gracefully shutdown the vLLM engine and background event loop. diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index dfdd5176d..d660ec55d 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -1,13 +1,15 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Server configuration package — aggregate root and per-deployment specs.""" -from .application_spec import ApplicationSpec, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, ServerArgs +from .application_spec import (ApplicationSpec, DataPlaneArgs, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, + ServerArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig from .telemetry import TelemetryConfig __all__ = [ 'ApplicationSpec', + 'DataPlaneArgs', 'HttpOptions', 'ModelArgs', 'PersistenceConfig', diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index 51015245c..539d756ca 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Per-deployment ``ApplicationSpec`` and typed argument schemas. -Each deployment kind (``server | model | sampler | processor``) carries its +Each deployment kind (``server | model | sampler | processor | data_plane``) carries its own ``args`` block with strict field validation. ``ApplicationSpec`` holds the routing metadata plus the deployment kind and validates ``args`` against the matching ``*Args`` schema in a model validator. @@ -58,6 +58,7 @@ class ModelArgs(_ArgsBase): queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) max_loras: int = 5 max_length: int | None = None + data_plane_url: str | None = None class SamplerArgs(_ArgsBase): @@ -73,6 +74,7 @@ class SamplerArgs(_ArgsBase): sampler_type: Literal['mock', 'vllm', 'torch'] engine_args: dict[str, Any] | None = None queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + data_plane_url: str | None = None class ServerStateArgs(_ArgsBase): @@ -111,11 +113,18 @@ class ProcessorArgs(_ArgsBase): queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) +class DataPlaneArgs(_ArgsBase): + """Args for the TransferQueue-backed client data plane.""" + + config: dict[str, Any] | None = None + + _ARGS_SCHEMA: dict[str, type[_ArgsBase]] = { 'server': ServerArgs, 'model': ModelArgs, 'sampler': SamplerArgs, 'processor': ProcessorArgs, + 'data_plane': DataPlaneArgs, } # ---------- ApplicationSpec ------------------------------------------------ # @@ -137,12 +146,12 @@ class ApplicationSpec(BaseModel): name: str route_prefix: str = '/' - import_path: Literal['server', 'model', 'sampler', 'processor'] + import_path: Literal['server', 'model', 'sampler', 'processor', 'data_plane'] # ``args`` is always populated by the ``mode='before'`` validator below # (which validates the raw block against the schema selected by # ``import_path`` and defaults a missing block to ``{}``), so the field is # required here — the validator runs first and fills it. - args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs + args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs | DataPlaneArgs deployments: list[dict[str, Any]] = Field(default_factory=list) @model_validator(mode='before') diff --git a/src/twinkle/server/data_plane/__init__.py b/src/twinkle/server/data_plane/__init__.py new file mode 100644 index 000000000..1d11b1ee2 --- /dev/null +++ b/src/twinkle/server/data_plane/__init__.py @@ -0,0 +1,4 @@ +from .app import DataPlaneManagement, build_data_plane_app +from .proxy import DataPlaneProxy + +__all__ = ['DataPlaneManagement', 'DataPlaneProxy', 'build_data_plane_app'] diff --git a/src/twinkle/server/data_plane/app.py b/src/twinkle/server/data_plane/app.py new file mode 100644 index 000000000..f37feb705 --- /dev/null +++ b/src/twinkle/server/data_plane/app.py @@ -0,0 +1,45 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from typing import Any + +from fastapi import FastAPI + +from twinkle.server.deployment import bind_deployment, build_deployment_app +from .store import TQDataRefStore + + +class DataPlaneManagement: + + def __init__(self, config: dict[str, Any] | None = None): + self.store = TQDataRefStore(config) + + +def build_data_plane_app( + deploy_options: dict[str, Any], + config: dict[str, Any] | None = None, +): + from .handlers import register_data_plane_routes + + deploy_options = dict(deploy_options) + autoscaling = deploy_options.get('autoscaling_config') + if autoscaling: + values = autoscaling.model_dump() if hasattr(autoscaling, 'model_dump') else autoscaling + if int(values.get('min_replicas', 1)) != 1 or int(values.get('max_replicas', 1)) != 1: + raise ValueError('data_plane must use exactly one replica') + else: + if int(deploy_options.get('num_replicas', 1)) != 1: + raise ValueError('data_plane must use exactly one replica') + deploy_options.setdefault('num_replicas', 1) + + def register(app: FastAPI, get_self: Any) -> None: + register_data_plane_routes(app, get_self) + + app = build_deployment_app('DataPlane', register) + return bind_deployment( + app, + DataPlaneManagement, + deploy_options, + deployment_name='DataPlaneManagement', + bind_kwargs={'config': config}, + ) diff --git a/src/twinkle/server/data_plane/handlers.py b/src/twinkle/server/data_plane/handlers.py new file mode 100644 index 000000000..ddef7284d --- /dev/null +++ b/src/twinkle/server/data_plane/handlers.py @@ -0,0 +1,52 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from fastapi import Depends, FastAPI + +import twinkle_client.types as types + +if TYPE_CHECKING: + from .app import DataPlaneManagement + + +def register_data_plane_routes(app: FastAPI, self_fn: Callable[[], 'DataPlaneManagement']) -> None: + + @app.post('/twinkle/put', response_model=types.DataRef) + async def put(body: types.DataPutRequest, + self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: + return await self.store.put( + body.rows, + kind=body.kind, + tags=body.tags, + ) + + @app.post('/twinkle/get', response_model=types.DataRowsResponse) + async def get(body: types.DataGetRequest, + self: DataPlaneManagement = Depends(self_fn)) -> types.DataRowsResponse: + rows = await self.store.get( + body.ref, + fields=body.fields, + ) + tags = ( + await self.store.get_tags(body.ref) + if body.include_tags else [] + ) + return types.DataRowsResponse(rows=rows, tags=tags) + + @app.post('/twinkle/append', response_model=types.DataRef) + async def append(body: types.DataAppendRequest, + self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: + return await self.store.append( + body.ref, + body.rows, + tags=body.tags, + ) + + @app.post('/twinkle/release') + async def release(body: types.DataReleaseRequest, + self: DataPlaneManagement = Depends(self_fn)) -> dict[str, str]: + await self.store.release(body.ref) + return {'status': 'ok'} diff --git a/src/twinkle/server/data_plane/proxy.py b/src/twinkle/server/data_plane/proxy.py new file mode 100644 index 000000000..3f5004b9f --- /dev/null +++ b/src/twinkle/server/data_plane/proxy.py @@ -0,0 +1,56 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Internal HTTP adapter used by Model and Sampler component deployments.""" +from __future__ import annotations + +from typing import Any + +import httpx + +from twinkle_client.http.headers import build_routing_headers +from twinkle_client.types.component import DataRef + + +class DataPlaneProxy: + + def __init__(self, base_url: str | None): + self.base_url = base_url.rstrip('/') if base_url else None + self.client = httpx.AsyncClient(timeout=None) if self.base_url else None + + @property + def enabled(self) -> bool: + return self.client is not None + + async def get( + self, + ref: DataRef, + ) -> list[dict[str, Any]]: + if self.client is None or self.base_url is None: + raise RuntimeError('data_plane_url is required when a component request uses input_ref') + response = await self.client.post( + f'{self.base_url}/twinkle/get', + json={'ref': ref.model_dump()}, + headers=build_routing_headers(f'data-ref-{ref.ref_id}'), + ) + response.raise_for_status() + return response.json()['rows'] + + async def put( + self, + rows: list[dict[str, Any]], + *, + kind: str, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if self.client is None or self.base_url is None: + raise RuntimeError('data_plane_url is required to store component output') + response = await self.client.post( + f'{self.base_url}/twinkle/put', + json={'rows': rows, 'kind': kind, 'tags': tags}, + headers=build_routing_headers(f'data-put-{kind}'), + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def close(self) -> None: + if self.client is not None: + await self.client.aclose() diff --git a/src/twinkle/server/data_plane/store.py b/src/twinkle/server/data_plane/store.py new file mode 100644 index 000000000..e1a367ddf --- /dev/null +++ b/src/twinkle/server/data_plane/store.py @@ -0,0 +1,133 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""TransferQueue KV storage behind opaque client DataRef values.""" +from __future__ import annotations + +import uuid +from typing import Any + +from twinkle.tq_utils import rows_to_tq_fields +from twinkle_client.common.json_utils import json_safe +from twinkle_client.types.component import DataRef + + +def _keys(ref: DataRef) -> list[str]: + return [str(index) for index in range(ref.size)] + + +def _partition(ref: DataRef) -> str: + """Resolve an opaque DataRef to its self-contained physical TQ partition.""" + return f'twinkle-client/{ref.ref_id}' + + +def _input_token_count(rows: list[dict[str, Any]]) -> int: + return sum( + len(row.get('input_ids', [])) + for row in rows + if isinstance(row.get('input_ids'), (list, tuple)) + ) + + +def _rows_from_tensordict(data: Any, size: int) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + fields = list(data.keys()) + for index in range(size): + rows.append({ + field: json_safe(data[field][index]) + for field in fields + }) + return rows + + +class TQDataRefStore: + + def __init__(self, config: dict[str, Any] | None = None): + import transfer_queue as tq + if config: + from omegaconf import OmegaConf + tq.init(OmegaConf.create(config)) + else: + tq.init() + + async def put( + self, + rows: list[dict[str, Any]], + *, + kind: str = 'data', + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if not rows: + raise ValueError('rows must not be empty') + if tags is not None and len(tags) != len(rows): + raise ValueError(f'tag count {len(tags)} does not match row count {len(rows)}') + import transfer_queue as tq + ref = DataRef( + ref_id=uuid.uuid4().hex, + size=len(rows), + fields=list(rows[0]), + kind=kind, + num_tokens=_input_token_count(rows), + ) + await tq.async_kv_batch_put( + keys=_keys(ref), + partition_id=_partition(ref), + fields=rows_to_tq_fields(rows), + tags=tags, + ) + return ref + + async def get( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> list[dict[str, Any]]: + import transfer_queue as tq + selected = fields if fields is not None else ref.fields + data = await tq.async_kv_batch_get( + keys=_keys(ref), + partition_id=_partition(ref), + select_fields=selected, + ) + return _rows_from_tensordict(data, ref.size) + + async def get_tags( + self, + ref: DataRef, + ) -> list[dict[str, Any]]: + """Return TQ sample tags in the same order as the rows in ``ref``.""" + import transfer_queue as tq + partition_id = _partition(ref) + partitions = await tq.async_kv_list(partition_id=partition_id) + partition = partitions.get(partition_id, {}) + return [dict(partition.get(key, {})) for key in _keys(ref)] + + async def append( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if len(rows) != ref.size: + raise ValueError(f'append row count {len(rows)} does not match DataRef size {ref.size}') + if not rows: + raise ValueError('rows must not be empty') + if tags is not None and len(tags) != len(rows): + raise ValueError(f'tag count {len(tags)} does not match row count {len(rows)}') + import transfer_queue as tq + await tq.async_kv_batch_put( + keys=_keys(ref), + partition_id=_partition(ref), + fields=rows_to_tq_fields(rows), + tags=tags, + ) + updates: dict[str, Any] = { + 'fields': list(dict.fromkeys([*ref.fields, *rows[0].keys()])), + } + if 'input_ids' in rows[0]: + updates['num_tokens'] = _input_token_count(rows) + return ref.model_copy(update=updates) + + async def release(self, ref: DataRef) -> None: + import transfer_queue as tq + await tq.async_kv_clear(keys=_keys(ref), partition_id=_partition(ref)) diff --git a/src/twinkle/server/deployment.py b/src/twinkle/server/deployment.py index bb4e3c1cd..ccd700964 100644 --- a/src/twinkle/server/deployment.py +++ b/src/twinkle/server/deployment.py @@ -1,11 +1,11 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Shared deployment-application construction. -Top-level, central deployment-construction infrastructure shared by all four -deployments (Gateway, Model, Sampler, Processor). It is intentionally NOT under +Top-level deployment-construction infrastructure shared by Gateway, Model, +Sampler, Processor, and DataPlane. It is intentionally NOT under ``utils/`` — it is core to how every deployment is built, not a generic helper. -It consolidates, in one place, the construction logic the four App_Builders +It consolidates, in one place, the construction logic the App_Builders used to repeat: - ``get_servable()`` — the single servable-object accessor; diff --git a/src/twinkle/server/launcher/builder_registry.py b/src/twinkle/server/launcher/builder_registry.py index 5b9eda698..a18be38b1 100644 --- a/src/twinkle/server/launcher/builder_registry.py +++ b/src/twinkle/server/launcher/builder_registry.py @@ -5,8 +5,8 @@ decomposition). No logic change. The operator-facing YAML ``import_path`` literals (``"server"``, ``"model"``, -``"sampler"``, ``"processor"``) are unchanged; only the internal builder -function the ``"server"`` literal resolves to was renamed +``"sampler"``, ``"processor"``, ``"data_plane"``) resolve to internal builder +functions. The function selected by the ``"server"`` literal was renamed (``build_server_app`` → ``build_gateway_app``). """ from __future__ import annotations @@ -19,6 +19,7 @@ 'model': 'build_model_app', 'sampler': 'build_sampler_app', 'processor': 'build_processor_app', + 'data_plane': 'build_data_plane_app', } @@ -28,6 +29,7 @@ def get_builders() -> dict[str, Callable]: Imported lazily so that importing the launcher package does not eagerly pull in every deployment module. """ + from twinkle.server.data_plane import build_data_plane_app from twinkle.server.gateway import build_gateway_app from twinkle.server.model import build_model_app from twinkle.server.processor import build_processor_app @@ -38,6 +40,7 @@ def get_builders() -> dict[str, Callable]: 'build_model_app': build_model_app, 'build_sampler_app': build_sampler_app, 'build_processor_app': build_processor_app, + 'build_data_plane_app': build_data_plane_app, } diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index bae7b2a93..c87d0bfa6 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -76,6 +76,7 @@ def __init__(self, backend: str, adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): self.backend = backend self.device_group = DeviceGroup(**device_group) @@ -100,6 +101,8 @@ def __init__(self, self.model = MODEL_SELECTOR.construct(backend, ctor_kwargs) self.state: ServerState = get_server_state() + from twinkle.server.data_plane import DataPlaneProxy + self.data_plane = DataPlaneProxy(data_plane_url) self._replica_registered = False # Initialize mixins @@ -146,6 +149,7 @@ async def shutdown(self) -> None: await self.state.unregister_replica(self.replica_id) except Exception: pass + await self.data_plane.close() def check_model_health(self) -> dict: """Probe model actors liveness via a lightweight ping. @@ -182,6 +186,7 @@ def build_model_app(model_id: str, backend: str, adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): """Build a unified model management application for distributed training. @@ -226,7 +231,8 @@ async def _on_shutdown(servable: Any) -> None: deploy_options, deployment_name='ModelManagement', request_router_config=RequestRouterConfig(request_router_class=StickyLoraRequestRouter), - bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config), + bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config, + data_plane_url), bind_kwargs=kwargs, ) diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 2074f5e40..0413c9144 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -26,6 +26,7 @@ validate_user_path) from twinkle.server.utils.validation import get_session_id_from_request from twinkle.utils.logger import get_logger +from twinkle_client.common.json_utils import json_safe from twinkle_client.common.serialize import deserialize_object logger = get_logger() @@ -55,6 +56,24 @@ def _get_twinkle_adapter_name(request: Request, adapter_name: str | None) -> str return owner_id + '-' + adapter_name +def _model_result_rows(result: Any, batch_size: int) -> list[dict[str, Any]]: + """Keep per-sample model outputs at TQ sample granularity.""" + if isinstance(result, list) and len(result) == batch_size and all(isinstance(item, dict) for item in result): + return result + if isinstance(result, dict): + batched = { + name + for name, value in result.items() + if isinstance(value, list) and len(value) == batch_size + } + if batched: + return [{ + name: value[index] if name in batched else value + for name, value in result.items() + } for index in range(batch_size)] + return [{'result': result}] + + def _register_twinkle_routes(app: FastAPI, self_fn: Callable[[], ModelManagement]) -> None: """Register all /twinkle/* routes on the given FastAPI app. @@ -122,6 +141,184 @@ async def _task(): task_type='forward', )) + @app.post('/twinkle/submit_forward') + async def submit_forward( + request: Request, + body: types.AsyncForwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> dict[str, Any]: + """Queue a forward pass and return immediately with a task id.""" + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs = ( + await self.data_plane.get(body.input_ref) + if body.input_ref is not None else body.inputs + ) + inputs = _parse_inputs(raw_inputs) + method = self.model.forward_only if body.forward_only else self.model.forward + ret = method(inputs=inputs, adapter_name=adapter_name, **body.forward_kwargs) + safe = json_safe(ret) + if self.data_plane.enabled: + rows = _model_result_rows(safe, len(inputs)) + output_ref = await self.data_plane.put(rows, kind='model-output') + return {'output_ref': output_ref.model_dump()} + return {'result': safe} + + raw_for_metrics = body.inputs or [] + inputs_list = raw_for_metrics if isinstance(raw_for_metrics, list) else [raw_for_metrics] + input_tokens = ( + body.input_ref.num_tokens + if body.input_ref is not None else + sum(len(item.get('input_ids', [])) if isinstance(item, dict) else 0 for item in inputs_list) + ) + batch_size = body.input_ref.size if body.input_ref is not None else len(inputs_list) + return await self.schedule_task( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='async_forward', + ) + + @app.post('/twinkle/submit_forward_backward') + async def submit_forward_backward( + request: Request, + body: types.AsyncForwardBackwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> dict[str, Any]: + """Queue the same forward/backward primitive exposed by the synchronous client.""" + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + def first_element(data): + while isinstance(data, list): + if len(data) == 0: + return None + data = data[0] + return data + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs = ( + await self.data_plane.get(body.input_ref) + if body.input_ref is not None else body.inputs + ) + inputs = _parse_inputs(raw_inputs) + for model_input in inputs: + for key in model_input: + if (isinstance(model_input[key], list) + and isinstance(first_element(model_input[key]), (int, float))): + model_input[key] = torch.tensor(model_input[key]) + ret = self.model.forward_backward(inputs=inputs, adapter_name=adapter_name, **body.kwargs) + return {'result': json_safe(ret)} + + raw_for_metrics = body.inputs or [] + inputs_list = raw_for_metrics if isinstance(raw_for_metrics, list) else [raw_for_metrics] + input_tokens = ( + body.input_ref.num_tokens + if body.input_ref is not None else + sum(len(item.get('input_ids', [])) for item in inputs_list if isinstance(item, dict)) + ) + batch_size = body.input_ref.size if body.input_ref is not None else len(inputs_list) + return await self.schedule_task( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='async_forward_backward', + ) + + @app.post('/twinkle/submit_clip_grad_and_step') + async def submit_clip_grad_and_step( + request: Request, + body: types.AsyncClipGradAndStepRequest, + self: ModelManagement = Depends(self_fn), + ) -> dict[str, Any]: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + self.model.clip_grad_and_step( + max_grad_norm=body.max_grad_norm, + norm_type=body.norm_type, + adapter_name=adapter_name, + **body.kwargs, + ) + return {'status': 'ok'} + + return await self.schedule_task( + _task, + model_id=adapter_name, + token=token, + task_type='async_clip_grad_and_step', + ) + + @app.post('/twinkle/submit_save') + async def submit_save( + request: Request, + body: types.AsyncSaveRequest, + self: ModelManagement = Depends(self_fn), + ) -> dict[str, Any]: + """Queue an adapter snapshot; used for explicit policy publication.""" + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + checkpoint_name = checkpoint_manager.get_ckpt_name(body.name) + save_dir = checkpoint_manager.get_save_dir(model_id=adapter_name, is_sampler=body.is_sampler) + twinkle_path = checkpoint_manager.save( + model_id=adapter_name, + name=checkpoint_name, + is_sampler=body.is_sampler, + ) + model_save_name = 'latest' if body.is_sampler else checkpoint_name + checkpoint_dir = self.model.save( + name=model_save_name, + output_dir=save_dir, + adapter_name=adapter_name, + save_optimizer=body.save_optimizer, + ) + return {'twinkle_path': twinkle_path, 'checkpoint_dir': checkpoint_dir} + + return await self.schedule_task( + _task, + model_id=adapter_name, + token=token, + task_type='async_save', + ) + + @app.post('/twinkle/remove_adapter') + async def remove_adapter( + request: Request, + body: types.AdapterRequest, + self: ModelManagement = Depends(self_fn), + ) -> dict[str, str]: + """Release a drained tenant's in-memory training adapter.""" + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + await self._cleanup_adapter(adapter_name) + return {'status': 'ok'} + + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + task_type='remove_adapter', + )) + @app.post('/twinkle/forward_only', response_model=types.ForwardResponse) async def forward_only( request: Request, @@ -467,7 +664,10 @@ async def _task(): async_upload=False, ) - future_ref = await self.schedule_background_task(_task, task_type='upload_to_hub') + future_ref = await self.schedule_background_task( + _task, + task_type='upload_to_hub', + ) request_id = future_ref.get('request_id') if request_id is None: raise HTTPException(status_code=500, detail=f'Upload task scheduling failed: {future_ref}') diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index d2cba54e9..9e8eaac4e 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -7,9 +7,11 @@ """ from __future__ import annotations +import asyncio +from typing import Any + from fastapi import FastAPI, Request from ray import serve -from typing import Any from twinkle import DeviceGroup from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app, init_twinkle_runtime @@ -19,6 +21,7 @@ from twinkle.server.utils.task_queue import TaskQueueConfig, TaskQueueMixin from twinkle.server.utils.validation import get_token_from_request from twinkle.utils.logger import get_logger + from .tinker_handlers import _register_tinker_sampler_routes from .twinkle_handlers import _register_twinkle_sampler_routes @@ -54,6 +57,20 @@ def _make_torch_sampler(kw: dict[str, Any]) -> Any: ) +def _construct_sampler_backend( + sampler_type: str, + sampler_kwargs: dict[str, Any], + data_plane_url: str | None, +) -> Any: + if sampler_type == 'vllm' and data_plane_url: + # Client-orchestrated async RL needs sampler admission to release the + # Ray actor immediately. VLLMSamplerTQ retains the inherited + # synchronous API for ordinary sample calls. + from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ + return VLLMSamplerTQ(**sampler_kwargs, context_manager=None) + return SAMPLER_SELECTOR.construct(sampler_type, sampler_kwargs) + + class SamplerManagement(LazyCleanupMixin, TaskQueueMixin): """Unified sampler management service. @@ -72,6 +89,7 @@ def __init__(self, sampler_type: str, engine_args: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): self.device_group = DeviceGroup(**device_group) self.device_mesh = init_twinkle_runtime( @@ -101,13 +119,21 @@ def __init__(self, ) else: sampler_kwargs.update(kwargs) - self.sampler = SAMPLER_SELECTOR.construct(sampler_type, sampler_kwargs) + self.sampler = _construct_sampler_backend(sampler_type, sampler_kwargs, data_plane_url) self.state: ServerState = get_server_state() + from twinkle.server.data_plane import DataPlaneProxy + self.data_plane = DataPlaneProxy(data_plane_url) # Initialize task queue mixin self._init_task_queue(queue_config, deployment_name='Sampler') + async def shutdown(self) -> None: + cancel_all = getattr(self.sampler, 'cancel_all_generations', None) + if callable(cancel_all): + await asyncio.to_thread(cancel_all) + await self.data_plane.close() + @serve.multiplexed(max_num_models_per_replica=5) async def _sticky_entry(self, sticky_key: str): return sticky_key @@ -131,6 +157,7 @@ def build_sampler_app(model_id: str, sampler_type: str, engine_args: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): """Build a unified sampler application for text generation inference. @@ -172,6 +199,7 @@ def register_routes(app: FastAPI, get_self: Any) -> None: 'version': '1.0.0', }, attach_replica_id_header=True, + on_shutdown=lambda servable: servable.shutdown(), ) return bind_deployment( @@ -179,7 +207,8 @@ def register_routes(app: FastAPI, get_self: Any) -> None: SamplerManagement, deploy_options, deployment_name='SamplerManagement', - bind_args=(model_id, nproc_per_node, device_group, device_mesh, sampler_type, engine_args, queue_config), + bind_args=(model_id, nproc_per_node, device_group, device_mesh, sampler_type, engine_args, queue_config, + data_plane_url), bind_kwargs=kwargs, ) diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index 300e5b0c3..4f6ccb520 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -84,6 +84,11 @@ def __init__( # ----- Sampler interface --------------------------------------------- # + @remote_function() + def unload_adapter_paths(self, adapter_paths: list[str]) -> None: + """Mirror the production cache-eviction API for control-plane tests.""" + return None + @remote_function() def sample( self, diff --git a/src/twinkle/server/sampler/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py index 91fdf7813..75e42723d 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -87,7 +87,8 @@ async def _do_sample(): stop=body.sampling_params.stop, ) - responses = self.sampler.sample( + sample_fn = getattr(self.sampler, 'sample_sync', self.sampler.sample) + responses = sample_fn( inputs=[prompt_inputs] * body.num_samples, sampling_params=sampling_params, adapter_path=adapter_uri, diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index 8a10f0b8d..bec3c74a2 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -9,6 +9,7 @@ import asyncio import json import traceback +import uuid from collections.abc import Callable from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse @@ -26,6 +27,7 @@ from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.server.utils.validation import get_session_id_from_request +from twinkle_client.common.json_utils import json_safe from twinkle.utils.logger import get_logger logger = get_logger() @@ -57,6 +59,120 @@ def _get_twinkle_sampler_adapter_name(request: Request, adapter_name: str | None return owner_id + '-' + adapter_name +def _sample_models_to_rows( + sample_models: list[types.SampleResponseModel], + *, + group_ids: list[str] | None, + policy_version: int | None, + adapter_uri: str | None, +) -> tuple[list[dict], list[dict]]: + """Flatten sampler output to one TQ row per generated sequence.""" + resolved_group_ids = group_ids or [uuid.uuid4().hex for _ in sample_models] + if len(resolved_group_ids) != len(sample_models): + raise ValueError( + f'group_ids contains {len(resolved_group_ids)} values for ' + f'{len(sample_models)} sampler inputs') + rows = [] + tags = [] + for prompt_index, (response, group_id) in enumerate(zip(sample_models, resolved_group_ids)): + for generation_idx, sequence in enumerate(response.sequences): + rows.append({ + **sequence.model_dump(), + 'prompt_logprobs': response.prompt_logprobs, + 'topk_prompt_logprobs': response.topk_prompt_logprobs, + }) + tags.append({ + 'record_type': 'sample', + 'group_id': group_id, + 'prompt_index': prompt_index, + 'generation_idx': generation_idx, + 'rollout_status': 'ROLLOUT_DONE', + 'rollout_policy_version': policy_version, + 'rollout_adapter_uri': adapter_uri, + }) + return rows, tags + + +def _responses_to_models(responses) -> list[types.SampleResponseModel]: + """Convert internal sampler responses to the HTTP response schema.""" + sample_models = [] + for response in responses: + sequences = [ + types.SampledSequenceModel( + stop_reason=sequence.stop_reason, + tokens=list(sequence.tokens), + logprobs=list(sequence.logprobs) if sequence.logprobs is not None else None, + decoded=sequence.decoded, + new_input_feature=( + _serialize_input_feature(sequence.new_input_feature) + if sequence.new_input_feature is not None else None + ), + ) + for sequence in response.sequences + ] + sample_models.append( + types.SampleResponseModel( + sequences=sequences, + prompt_logprobs=response.prompt_logprobs, + topk_prompt_logprobs=response.topk_prompt_logprobs, + )) + return sample_models + + +def _submission_states(value) -> list[dict]: + """Normalize Twinkle's single-worker unwrapping to a list of states.""" + return value if isinstance(value, list) else [value] + + +async def _await_generation( + sampler, + submission_id: str, + inputs, + params: SamplingParams, + *, + adapter_name: str, + adapter_path: str | None, +): + """Submit a generation and poll without occupying a worker thread.""" + submitted = False + collected = False + try: + # Mark before dispatch so a partial multi-DP admission is still rolled + # back if one actor rejects while another has already registered it. + submitted = True + await asyncio.to_thread( + sampler.submit_generation, + submission_id, + inputs, + params, + adapter_name=adapter_name, + adapter_path=adapter_path, + ) + poll_interval = 0.01 + while True: + states = _submission_states( + await asyncio.to_thread(sampler.get_generation_status, submission_id)) + failed = next( + (state for state in states if state.get('status') not in ('running', 'completed')), + None, + ) + if failed is not None: + error = failed.get('error') or failed.get('status', 'unknown failure') + raise RuntimeError(f'generation {submission_id} failed: {error}') + if states and all(state.get('status') == 'completed' for state in states): + responses = await asyncio.to_thread(sampler.collect_generation, submission_id) + collected = True + return responses + await asyncio.sleep(poll_interval) + poll_interval = min(poll_interval * 1.5, 0.25) + finally: + if submitted and not collected: + try: + await asyncio.to_thread(sampler.cancel_generation, submission_id) + except Exception: + logger.warning('Failed to cancel generation %s', submission_id, exc_info=True) + + def _register_twinkle_sampler_routes(app: FastAPI, self_fn: Callable[[], SamplerManagement]) -> None: """Register all /twinkle/* sampler routes on the given FastAPI app. @@ -127,32 +243,14 @@ async def _task(): params = SamplingParams.from_dict(body.sampling_params) # Sample - responses = self.sampler.sample( + sample_fn = getattr(self.sampler, 'sample_sync', self.sampler.sample) + responses = sample_fn( inputs, params, adapter_name=full_adapter_name, adapter_path=adapter_path, ) - - sample_models = [] - for response in responses: - sequences = [ - types.SampledSequenceModel( - stop_reason=seq.stop_reason, - tokens=list(seq.tokens), - logprobs=list(seq.logprobs) if seq.logprobs is not None else None, - decoded=seq.decoded, - new_input_feature=_serialize_input_feature(seq.new_input_feature) - if seq.new_input_feature is not None else None, - ) for seq in response.sequences - ] - sample_models.append( - types.SampleResponseModel( - sequences=sequences, - prompt_logprobs=response.prompt_logprobs, - topk_prompt_logprobs=response.topk_prompt_logprobs, - )) - return types.SampleResponseModelList(samples=sample_models) + return types.SampleResponseModelList(samples=_responses_to_models(responses)) # Calculate metrics for queue scheduling inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] @@ -165,6 +263,103 @@ async def _task(): task_type='sample', )) + @app.post('/twinkle/submit_sample') + async def submit_sample( + request: Request, + body: types.AsyncSampleRequest, + self: SamplerManagement = Depends(self_fn), + ) -> dict: + """Queue sampling and return immediately for client-side orchestration.""" + token = await self._on_request_start(request) + + async def _task(): + adapter_path = None + full_adapter_name = _get_twinkle_sampler_adapter_name(request, body.adapter_name) or '' + if body.adapter_uri: + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + _, adapter_path = checkpoint_manager.parse_adapter_uri(body.adapter_uri) + + inputs = ( + await self.data_plane.get(body.input_ref) + if body.input_ref is not None else body.inputs + ) + if isinstance(inputs, list) and inputs: + first = inputs[0] + if isinstance(first, dict) and 'input_ids' in first: + inputs = [InputFeature(**item) for item in inputs] + else: + inputs = [Trajectory(**item) for item in inputs] + elif isinstance(inputs, dict): + inputs = [InputFeature(**inputs)] if 'input_ids' in inputs else [Trajectory(**inputs)] + + params_dict = dict(body.sampling_params or {}) + params_dict['num_samples'] = body.num_samples + params = SamplingParams.from_dict(params_dict) + if callable(getattr(self.sampler, 'submit_generation', None)): + responses = await _await_generation( + self.sampler, + uuid.uuid4().hex, + inputs, + params, + adapter_name=full_adapter_name, + adapter_path=adapter_path, + ) + else: + # Mock/Torch and vLLM deployments without a DataPlane retain + # the compatibility path. DataPlane-enabled vLLM deployments + # are constructed with VLLMSamplerTQ and never wait here. + responses = await asyncio.to_thread( + self.sampler.sample, + inputs, + params, + adapter_name=full_adapter_name, + adapter_path=adapter_path, + ) + sample_models = _responses_to_models(responses) + payload = types.SampleResponseModelList(samples=sample_models).model_dump() + if self.data_plane.enabled: + rows, tags = _sample_models_to_rows( + sample_models, + group_ids=body.group_ids, + policy_version=body.policy_version, + adapter_uri=body.adapter_uri, + ) + output_ref = await self.data_plane.put( + [json_safe(item) for item in rows], + kind='rollout', + tags=tags, + ) + return {'output_ref': output_ref.model_dump()} + return payload + + return await self.schedule_background_task( + _task, + model_id=full_adapter_name if (full_adapter_name := _get_twinkle_sampler_adapter_name( + request, body.adapter_name)) else None, + task_type='async_sample', + ) + + @app.post('/twinkle/unload_adapter_paths') + async def unload_adapter_paths( + request: Request, + body: types.UnloadAdapterPathsRequest, + self: SamplerManagement = Depends(self_fn), + ) -> dict[str, str]: + """Best-effort eviction of published LoRA snapshots from sampler caches.""" + token = await self._on_request_start(request) + resolved_paths = [] + for adapter_path in body.adapter_paths: + if adapter_path.startswith('twinkle://'): + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + _, adapter_path = checkpoint_manager.parse_adapter_uri(adapter_path) + resolved_paths.append(adapter_path) + unload = getattr(self.sampler, 'unload_adapter_paths', None) + if unload is not None: + unload(resolved_paths) + return {'status': 'ok'} + @app.post('/twinkle/set_template', response_model=types.SetTemplateResponse) async def set_template( request: Request, diff --git a/src/twinkle/server/utils/task_queue/mixin.py b/src/twinkle/server/utils/task_queue/mixin.py index 7b6f895e0..05aa45fab 100644 --- a/src/twinkle/server/utils/task_queue/mixin.py +++ b/src/twinkle/server/utils/task_queue/mixin.py @@ -1,10 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -TaskQueueMixin: serial compute queue + background-task execution. +"""TaskQueueMixin: serial compute queue plus admitted concurrent tasks. -Two execution paths: - schedule_task() / schedule_task_and_wait() -> serial compute queue (GPU ops) - schedule_background_task() -> fire-and-forget asyncio Task (I/O ops) +``schedule_task`` serializes stateful Model operations. Detached tasks are +used for I/O and for engines such as vLLM that own their compute concurrency +and continuous batching internally. """ from __future__ import annotations @@ -39,8 +38,8 @@ class TaskQueueMixin: Use for GPU operations: forward, backward, step, save, load, etc. 2. Background task (schedule_background_task): - asyncio.create_task, runs concurrently with compute queue. - Use for pure I/O: upload_to_hub, etc. + asyncio.create_task, runs concurrently with compute queue. Use for I/O + or an engine that provides its own safe concurrency and batching. Status is still tracked; clients can poll the same status endpoints. Requirements @@ -200,8 +199,15 @@ async def schedule_task( """ request_id = f'req_{uuid.uuid4().hex}' - preflight_result = await self._perform_preflight_checks(request_id, model_id, token, input_tokens, batch_size, - data_world_size, batch_size_multiple) + preflight_result = await self._perform_preflight_checks( + request_id=request_id, + model_id=model_id, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=data_world_size, + batch_size_multiple=batch_size_multiple, + ) if preflight_result is not None: return preflight_result @@ -209,7 +215,11 @@ async def schedule_task( self._event_loop = asyncio.get_running_loop() await self.state.store_future_status( - request_id, TaskStatus.PENDING.value, model_id, queue_state=QueueState.ACTIVE.value) + request_id, + TaskStatus.PENDING.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) queue_key = self._queue_key(model_id=model_id, token=token) self._compute_worker.ensure_queue_registered(queue_key) @@ -227,7 +237,11 @@ async def schedule_task( created_at=time.monotonic(), )) await self.state.store_future_status( - request_id, TaskStatus.QUEUED.value, model_id, queue_state=QueueState.ACTIVE.value) + request_id, + TaskStatus.QUEUED.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) logger.info(f'[TaskQueue] Task {request_id} queued, type={task_type or "unknown"}, ' f'model_id={model_id}, queue_key={queue_key}, ' f'queue_depth={q.qsize()}, input_tokens={input_tokens}') @@ -294,12 +308,12 @@ async def schedule_background_task( model_id: str | None = None, task_type: str | None = None, ) -> dict[str, Any]: - """Schedule a fire-and-forget background task (bypasses compute queue). + """Schedule a fire-and-forget task outside the serial queue. - Designed for pure I/O operations such as upload_to_hub that do not - require GPU serialization. The task is launched immediately as an - asyncio.create_task so it runs concurrently with the compute queue - without blocking any other user's training operations. + The task is launched immediately as an asyncio task. This is suitable + for pure I/O and for an inference engine such as vLLM that performs its + own safe request concurrency and continuous batching. Stateful Model + operations must continue to use :meth:`schedule_task`. Status is tracked via state.store_future_status so clients can poll progress through the same status endpoints as schedule_task(). @@ -317,7 +331,11 @@ async def schedule_background_task( f'type={task_type or "unknown"}, model_id={model_id}') await self.state.store_future_status( - request_id, TaskStatus.RUNNING.value, model_id, queue_state=QueueState.ACTIVE.value) + request_id, + TaskStatus.RUNNING.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) async def _run() -> None: try: @@ -327,7 +345,8 @@ async def _run() -> None: TaskStatus.COMPLETED.value, model_id, result=result, - queue_state=QueueState.ACTIVE.value) + queue_state=QueueState.ACTIVE.value, + ) logger.info(f'[TaskQueue] Background task {request_id} completed, type={task_type or "unknown"}') except Exception: error_payload = task_error_payload(traceback.format_exc()) @@ -336,7 +355,8 @@ async def _run() -> None: TaskStatus.FAILED.value, model_id, result=error_payload, - queue_state=QueueState.ACTIVE.value) + queue_state=QueueState.ACTIVE.value, + ) logger.error(f'[TaskQueue] Background task {request_id} FAILED, type={task_type or "unknown"}:\n' f'{traceback.format_exc(limit=3)}') diff --git a/src/twinkle/server/utils/task_queue/worker.py b/src/twinkle/server/utils/task_queue/worker.py index a26888cfe..4084febe2 100644 --- a/src/twinkle/server/utils/task_queue/worker.py +++ b/src/twinkle/server/utils/task_queue/worker.py @@ -192,7 +192,11 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue q.task_done() in the finally block. """ await self._state.store_future_status( - task.request_id, TaskStatus.RUNNING.value, task.model_id, queue_state=QueueState.ACTIVE.value) + task.request_id, + TaskStatus.RUNNING.value, + task.model_id, + queue_state=QueueState.ACTIVE.value, + ) task_type = task.task_type or 'unknown' exec_start = time.monotonic() @@ -226,7 +230,8 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue TaskStatus.COMPLETED.value, task.model_id, result=result, - queue_state=QueueState.ACTIVE.value) + queue_state=QueueState.ACTIVE.value, + ) except asyncio.TimeoutError: task_status = 'timeout' exec_time = time.monotonic() - exec_start diff --git a/src/twinkle/tq_utils.py b/src/twinkle/tq_utils.py new file mode 100644 index 000000000..f7c34f073 --- /dev/null +++ b/src/twinkle/tq_utils.py @@ -0,0 +1,43 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Small TransferQueue packing helpers shared by both async-RL modes.""" +from __future__ import annotations + +from numbers import Number +from typing import Any + + +def rows_to_tq_fields(rows: list[dict[str, Any]]): + from tensordict import TensorDict + + if not rows: + return TensorDict({}, batch_size=[0]) + field_names = tuple(rows[0].keys()) + expected = set(field_names) + for row_index, row in enumerate(rows): + actual = set(row) + if actual != expected: + missing = sorted(expected - actual) + extra = sorted(actual - expected) + raise ValueError(f'TQ row {row_index} fields mismatch: missing={missing}, extra={extra}') + columns = {field_name: [row[field_name] for row in rows] for field_name in field_names} + return columns_to_tq_fields(columns, len(rows)) + + +def columns_to_tq_fields(columns: dict[str, list[Any]], size: int): + import torch + from tensordict import TensorDict + from tensordict.tensorclass import NonTensorStack + + if size < 0: + raise ValueError(f'TQ field size must be non-negative, got {size}') + packed = {} + for field_name, values in columns.items(): + if not isinstance(values, list): + raise TypeError(f'TQ field {field_name!r} must be a list, got {type(values)!r}') + if len(values) != size: + raise ValueError(f'TQ field {field_name!r} must contain {size} values, got {len(values)}') + if all(isinstance(item, Number) and not isinstance(item, bool) for item in values): + packed[field_name] = torch.tensor(values) + else: + packed[field_name] = NonTensorStack(*values) + return TensorDict(packed, batch_size=[size]) diff --git a/src/twinkle_agentic/async_rl/__init__.py b/src/twinkle_agentic/async_rl/__init__.py new file mode 100644 index 000000000..0f79590b0 --- /dev/null +++ b/src/twinkle_agentic/async_rl/__init__.py @@ -0,0 +1,33 @@ +"""Native TransferQueue building blocks for YAML-driven async multi-LoRA RL.""" + +from .context_manager import ContextStatus, LoraContextManager +from .data_plane import TQDataPlane +from .native_tq import ContextGRPOGroupNSampler +from .pipeline import AsyncMultiLoraGRPOConfig, AsyncMultiLoraGRPOPipeline, create_cpu_actor +from .scheduler import ContextSchedulePolicy, ContextScheduler, ScheduleCandidate, SchedulerConfig +from .types import LoraContext, PartitionAdmission, PreparedPartition, PromptGroup, RolloutPolicy +from .vllm_sampler_tq import VLLMSamplerTQ +from .workers import AdvantageWorker, RolloutWorker, TrainerWorker + +__all__ = [ + 'AdvantageWorker', + 'AsyncMultiLoraGRPOConfig', + 'AsyncMultiLoraGRPOPipeline', + 'ContextSchedulePolicy', + 'ContextScheduler', + 'ContextStatus', + 'ContextGRPOGroupNSampler', + 'LoraContext', + 'LoraContextManager', + 'PartitionAdmission', + 'PreparedPartition', + 'PromptGroup', + 'RolloutPolicy', + 'RolloutWorker', + 'ScheduleCandidate', + 'SchedulerConfig', + 'TQDataPlane', + 'TrainerWorker', + 'create_cpu_actor', + 'VLLMSamplerTQ', +] diff --git a/src/twinkle_agentic/async_rl/context_manager.py b/src/twinkle_agentic/async_rl/context_manager.py new file mode 100644 index 000000000..5933b6dc6 --- /dev/null +++ b/src/twinkle_agentic/async_rl/context_manager.py @@ -0,0 +1,300 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-owner control plane for async multi-LoRA RL.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import StrEnum + +from .types import LoraContext, PartitionAdmission, RolloutPolicy + + +class ContextStatus(StrEnum): + ADDING = 'ADDING' + ACTIVE = 'ACTIVE' + DRAINING = 'DRAINING' + EXHAUSTED = 'EXHAUSTED' + FINISHED = 'FINISHED' + REMOVED = 'REMOVED' + FAILED = 'FAILED' + + +@dataclass +class _ContextState: + context: LoraContext + policy: RolloutPolicy + policy_history: list[RolloutPolicy] = field(default_factory=list) + next_step: int = 0 + completed_partitions: int = 0 + status: ContextStatus = ContextStatus.ACTIVE + dataset_exhausted: bool = False + training_partition_id: str | None = None + live_partitions: dict[str, PartitionAdmission] = field(default_factory=dict) + rollout_policy_references: dict[str, int] = field(default_factory=dict) + max_steps: int | None = None + + +class LoraContextManager: + """Ray-safe control plane; it deliberately knows no TQ sample readiness.""" + + def __init__(self, *, max_staleness: int = 0, max_steps: int | None = None): + if max_staleness < 0: + raise ValueError('max_staleness must be non-negative') + if max_steps is not None and max_steps < 0: + raise ValueError('max_steps must be non-negative') + self.max_staleness = int(max_staleness) + self.max_steps = max_steps + self._contexts: dict[str, _ContextState] = {} + self._creation_order = 0 + self._stop_requested = max_steps == 0 + + def register_context(self, + context: LoraContext, + *, + adapter_path: str | None = None, + policy_version: int = 0, + max_steps: int | None = None, + status: ContextStatus = ContextStatus.ACTIVE) -> None: + if context.key in self._contexts: + return + if max_steps is not None and max_steps < 0: + raise ValueError('max_steps must be non-negative') + self._contexts[context.key] = _ContextState( + context=context, + policy=RolloutPolicy(context.key, context.adapter_name, int(policy_version), adapter_path), + policy_history=[RolloutPolicy(context.key, context.adapter_name, int(policy_version), adapter_path)], + status=status, + max_steps=max_steps, + ) + + def reserve_context(self, context: LoraContext, *, max_steps: int | None = None) -> None: + if context.key in self._contexts: + raise KeyError(f'context already exists: {context.key}') + if max_steps is not None and max_steps < 0: + raise ValueError('max_steps must be non-negative') + self.register_context(context, status=ContextStatus.ADDING, max_steps=max_steps) + + def activate_context(self, + context: LoraContext | str, + *, + adapter_path: str | None = None, + policy_version: int = 0) -> None: + state = self._state(context) + if state.status is not ContextStatus.ADDING: + raise RuntimeError(f'{state.context.key} cannot activate from {state.status}') + policy = RolloutPolicy(state.context.key, state.context.adapter_name, int(policy_version), adapter_path) + state.policy = policy + state.policy_history = [policy] + state.status = ContextStatus.ACTIVE + + def request_context_drain(self, context: LoraContext | str) -> None: + state = self._state(context) + if state.status in (ContextStatus.REMOVED, ContextStatus.FAILED): + return + if state.status is ContextStatus.ADDING: + raise RuntimeError(f'{state.context.key} is still being added') + state.status = ContextStatus.DRAINING + + def context_is_drained(self, context: LoraContext | str) -> bool: + return not self._state(context).live_partitions + + def fail_context(self, context: LoraContext | str) -> None: + state = self._state(context) + if state.live_partitions: + raise RuntimeError(f'{state.context.key} still has live partitions') + state.status = ContextStatus.FAILED + + def unregister_context(self, context: LoraContext | str) -> None: + state = self._state(context) + if state.live_partitions: + raise RuntimeError(f'{state.context.key} still has live partitions') + state.status = ContextStatus.REMOVED + self._contexts.pop(state.context.key) + + def context_snapshot(self, context: LoraContext | str) -> dict[str, object]: + state = self._state(context) + return { + 'context': state.context, + 'status': state.status, + 'policy_version': state.policy.version, + 'adapter_path': state.policy.adapter_path, + 'next_step': state.next_step, + 'completed_partitions': state.completed_partitions, + 'live_partitions': len(state.live_partitions), + 'dataset_exhausted': state.dataset_exhausted, + 'max_steps': state.max_steps, + } + + def list_context_snapshots(self) -> list[dict[str, object]]: + return [self.context_snapshot(key) for key in self._contexts] + + def context_adapter_paths(self, context: LoraContext | str) -> list[str]: + return [ + policy.adapter_path + for policy in self._state(context).policy_history + if policy.adapter_path is not None + ] + + def get_rollout_policy(self, context: LoraContext | str) -> RolloutPolicy: + return self._state(context).policy + + def acquire_rollout_policy(self, context: LoraContext | str) -> RolloutPolicy: + """Pin the current policy while one sampler request is using it.""" + state = self._state(context) + policy = state.policy + if policy.adapter_path is not None: + references = state.rollout_policy_references + references[policy.adapter_path] = references.get(policy.adapter_path, 0) + 1 + return policy + + def release_rollout_policy(self, policy: RolloutPolicy) -> None: + """Release a policy previously returned by :meth:`acquire_rollout_policy`.""" + if policy.adapter_path is None: + return + state = self._state(policy.context_key) + references = state.rollout_policy_references + count = references.get(policy.adapter_path, 0) + if count <= 0: + raise RuntimeError(f'rollout policy was not acquired: {policy.adapter_path}') + if count == 1: + references.pop(policy.adapter_path) + else: + references[policy.adapter_path] = count - 1 + + def request_rollout_partition(self, context: LoraContext | str, *, target_groups: int, + num_generations: int) -> PartitionAdmission | None: + state = self._state(context) + if target_groups <= 0 or num_generations <= 0: + raise ValueError('target_groups and num_generations must be positive') + if state.status is not ContextStatus.ACTIVE or self._stop_requested: + return None + if state.max_steps is not None and state.completed_partitions + len(state.live_partitions) >= state.max_steps: + state.dataset_exhausted = True + state.status = ContextStatus.EXHAUSTED + self._finish_if_drained(state) + return None + if self.max_steps is not None and (self.completed_partitions + len(self.list_live_partitions()) + >= self.max_steps): + return None + oldest_unreleased_step = ( + min(admission.step + for admission in state.live_partitions.values()) if state.live_partitions else state.next_step) + if state.next_step - oldest_unreleased_step > self.max_staleness: + return None + admission = PartitionAdmission( + context=state.context, + partition_id=state.context.partition_id(state.next_step), + step=state.next_step, + target_groups=target_groups, + num_generations=num_generations, + created_order=self._creation_order, + ) + self._creation_order += 1 + state.next_step += 1 + state.live_partitions[admission.partition_id] = admission + return admission + + def list_live_partitions(self) -> list[PartitionAdmission]: + return sorted( + (admission for state in self._contexts.values() for admission in state.live_partitions.values()), + key=lambda admission: admission.created_order, + ) + + def list_trainable_partitions(self) -> list[PartitionAdmission]: + """Return at most one partition per context, in partition step order.""" + partitions = [] + for state in self._contexts.values(): + if not state.live_partitions: + continue + if state.training_partition_id is not None: + partitions.append(state.live_partitions[state.training_partition_id]) + continue + partitions.append(min(state.live_partitions.values(), key=lambda admission: admission.step)) + return sorted(partitions, key=lambda admission: admission.created_order) + + def on_dataset_exhausted(self, context: LoraContext | str) -> None: + state = self._state(context) + state.dataset_exhausted = True + if state.status is ContextStatus.ACTIVE: + state.status = ContextStatus.EXHAUSTED + self._finish_if_drained(state) + + def on_partition_training_started(self, admission: PartitionAdmission) -> None: + state = self._state(admission.context) + if state.training_partition_id not in (None, admission.partition_id): + raise RuntimeError(f'{state.context.key} already trains {state.training_partition_id}') + self._require_live(state, admission) + oldest_partition = min(state.live_partitions.values(), key=lambda candidate: candidate.step) + if oldest_partition.partition_id != admission.partition_id: + raise RuntimeError(f'{admission.partition_id} cannot train before {oldest_partition.partition_id}') + state.training_partition_id = admission.partition_id + + def on_partition_trained(self, admission: PartitionAdmission, *, adapter_path: str) -> RolloutPolicy: + state = self._state(admission.context) + self._require_live(state, admission) + next_policy = RolloutPolicy(state.context.key, state.context.adapter_name, state.policy.version + 1, + adapter_path) + state.policy = next_policy + state.policy_history.append(next_policy) + return next_policy + + def on_partition_cleared(self, admission: PartitionAdmission) -> None: + state = self._state(admission.context) + self._require_live(state, admission) + state.live_partitions.pop(admission.partition_id) + if state.training_partition_id == admission.partition_id: + state.training_partition_id = None + state.completed_partitions += 1 + if state.max_steps is not None and state.completed_partitions >= state.max_steps: + state.dataset_exhausted = True + if state.status is ContextStatus.ACTIVE: + state.status = ContextStatus.EXHAUSTED + if self.max_steps is not None and self.completed_partitions >= self.max_steps: + self._stop_requested = True + self._finish_if_drained(state) + + @property + def completed_partitions(self) -> int: + return sum(state.completed_partitions for state in self._contexts.values()) + + def get_completed_partitions(self) -> int: + return self.completed_partitions + + def is_run_finished(self) -> bool: + if self._stop_requested: + return not self.list_live_partitions() + return bool(self._contexts) and all(state.status is ContextStatus.FINISHED and not state.live_partitions + for state in self._contexts.values()) + + def is_rollout_admission_closed(self) -> bool: + """Whether the producer must stop reading new prompts. + + A global train limit closes admission only. Existing live partitions + remain available to AdvantageWorker and TrainerWorker until drained. + """ + return self._stop_requested or all(state.status is not ContextStatus.ACTIVE + for state in self._contexts.values()) + + def adapter_paths_to_keep(self) -> set[str]: + paths: set[str] = set() + for state in self._contexts.values(): + if state.policy.adapter_path is not None: + paths.add(state.policy.adapter_path) + paths.update(state.rollout_policy_references) + return paths + + def context_status(self, context: LoraContext | str) -> ContextStatus: + return self._state(context).status + + def _finish_if_drained(self, state: _ContextState) -> None: + if state.dataset_exhausted and not state.live_partitions and state.status is ContextStatus.EXHAUSTED: + state.status = ContextStatus.FINISHED + + def _state(self, context: LoraContext | str) -> _ContextState: + key = context if isinstance(context, str) else context.key + return self._contexts[key] + + @staticmethod + def _require_live(state: _ContextState, admission: PartitionAdmission) -> None: + if state.live_partitions.get(admission.partition_id) != admission: + raise KeyError(f'unknown live partition {admission.partition_id}') diff --git a/src/twinkle_agentic/async_rl/data_plane.py b/src/twinkle_agentic/async_rl/data_plane.py new file mode 100644 index 000000000..0bc38f6ca --- /dev/null +++ b/src/twinkle_agentic/async_rl/data_plane.py @@ -0,0 +1,273 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The only async-RL layer that speaks native TransferQueue BatchMeta.""" + +from __future__ import annotations + +from typing import Any, Sequence + +from .native_tq import (AsyncTQClient, append_fields, batch_size_for_groups, clear_partition, fetch_ready_batch, + metadata_size, preallocate_partition, set_sample_tags, split_batch_meta) +from .tq_utils import REQUIRED_MODEL_INPUT_FIELDS, ROLLOUT_TRAIN_FIELDS, columns_to_tq_fields, rows_to_tq_fields +from .types import ClaimedBatch, LoraContext, PartitionAdmission, PreparedPartition, PromptGroup, RolloutOutput + +_REQUIRED_ROLLOUT_FIELDS = frozenset((*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs', 'rewards')) + + +def build_rollout_group_sample_write( + group: PromptGroup, + samples: Sequence[RolloutOutput], + *, + rewards: list[float] | None = None, + expected_num_generations: int, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + group_samples = [dict(sample) for sample in samples] + if expected_num_generations <= 0: + raise ValueError(f'expected_num_generations must be positive, got {expected_num_generations}') + if len(group_samples) != expected_num_generations: + raise ValueError(f'group {group.group_id} expected {expected_num_generations} rollout samples, ' + f'got {len(group_samples)}') + if rewards is not None and len(rewards) != len(group_samples): + raise ValueError(f'reward count {len(rewards)} does not match sample count {len(group_samples)}') + + sample_fields: list[dict[str, Any]] = [] + sample_tags: list[dict[str, Any]] = [] + generation_indices: list[int] = [] + reward_iter = iter(rewards or []) + for sample_index, trajectory in enumerate(group_samples): + sample = dict(trajectory) + if rewards is not None: + sample['rewards'] = float(next(reward_iter)) + generation_idx = int(sample.get('generation_idx', sample_index)) + sample_key = f'samples/{group.group_id}/{generation_idx}' + generation_indices.append(generation_idx) + logprobs = _require_rollout_logprobs(sample, sample_key=sample_key) + sample['logprobs'] = logprobs + sample_fields.append(_rollout_sample_fields(sample)) + sample_tags.append( + _sample_tag( + context=group.context, + group=group, + sample=sample, + sample_key=sample_key, + generation_idx=generation_idx, + logprobs=logprobs, + )) + + expected_indices = list(range(expected_num_generations)) + if generation_indices != expected_indices: + raise ValueError(f'group {group.group_id} generation_idx must be 0..{expected_num_generations - 1} ' + f'in order, got {generation_indices}') + return sample_fields, sample_tags + + +def _require_rollout_logprobs(sample: dict[str, Any], *, sample_key: str) -> list[float]: + logprobs = sample.get('logprobs') + if not isinstance(logprobs, list): + raise TypeError(f'rollout sample {sample_key!r} logprobs must be list[float], got {type(logprobs)!r}') + values: list[float] = [] + for index, value in enumerate(logprobs): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f'rollout sample {sample_key!r} logprobs[{index}] must be a float, got {type(value)!r}') + values.append(float(value)) + labels = sample.get('labels') + if labels is not None: + trainable_tokens = sum(1 for label in labels if label != -100) + if len(values) != trainable_tokens: + raise ValueError(f'rollout sample {sample_key!r} logprobs length must match trainable labels: ' + f'{len(values)} != {trainable_tokens}') + return values + + +def _rollout_sample_fields(sample: dict[str, Any]) -> dict[str, Any]: + return {field_name: sample[field_name] for field_name in ROLLOUT_TRAIN_FIELDS if field_name in sample} + + +def _sample_tag( + *, + context: LoraContext, + group: PromptGroup, + sample: dict[str, Any], + sample_key: str, + generation_idx: int, + logprobs: list[float], +) -> dict[str, Any]: + tag = { + 'record_type': 'sample', + 'sample_status': 'success', + 'context_key': context.key, + 'tenant_id': context.tenant_id, + 'training_run_id': context.training_run_id, + 'adapter_name': context.adapter_name, + 'sample_id': sample.get('sample_id', sample_key), + 'group_id': group.group_id, + 'generation_idx': generation_idx, + 'rollout_policy_version': int(sample['rollout_policy_version']), + 'rollout_adapter_path': sample.get('rollout_adapter_path'), + 'logprobs_length': len(logprobs), + } + for field_name in ('rollout_policy_versions', 'initial_policy_version', 'final_policy_version', + 'policy_version_span'): + if field_name in sample: + tag[field_name] = sample[field_name] + trainable_tokens = _trainable_token_count(sample.get('labels')) + if trainable_tokens is not None: + tag['trainable_tokens'] = trainable_tokens + for sample_field, tag_field in ( + ('input_ids', 'input_length'), + ('labels', 'label_length'), + ('attention_mask', 'attention_length'), + ): + length = _safe_len(sample.get(sample_field)) + if length is not None: + tag[tag_field] = length + for field_name in ('stop_reason', 'truncated', 'turns'): + if field_name in sample: + tag[field_name] = sample[field_name] + if 'completion_length' in sample: + tag['completion_length'] = int(sample['completion_length']) + return tag + + +def _trainable_token_count(labels: Any) -> int | None: + if labels is None: + return None + return sum(1 for label in labels if label != -100) + + +def _safe_len(value: Any) -> int | None: + if value is None: + return None + try: + return len(value) + except TypeError: + return None + + +class TQDataPlane: + """Maps partition-level training operations to native TQ operations.""" + + def __init__(self, client: AsyncTQClient | None = None): + self._client = client + + @property + def client(self) -> AsyncTQClient: + if self._client is None: + import transfer_queue as tq + tq.init() + self._client = tq.get_client() + return self._client + + async def prepare_rollout_partition( + self, + admission: PartitionAdmission, + prompts: Sequence[dict[str, Any]], + sampling_params: Any, + ) -> PreparedPartition: + if len(prompts) != admission.target_groups: + raise ValueError(f'{admission.partition_id} expected {admission.target_groups} prompts, got {len(prompts)}') + rows = [dict(prompt) for prompt in prompts for _ in range(admission.num_generations)] + metadata = await preallocate_partition( + self.client, partition_id=admission.partition_id, prompt_fields=rows_to_tq_fields(rows)) + group_batch_metas = split_batch_meta(metadata, admission.num_generations) + groups = [] + for index, (prompt, batch_meta) in enumerate(zip(prompts, group_batch_metas)): + group_id = f'{admission.partition_id}/group_{index}' + await set_sample_tags(self.client, batch_meta, [{ + 'group_id': group_id, + 'generation_idx': generation_idx, + 'rollout_status': 'PENDING', + } for generation_idx in range(admission.num_generations)]) + groups.append(PromptGroup( + context=admission.context, + partition=admission, + group_id=group_id, + prompt=dict(prompt), + batch_meta=batch_meta, + )) + return PreparedPartition(admission, tuple(groups), sampling_params) + + async def complete_rollout_group( + self, + group: PromptGroup, + *, + rollout_rows: Sequence[RolloutOutput], + rewards: Sequence[float], + submission_id: str, + tag_metrics: dict[str, Any] | None = None, + ) -> None: + expected = group.partition.num_generations + sample_fields, sample_tags = build_rollout_group_sample_write( + group, + rollout_rows, + rewards=list(rewards), + expected_num_generations=expected, + ) + for index, fields in enumerate(sample_fields): + missing = sorted(_REQUIRED_ROLLOUT_FIELDS - set(fields)) + if missing: + raise ValueError(f'rollout sample {group.group_id}/{index} is missing training fields {missing}') + metrics = dict(tag_metrics or {}) + completed_tags = [] + for tag in sample_tags: + completed_tag = dict(tag) + completed_tag.update(metrics) + completed_tag.update({'rollout_status': 'ROLLOUT_DONE', 'submission_id': submission_id}) + completed_tags.append(completed_tag) + await set_sample_tags(self.client, group.batch_meta, completed_tags) + await append_fields(self.client, rows_to_tq_fields(sample_fields), group.batch_meta) + + async def claim_advantage_batch(self, admission: PartitionAdmission, group_count: int) -> ClaimedBatch | None: + metadata = await self._claim(admission, group_count, ['input_ids', 'logprobs', 'rewards'], + self._advantage_task(admission)) + if metadata is None: + return None + return ClaimedBatch( + admission=admission, + data=await self.client.async_get_data(metadata.select_fields(['rewards'])), + batch_meta=metadata, + ) + + async def write_advantages(self, batch: ClaimedBatch, *, advantages: Any, returns: Any) -> None: + size = metadata_size(batch.batch_meta) + fields = columns_to_tq_fields({'advantages': list(advantages), 'returns': list(returns)}, size) + await append_fields(self.client, fields, batch.batch_meta) + + async def claim_training_batch(self, admission: PartitionAdmission, group_count: int) -> ClaimedBatch | None: + metadata = await self._claim( + admission, + group_count, + [*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs', 'rewards', 'advantages', 'returns'], + self._trainer_task(admission), + ) + if metadata is None: + return None + return ClaimedBatch( + admission=admission, + data=await self.client.async_get_data(metadata), + batch_meta=metadata, + sample_tags=tuple(metadata.get_all_custom_meta()), + ) + + async def is_training_consumed(self, admission: PartitionAdmission) -> bool: + return await self.client.async_check_consumption_status(self._trainer_task(admission), admission.partition_id) + + async def clear_partition(self, admission: PartitionAdmission) -> None: + await clear_partition(self.client, admission.partition_id) + + async def _claim(self, admission: PartitionAdmission, groups: int, fields: list[str], task: str) -> Any | None: + return await fetch_ready_batch( + self.client, + data_fields=fields, + batch_size=batch_size_for_groups(groups, admission.num_generations), + partition_id=admission.partition_id, + task_name=task, + num_generations=admission.num_generations, + ) + + @staticmethod + def _advantage_task(admission: PartitionAdmission) -> str: + return f'async_rl/advantage/{admission.context.key}' + + @staticmethod + def _trainer_task(admission: PartitionAdmission) -> str: + return f'async_rl/trainer/{admission.context.key}' diff --git a/src/twinkle_agentic/async_rl/metrics.py b/src/twinkle_agentic/async_rl/metrics.py new file mode 100644 index 000000000..abb4696d1 --- /dev/null +++ b/src/twinkle_agentic/async_rl/metrics.py @@ -0,0 +1,129 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stateless metrics specific to async RL policy and advantage semantics.""" + +from __future__ import annotations + +import math +import statistics +from collections.abc import Mapping, Sequence +from typing import Any + + +def _p95(values: list[float]) -> float: + ordered = sorted(values) + return ordered[max(0, (95 * len(ordered) + 99) // 100 - 1)] + + +def rollout_metrics( + *, + rewards: Mapping[str, Sequence[float]] | None = None, + completion_lengths: Sequence[int] = (), + stop_reasons: Sequence[str | None] = (), + rollout_latency_s: float | None = None, +) -> dict[str, float | int]: + """Summarize one RL rollout collection without retaining state.""" + metrics: dict[str, float | int] = {} + reward_counts = [len(values) for values in (rewards or {}).values() if values] + if len(set(reward_counts)) > 1: + raise ValueError(f'reward metric lengths must match, got {reward_counts}') + sample_count = len(completion_lengths) or (reward_counts[0] if reward_counts else 0) + if completion_lengths and reward_counts and any(count != sample_count for count in reward_counts): + raise ValueError( + f'reward and completion metric lengths must match: {reward_counts} != {sample_count}') + if stop_reasons and len(stop_reasons) != len(completion_lengths): + raise ValueError( + 'stop reason and completion metric lengths must match: ' + f'{len(stop_reasons)} != {len(completion_lengths)}') + if sample_count: + metrics['sample_count'] = sample_count + if completion_lengths: + lengths = [int(value) for value in completion_lengths] + output_tokens = sum(lengths) + truncated_count = sum(reason == 'length' for reason in stop_reasons) + metrics.update({ + 'completion_length_mean': output_tokens / sample_count, + 'completion_length_p95': _p95(lengths), + 'completion_length_max': max(lengths), + 'completion_truncated_count': truncated_count, + 'completion_truncated_ratio': truncated_count / sample_count, + 'output_tokens': output_tokens, + }) + if rollout_latency_s is not None: + latency = float(rollout_latency_s) + metrics['rollout_latency_s'] = latency + metrics['output_tokens_per_s'] = output_tokens / latency if latency > 0 else 0.0 + elif rollout_latency_s is not None: + metrics['rollout_latency_s'] = float(rollout_latency_s) + + for name, raw_values in (rewards or {}).items(): + values = [float(value) for value in raw_values] + if not values: + continue + prefix = 'reward' if name == 'reward' else f'{name}_reward' + metrics[prefix] = sum(values) / len(values) + metrics[f'{prefix}_std'] = statistics.stdev(values) if len(values) > 1 else 0.0 + return metrics + + +def training_policy_metrics( + sample_tags: tuple[dict[str, Any], ...], + train_policy_version: int, +) -> dict[str, float | int]: + if not sample_tags: + raise ValueError('training batch must contain sample policy tags') + final_versions = [int(tag['final_policy_version']) for tag in sample_tags] + spans = [int(tag['policy_version_span']) for tag in sample_tags] + gaps = [int(train_policy_version) - version for version in final_versions] + if any(gap < 0 for gap in gaps): + raise ValueError( + f'training policy version {train_policy_version} is older than rollout versions {final_versions}') + return { + 'policy_version_gap_mean': sum(gaps) / len(gaps), + 'policy_version_gap_p95': _p95(gaps), + 'policy_version_gap_max': max(gaps), + 'rollout_policy_span_mean': sum(spans) / len(spans), + 'rollout_policy_span_max': max(spans), + } + + +def advantage_signal_metrics( + rewards: Sequence[float], + advantages: Sequence[float], + *, + num_generations: int, + zero_tolerance: float = 1e-8, +) -> dict[str, float | int]: + """Summarize whether GRPO groups provide a useful learning signal.""" + if num_generations <= 0: + raise ValueError(f'num_generations must be positive, got {num_generations}') + if len(rewards) != len(advantages): + raise ValueError(f'rewards and advantages must have equal length: {len(rewards)} != {len(advantages)}') + if len(rewards) == 0 or len(rewards) % num_generations: + raise ValueError( + f'advantage metrics require complete groups: sample_count={len(rewards)}, ' + f'num_generations={num_generations}') + + reward_values = [float(value) for value in rewards] + advantage_values = [float(value) for value in advantages] + group_reward_stds: list[float] = [] + zero_advantage_groups = 0 + for start in range(0, len(reward_values), num_generations): + group_rewards = reward_values[start:start + num_generations] + group_advantages = advantage_values[start:start + num_generations] + reward_mean = sum(group_rewards) / num_generations + group_reward_stds.append( + math.sqrt(sum((value - reward_mean)**2 for value in group_rewards) / num_generations)) + if max(abs(value) for value in group_advantages) <= zero_tolerance: + zero_advantage_groups += 1 + + advantage_mean = sum(advantage_values) / len(advantage_values) + group_count = len(group_reward_stds) + return { + 'group_count': group_count, + 'group_reward_std_mean': sum(group_reward_stds) / group_count, + 'zero_advantage_group_ratio': zero_advantage_groups / group_count, + 'positive_advantage_ratio': sum(value > zero_tolerance for value in advantage_values) / len(advantage_values), + 'advantage_mean': advantage_mean, + 'advantage_std': math.sqrt( + sum((value - advantage_mean)**2 for value in advantage_values) / len(advantage_values)), + } diff --git a/src/twinkle_agentic/async_rl/native_tq.py b/src/twinkle_agentic/async_rl/native_tq.py new file mode 100644 index 000000000..e42f68739 --- /dev/null +++ b/src/twinkle_agentic/async_rl/native_tq.py @@ -0,0 +1,187 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Adapters and helpers for the native TransferQueue client API. + +The async RL data path deliberately uses ``BatchMeta`` as its descriptor. A +``kv_list`` result is a diagnostic snapshot, not a queue cursor, and therefore +must not be used to drive the hot path. +""" + +from __future__ import annotations + +from typing import Any, Protocol, Sequence + +from transfer_queue import GRPOGroupNSampler + + +class AsyncTQClient(Protocol): + + async def async_get_meta(self, + *, + data_fields: list[str], + batch_size: int, + partition_id: str, + mode: str = 'fetch', + task_name: str | None = None, + sampling_config: dict[str, Any] | None = None) -> Any: + ... + + async def async_get_data(self, metadata: Any) -> Any: + ... + + async def async_put(self, data: Any, metadata: Any | None = None, partition_id: str | None = None) -> Any: + ... + + async def async_clear_partition(self, partition_id: str) -> Any: + ... + + async def async_check_consumption_status(self, task_name: str, partition_id: str) -> bool: + ... + + async def async_set_custom_meta(self, metadata: Any) -> Any: + ... + + +class ContextGRPOGroupNSampler(GRPOGroupNSampler): + """Select complete prompt groups using the request's generation count.""" + + def sample( + self, + ready_indexes: list[int], + batch_size: int, + task_name: str = '', + partition_id: str = '', + *args: Any, + **kwargs: Any, + ) -> tuple[list[int], list[int]]: + group_size = int(kwargs['n_samples_per_prompt']) + if group_size <= 0: + raise ValueError(f'n_samples_per_prompt must be positive, got {group_size}') + if batch_size % group_size: + raise ValueError(f'batch_size ({batch_size}) must be a multiple of n_samples_per_prompt ({group_size})') + + states = self._states.get(partition_id, {}).get(task_name, {}) + dp_rank = kwargs.get('dp_rank') + batch_index = kwargs.get('batch_index') + if dp_rank in states and batch_index in states[dp_rank]: + return states[dp_rank][batch_index] + + ready = sorted(ready_indexes) + selected: list[int] = [] + offset = 0 + while offset <= len(ready) - group_size and len(selected) < batch_size: + group = ready[offset:offset + group_size] + if all(right - left == 1 for left, right in zip(group, group[1:])): + selected.extend(group) + offset += group_size + else: + offset += 1 + + if len(selected) != batch_size: + return [], [] + + result = (selected, selected.copy()) + if dp_rank is not None: + states.setdefault(dp_rank, {})[batch_index] = result + self._states.setdefault(partition_id, {})[task_name] = states + return result + + +def batch_size_for_groups(groups: int, num_generations: int) -> int: + if groups <= 0: + raise ValueError(f'groups must be positive, got {groups}') + if num_generations <= 0: + raise ValueError(f'num_generations must be positive, got {num_generations}') + return groups * num_generations + + +def validate_group_batch_size(batch_size: int, num_generations: int) -> None: + if batch_size <= 0: + raise ValueError(f'batch_size must be positive, got {batch_size}') + if num_generations <= 0: + raise ValueError(f'num_generations must be positive, got {num_generations}') + if batch_size % num_generations: + raise ValueError(f'batch_size={batch_size} must be divisible by num_generations={num_generations}') + + +def metadata_size(metadata: Any) -> int: + """Return the native BatchMeta size.""" + if metadata is None: + return 0 + return int(metadata.size) + + +async def fetch_ready_batch( + client: AsyncTQClient, + *, + data_fields: list[str], + batch_size: int, + partition_id: str, + task_name: str, + num_generations: int, + sampling_config: dict[str, Any] | None = None, +) -> Any | None: + """Fetch one complete group-aligned batch using TQ production status. + + The caller owns the outer service loop. This helper performs one request + only, which keeps shutdown and failure handling explicit and avoids hiding + an unbounded wait in a data-plane utility. + """ + + validate_group_batch_size(batch_size, num_generations) + config = dict(sampling_config or {}) + config['n_samples_per_prompt'] = num_generations + metadata = await client.async_get_meta( + data_fields=list(data_fields), + batch_size=batch_size, + partition_id=partition_id, + mode='fetch', + task_name=task_name, + sampling_config=config, + ) + return metadata if metadata_size(metadata) else None + + +async def append_fields(client: AsyncTQClient, data: Any, metadata: Any) -> Any: + """Append fields to exactly the samples described by ``metadata``.""" + + if metadata_size(metadata) == 0: + raise ValueError('cannot append fields to an empty BatchMeta') + return await client.async_put(data=data, metadata=metadata) + + +async def set_sample_tags(client: AsyncTQClient, metadata: Any, tags: Sequence[dict[str, Any]]) -> None: + """Persist tags through the native metadata API in one controller request.""" + + if metadata_size(metadata) != len(tags): + raise ValueError(f'metadata size {metadata_size(metadata)} does not match tags {len(tags)}') + metadata.update_custom_meta([dict(tag) for tag in tags]) + await client.async_set_custom_meta(metadata) + + +def split_batch_meta(metadata: Any, group_size: int) -> list[Any]: + """Split a preallocated BatchMeta into contiguous prompt-group views.""" + + size = metadata_size(metadata) + if group_size <= 0: + raise ValueError(f'group_size must be positive, got {group_size}') + if size % group_size: + raise ValueError(f'metadata size {size} is not divisible by group_size {group_size}') + return [metadata.select_samples(list(range(start, start + group_size))) for start in range(0, size, group_size)] + + +async def preallocate_partition( + client: AsyncTQClient, + *, + partition_id: str, + prompt_fields: Any, +) -> Any: + """Insert prompt rows once and return the native BatchMeta descriptor.""" + + batch_size = getattr(prompt_fields, 'batch_size', None) + if batch_size is None or len(batch_size) == 0 or int(batch_size[0]) <= 0: + raise ValueError('prompt_fields must be a non-empty batched TensorDict') + return await client.async_put(data=prompt_fields, partition_id=partition_id) + + +async def clear_partition(client: AsyncTQClient, partition_id: str) -> None: + await client.async_clear_partition(partition_id) diff --git a/src/twinkle_agentic/async_rl/pipeline.py b/src/twinkle_agentic/async_rl/pipeline.py new file mode 100644 index 000000000..012372bea --- /dev/null +++ b/src/twinkle_agentic/async_rl/pipeline.py @@ -0,0 +1,687 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Driver for the independent async-RL Ray workers.""" + +from __future__ import annotations + +import asyncio +import time +from dataclasses import dataclass +from functools import partial +from pydoc import locate +from typing import Any, Sequence + +from twinkle.metric import MetricRecord, MetricsReporter, create_metrics_reporter +from .context_manager import LoraContextManager +from .data_plane import TQDataPlane +from .scheduler import ContextSchedulePolicy, SchedulerConfig +from .types import LoraContext, PartitionAdmission +from .utils import (TrainBatchConfig, build_native_fsdp_model_kwargs, + configure_lora_lr_scheduler, resolve_context_learning_rate, + resolve_context_lora_target_modules, resolve_context_loss_config, + resolve_model_attention_implementation, sampler_data_parallel_size, + resolve_sequence_parallel_size, validate_context_batch_config) +from .workers import AdvantageWorker, RolloutWorker, TrainerWorker + + +@dataclass(frozen=True) +class AsyncMultiLoraGRPOConfig: + metrics_drain_interval_s: float = 1.0 + + +class AsyncMultiLoraGRPOPipeline: + """Owns the production async-RL runtime and drives its worker services. + + ``from_config`` is the real training construction path. The explicit + constructor remains available to fake-TQ tests, where injecting a fake + sampler/model is the point of the test. + """ + + def __init__(self, + *, + context_manager: LoraContextManager, + rollout_worker: RolloutWorker, + advantage_worker: AdvantageWorker, + trainer_worker: TrainerWorker, + metrics: MetricsReporter | None = None, + config: AsyncMultiLoraGRPOConfig = AsyncMultiLoraGRPOConfig(), + sampler: Any | None = None, + model: Any | None = None, + contexts: Sequence[LoraContext] = ()): + self.context_manager = context_manager + self.rollout_worker = rollout_worker + self.advantage_worker = advantage_worker + self.trainer_worker = trainer_worker + self.sampler = sampler + self.model = model + self.contexts = tuple(contexts) + self.metrics = metrics + self.config = config + + @classmethod + def from_config( + cls, + raw_config: dict[str, Any], + *, + persistent: bool = False, + ) -> AsyncMultiLoraGRPOPipeline: + """Build the complete Ray/TQ runtime from the async-RL YAML mapping.""" + from omegaconf import OmegaConf + + raw_config = OmegaConf.to_container(OmegaConf.create(raw_config), resolve=True) + if not isinstance(raw_config, dict): + raise TypeError('async-RL config must resolve to a mapping') + + import ray + import transfer_queue as tq + from peft import LoraConfig + + import twinkle + from twinkle import DeviceGroup, DeviceMesh + from twinkle.data_format import SamplingParams + from twinkle.model import MultiLoraTransformersModel + from twinkle.processor import InputProcessor + from .native_tq import ContextGRPOGroupNSampler + + runtime = raw_config['runtime'] + model_config = raw_config['model'] + lora_config_data = raw_config['lora'] + loss_config_data = raw_config.get('loss') + template_config = raw_config.get('template', {}) + template_cls = template_config.get('cls', 'Qwen3_5Template') + enable_thinking = bool(template_config.get('enable_thinking', False)) + rollout_output_config = dict(raw_config.get('rollout_output') or {}) + sampler_gpus = int(runtime['sampler_gpus']) + sampler_tp = int(runtime['sampler_tp']) + sampler_dp = sampler_data_parallel_size(sampler_gpus, sampler_tp) + model_dp = int(runtime['model_gpus']) + sequence_parallel_size = resolve_sequence_parallel_size( + model_dp, + int(model_config['sequence_parallel_size']), + ) + padding_free = bool(model_config['padding_free']) + attn_implementation = resolve_model_attention_implementation( + model_config, + padding_free=padding_free, + sequence_parallel_size=sequence_parallel_size, + ) + model_max_length = int(model_config['max_length']) + sampler_config = raw_config['sampler'] + total_gpus = model_dp + sampler_gpus + device_groups = [ + DeviceGroup('model', list(range(int(runtime['model_gpus']))), device_type='GPU'), + DeviceGroup( + 'sampler', + list(range(int(runtime['model_gpus']), total_gpus)), + device_type='GPU', + gpus_per_worker=sampler_tp, + ), + ] + twinkle.initialize(mode='ray', nproc_per_node=total_gpus, groups=device_groups, lazy_collect=False) + tq.init( + OmegaConf.create( + { + 'controller': { + 'sampler': ContextGRPOGroupNSampler, + 'polling_mode': bool(raw_config['tq'].get('polling_mode', True)), + }, + 'backend': { + 'SimpleStorage': { + 'num_data_storage_units': raw_config['tq']['storage_units'] + } + }, + }, + flags={'allow_objects': True})) + + model_mesh = DeviceMesh.from_sizes( + world_size=model_dp, + dp_size=model_dp, + ulysses_size=sequence_parallel_size, + ) + model_data_parallel_size = model_mesh.data_world_size + sampler_mesh = DeviceMesh.from_sizes(world_size=sampler_gpus, dp_size=sampler_dp, tp_size=sampler_tp) + model_kwargs = build_native_fsdp_model_kwargs(model_config) + if attn_implementation is not None: + model_kwargs['attn_implementation'] = attn_implementation + model = MultiLoraTransformersModel( + model_id=runtime['model_id'], + device_mesh=model_mesh, + remote_group='model', + max_length=model_max_length, + **model_kwargs, + ) + contexts: list[LoraContext] = [] + prompt_sources: dict[str, Any] = {} + rollout_config: dict[str, dict[str, Any]] = {} + train_batch_configs: dict[str, TrainBatchConfig] = {} + rewards: dict[str, Any] = {} + evaluation_config: dict[str, dict[str, Any]] = {} + evaluation_rewards: dict[str, Any] = {} + initial_paths: dict[str, str] = {} + global_evaluation = dict(raw_config.get('evaluation') or {}) + for item in raw_config['lora_contexts']: + train = item['train'] + context = LoraContext( + item['tenant_id'], + item['training_run_id'], + runtime['model_id'], + item['adapter_name'], + ) + contexts.append(context) + adapter_lora_config = LoraConfig( + target_modules=resolve_context_lora_target_modules(item, lora_config_data), + r=lora_config_data['r'], + lora_alpha=lora_config_data['alpha'], + lora_dropout=lora_config_data['dropout'], + ) + model.add_adapter_to_model( + context.adapter_name, + adapter_lora_config, + gradient_accumulation_steps=1, + ) + model.set_optimizer( + 'AdamW', + lr=resolve_context_learning_rate(train, lora_config_data), + adapter_name=context.adapter_name, + ) + configure_lora_lr_scheduler(model, context.adapter_name, lora_config_data) + loss_cls, loss_kwargs = resolve_context_loss_config(item, loss_config_data) + model.set_loss( + loss_cls, + adapter_name=context.adapter_name, + **loss_kwargs, + ) + model.set_processor( + InputProcessor, + adapter_name=context.adapter_name, + padding_free=padding_free, + ) + model.set_template( + template_cls, + model_id=runtime['model_id'], + adapter_name=context.adapter_name, + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + + rollout = item['rollout'] + rollout_batch_size = int(rollout['batch_size']) + num_generations = int(rollout['num_generations']) + train_batch_config = TrainBatchConfig( + mini_batch_size=int(train['mini_batch_size']), + micro_batch_size=int(train['micro_batch_size']), + dynamic_batching=bool(train.get('dynamic_batching', False)), + max_tokens_per_micro_batch=( + int(train['max_tokens_per_micro_batch']) + if train.get('max_tokens_per_micro_batch') is not None else None + ), + packing_algorithm=str(train.get('packing_algorithm', 'ffd')), + ) + validate_context_batch_config( + context.key, + rollout_groups=rollout_batch_size, + num_generations=num_generations, + train=train_batch_config, + sampler_dp=sampler_dp, + model_dp=model_data_parallel_size, + ) + prompt_sources[context.key] = partial( + _prompt_batches, + item['dataset'], + model_id=runtime['model_id'], + batch_size=rollout_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + ) + rollout_config[context.key] = { + 'context': + context, + 'batch_size': + rollout_batch_size, + 'num_generations': + num_generations, + 'sampling_params': + SamplingParams( + max_tokens=rollout['max_tokens'], + temperature=rollout['temperature'], + top_p=rollout['top_p'], + repetition_penalty=float(rollout.get('repetition_penalty', 1.0)), + logprobs=1, + num_samples=1, + ), + } + train_batch_configs[context.key] = train_batch_config + rewards[context.key] = _reward_for_context( + item.get('reward'), + context_key=context.key, + ) + if bool(global_evaluation.get('enabled', False)): + eval_dataset = item.get('eval_dataset') + if eval_dataset is None: + raise ValueError(f'eval_dataset is required for periodic evaluation of {context.key}') + eval_batch_size = int(global_evaluation.get('batch_size', 16)) + eval_interval = int(global_evaluation.get('interval', 1)) + if eval_batch_size <= 0 or eval_interval <= 0: + raise ValueError('evaluation.batch_size and evaluation.interval must be positive') + eval_sampling = dict(global_evaluation.get('sampling_params') or {}) + evaluation_config[context.key] = { + 'interval': eval_interval, + 'dataset_name': eval_dataset.get('name', eval_dataset['dataset_id']), + 'prompt_batches': partial( + _prompt_batches, + eval_dataset, + model_id=runtime['model_id'], + batch_size=eval_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + full_batches_only=False, + ), + 'sampling_params': SamplingParams( + max_tokens=int(eval_sampling.get('max_tokens', rollout['max_tokens'])), + temperature=float(eval_sampling.get('temperature', 0.0)), + top_p=float(eval_sampling.get('top_p', 1.0)), + repetition_penalty=float(eval_sampling.get('repetition_penalty', 1.0)), + logprobs=0, + num_samples=1, + ), + } + evaluation_rewards[context.key] = _reward_for_context( + eval_dataset.get('reward'), + context_key=f'{context.key} evaluation', + ) + initial_paths[context.key] = model.save( + f'async-{context.adapter_name}-initial', + output_dir=runtime['output_dir'], + adapter_name=context.adapter_name, + ) + + manager = create_cpu_actor( + LoraContextManager, + max_staleness=runtime['max_staleness'], + max_steps=runtime['max_steps'], + ) + for context in contexts: + ray.get(manager.register_context.remote(context, adapter_path=initial_paths[context.key])) + + from .vllm_sampler_tq import VLLMSamplerTQ + sampler_engine_args = { + 'tensor_parallel_size': sampler_tp, + 'enable_lora': True, + 'max_loras': int(runtime['sampler_max_loras']), + 'max_lora_rank': lora_config_data['r'], + 'max_model_len': int(sampler_config['max_model_len']), + 'gpu_memory_utilization': float(sampler_config['gpu_memory_utilization']), + 'max_num_seqs': int(sampler_config['max_num_seqs']), + 'enforce_eager': bool(sampler_config['enforce_eager']), + 'seed': int(runtime.get('seed', 1)), + } + if sampler_config.get('max_num_batched_tokens') is not None: + sampler_engine_args['max_num_batched_tokens'] = int(sampler_config['max_num_batched_tokens']) + sampler = VLLMSamplerTQ( + model_id=runtime['model_id'], + remote_group='sampler', + device_mesh=sampler_mesh, + engine_args=sampler_engine_args, + reward_registry=rewards, + context_manager=manager, + rollout_max_retries=int(runtime.get('rollout_max_retries', 2)), + rollout_retry_delay_s=float(runtime.get('rollout_retry_delay_s', 0.5)), + rollout_output_dir=( + rollout_output_config.get('output_dir') + if bool(rollout_output_config.get('enabled', False)) else None + ), + rollout_output_include_token_ids=bool( + rollout_output_config.get('include_token_ids', False) + ), + ) + sampler.set_template( + template_cls, + model_id=runtime['model_id'], + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + + rollout_worker = create_cpu_actor( + RolloutWorker, + context_manager=manager, + data_plane=TQDataPlane(), + sampler=sampler, + prompt_batches=prompt_sources, + rollout_config=rollout_config, + scheduler=_scheduler(raw_config['scheduler']['rollout']), + allow_partial_rollout=runtime['allow_partial_rollout'], + persistent=persistent, + ) + advantage_worker = create_cpu_actor( + AdvantageWorker, + context_manager=manager, + data_plane=TQDataPlane(), + advantage_fn=_compute_advantages, + scheduler=_scheduler(raw_config['scheduler']['advantage']), + persistent=persistent, + ) + trainer_worker = create_cpu_actor( + TrainerWorker, + context_manager=manager, + data_plane=TQDataPlane(), + train_fn=partial( + _train_batch, + model, + train_batch_configs, + model_data_parallel_size=model_data_parallel_size, + ), + train_with_config_fn=partial( + _train_batch_with_config, + model, + model_data_parallel_size=model_data_parallel_size, + ), + train_batch_configs=train_batch_configs, + save_adapter=partial(_save_adapter, model, runtime['output_dir']), + mini_batch_sizes={ + key: config.mini_batch_size for key, config in train_batch_configs.items() + }, + scheduler=_scheduler(raw_config['scheduler']['train']), + keep_adapter_versions=runtime['keep_adapter_versions'], + initial_adapter_paths=initial_paths, + remove_adapter=partial(_remove_adapter_snapshot, sampler), + evaluation_config=evaluation_config, + evaluate_batch=partial(_evaluate_batch, sampler, evaluation_rewards) if evaluation_config else None, + evaluate_with_reward_fn=partial(_evaluate_batch_with_reward, sampler), + evaluation_rewards=evaluation_rewards, + persistent=persistent, + ) + raw_metrics_config = raw_config.get('metrics') + metrics_config = dict(raw_metrics_config or {}) + metrics = create_metrics_reporter( + raw_metrics_config, + run_id=str(runtime.get('run_id', 'async_multi_lora_grpo')), + ) + return cls( + context_manager=manager, + rollout_worker=rollout_worker, + advantage_worker=advantage_worker, + trainer_worker=trainer_worker, + sampler=sampler, + metrics=metrics, + config=AsyncMultiLoraGRPOConfig( + metrics_drain_interval_s=float(metrics_config.get('drain_interval_s', 1.0)), + ), + model=model, + contexts=contexts, + ) + + async def run_async(self) -> dict[str, Any]: + started = time.perf_counter() + workers = [self.rollout_worker, self.advantage_worker, self.trainer_worker] + await asyncio.gather(*(worker.start.remote() for worker in workers)) + try: + while True: + await self._drain_metrics() + states = await asyncio.gather(*(worker.get_service_state.remote() for worker in workers)) + failures = [state['failure'] for state in states if state['failure']] + if failures: + raise RuntimeError(f'async RL worker failed: {failures[0]}') + if self.sampler is not None: + await asyncio.to_thread(self.sampler.check_health) + running = any(bool(state['running']) for state in states) + if not running: + if await self.context_manager.is_run_finished.remote(): + break + raise RuntimeError('async RL workers stopped before all contexts were drained') + await asyncio.sleep(self.config.metrics_drain_interval_s) + except Exception as exc: + if self.metrics is not None: + self.metrics.record(MetricRecord( + stage='run', + status='failed', + values={'wall_time_s': time.perf_counter() - started}, + attributes={'error': f'{type(exc).__name__}: {exc}'}, + )) + self.metrics.flush() + raise + finally: + await asyncio.gather(*(worker.stop.remote() for worker in workers), return_exceptions=True) + await self._drain_metrics() + result = { + 'trained_partitions': await self.context_manager.get_completed_partitions.remote(), + 'wall_time_s': time.perf_counter() - started, + } + if self.metrics is not None: + self.metrics.record(MetricRecord(stage='run', values=result)) + self.metrics.flush() + result['metrics_health'] = self.metrics.health() + return result + + def run(self) -> dict[str, Any]: + try: + return asyncio.run(self.run_async()) + finally: + if self.metrics is not None: + self.metrics.close() + + async def _drain_metrics(self) -> None: + workers = [self.rollout_worker, self.advantage_worker, self.trainer_worker] + for worker in workers: + records = await worker.drain_metric_records.remote() + if self.metrics is not None: + self.metrics.record_many(records) + if self.sampler is not None: + records = await asyncio.to_thread(self.sampler.drain_metric_records) + if self.metrics is not None: + self.metrics.record_many(records) + + +def create_cpu_actor(cls: type, *args: Any, **kwargs: Any) -> Any: + """Deploy a CPU service as one raw Ray actor; local tests use the class directly.""" + + import ray + actor_class = ray.remote( + num_cpus=1, + runtime_env={'env_vars': { + 'TWINKLE_MODE': 'ray' + }}, + )( + cls) + return actor_class.remote(*args, **kwargs) + + +def _scheduler(config: dict[str, Any]) -> SchedulerConfig: + return SchedulerConfig(ContextSchedulePolicy(config['policy']), config.get('max_consecutive_units')) + + +def _prompt_batches( + dataset_config: dict[str, Any], + *, + model_id: str, + batch_size: int, + template_cls: str, + enable_thinking: bool, + full_batches_only: bool = True, +): + """Create a lazy, full-batch-only prompt source for one context.""" + from twinkle.dataloader import DataLoader + from twinkle.dataset import Dataset, DatasetMeta + from twinkle.preprocessor import llm as llm_processors + + def batches(): + data_num = dataset_config.get('data_num') + dataset = Dataset( + DatasetMeta( + dataset_config['dataset_id'], + subset_name=dataset_config.get('subset_name'), + split=dataset_config.get('split', 'train'), + data_slice=range(int(data_num)) if data_num is not None else None, + )) + dataset.set_template( + template_cls, + model_id=model_id, + max_length=dataset_config['max_length'], + enable_thinking=enable_thinking, + ) + processor_name = dataset_config.get('processor', 'GSM8KProcessor') + processor_cls = getattr(llm_processors, processor_name) + if processor_name == 'GSM8KProcessor': + processor = processor_cls(system=dataset_config['system_prompt']) + else: + processor = processor_cls() + dataset.map(processor) + dataset.encode(add_generation_prompt=True) + loader = DataLoader( + dataset=dataset, + batch_size=batch_size, + min_batch_size=batch_size if full_batches_only else 1, + ) + remaining = data_num + remaining = None if remaining is None else int(remaining) + for batch in loader: + if full_batches_only and (len(batch) != batch_size or (remaining is not None and remaining < batch_size)): + return + yield batch + if remaining is not None: + remaining -= batch_size + + return batches() + + +def _evaluate_batch( + sampler: Any, + reward_registry: dict[str, Any], + prompts: Sequence[dict[str, Any]], + admission: PartitionAdmission, + adapter_path: str, + policy_version: int, + sampling_params: Any, +) -> dict[str, Any]: + reward_fn = reward_registry[admission.context.key] + return _evaluate_batch_with_reward( + sampler, + prompts, + admission, + adapter_path, + policy_version, + sampling_params, + reward_fn, + ) + + +def _evaluate_batch_with_reward( + sampler: Any, + prompts: Sequence[dict[str, Any]], + admission: PartitionAdmission, + adapter_path: str, + policy_version: int, + sampling_params: Any, + reward_fn: Any, +) -> dict[str, Any]: + from .utils import sample_responses_to_rollout_rows + + responses = sampler.evaluate( + list(prompts), + sampling_params, + admission.context.adapter_name, + adapter_path, + ) + rows = sample_responses_to_rollout_rows(list(prompts), responses, policy_version=policy_version) + rewards = list(reward_fn(rows, context=admission.context)) + return { + 'rewards': rewards, + 'completion_lengths': [int(row['completion_length']) for row in rows], + } + + +def _reward_for_context( + reward_config: dict[str, Any] | None = None, + *, + context_key: str, +) -> Any: + from twinkle.reward import Reward + + config = dict(reward_config or {}) + class_path = config.get('class_path', '') + reward_cls = locate(class_path) + if not isinstance(reward_cls, type) or not issubclass(reward_cls, Reward): + raise TypeError(f'reward.class_path {class_path!r} for {context_key} must reference a Reward subclass') + return reward_cls(**dict(config.get('kwargs') or {})) + + +def _compute_advantages(data: Any, admission: PartitionAdmission) -> tuple[list[float], list[float]]: + from twinkle.advantage import GRPOAdvantage + rewards = [float(value) for value in data['rewards']] + advantages = GRPOAdvantage()(rewards, num_generations=admission.num_generations, scale='group').tolist() + return advantages, rewards + + +def _train_batch( + model: Any, + train_batch_configs: dict[str, TrainBatchConfig], + data: Any, + admission: PartitionAdmission, + *, + model_data_parallel_size: int = 1, +) -> dict[str, Any]: + config = train_batch_configs[admission.context.key] + return _train_batch_with_config( + model, + data, + admission, + config, + model_data_parallel_size=model_data_parallel_size, + ) + + +def _train_batch_with_config( + model: Any, + data: Any, + admission: PartitionAdmission, + config: TrainBatchConfig, + *, + model_data_parallel_size: int = 1, +) -> dict[str, Any]: + from .tq_utils import REQUIRED_MODEL_INPUT_FIELDS + + size = int(data.batch_size[0]) + inputs = [{name: data[name][index] for name in REQUIRED_MODEL_INPUT_FIELDS} for index in range(size)] + old_logps = list(data['logprobs']) + advantages = list(data['advantages']) + if size != config.mini_batch_size: + raise ValueError( + f'train batch for {admission.context.key} has {size} samples; ' + f'expected mini_batch_size={config.mini_batch_size}') + + if size % model_data_parallel_size: + raise ValueError(f'train batch size {size} must be divisible by model DP size ' + f'{model_data_parallel_size}') + model.forward_backward( + inputs=inputs, + old_logps=old_logps, + advantages=advantages, + adapter_name=admission.context.adapter_name, + micro_batch_size=config.micro_batch_size, + dynamic_batching=config.dynamic_batching, + max_tokens_per_micro_batch=config.max_tokens_per_micro_batch, + packing_algorithm=config.packing_algorithm, + sync_gradients=True, + loss_scale=1.0, + ) + + model.clip_grad_and_step(adapter_name=admission.context.adapter_name) + metrics = dict(model.calculate_metric(is_training=True, adapter_name=admission.context.adapter_name)) + metrics['mini_batch_size'] = config.mini_batch_size + metrics['micro_batch_size_per_rank'] = config.micro_batch_size + metrics['dynamic_batching'] = config.dynamic_batching + return metrics + + +def _save_adapter(model: Any, output_dir: str, admission: PartitionAdmission) -> str: + return model.save( + f'async-{admission.context.adapter_name}-v{admission.step + 1}', + output_dir=output_dir, + adapter_name=admission.context.adapter_name, + ) + + +def _remove_adapter_snapshot(sampler: Any, adapter_path: str) -> None: + """Unload an unreferenced policy from vLLM before deleting its checkpoint.""" + from .workers import _remove_local_adapter + + sampler.unload_lora_paths([adapter_path]) + _remove_local_adapter(adapter_path) diff --git a/src/twinkle_agentic/async_rl/scheduler.py b/src/twinkle_agentic/async_rl/scheduler.py new file mode 100644 index 000000000..e77c53a93 --- /dev/null +++ b/src/twinkle_agentic/async_rl/scheduler.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Context selection policies used independently by rollout/advantage/training.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Sequence + +from .types import LoraContext, PartitionAdmission + + +class ContextSchedulePolicy(StrEnum): + ROUND_ROBIN = 'round_robin' + STICKY = 'sticky' + OLDEST_PARTITION = 'oldest_partition' + + +@dataclass(frozen=True) +class SchedulerConfig: + policy: ContextSchedulePolicy = ContextSchedulePolicy.ROUND_ROBIN + max_consecutive_units: int | None = 1 + + +@dataclass(frozen=True) +class ScheduleCandidate: + context: LoraContext + partition: PartitionAdmission | None = None + + +class ContextScheduler: + + def __init__(self, config: SchedulerConfig): + self.config = config + self._cursor = 0 + self._sticky_key: str | None = None + self._consecutive = 0 + + def choose(self, candidates: Sequence[ScheduleCandidate]) -> ScheduleCandidate | None: + if not candidates: + return None + if self.config.policy is ContextSchedulePolicy.OLDEST_PARTITION: + return min( + candidates, + key=lambda item: (item.partition.created_order if item.partition else float('inf'), item.context.key)) + if self.config.policy is ContextSchedulePolicy.STICKY and self._sticky_key is not None: + cap = self.config.max_consecutive_units + if cap is None or self._consecutive < cap: + for candidate in candidates: + if candidate.context.key == self._sticky_key: + return candidate + else: + for candidate in candidates: + if candidate.context.key != self._sticky_key: + return candidate + index = self._cursor % len(candidates) + return candidates[index] + + def on_success(self, candidate: ScheduleCandidate) -> None: + if self.config.policy is ContextSchedulePolicy.STICKY and candidate.context.key == self._sticky_key: + self._consecutive += 1 + return + self._sticky_key = candidate.context.key + self._consecutive = 1 + if self.config.policy is ContextSchedulePolicy.ROUND_ROBIN: + self._cursor += 1 + + def on_blocked(self, candidate: ScheduleCandidate) -> None: + if candidate.context.key == self._sticky_key: + self._sticky_key = None + self._consecutive = 0 + if self.config.policy is ContextSchedulePolicy.ROUND_ROBIN: + self._cursor += 1 diff --git a/src/twinkle_agentic/async_rl/tq_utils.py b/src/twinkle_agentic/async_rl/tq_utils.py new file mode 100644 index 000000000..55284567e --- /dev/null +++ b/src/twinkle_agentic/async_rl/tq_utils.py @@ -0,0 +1,29 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from twinkle.tq_utils import columns_to_tq_fields, rows_to_tq_fields + +TRANSFORMERS_INPUT_FIELDS = ( + 'input_ids', + 'labels', + 'attention_mask', + 'position_ids', + 'cu_seqlens', + 'completion_mask', + 'pixel_values', + 'image_grid_thw', + 'video_pixel_values', + 'video_grid_thw', + 'input_features', + 'feature_attention_mask', +) +REQUIRED_MODEL_INPUT_FIELDS = ('input_ids', 'labels', 'attention_mask', 'position_ids') +ROLLOUT_TRAIN_FIELDS = (*TRANSFORMERS_INPUT_FIELDS, 'logprobs', 'rewards', 'advantages', 'returns') + +__all__ = [ + 'ROLLOUT_TRAIN_FIELDS', + 'REQUIRED_MODEL_INPUT_FIELDS', + 'TRANSFORMERS_INPUT_FIELDS', + 'columns_to_tq_fields', + 'rows_to_tq_fields', +] diff --git a/src/twinkle_agentic/async_rl/types.py b/src/twinkle_agentic/async_rl/types.py new file mode 100644 index 000000000..e31352f6c --- /dev/null +++ b/src/twinkle_agentic/async_rl/types.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training-domain descriptors for native-TQ async multi-LoRA RL.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TypedDict + + +@dataclass(frozen=True) +class LoraContext: + tenant_id: str + training_run_id: str + base_model_id: str + adapter_name: str + tool_profile: str = 'default' + + @property + def key(self) -> str: + return f'{self.tenant_id}/{self.training_run_id}/{self.adapter_name}' + + def partition_id(self, step: int) -> str: + return f'{self.key}/train_{step}' + + +@dataclass(frozen=True) +class RolloutPolicy: + """The immutable policy snapshot used by one rollout group.""" + + context_key: str + adapter_name: str + version: int + adapter_path: str | None + + +@dataclass(frozen=True) +class PartitionAdmission: + """Control-plane admission result for one complete training-data partition. + + A partition is only a training-data batch. Its prompt groups may be + generated by different policy snapshots. + """ + + context: LoraContext + partition_id: str + step: int + target_groups: int + num_generations: int + created_order: int + + @property + def sample_count(self) -> int: + return self.target_groups * self.num_generations + + +@dataclass +class PromptGroup: + """A prompt and the BatchMeta descriptor reserved for its generated samples.""" + + context: LoraContext + partition: PartitionAdmission + group_id: str + prompt: dict[str, Any] + batch_meta: Any + + @property + def partition_id(self) -> str: + return self.partition.partition_id + + @property + def num_samples(self) -> int: + return self.partition.num_generations + +@dataclass +class PreparedPartition: + """Partition prepared in TQ and ready for sampler submission.""" + + admission: PartitionAdmission + groups: tuple[PromptGroup, ...] + sampling_params: Any + + +@dataclass +class ClaimedBatch: + """A TQ-consumed batch and the BatchMeta descriptor selecting its samples.""" + + admission: PartitionAdmission + data: Any + batch_meta: Any + sample_tags: tuple[dict[str, Any], ...] = () + + +class RolloutOutput(TypedDict, total=False): + logprobs: list[float] + rewards: float + completion_length: int + generation_idx: int + rollout_policy_version: int + rollout_adapter_path: str | None + rollout_policy_versions: list[int] + initial_policy_version: int + final_policy_version: int + policy_version_span: int diff --git a/src/twinkle_agentic/async_rl/utils.py b/src/twinkle_agentic/async_rl/utils.py new file mode 100644 index 000000000..bd2a176de --- /dev/null +++ b/src/twinkle_agentic/async_rl/utils.py @@ -0,0 +1,218 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared configuration helpers for synchronous and asynchronous RL runners.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal + +from twinkle.data_format import SampleResponse + +from .types import RolloutOutput + + +@dataclass(frozen=True) +class TrainBatchConfig: + mini_batch_size: int + micro_batch_size: int + dynamic_batching: bool = False + max_tokens_per_micro_batch: int | None = None + packing_algorithm: Literal['ffd', 'kk'] = 'ffd' + + +def _extract_sampled_token_logps(logprobs: Any) -> list[float]: + return [0.0 if not item else float(item[0][1]) for item in logprobs or []] + + +def sample_responses_to_rollout_rows( + sources: list[dict[str, Any]], + responses: list[SampleResponse], + *, + policy_version: int | None, +) -> list[RolloutOutput]: + rows: list[RolloutOutput] = [] + for source, response in zip(sources, responses): + for sequence in response.sequences: + row = dict(source) + row.update(sequence.new_input_feature or {}) + row.setdefault('group_id', source['group_id']) + row.setdefault('generation_idx', source['generation_idx']) + row['logprobs'] = _extract_sampled_token_logps(sequence.logprobs) + row['stop_reason'] = sequence.stop_reason + row['completion_length'] = len(sequence.tokens) + row['rollout_policy_version'] = policy_version + rows.append(row) + return rows + + +def resolve_adapter_path(adapter_path: str) -> str: + path = os.path.abspath(os.path.expanduser(str(adapter_path))) + if not os.path.exists(path): + raise FileNotFoundError(f'local LoRA adapter path does not exist: {path}') + return path + + +def sampler_data_parallel_size(sampler_gpus: int, sampler_tp: int) -> int: + if sampler_gpus <= 0: + raise ValueError(f'sampler_gpus must be positive, got {sampler_gpus}') + if sampler_tp <= 0: + raise ValueError(f'sampler_tp must be positive, got {sampler_tp}') + if sampler_gpus % sampler_tp != 0: + raise ValueError(f'sampler_gpus ({sampler_gpus}) must be divisible by sampler_tp ({sampler_tp})') + return sampler_gpus // sampler_tp + + +def resolve_sequence_parallel_size(model_gpus: int, configured_size: int) -> int: + if configured_size <= 0: + raise ValueError(f'model.sequence_parallel_size must be positive, got {configured_size}') + if model_gpus % configured_size: + raise ValueError(f'runtime.model_gpus ({model_gpus}) must be divisible by ' + f'model.sequence_parallel_size ({configured_size})') + return configured_size + + +def resolve_model_attention_implementation( + model_config: Mapping[str, Any], + *, + padding_free: bool, + sequence_parallel_size: int, +) -> str | None: + implementation = model_config.get('attn_implementation') + if implementation is not None: + implementation = str(implementation) + if padding_free and sequence_parallel_size > 1 and implementation != 'flash_attention_2': + raise ValueError( + 'model.attn_implementation must be flash_attention_2 when ' + 'model.padding_free=true and model.sequence_parallel_size>1') + return implementation + + +def build_native_fsdp_model_kwargs(model_config: Mapping[str, Any]) -> dict[str, Any]: + strategy = str(model_config.get('strategy', 'native_fsdp')) + if strategy != 'native_fsdp': + raise ValueError(f'model.strategy must be native_fsdp for RL training, got {strategy!r}') + return { + 'strategy': strategy, + 'fsdp_config': dict(model_config.get('fsdp_config') or {}), + } + + +def validate_context_batch_config( + context_key: str, + *, + rollout_groups: int, + num_generations: int, + train: TrainBatchConfig, + sampler_dp: int, + model_dp: int, +) -> None: + values = { + 'rollout.batch_size': rollout_groups, + 'rollout.num_generations': num_generations, + 'train.mini_batch_size': train.mini_batch_size, + 'train.micro_batch_size': train.micro_batch_size, + } + for name, value in values.items(): + if value <= 0: + raise ValueError(f'{name} for {context_key} must be positive, got {value}') + if rollout_groups % sampler_dp: + raise ValueError(f'rollout.batch_size for {context_key} must be divisible by sampler DP size ' + f'({sampler_dp}), got {rollout_groups}') + partition_samples = rollout_groups * num_generations + if partition_samples % train.mini_batch_size: + raise ValueError( + f'partition for {context_key} has {partition_samples} samples and must be divisible by ' + f'train.mini_batch_size={train.mini_batch_size}') + if train.mini_batch_size % num_generations: + raise ValueError( + f'train.mini_batch_size for {context_key} must preserve complete prompt groups: ' + f'{train.mini_batch_size} % {num_generations} != 0') + if train.mini_batch_size % model_dp: + raise ValueError(f'train.mini_batch_size for {context_key} must be divisible by ' + f'model DP size {model_dp}') + samples_per_rank = train.mini_batch_size // model_dp + if train.micro_batch_size > samples_per_rank: + raise ValueError(f'train.micro_batch_size for {context_key} must not exceed the per-rank train batch ' + f'({samples_per_rank}), got {train.micro_batch_size}') + if train.dynamic_batching: + if train.max_tokens_per_micro_batch is None or train.max_tokens_per_micro_batch <= 0: + raise ValueError( + f'train.max_tokens_per_micro_batch for {context_key} must be positive when ' + 'train.dynamic_batching=true') + if train.packing_algorithm not in ('ffd', 'kk'): + raise ValueError( + f'train.packing_algorithm for {context_key} must be ffd or kk, ' + f'got {train.packing_algorithm!r}') + + +def configure_lora_lr_scheduler( + model: Any, + adapter_name: str, + lora_config: Mapping[str, Any], +) -> None: + scheduler_config = lora_config.get('lr_scheduler') + if scheduler_config is None: + return + scheduler_config = dict(scheduler_config) + scheduler_cls = scheduler_config.pop('cls') + model.set_lr_scheduler( + scheduler_cls, + adapter_name=adapter_name, + **scheduler_config, + ) + + +def resolve_context_learning_rate( + train_config: Mapping[str, Any], + lora_defaults: Mapping[str, Any], +) -> float: + configured = train_config.get('learning_rate', lora_defaults.get('learning_rate')) + if configured is None: + raise ValueError('train.learning_rate or lora.learning_rate must be configured') + learning_rate = float(configured) + if not math.isfinite(learning_rate) or learning_rate <= 0: + raise ValueError(f'train.learning_rate must be a positive finite value, got {configured!r}') + return learning_rate + + +def resolve_context_lora_target_modules( + context_config: Mapping[str, Any], + lora_defaults: Mapping[str, Any], +) -> str | list[str]: + context_lora_config = dict(context_config.get('lora') or {}) + target_modules = context_lora_config.get( + 'target_modules', + lora_defaults.get('target_modules', 'all-linear'), + ) + if isinstance(target_modules, str): + if not target_modules: + raise ValueError('lora.target_modules must not be empty') + return target_modules + if isinstance(target_modules, Sequence) and target_modules: + modules = list(target_modules) + if all(isinstance(module, str) and module for module in modules): + return modules + raise ValueError( + 'lora.target_modules must be a non-empty string or sequence of module names, ' + f'got {target_modules!r}') + + +def resolve_context_loss_config( + context_config: Mapping[str, Any], + loss_defaults: Mapping[str, Any] | None = None, +) -> tuple[str, dict[str, Any]]: + loss_config: dict[str, Any] = { + 'cls': 'GRPOLoss', + 'epsilon': 0.2, + 'normalization': 'sequence_mean', + } + loss_config.update(dict(loss_defaults or {})) + loss_config.update(dict(context_config.get('loss') or {})) + + loss_cls = loss_config.pop('cls', None) + if not isinstance(loss_cls, str) or not loss_cls: + raise ValueError(f'loss.cls must be a non-empty string, got {loss_cls!r}') + return loss_cls, loss_config diff --git a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py new file mode 100644 index 000000000..141fbe821 --- /dev/null +++ b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py @@ -0,0 +1,805 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +import json +import os +import re +import time +import uuid +from concurrent.futures import Future +from copy import copy +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from twinkle import DeviceMesh, get_logger, remote_class, remote_function +from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams, user_data_get +from twinkle.hub import HubOperation +from twinkle.metric import MetricBuffer, MetricRecord +from twinkle.sampler.vllm_sampler import vLLMSampler + +from .data_plane import TQDataPlane +from .metrics import rollout_metrics +from .types import LoraContext, PromptGroup, RolloutOutput, RolloutPolicy +from .utils import resolve_adapter_path, sample_responses_to_rollout_rows + +logger = get_logger() + + +def _dispatch_generation( + worker_count: int, + worker_index: int, + args: tuple[Any, ...], + kwargs: dict[str, Any], + **_dispatch_kwargs, +) -> tuple[tuple[Any, ...], dict[str, Any]]: + """Slice CS inputs while allowing a prompt count smaller than DP size.""" + sliced_args = list(args) + sliced_kwargs = dict(kwargs) + if len(sliced_args) > 1: + inputs = sliced_args[1] + target = ('args', 1) + elif 'inputs' in sliced_kwargs: + inputs = sliced_kwargs['inputs'] + target = ('kwargs', 'inputs') + else: + raise ValueError('submit_generation requires inputs') + + input_list = list(inputs) if isinstance(inputs, (list, tuple)) else [inputs] + size, remainder = divmod(len(input_list), worker_count) + start = worker_index * size + min(worker_index, remainder) + stop = (worker_index + 1) * size + min(worker_index + 1, remainder) + shard = input_list[start:stop] + if target[0] == 'args': + sliced_args[target[1]] = shard + else: + sliced_kwargs[target[1]] = shard + return tuple(sliced_args), sliced_kwargs + + +def _path_component(value: str) -> str: + return re.sub(r'[^A-Za-z0-9._-]+', '_', value).strip('._') or 'unknown' + + +def _compute_rewards( + reward_registry: dict[str, Any], + context: LoraContext, + rollout_rows: list[RolloutOutput], +) -> list[float] | None: + reward_fn = reward_registry.get(context.key) + if reward_fn is None: + return None + return list(reward_fn(rollout_rows, context=context)) + + +def _compute_reward_metrics( + reward_registry: dict[str, Any], + context: LoraContext, + rollout_rows: list[RolloutOutput], + rewards: list[float], +) -> dict[str, Any]: + reward_fn = reward_registry.get(context.key) + metric_payload = getattr(reward_fn, 'metric_payload', None) + if metric_payload is None: + return {} + return dict(metric_payload(rollout_rows, rewards=rewards, context=context)) + + +@dataclass(frozen=True) +class _GeneratedSample: + response: SampleResponse + policies: tuple[RolloutPolicy, ...] + attempts: int + was_aborted: bool + resumed_partial_output: bool + + @property + def initial_policy(self) -> RolloutPolicy: + return self.policies[0] + + @property + def final_policy(self) -> RolloutPolicy: + return self.policies[-1] + + @property + def retry_count(self) -> int: + return self.attempts - 1 + + +@dataclass(frozen=True) +class _PromptGroupRolloutStats: + completion_lengths: tuple[int, ...] + stop_reasons: tuple[str | None, ...] + policy_versions: tuple[int, ...] + + +@remote_class() +class VLLMSamplerTQ(vLLMSampler): + """vLLM sampler that writes async RL rollout results directly to TransferQueue. + + ``sample()`` is intentionally fire-and-forget: it schedules generation work + on the sampler actor's vLLM event loop and returns submission metadata + without waiting for any prompt group to finish. + """ + + def __init__( + self, + model_id: str, + engine_args: dict[str, Any] | None = None, + device_mesh: DeviceMesh | None = None, + *, + context_manager: Any | None = None, + reward_registry: dict[str, Any] | None = None, + rollout_max_retries: int = 2, + rollout_retry_delay_s: float = 0.5, + rollout_output_dir: str | None = None, + rollout_output_include_token_ids: bool = False, + **kwargs, + ): + self.context_manager = context_manager + super().__init__(model_id=model_id, engine_args=engine_args, device_mesh=device_mesh, **kwargs) + self.data_plane = TQDataPlane() + self.reward_registry = dict(reward_registry or {}) + self.rollout_max_retries = int(rollout_max_retries) + self.rollout_retry_delay_s = float(rollout_retry_delay_s) + self.rollout_output_dir = ( + Path(rollout_output_dir).expanduser().resolve() + if rollout_output_dir is not None else None + ) + self.rollout_output_include_token_ids = bool(rollout_output_include_token_ids) + if self.rollout_max_retries < 0: + raise ValueError(f'rollout_max_retries must be non-negative, got {self.rollout_max_retries}') + if self.rollout_retry_delay_s < 0: + raise ValueError(f'rollout_retry_delay_s must be non-negative, got {self.rollout_retry_delay_s}') + self._background_submissions: dict[str, Future] = {} + # Generation submissions are used by the client-orchestrated server + # path. Unlike ``_background_submissions`` above, their results must + # remain available until SamplerManagement collects them and writes + # them to the opaque client DataPlane. + self._generation_submissions: dict[str, Future[list[SampleResponse]]] = {} + self.metric_buffer = MetricBuffer() + self._failure: str | None = None + + def _record_metrics( + self, + group: PromptGroup, + values: dict[str, Any], + *, + status: str = 'completed', + attributes: dict[str, Any] | None = None, + policy_version: int | None = None, + ) -> None: + self.metric_buffer.record(MetricRecord( + stage='rollout', + values=dict(values), + context_key=group.context.key, + partition_id=group.partition_id, + partition_index=group.partition.step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) + + @remote_function(dispatch='all', collect='flatten', lazy_collect=False) + def drain_metric_records(self) -> list[MetricRecord]: + return self.metric_buffer.drain() + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def check_health(self) -> None: + if self._failure is not None: + raise RuntimeError(self._failure) + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def register_reward(self, context_key: str, reward: Any) -> None: + if context_key in self.reward_registry: + raise KeyError(f'reward already registered for {context_key}') + self.reward_registry[context_key] = reward + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def unregister_reward(self, context_key: str) -> None: + self.reward_registry.pop(context_key, None) + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def unload_lora_paths(self, adapter_paths: list[str]) -> None: + # Unloading is keyed by the normalized path stored in VLLMEngine's + # request cache. The checkpoint itself may already have been pruned, + # so unlike loading this must not require the path to still exist. + local_paths = [ + os.path.abspath(os.path.expanduser(str(path))) + for path in adapter_paths + ] + self._submit_in_loop(self.engine.unload_lora_paths(local_paths)).result() + + @remote_function(dispatch='slice_dp', collect='none', lazy_collect=False) + def sample( + self, + groups: list[PromptGroup], + sampling_params: SamplingParams, + allow_partial_rollout: bool = False, + ) -> dict[str, Any]: + """Schedule this DP worker's complete prompt groups and return immediately.""" + if self.context_manager is None: + raise RuntimeError('context_manager is required for native TQ prompt-group sampling') + submission_id = str(uuid.uuid4()) + submitted_at = time.perf_counter() + future = self._submit_in_loop( + self._sample_prompt_groups( + submission_id, + groups, + sampling_params, + bool(allow_partial_rollout), + submitted_at, + )) + self._background_submissions[submission_id] = future + future.add_done_callback(self._on_submission_done(submission_id)) + return { + 'submission_id': submission_id, + 'submitted_prompt_groups': len(groups), + 'submitted_samples': sum(group.num_samples for group in groups), + } + + @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) + def sample_sync( + self, + inputs: Any, + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = '', + adapter_path: str | None = None, + *, + return_encoded: bool = False, + use_base_model: bool = False, + ) -> list[SampleResponse]: + """Run the inherited blocking sampler API for synchronous CS calls.""" + return super().sample( + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_path=adapter_path, + return_encoded=return_encoded, + use_base_model=use_base_model, + ) + + @remote_function(dispatch=_dispatch_generation, collect='none', lazy_collect=False) + def submit_generation( + self, + submission_id: str, + inputs: Any, + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = '', + adapter_path: str | None = None, + *, + use_base_model: bool = False, + ) -> dict[str, Any]: + """Submit a CS sampling shard without blocking the Ray actor. + + The generated responses stay local to this DP worker until + :meth:`collect_generation` consumes them. This gives the HTTP + service the same fast-admission property as the native TQ rollout path + without exposing PromptGroup or BatchMeta to the client. + """ + if submission_id in self._generation_submissions: + raise KeyError(f'generation submission already exists: {submission_id}') + future = self._submit_in_loop( + self._generate_inputs( + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_path=adapter_path, + use_base_model=use_base_model, + )) + self._generation_submissions[submission_id] = future + return {'submission_id': submission_id, 'status': 'running'} + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def get_generation_status(self, submission_id: str) -> dict[str, Any]: + """Return this DP worker's submission state without waiting.""" + future = self._generation_submissions.get(submission_id) + if future is None: + return { + 'submission_id': submission_id, + 'status': 'missing', + 'error': f'unknown generation submission: {submission_id}', + } + if future.cancelled(): + return {'submission_id': submission_id, 'status': 'cancelled'} + if not future.done(): + return {'submission_id': submission_id, 'status': 'running'} + error = future.exception() + if error is not None: + return { + 'submission_id': submission_id, + 'status': 'failed', + 'error': f'{type(error).__name__}: {error}', + } + return {'submission_id': submission_id, 'status': 'completed'} + + @remote_function(dispatch='all', collect='flatten', lazy_collect=False) + def collect_generation(self, submission_id: str) -> list[SampleResponse]: + """Consume completed responses from every DP worker.""" + future = self._generation_submissions.get(submission_id) + if future is None: + raise KeyError(f'unknown generation submission: {submission_id}') + if not future.done(): + raise RuntimeError(f'generation submission is still running: {submission_id}') + try: + return future.result() + finally: + self._generation_submissions.pop(submission_id, None) + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def cancel_generation(self, submission_id: str) -> dict[str, Any]: + """Cancel and forget one generation submission on every DP worker.""" + future = self._generation_submissions.pop(submission_id, None) + if future is None: + return {'submission_id': submission_id, 'status': 'missing'} + was_done = future.done() + cancelled = future.cancel() + if cancelled: + status = 'cancelled' + elif was_done: + status = 'completed' + else: + status = 'cancellation_requested' + return { + 'submission_id': submission_id, + 'status': status, + } + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def cancel_all_generations(self) -> dict[str, int]: + """Cancel all retained CS submissions during replica shutdown.""" + submissions = list(self._generation_submissions.values()) + self._generation_submissions.clear() + cancelled = sum(future.cancel() for future in submissions if not future.done()) + return {'submissions': len(submissions), 'cancelled': cancelled} + + @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) + def evaluate( + self, + inputs: list[dict[str, Any]], + sampling_params: SamplingParams, + adapter_name: str, + adapter_path: str, + ) -> list[SampleResponse]: + """Synchronously evaluate one adapter without writing results to TQ.""" + return super().sample( + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_path=adapter_path, + ) + + def _submit_in_loop(self, coro) -> Future: + return asyncio.run_coroutine_threadsafe(coro, self._async_loop) + + async def _generate_inputs( + self, + inputs: Any, + sampling_params: SamplingParams | dict[str, Any] | None, + *, + adapter_name: str, + adapter_path: str | None, + use_base_model: bool, + ) -> list[SampleResponse]: + """Asynchronous counterpart of ``vLLMSampler.sample`` for CS use.""" + if sampling_params is None: + sampling_params = SamplingParams() + elif isinstance(sampling_params, dict): + sampling_params = SamplingParams.from_dict(sampling_params) + + inputs_list = self._normalize_inputs(inputs) + if not inputs_list: + return [] + + is_trajectory = 'input_ids' not in inputs_list[0] + logprobs_only = False + if sampling_params.max_tokens == 0: + sampling_params = copy(sampling_params) + sampling_params.max_tokens = 1 + logprobs_only = True + + multi_modal_data_list = [self._extract_multi_modal_data(feat) for feat in inputs_list] + if is_trajectory: + if self.template is None: + raise ValueError('Use set_template to add a template when trying to input Trajectory') + encoded_inputs = [ + self.encode_trajectory_for_vllm(trajectory, adapter_name, not logprobs_only) + for trajectory in inputs_list + ] + else: + encoded_inputs = inputs_list + + lora_request = None + if adapter_path is not None: + logger.info(f'Loading LoRA from {adapter_path}') + local_adapter_path = HubOperation.download_model(model_id_or_path=adapter_path) + lora_request = await self.engine._get_or_load_lora(local_adapter_path) + if lora_request is None: + logger.warning(f'Failed to pre-load LoRA from {local_adapter_path}, ' + 'sampling will proceed without LoRA') + + return await asyncio.gather(*( + self._sample_single( + feat, + sampling_params, + lora_request=lora_request, + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + disable_lora=use_base_model, + ) + for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list) + )) + + def _on_submission_done(self, submission_id: str): + + def callback(future: Future) -> None: + self._background_submissions.pop(submission_id, None) + error = future.exception() + if error is not None: + self._failure = f'{type(error).__name__}: {error}' + logger.warning('VLLMSamplerTQ background submission failed: submission=%s error=%s', submission_id, + error) + + return callback + + async def _sample_prompt_groups( + self, + submission_id: str, + groups: list[PromptGroup], + sampling_params: SamplingParams, + allow_partial_rollout: bool, + submitted_at: float, + ) -> None: + results = await asyncio.gather( + *(self._run_prompt_group( + submission_id=submission_id, + group=group, + sampling_params=sampling_params, + allow_partial_rollout=allow_partial_rollout, + ) for group in groups), + return_exceptions=True) + failed_group = next( + ((group, result) for group, result in zip(groups, results) if isinstance(result, Exception)), None) + if failed_group is not None: + group, error = failed_group + self._record_metrics( + group, + {}, + status='failed', + attributes={'scope': 'group', 'group_id': group.group_id, 'error': str(error)}, + ) + raise RuntimeError(f'rollout failed for {group.group_id}: {error}') from error + + rollout_stats = [result for result in results if isinstance(result, _PromptGroupRolloutStats)] + metric_rows = [ + { + 'completion_length': completion_length, + 'stop_reason': stop_reason, + } + for stats in rollout_stats + for completion_length, stop_reason in zip(stats.completion_lengths, stats.stop_reasons) + ] + policy_versions = [version for stats in rollout_stats for version in stats.policy_versions] + first_group = groups[0] + dp_size = self.device_mesh.dp_world_size or 1 + self._record_metrics( + first_group, + { + 'prompt_group_count': len(groups), + **rollout_metrics( + completion_lengths=[row['completion_length'] for row in metric_rows], + stop_reasons=[row['stop_reason'] for row in metric_rows], + rollout_latency_s=time.perf_counter() - submitted_at, + ), + 'policy_version_min': min(policy_versions), + 'policy_version_max': max(policy_versions), + 'sampler_dp_size': dp_size, + }, + attributes={'scope': 'partition' if dp_size == 1 else 'shard'}, + policy_version=max(policy_versions), + ) + + async def _run_prompt_group( + self, + *, + submission_id: str, + group: PromptGroup, + sampling_params: SamplingParams, + allow_partial_rollout: bool, + ) -> _PromptGroupRolloutStats: + """Sample all generations for one group, then write that group once.""" + started = asyncio.get_running_loop().time() + num_generations = group.num_samples + sources = [{ + **group.prompt, 'group_id': group.group_id, + 'generation_idx': generation_idx + } for generation_idx in range(num_generations)] + generated_samples = await self._generate_group_samples( + group.context, sources, sampling_params, allow_partial_rollout=allow_partial_rollout) + rows = [] + for source, generated in zip(sources, generated_samples): + sample_rows = sample_responses_to_rollout_rows([source], [generated.response], + policy_version=generated.final_policy.version) + if len(sample_rows) != 1: + raise ValueError(f'generation {source["generation_idx"]} produced {len(sample_rows)} samples') + row = sample_rows[0] + versions = [policy.version for policy in generated.policies] + row.update({ + 'rollout_policy_version': generated.final_policy.version, + 'rollout_adapter_path': generated.final_policy.adapter_path, + 'rollout_policy_versions': versions, + 'initial_policy_version': generated.initial_policy.version, + 'final_policy_version': generated.final_policy.version, + 'policy_version_span': generated.final_policy.version - generated.initial_policy.version, + }) + rows.append(row) + if len(rows) != num_generations: + raise ValueError(f'group {group.group_id} expected {num_generations} rollout samples, got {len(rows)}') + + rewards = _compute_rewards(self.reward_registry, group.context, rows) + if rewards is None: + raise ValueError(f'no reward function registered for context {group.context.key}') + reward_metrics = _compute_reward_metrics(self.reward_registry, group.context, rows, rewards) + await self.data_plane.complete_rollout_group( + group, + rollout_rows=rows, + rewards=rewards, + submission_id=submission_id, + tag_metrics=reward_metrics, + ) + + rollout_latency_s = asyncio.get_running_loop().time() - started + policy_versions = [policy.version for sample in generated_samples for policy in sample.policies] + if self.rollout_output_dir is not None: + try: + await asyncio.to_thread( + self._write_rollout_group, + submission_id, + group, + generated_samples, + rows, + rewards, + ) + except Exception as error: + logger.warning('Failed to write rollout output for %s: %s', group.group_id, error) + self._record_metrics( + group, + { + **rollout_metrics( + rewards={'reward': rewards}, + completion_lengths=[int(row['completion_length']) for row in rows], + stop_reasons=[row.get('stop_reason') for row in rows], + rollout_latency_s=rollout_latency_s, + ), + 'retry_count': + sum(sample.retry_count for sample in generated_samples), + 'aborted_sample_count': + sum(sample.was_aborted for sample in generated_samples), + 'partial_resumed_sample_count': + sum(sample.resumed_partial_output for sample in generated_samples), + 'policy_version_min': + min(policy_versions), + 'policy_version_max': + max(policy_versions), + **reward_metrics, + }, + attributes={'scope': 'group', 'group_id': group.group_id}, + policy_version=max(policy_versions), + ) + return _PromptGroupRolloutStats( + completion_lengths=tuple(int(row['completion_length']) for row in rows), + stop_reasons=tuple(row.get('stop_reason') for row in rows), + policy_versions=tuple(policy_versions), + ) + + def _write_rollout_group( + self, + submission_id: str, + group: PromptGroup, + generated_samples: list[_GeneratedSample], + rows: list[RolloutOutput], + rewards: list[float], + ) -> None: + policy_version = max(int(row['rollout_policy_version']) for row in rows) + partition_name = _path_component(group.partition_id.rsplit('/', 1)[-1]) + group_name = _path_component(group.group_id.rsplit('/', 1)[-1]) + output_dir = self.rollout_output_dir.joinpath( + _path_component(group.context.tenant_id), + _path_component(group.context.training_run_id), + _path_component(group.context.adapter_name), + f'policy_{policy_version}', + ) + output_dir.mkdir(parents=True, exist_ok=True) + output_path = output_dir / f'{partition_name}-{group_name}.jsonl' + temporary_path = output_path.with_suffix(f'.jsonl.{uuid.uuid4().hex}.tmp') + ground_truth = user_data_get(group.prompt.get('user_data'), 'ground_truth') + + with temporary_path.open('w', encoding='utf-8') as stream: + for generated, row, reward in zip(generated_samples, rows, rewards): + response = generated.response + sequence = response.sequences[0] + prompt_token_ids = list(response.prompt_token_ids or []) + completion_token_ids = list(sequence.tokens) + record = { + 'submission_id': submission_id, + 'context_key': group.context.key, + 'tenant_id': group.context.tenant_id, + 'training_run_id': group.context.training_run_id, + 'adapter_name': group.context.adapter_name, + 'partition_id': group.partition_id, + 'group_id': group.group_id, + 'sample_idx': int(row['generation_idx']), + 'seqlen': len(prompt_token_ids) + len(completion_token_ids), + 'prompt_len': len(prompt_token_ids), + 'completion_len': len(completion_token_ids), + 'head_version': int(row['initial_policy_version']), + 'tail_version': int(row['final_policy_version']), + 'policy_versions': list(row['rollout_policy_versions']), + 'adapter_path': row.get('rollout_adapter_path'), + 'reward': float(reward), + 'ground_truth': ground_truth, + 'stop_reason': row.get('stop_reason'), + 'retry_count': generated.retry_count, + 'was_aborted': generated.was_aborted, + 'resumed_partial_output': generated.resumed_partial_output, + 'prompt': self.template.decode(prompt_token_ids, skip_special_tokens=False), + 'completion': self.template.decode(completion_token_ids, skip_special_tokens=False), + } + if self.rollout_output_include_token_ids: + record.update({ + 'prompt_token_ids': prompt_token_ids, + 'completion_token_ids': completion_token_ids, + 'logprobs': list(row['logprobs']), + }) + stream.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + os.replace(temporary_path, output_path) + + async def _load_lora_for_policy(self, policy: RolloutPolicy) -> Any: + """Load the adapter selected for one group's rollout snapshot.""" + if policy.adapter_path is None: + return None + local_path = await asyncio.to_thread(resolve_adapter_path, policy.adapter_path) + lora_request = await self.engine._get_or_load_lora(local_path) + if lora_request is None: + raise RuntimeError(f'failed to load LoRA adapter from {local_path}') + return lora_request + + async def _generate_group_samples( + self, + context: LoraContext, + sources: list[dict[str, Any]], + sampling_params: SamplingParams, + *, + allow_partial_rollout: bool, + ) -> list[_GeneratedSample]: + logprobs_only = False + if sampling_params.max_tokens == 0: + sampling_params = copy(sampling_params) + sampling_params.max_tokens = 1 + logprobs_only = True + + is_trajectory = 'input_ids' not in sources[0] + multi_modal_data_list = [self._extract_multi_modal_data(source) for source in sources] + if is_trajectory: + template = self.template + assert template is not None, 'Use set_template before sampling trajectories' + encoded_inputs = [ + self.encode_trajectory_for_vllm(source, context.adapter_name, not logprobs_only) for source in sources + ] + else: + encoded_inputs = sources + tasks = [ + self._generate_sample( + context, + feat, + sampling_params, + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + allow_partial_rollout=allow_partial_rollout, + ) for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list) + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + failures = [result for result in results if isinstance(result, Exception)] + if failures: + raise failures[0] + return results + + async def _generate_sample( + self, + context: LoraContext, + original_input: dict[str, Any], + sampling_params: SamplingParams, + *, + multi_modal_data: dict[str, Any] | None, + logprobs_only: bool, + allow_partial_rollout: bool, + ) -> _GeneratedSample: + current_input = original_input + partial_responses: list[SampleResponse] = [] + partial_policies: list[RolloutPolicy] = [] + generated_tokens = 0 + last_error: Exception | None = None + was_aborted = False + resumed_partial_output = False + + for attempt in range(self.rollout_max_retries + 1): + policy = await self.context_manager.acquire_rollout_policy.remote(context) + attempt_params = copy(sampling_params) + if allow_partial_rollout and attempt_params.max_tokens is not None: + attempt_params.max_tokens -= generated_tokens + try: + try: + response = await self._sample_single( + current_input, + attempt_params, + lora_request=await self._load_lora_for_policy(policy), + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + ) + sequence = response.sequences[0] + except Exception as exc: + last_error = exc + else: + if sequence.stop_reason not in {'abort', 'error'}: + if not allow_partial_rollout or not partial_responses: + return _GeneratedSample( + response, (policy, ), attempt + 1, was_aborted, resumed_partial_output) + partial_responses.append(response) + partial_policies.append(policy) + return _GeneratedSample( + self._merge_partial_responses(partial_responses), tuple(partial_policies), attempt + 1, + was_aborted, resumed_partial_output) + + last_error = RuntimeError(f'generation stopped with {sequence.stop_reason}') + was_aborted = was_aborted or sequence.stop_reason == 'abort' + if allow_partial_rollout and sequence.tokens: + resumed_partial_output = True + partial_responses.append(response) + partial_policies.append(policy) + generated_tokens += len(sequence.tokens) + current_input = sequence.new_input_feature + if sampling_params.max_tokens is not None and generated_tokens >= sampling_params.max_tokens: + return _GeneratedSample( + self._merge_partial_responses(partial_responses, stop_reason='length'), + tuple(partial_policies), + attempt + 1, + was_aborted, + resumed_partial_output, + ) + elif not allow_partial_rollout: + current_input = original_input + finally: + await self.context_manager.release_rollout_policy.remote(policy) + + if attempt < self.rollout_max_retries: + await asyncio.sleep(self.rollout_retry_delay_s) + + error_detail = f'{type(last_error).__name__}: {last_error}' + error = RuntimeError( + f'generation failed after {self.rollout_max_retries + 1} attempts; last error: {error_detail}') + raise error from last_error + + def _merge_partial_responses( + self, + responses: list[SampleResponse], + *, + stop_reason: str | None = None, + ) -> SampleResponse: + sequences = [response.sequences[0] for response in responses] + tokens = [token for sequence in sequences for token in sequence.tokens] + logprobs = [logprob for sequence in sequences for logprob in (sequence.logprobs or [])] + final_sequence = sequences[-1] + return SampleResponse( + prompt_token_ids=responses[0].prompt_token_ids, + sequences=[ + SampledSequence( + stop_reason=stop_reason or final_sequence.stop_reason, + tokens=tokens, + logprobs=logprobs, + decoded=self.template.decode(tokens), + new_input_feature=final_sequence.new_input_feature, + routed_experts=final_sequence.routed_experts, + ) + ], + ) diff --git a/src/twinkle_agentic/async_rl/workers.py b/src/twinkle_agentic/async_rl/workers.py new file mode 100644 index 000000000..905b63533 --- /dev/null +++ b/src/twinkle_agentic/async_rl/workers.py @@ -0,0 +1,666 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Long-lived, context-scheduling async-RL workers.""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import time +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence +from typing import Any + +from twinkle.metric import MetricBuffer, MetricRecord +from .context_manager import ContextStatus, LoraContextManager +from .data_plane import TQDataPlane +from .metrics import advantage_signal_metrics, training_policy_metrics +from .scheduler import ContextScheduler, ScheduleCandidate, SchedulerConfig +from .types import LoraContext, PartitionAdmission + + +class _Worker: + """A long-lived Ray service with one privately owned background loop.""" + + def __init__(self): + self._service_task: asyncio.Task[None] | None = None + self._stop_requested = False + self._failure: str | None = None + self.metric_buffer = MetricBuffer() + + async def start(self) -> None: + if self._service_task is not None and not self._service_task.done(): + return + self._stop_requested = False + self._failure = None + self._service_task = asyncio.create_task(self._run_service()) + + async def stop(self) -> None: + self._stop_requested = True + if self._service_task is None or self._service_task.done(): + return + self._service_task.cancel() + try: + await self._service_task + except asyncio.CancelledError: + pass + + async def get_service_state(self) -> dict[str, str | bool | None]: + return { + 'running': self._service_task is not None and not self._service_task.done(), + 'failure': self._failure, + } + + def drain_metric_records(self) -> list[MetricRecord]: + return self.metric_buffer.drain() + + def _record_metric( + self, + stage: str, + *, + context: LoraContext | None = None, + admission: PartitionAdmission | None = None, + partition_id: str | None = None, + values: dict[str, Any] | None = None, + status: str = 'completed', + attributes: dict[str, Any] | None = None, + optimizer_step: int | None = None, + policy_version: int | None = None, + ) -> None: + self.metric_buffer.record(MetricRecord( + stage=stage, + values=dict(values or {}), + context_key=context.key if context is not None else None, + partition_id=admission.partition_id if admission is not None else partition_id, + partition_index=admission.step if admission is not None else None, + optimizer_step=optimizer_step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) + + async def _run_service(self) -> None: + try: + await self._serve() + except asyncio.CancelledError: + return + except Exception as exc: + self._failure = f'{type(exc).__name__}: {exc}' + + async def _serve(self) -> None: + raise NotImplementedError + + +class RolloutWorker(_Worker): + """Admits full prompt batches and submits them to the sampler without waiting.""" + + def __init__(self, + *, + context_manager: LoraContextManager, + data_plane: TQDataPlane, + sampler: Any, + prompt_batches: dict[str, Iterable[Sequence[dict[str, Any]]] + | Callable[[], Iterable[Sequence[dict[str, Any]]]]], + rollout_config: dict[str, dict[str, Any]], + scheduler: SchedulerConfig, + allow_partial_rollout: bool = False, + persistent: bool = False, + idle_delay_s: float = 0.05): + super().__init__() + self.data_plane = data_plane + self.sampler = sampler + self.context_manager = context_manager + self.idle_delay_s = idle_delay_s + self.rollout_config = rollout_config + self.scheduler = ContextScheduler(scheduler) + self.allow_partial_rollout = allow_partial_rollout + self.persistent = persistent + self._prompt_batch_iterators = { + key: iter(value() if callable(value) else value) + for key, value in prompt_batches.items() + } + self._next_batch_tasks: dict[str, asyncio.Task[Sequence[dict[str, Any]] | None]] = {} + self._exhausted: set[str] = set() + self._contexts_changed = asyncio.Event() + + async def register_context( + self, + context: LoraContext, + prompt_batches: Iterable[Sequence[dict[str, Any]]] | Callable[[], Iterable[Sequence[dict[str, Any]]]], + rollout_config: dict[str, Any], + ) -> None: + key = context.key + if key in self._prompt_batch_iterators: + raise KeyError(f'rollout context already exists: {key}') + self._prompt_batch_iterators[key] = iter(prompt_batches() if callable(prompt_batches) else prompt_batches) + self.rollout_config[key] = dict(rollout_config) + self._exhausted.discard(key) + if self._service_task is not None and not self._service_task.done(): + self._start_next_batch(key) + self._contexts_changed.set() + + async def unregister_context(self, context: LoraContext | str) -> None: + key = context if isinstance(context, str) else context.key + task = self._next_batch_tasks.pop(key, None) + if task is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + iterator = self._prompt_batch_iterators.pop(key, None) + self.rollout_config.pop(key, None) + self._exhausted.discard(key) + close = getattr(iterator, 'close', None) + if callable(close): + await asyncio.to_thread(close) + self._contexts_changed.set() + + def _start_next_batch(self, key: str) -> None: + if key in self._prompt_batch_iterators and key not in self._next_batch_tasks: + self._next_batch_tasks[key] = asyncio.create_task( + asyncio.to_thread(next, self._prompt_batch_iterators[key], None)) + + async def stop(self) -> None: + await super().stop() + pending = list(self._next_batch_tasks.values()) + self._next_batch_tasks.clear() + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + async def _serve(self) -> None: + for key in self._prompt_batch_iterators: + self._start_next_batch(key) + while not self._stop_requested: + if not self.persistent and await self.context_manager.is_rollout_admission_closed.remote(): + return + candidates = [] + for key in list(self._prompt_batch_iterators): + if key in self._exhausted: + continue + status = await self.context_manager.context_status.remote(key) + if status is not ContextStatus.ACTIVE: + if status in (ContextStatus.EXHAUSTED, ContextStatus.FINISHED): + self._exhausted.add(key) + continue + config = self.rollout_config[key] + task = self._next_batch_tasks.get(key) + if task is None: + self._start_next_batch(key) + task = self._next_batch_tasks.get(key) + if task is not None and task.done(): + candidates.append(ScheduleCandidate(config['context'])) + candidate = self.scheduler.choose(candidates) + if candidate is None: + if not self.persistent and len(self._exhausted) == len(self._prompt_batch_iterators): + return + self._contexts_changed.clear() + try: + await asyncio.wait_for(self._contexts_changed.wait(), timeout=self.idle_delay_s) + except TimeoutError: + pass + continue + key = candidate.context.key + config = self.rollout_config[key] + batch_task = self._next_batch_tasks[key] + try: + batch = batch_task.result() + except Exception as exc: + self._next_batch_tasks.pop(key) + self._record_metric( + 'rollout', + context=candidate.context, + status='failed', + attributes={'error': f'prompt loading failed: {exc}'}, + ) + raise RuntimeError(f'prompt loading failed for {key}: {exc}') from exc + if batch is None or len(batch) != int(config['batch_size']): + self._next_batch_tasks.pop(key) + self._exhausted.add(key) + await self.context_manager.on_dataset_exhausted.remote(candidate.context) + self.scheduler.on_blocked(candidate) + continue + admission = await self.context_manager.request_rollout_partition.remote( + candidate.context, + target_groups=len(batch), + num_generations=int(config['num_generations']), + ) + if admission is None: + self.scheduler.on_blocked(candidate) + await asyncio.sleep(self.idle_delay_s) + continue + self._next_batch_tasks.pop(key) + submission_started = time.perf_counter() + try: + prepared = await self.data_plane.prepare_rollout_partition( + admission, + list(batch), + config['sampling_params'], + ) + await asyncio.to_thread( + self.sampler.sample, + list(prepared.groups), + prepared.sampling_params, + self.allow_partial_rollout, + ) + except Exception as exc: + self._record_metric( + 'rollout', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError(f'rollout submission failed for {admission.partition_id}: {exc}') from exc + self._start_next_batch(key) + self.scheduler.on_success(candidate) + self._record_metric( + 'rollout', + context=admission.context, + admission=admission, + status='submitted', + values={ + 'prompt_count': admission.target_groups, + 'sample_count': admission.sample_count, + 'num_generations': admission.num_generations, + 'rollout_submission_latency_s': time.perf_counter() - submission_started, + }, + attributes={'scope': 'partition'}, + ) + + +class AdvantageWorker(_Worker): + + def __init__(self, + *, + context_manager: LoraContextManager, + data_plane: TQDataPlane, + advantage_fn: Callable[[Any, PartitionAdmission], tuple[Sequence[float], Sequence[float]]], + scheduler: SchedulerConfig, + persistent: bool = False, + idle_delay_s: float = 0.05): + super().__init__() + self.data_plane = data_plane + self.context_manager = context_manager + self.idle_delay_s = idle_delay_s + self.advantage_fn = advantage_fn + self.scheduler = ContextScheduler(scheduler) + self.persistent = persistent + + async def _serve(self) -> None: + while not self._stop_requested: + if not self.persistent and await self.context_manager.is_run_finished.remote(): + return + admissions = await self.context_manager.list_live_partitions.remote() + blocked: set[str] = set() + progressed = False + for _ in range(len(admissions)): + candidates = [ + ScheduleCandidate(admission.context, admission) for admission in admissions + if admission.partition_id not in blocked + ] + candidate = self.scheduler.choose(candidates) + if candidate is None: + break + admission = candidate.partition + batch = await self.data_plane.claim_advantage_batch(admission, 1) + if batch is None: + blocked.add(admission.partition_id) + self.scheduler.on_blocked(candidate) + continue + started = time.perf_counter() + try: + advantages, returns = self.advantage_fn(batch.data, admission) + await self.data_plane.write_advantages(batch, advantages=advantages, returns=returns) + except Exception as exc: + self._record_metric( + 'advantage', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError(f'advantage failed for {admission.partition_id}: {exc}') from exc + self.scheduler.on_success(candidate) + policy = await self.context_manager.get_rollout_policy.remote(admission.context) + advantage_metrics = advantage_signal_metrics( + batch.data['rewards'], + advantages, + num_generations=admission.num_generations, + ) + advantage_metrics.update({ + 'sample_count': len(advantages), + 'advantage_latency_s': time.perf_counter() - started, + }) + self._record_metric( + 'advantage', + context=admission.context, + admission=admission, + values=advantage_metrics, + policy_version=policy.version, + ) + progressed = True + break + if not progressed: + await asyncio.sleep(self.idle_delay_s) + + +class TrainerWorker(_Worker): + + def __init__(self, + *, + context_manager: LoraContextManager, + data_plane: TQDataPlane, + train_fn: Callable[[Any, PartitionAdmission], dict[str, Any] | None], + save_adapter: Callable[[PartitionAdmission], str], + mini_batch_sizes: dict[str, int], + scheduler: SchedulerConfig, + train_with_config_fn: Callable[[Any, PartitionAdmission, Any], dict[str, Any] | None] | None = None, + train_batch_configs: dict[str, Any] | None = None, + keep_adapter_versions: int = 0, + initial_adapter_paths: dict[str, str] | None = None, + remove_adapter: Callable[[str], None] | None = None, + evaluation_config: dict[str, dict[str, Any]] | None = None, + evaluate_batch: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any], + dict[str, Any]] | None = None, + evaluate_with_reward_fn: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any, + Any], dict[str, Any]] | None = None, + evaluation_rewards: dict[str, Any] | None = None, + persistent: bool = False, + idle_delay_s: float = 0.05): + super().__init__() + self.data_plane = data_plane + self.context_manager = context_manager + self.idle_delay_s = idle_delay_s + self.train_fn = train_fn + self.train_with_config_fn = train_with_config_fn + self.train_batch_configs = dict(train_batch_configs or {}) + self.save_adapter = save_adapter + self.mini_batch_sizes = mini_batch_sizes + self.scheduler = ContextScheduler(scheduler) + self.keep_adapter_versions = max(0, int(keep_adapter_versions)) + self._adapter_history: dict[str, list[str]] = defaultdict(list) + for context_key, path in (initial_adapter_paths or {}).items(): + if path: + self._adapter_history[context_key].append(path) + self.remove_adapter = remove_adapter or _remove_local_adapter + self._adapter_removal_tasks: set[asyncio.Task[None]] = set() + self.evaluation_config = dict(evaluation_config or {}) + self.evaluate_batch = evaluate_batch + self.evaluate_with_reward_fn = evaluate_with_reward_fn + self.evaluation_rewards = dict(evaluation_rewards or {}) + self._evaluation_batches: dict[str, list[Sequence[dict[str, Any]]]] = {} + self._optimizer_steps: dict[str, int] = defaultdict(int) + self.persistent = persistent + + async def register_context( + self, + context: LoraContext, + *, + mini_batch_size: int, + train_batch_config: Any | None = None, + initial_adapter_path: str | None = None, + evaluation_config: dict[str, Any] | None = None, + evaluation_reward: Any | None = None, + ) -> None: + key = context.key + if key in self.mini_batch_sizes: + raise KeyError(f'trainer context already exists: {key}') + self.mini_batch_sizes[key] = int(mini_batch_size) + if train_batch_config is not None: + self.train_batch_configs[key] = train_batch_config + if initial_adapter_path: + self._adapter_history[key].append(initial_adapter_path) + if evaluation_config is not None: + self.evaluation_config[key] = dict(evaluation_config) + if evaluation_reward is not None: + self.evaluation_rewards[key] = evaluation_reward + + async def unregister_context(self, context: LoraContext | str) -> None: + key = context if isinstance(context, str) else context.key + self.mini_batch_sizes.pop(key, None) + self.train_batch_configs.pop(key, None) + self.evaluation_config.pop(key, None) + self.evaluation_rewards.pop(key, None) + self._evaluation_batches.pop(key, None) + self._adapter_history.pop(key, None) + self._optimizer_steps.pop(key, None) + + async def stop(self) -> None: + await super().stop() + pending = tuple(self._adapter_removal_tasks) + if pending: + await asyncio.gather(*pending) + + async def _serve(self) -> None: + while not self._stop_requested: + if not self.persistent and await self.context_manager.is_run_finished.remote(): + return + admissions = await self.context_manager.list_trainable_partitions.remote() + blocked: set[str] = set() + progressed = False + for _ in range(len(admissions)): + candidates = [ + ScheduleCandidate(admission.context, admission) for admission in admissions + if admission.partition_id not in blocked + ] + candidate = self.scheduler.choose(candidates) + if candidate is None: + break + admission = candidate.partition + mini_batch_size = self.mini_batch_sizes[admission.context.key] + batch = await self.data_plane.claim_training_batch( + admission, + mini_batch_size // admission.num_generations, + ) + if batch is None: + if await self.data_plane.is_training_consumed(admission): + try: + await self.context_manager.on_partition_training_started.remote(admission) + await self._finish_partition(admission) + self.scheduler.on_success(candidate) + progressed = True + break + except Exception as exc: + self._record_metric( + 'train', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError( + f'training completion failed for {admission.partition_id}: {exc}') from exc + blocked.add(admission.partition_id) + self.scheduler.on_blocked(candidate) + continue + try: + await self.context_manager.on_partition_training_started.remote(admission) + policy = await self.context_manager.get_rollout_policy.remote(admission.context) + sample_count = len(batch.data['input_ids']) + if sample_count != mini_batch_size: + raise RuntimeError( + f'training claim for {admission.partition_id} returned {sample_count} samples; ' + f'expected mini_batch_size={mini_batch_size}') + started = time.perf_counter() + if self.train_with_config_fn is not None: + config = self.train_batch_configs[admission.context.key] + metrics = dict(self.train_with_config_fn(batch.data, admission, config) or {}) + else: + metrics = dict(self.train_fn(batch.data, admission) or {}) + except Exception as exc: + self._record_metric( + 'train', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError(f'training failed for {admission.partition_id}: {exc}') from exc + self.scheduler.on_success(candidate) + metrics['sample_count'] = sample_count + metrics['reward'] = ( + sum(float(value) for value in batch.data['rewards']) / sample_count) + metrics['train_latency_s'] = time.perf_counter() - started + metrics.update(training_policy_metrics(batch.sample_tags, policy.version)) + context_key = admission.context.key + self._optimizer_steps[context_key] += 1 + optimizer_step = self._optimizer_steps[context_key] + self._record_metric( + 'train', + context=admission.context, + admission=admission, + values=metrics, + optimizer_step=optimizer_step, + policy_version=policy.version, + ) + progressed = True + break + if not progressed: + await asyncio.sleep(self.idle_delay_s) + + async def _finish_partition(self, admission: PartitionAdmission) -> None: + finalize_started = time.perf_counter() + save_started = time.perf_counter() + adapter_path = self.save_adapter(admission) + adapter_save_latency_s = time.perf_counter() - save_started + publish_started = time.perf_counter() + policy = await self.context_manager.on_partition_trained.remote(admission, adapter_path=adapter_path) + policy_publish_latency_s = time.perf_counter() - publish_started + self._record_metric( + 'policy', + context=admission.context, + admission=admission, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + }, + attributes={'operation': 'publish', 'adapter_path': adapter_path}, + optimizer_step=self._optimizer_steps[admission.context.key], + policy_version=policy.version, + ) + await self._evaluate_policy(admission, adapter_path, policy.version) + clear_started = time.perf_counter() + await self.data_plane.clear_partition(admission) + tq_clear_latency_s = time.perf_counter() - clear_started + release_started = time.perf_counter() + await self.context_manager.on_partition_cleared.remote(admission) + partition_release_latency_s = time.perf_counter() - release_started + self._adapter_history[admission.context.key].append(adapter_path) + prune_started = time.perf_counter() + await self._prune_adapter_history(admission.context) + adapter_prune_schedule_latency_s = time.perf_counter() - prune_started + self._record_metric( + 'partition', + context=admission.context, + admission=admission, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + 'tq_clear_latency_s': tq_clear_latency_s, + 'partition_release_latency_s': partition_release_latency_s, + 'adapter_prune_schedule_latency_s': adapter_prune_schedule_latency_s, + 'partition_finalize_latency_s': time.perf_counter() - finalize_started, + }, + optimizer_step=self._optimizer_steps[admission.context.key], + policy_version=policy.version, + ) + + async def _evaluate_policy(self, admission: PartitionAdmission, adapter_path: str, policy_version: int) -> None: + config = self.evaluation_config.get(admission.context.key) + if config is None or (self.evaluate_batch is None and self.evaluate_with_reward_fn is None): + return + interval = int(config['interval']) + if policy_version % interval: + return + + context_key = admission.context.key + if context_key not in self._evaluation_batches: + source = config['prompt_batches'] + self._evaluation_batches[context_key] = list(source() if callable(source) else source) + batches = self._evaluation_batches[context_key] + started = time.perf_counter() + rewards: list[float] = [] + completion_lengths: list[int] = [] + prompt_count = 0 + for batch in batches: + if self.evaluate_with_reward_fn is not None: + result = await asyncio.to_thread( + self.evaluate_with_reward_fn, + batch, + admission, + adapter_path, + policy_version, + config['sampling_params'], + self.evaluation_rewards[context_key], + ) + else: + result = await asyncio.to_thread( + self.evaluate_batch, + batch, + admission, + adapter_path, + policy_version, + config['sampling_params'], + ) + rewards.extend(float(value) for value in result['rewards']) + completion_lengths.extend(int(value) for value in result['completion_lengths']) + prompt_count += len(batch) + if not rewards: + raise ValueError(f'evaluation dataset is empty for {context_key}') + self._record_metric( + 'evaluation', + context=admission.context, + admission=admission, + values={ + 'accuracy': sum(rewards) / len(rewards), + 'sample_count': len(rewards), + 'prompt_count': prompt_count, + 'completion_length': sum(completion_lengths) / len(completion_lengths), + 'eval_latency_s': time.perf_counter() - started, + }, + attributes={'eval_dataset': config['dataset_name']}, + optimizer_step=self._optimizer_steps[context_key], + policy_version=policy_version, + ) + + async def _prune_adapter_history(self, context: LoraContext) -> None: + protected = set(await self.context_manager.adapter_paths_to_keep.remote()) + context_key = context.key + history = self._adapter_history[context_key] + retained_history = set(history[-self.keep_adapter_versions:]) if self.keep_adapter_versions else set() + retained = protected | retained_history + stale = [path for path in history if path not in retained] + self._adapter_history[context_key] = [path for path in history if path in retained] + for path in stale: + task = asyncio.create_task(self._remove_adapter(context, path)) + self._adapter_removal_tasks.add(task) + task.add_done_callback(self._adapter_removal_tasks.discard) + + async def _remove_adapter(self, context: LoraContext, path: str) -> None: + started = time.perf_counter() + try: + await asyncio.to_thread(self.remove_adapter, path) + except OSError as exc: + self._record_metric( + 'policy', + context=context, + status='failed', + values={ + 'adapter_prune_latency_s': time.perf_counter() - started, + }, + attributes={'operation': 'adapter_prune', 'adapter_path': path, 'error': str(exc)}, + ) + return + self._record_metric( + 'policy', + context=context, + values={ + 'adapter_prune_latency_s': time.perf_counter() - started, + }, + attributes={'operation': 'adapter_prune', 'adapter_path': path}, + ) + + +def _remove_local_adapter(path: str) -> None: + if os.path.isdir(path): + shutil.rmtree(path) diff --git a/src/twinkle_client/__init__.py b/src/twinkle_client/__init__.py index a5105d497..a3d281af1 100644 --- a/src/twinkle_client/__init__.py +++ b/src/twinkle_client/__init__.py @@ -72,4 +72,7 @@ def init_twinkle_client( ) -__all__ = ['init_tinker_client', 'init_twinkle_client'] +from .remote_task import RemoteTask, RemoteTaskError +from .data_plane import DataPlaneClient + +__all__ = ['DataPlaneClient', 'RemoteTask', 'RemoteTaskError', 'init_tinker_client', 'init_twinkle_client'] diff --git a/src/twinkle_client/async_rl/__init__.py b/src/twinkle_client/async_rl/__init__.py new file mode 100644 index 000000000..f26758240 --- /dev/null +++ b/src/twinkle_client/async_rl/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .workers import Worker, WorkerPipeline + +__all__ = ['Worker', 'WorkerPipeline'] diff --git a/src/twinkle_client/async_rl/workers.py b/src/twinkle_client/async_rl/workers.py new file mode 100644 index 000000000..b98a9fb1f --- /dev/null +++ b/src/twinkle_client/async_rl/workers.py @@ -0,0 +1,61 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Small client-side lifecycle primitives for composable async-RL workers. + +Workers remain concrete long-running roles. This module deliberately does not +introduce an algorithm graph or a data-dependency DSL; role implementations +coordinate through ordinary asyncio queues and server-side DataRefs. +""" +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Sequence + + +class Worker(ABC): + """One long-running client-side computation role.""" + + def __init__(self, name: str) -> None: + if not name: + raise ValueError('worker name must not be empty') + self.name = name + + @abstractmethod + async def run(self) -> None: + """Run until this role has drained its input or fails.""" + + +class WorkerPipeline: + """Run a concrete set of worker roles and propagate failures as one unit.""" + + def __init__(self, workers: Sequence[Worker]) -> None: + self.workers = tuple(workers) + if not self.workers: + raise ValueError('at least one worker is required') + names = [worker.name for worker in self.workers] + if len(names) != len(set(names)): + raise ValueError(f'worker names must be unique, got {names}') + + async def run(self) -> None: + tasks = { + asyncio.create_task(worker.run(), name=worker.name): worker + for worker in self.workers + } + try: + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) + failure = next( + (task.exception() for task in done if not task.cancelled() and task.exception() is not None), + None, + ) + if failure is not None: + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + raise failure + await asyncio.gather(*pending) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + diff --git a/src/twinkle_client/common/json_utils.py b/src/twinkle_client/common/json_utils.py new file mode 100644 index 000000000..51c039c15 --- /dev/null +++ b/src/twinkle_client/common/json_utils.py @@ -0,0 +1,33 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Lightweight JSON conversion helpers shared by component clients and servers.""" +from __future__ import annotations + +from collections.abc import Mapping +from numbers import Number +from typing import Any + +from pydantic import BaseModel + + +_PRIMITIVE_TYPES = (str, Number, bool, bytes, type(None)) + + +def json_safe(obj: Any) -> Any: + """Recursively convert models, tensors, and arrays into JSON-compatible values. + + This module intentionally does not import :mod:`twinkle`: component protocol + types are imported while the top-level package is still being initialized. + """ + if isinstance(obj, BaseModel): + return json_safe(obj.model_dump()) + if isinstance(obj, Mapping): + return {key: json_safe(value) for key, value in obj.items()} + if isinstance(obj, (list, tuple, set, frozenset)): + return [json_safe(value) for value in obj] + tolist = getattr(obj, 'tolist', None) + if callable(tolist) and not isinstance(obj, _PRIMITIVE_TYPES): + try: + return json_safe(tolist()) + except Exception: + pass + return obj diff --git a/src/twinkle_client/common/serialize.py b/src/twinkle_client/common/serialize.py index 42a27beb3..0f0a49096 100644 --- a/src/twinkle_client/common/serialize.py +++ b/src/twinkle_client/common/serialize.py @@ -7,6 +7,7 @@ from typing import Any, Mapping from twinkle.dataset import DatasetMeta +from .json_utils import json_safe supported_types = { DatasetMeta, diff --git a/src/twinkle_client/data_plane.py b/src/twinkle_client/data_plane.py new file mode 100644 index 000000000..281df8315 --- /dev/null +++ b/src/twinkle_client/data_plane.py @@ -0,0 +1,129 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Client for the server-side TransferQueue DataRef service.""" +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any, TypeVar + +from twinkle_client.common.json_utils import json_safe +from twinkle_client.http import get_base_url, http_post +from twinkle_client.types.component import DataRef, DataRowsResponse + + +_T = TypeVar('_T') + + +async def _call_in_thread(func: Callable[..., _T], /, *args: Any, **kwargs: Any) -> _T: + """Run one synchronous DataPlane operation without blocking the event loop.""" + return await asyncio.to_thread(func, *args, **kwargs) + + +class DataPlaneClient: + + def __init__(self, server_url: str | None = None): + self.server_url = (server_url or f'{get_base_url()}/data-plane').rstrip('/') + + def put( + self, + rows: list[dict[str, Any]], + *, + kind: str = 'data', + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + response = http_post( + f'{self.server_url}/twinkle/put', + json_data={'rows': json_safe(rows), 'kind': kind, 'tags': json_safe(tags)}, + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def aput( + self, + rows: list[dict[str, Any]], + *, + kind: str = 'data', + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + """Asynchronously store rows while preserving :meth:`put` semantics.""" + if tags is None: + return await _call_in_thread(self.put, rows, kind=kind) + return await _call_in_thread(self.put, rows, kind=kind, tags=tags) + + def get(self, ref: DataRef, *, fields: list[str] | None = None) -> list[dict[str, Any]]: + response = http_post( + f'{self.server_url}/twinkle/get', + json_data={'ref': ref.model_dump(), 'fields': fields}, + ) + response.raise_for_status() + return DataRowsResponse(**response.json()).rows + + def get_batch( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> DataRowsResponse: + response = http_post( + f'{self.server_url}/twinkle/get', + json_data={'ref': ref.model_dump(), 'fields': fields, 'include_tags': True}, + ) + response.raise_for_status() + return DataRowsResponse(**response.json()) + + async def aget(self, ref: DataRef, *, fields: list[str] | None = None) -> list[dict[str, Any]]: + """Asynchronously fetch rows while preserving :meth:`get` semantics.""" + if fields is None: + return await _call_in_thread(self.get, ref) + return await _call_in_thread(self.get, ref, fields=fields) + + async def aget_batch( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> DataRowsResponse: + if fields is None: + return await _call_in_thread(self.get_batch, ref) + return await _call_in_thread(self.get_batch, ref, fields=fields) + + def append( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + response = http_post( + f'{self.server_url}/twinkle/append', + json_data={ + 'ref': ref.model_dump(), + 'rows': json_safe(rows), + 'tags': json_safe(tags), + }, + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def aappend( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + """Asynchronously append rows while preserving :meth:`append` semantics.""" + if tags is None: + return await _call_in_thread(self.append, ref, rows) + return await _call_in_thread(self.append, ref, rows, tags=tags) + + def release(self, ref: DataRef) -> None: + response = http_post( + f'{self.server_url}/twinkle/release', + json_data={'ref': ref.model_dump()}, + ) + response.raise_for_status() + + async def arelease(self, ref: DataRef) -> None: + """Asynchronously release a reference while preserving :meth:`release` semantics.""" + await _call_in_thread(self.release, ref) diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index c628c353b..c620269ff 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -2,6 +2,9 @@ from pathlib import Path import time from twinkle_client.http import http_get, http_post +from twinkle_client.remote_task import RemoteTask +from twinkle_client.common.json_utils import json_safe +from twinkle_client.types.component import ComponentTaskRef, DataRef from twinkle_client.types.model import ( CalculateLossResponse, CalculateMetricResponse, @@ -26,6 +29,8 @@ def __init__(self, model_id: str, **kwargs): """Initialize model client.""" from twinkle_client.http import get_base_url self.server_url = get_base_url() + from twinkle_client.data_plane import DataPlaneClient + self.data_plane = DataPlaneClient(kwargs.pop('data_plane_url', None)) if '://' in model_id: model_id = model_id.split('://')[1] @@ -37,6 +42,85 @@ def __init__(self, model_id: str, **kwargs): ) response.raise_for_status() + def submit_forward( + self, + inputs: Any | DataRef, + *, + forward_only: bool = False, + **forward_kwargs, + ) -> RemoteTask: + """Submit the Model component's forward primitive directly.""" + body = { + 'adapter_name': self.adapter_name or '', + 'forward_only': forward_only, + 'forward_kwargs': forward_kwargs, + } + body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( + inputs.model_dump() if isinstance(inputs, DataRef) else json_safe(inputs)) + response = http_post( + url=f'{self.server_url}/submit_forward', + json_data=body, + ) + response.raise_for_status() + return RemoteTask(ComponentTaskRef(**response.json())) + + def submit_forward_only( + self, + inputs: Any | DataRef, + **forward_kwargs, + ) -> RemoteTask: + return self.submit_forward( + inputs, + forward_only=True, + **forward_kwargs, + ) + + def submit_forward_backward( + self, + inputs: Any | DataRef, + **kwargs, + ) -> RemoteTask: + """Submit forward/backward without prescribing the surrounding train loop.""" + body = {'adapter_name': self.adapter_name or '', 'kwargs': json_safe(kwargs)} + body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( + inputs.model_dump() if isinstance(inputs, DataRef) else json_safe(inputs)) + response = http_post( + url=f'{self.server_url}/submit_forward_backward', + json_data=body, + ) + response.raise_for_status() + return RemoteTask(ComponentTaskRef(**response.json())) + + def submit_clip_grad_and_step( + self, + max_grad_norm: float = 1.0, + norm_type: int = 2, + **kwargs, + ) -> RemoteTask: + response = http_post( + url=f'{self.server_url}/submit_clip_grad_and_step', + json_data={ + 'adapter_name': self.adapter_name or '', + 'max_grad_norm': max_grad_norm, + 'norm_type': norm_type, + 'kwargs': kwargs, + }, + ) + response.raise_for_status() + return RemoteTask(ComponentTaskRef(**response.json())) + + def submit_save(self, name: str, *, save_optimizer: bool = False) -> RemoteTask: + response = http_post( + url=f'{self.server_url}/submit_save', + json_data={ + 'adapter_name': self.adapter_name or '', + 'name': name, + 'save_optimizer': save_optimizer, + }, + ) + response.raise_for_status() + return RemoteTask(ComponentTaskRef(**response.json())) + def add_adapter_to_model(self, adapter_name: str, config: Dict[str, Any], **kwargs) -> None: """Add a new adapter to the model.""" save_dir = kwargs.get('save_dir') @@ -49,6 +133,17 @@ def add_adapter_to_model(self, adapter_name: str, config: Dict[str, Any], **kwar response.raise_for_status() self.adapter_name = adapter_name + def remove_adapter(self, adapter_name: str | None = None) -> None: + """Release one client-owned adapter from the training component.""" + name = adapter_name or self.adapter_name + response = http_post( + url=f'{self.server_url}/remove_adapter', + json_data={'adapter_name': name}, + ) + response.raise_for_status() + if name == self.adapter_name: + self.adapter_name = None + def forward(self, inputs: Any, **kwargs) -> ForwardResponse: """Execute forward pass on the model.""" response = http_post( diff --git a/src/twinkle_client/remote_task.py b/src/twinkle_client/remote_task.py new file mode 100644 index 000000000..d3ff36cec --- /dev/null +++ b/src/twinkle_client/remote_task.py @@ -0,0 +1,94 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Future handle returned by an individual Model or Sampler component.""" +from __future__ import annotations + +import asyncio +import time +from typing import Any + +from twinkle_client.http import get_base_url, http_post +from twinkle_client.http.http_utils import _build_headers +from twinkle_client.types.component import ComponentTaskRef + + +class RemoteTaskError(RuntimeError): + pass + + +class RemoteTask: + """A wrapper over the server's existing component future registry.""" + + def __init__(self, task: ComponentTaskRef | str): + self.request_id = task if isinstance(task, str) else task.request_id + self.model_id = None if isinstance(task, str) else task.model_id + self._url = f'{get_base_url()}/retrieve_future' + + def poll(self, timeout: float | None = None) -> Any | None: + request_kwargs = {} if timeout is None else {'timeout': timeout} + response = http_post( + self._url, + json_data={'request_id': self.request_id}, + **request_kwargs, + ) + response.raise_for_status() + return self._resolve_payload(response.json()) + + @staticmethod + def _resolve_payload(payload: Any) -> Any | None: + if isinstance(payload, dict) and payload.get('type') == 'try_again': + return None + if isinstance(payload, dict) and 'error' in payload: + raise RemoteTaskError(payload['error']) + return payload + + def result(self, timeout: float | None = None) -> Any: + import requests + + deadline = None if timeout is None else time.monotonic() + timeout + while True: + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError(f'component task {self.request_id} did not finish within {timeout}s') + try: + result = self.poll(timeout=remaining) + except requests.Timeout as exc: + raise TimeoutError( + f'component task {self.request_id} did not finish within {timeout}s') from exc + if result is not None: + return result + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError(f'component task {self.request_id} did not finish within {timeout}s') + + async def aresult(self, timeout: float | None = None) -> Any: + import httpx + deadline = None if timeout is None else time.monotonic() + timeout + poll_interval = 0.05 + async with httpx.AsyncClient(timeout=600) as client: + while True: + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError( + f'component task {self.request_id} did not finish within {timeout}s') + try: + response = await client.post( + self._url, + headers=_build_headers(), + json={'request_id': self.request_id}, + timeout=remaining if remaining is not None else 600, + ) + except httpx.TimeoutException as exc: + raise TimeoutError( + f'component task {self.request_id} did not finish within {timeout}s') from exc + response.raise_for_status() + result = self._resolve_payload(response.json()) + if result is not None: + return result + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError(f'component task {self.request_id} did not finish within {timeout}s') + remaining = None if deadline is None else deadline - time.monotonic() + await asyncio.sleep( + poll_interval if remaining is None else min(poll_interval, max(remaining, 0.0))) + poll_interval = min(poll_interval * 1.5, 1.0) + + def __await__(self): + return self.aresult().__await__() diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index 55c5800b7..54c6e450b 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -4,8 +4,7 @@ This module hosts :class:`ClientMultiTurnRollout`, a hand-maintained multi-turn rollout orchestrator whose algorithmic structure mirrors ``twinkle_agentic.rollout.multi_turn.MultiTurnRollout`` but issues sampling over -HTTP via ``twinkle_client.sampler.vLLMSampler.sample()`` instead of holding a -Ray actor handle. +HTTP via the client sampler instead of holding a Ray actor handle. Design notes: * It deliberately does NOT subclass ``MultiTurnRollout``. That class is @@ -14,9 +13,12 @@ not match the HTTP-client semantics here. * The ``tool_manager`` type is reused directly from ``twinkle_agentic.tools.tool_manager.ToolManager`` (imported, not copied). + * ``arun`` uses the Sampler component's asynchronous future API; ``__call__`` + remains a synchronous compatibility wrapper. * Bridge-token stitching is reused from ``twinkle_agentic.rollout.bridge.extend_with_bridge``. """ +import asyncio import dataclasses from typing import Any, Dict, List, Optional @@ -33,8 +35,7 @@ class ClientMultiTurnRollout: Mirrors the per-trajectory state machine of ``twinkle_agentic.rollout.multi_turn.MultiTurnRollout`` but issues sampling - via ``vLLMSampler.sample()`` (an HTTP call to ``/twinkle/sample``) rather - than a Ray actor call. + through the HTTP Sampler component rather than a Ray actor call. """ def __init__( @@ -67,11 +68,24 @@ def __init__( f'got {self.sampling_params.num_samples}') def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - """Run the batched multi-turn rollout state machine over HTTP. + """Synchronous wrapper for :meth:`arun`. + + Async client orchestrators should call ``await rollout.arun(...)`` so + independent rollout pipelines can overlap with Model training. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(self.arun(trajectories, **kwargs)) + raise RuntimeError('ClientMultiTurnRollout.__call__ cannot run inside an event loop; ' + 'use `await rollout.arun(...)` instead') + + async def arun(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: + """Asynchronously run the batched multi-turn state machine over HTTP. Structurally mirrors ``MultiTurnRollout.__call__`` but issues each - round's sampling through ``vLLMSampler.sample()`` (an HTTP POST to - ``/twinkle/sample``) rather than a Ray actor call. Every round makes a + round's sampling through ``vLLMSampler.asample()`` rather than a Ray + actor call. Every round makes a SINGLE batched HTTP call for all currently-live trajectories so the sampler can run them in parallel; finished trajectories are parked and excluded from later batches. @@ -138,7 +152,19 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] # upstream concern (retry/backoff) and failures are never # silently swallowed. batch_pifs = [pifs[i] for i in active] - resps = self.sampler.sample(batch_pifs, sampling_params=sampling_params) + sample_kwargs = {'sampling_params': sampling_params} + for name in ('adapter_name', 'adapter_uri'): + if name in kwargs: + sample_kwargs[name] = kwargs[name] + async_sample = getattr(self.sampler, 'asample', None) + if callable(async_sample): + resps = await async_sample(batch_pifs, **sample_kwargs) + else: + resps = await asyncio.to_thread( + self.sampler.sample, + batch_pifs, + **sample_kwargs, + ) pending_bridges: List[tuple] = [] # (global_idx, tool_messages) for local_idx, global_idx in enumerate(active): diff --git a/src/twinkle_client/sampler/vllm_sampler.py b/src/twinkle_client/sampler/vllm_sampler.py index 0d553bb32..c0c865ba8 100644 --- a/src/twinkle_client/sampler/vllm_sampler.py +++ b/src/twinkle_client/sampler/vllm_sampler.py @@ -3,6 +3,9 @@ from twinkle_client.types.sampler import AddAdapterResponse, SampleResponseModel, SetTemplateResponse from peft import PeftConfig from twinkle.data_format import Trajectory, InputFeature +from twinkle_client.common.json_utils import json_safe +from twinkle_client.remote_task import RemoteTask +from twinkle_client.types.component import ComponentTaskRef, DataRef # Intentionally does NOT subclass ``twinkle.sampler.base.Sampler``: importing @@ -17,17 +20,7 @@ def _json_safe(obj: Any) -> Any: duck-typing (``.tolist()``) so this stays free of a hard torch/numpy import, honouring the CPU-only client contract noted above. """ - if isinstance(obj, dict): - return {k: _json_safe(v) for k, v in obj.items()} - if isinstance(obj, (list, tuple)): - return [_json_safe(x) for x in obj] - tolist = getattr(obj, 'tolist', None) - if callable(tolist) and not isinstance(obj, (str, bytes, int, float, bool)): - try: - return _json_safe(tolist()) - except Exception: - return obj - return obj + return json_safe(obj) class vLLMSampler: @@ -41,6 +34,8 @@ def __init__(self, model_id: str, **kwargs): """Create the sampler instance on server.""" from twinkle_client.http import get_base_url self.server_url = get_base_url() + from twinkle_client.data_plane import DataPlaneClient + self.data_plane = DataPlaneClient(kwargs.pop('data_plane_url', None)) self.adapter_name = None if '://' in model_id: @@ -85,6 +80,8 @@ def sample( Returns: SampleResponseModel with 'sequences' list, each containing tokens, logprobs, stop_reason. """ + sampling_params = dict(sampling_params or {}) + sampling_params['num_samples'] = num_samples json_data = { 'inputs': _json_safe(inputs), 'sampling_params': sampling_params, @@ -101,6 +98,98 @@ def sample( response.raise_for_status() return [SampleResponseModel(**r) for r in response.json()['samples']] + def submit_sample( + self, + inputs: Union[List[Trajectory], List[InputFeature], DataRef], + sampling_params: Optional[Dict[str, Any]] = None, + *, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + policy_version: int | None = None, + group_ids: list[str] | None = None, + num_samples: int = 1, + ) -> RemoteTask: + """Submit directly to the Sampler component and return immediately.""" + body = { + 'sampling_params': sampling_params, + 'adapter_name': adapter_name, + 'adapter_uri': adapter_uri, + 'policy_version': policy_version, + 'group_ids': group_ids, + 'num_samples': num_samples, + } + body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( + inputs.model_dump() if isinstance(inputs, DataRef) else _json_safe(inputs)) + response = http_post( + url=f'{self.server_url}/submit_sample', + json_data=json_safe(body), + ) + response.raise_for_status() + return RemoteTask(ComponentTaskRef(**response.json())) + + async def asample( + self, + inputs: Union[List[Trajectory], List[InputFeature], DataRef], + sampling_params: Optional[Dict[str, Any]] = None, + *, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + policy_version: int | None = None, + group_ids: list[str] | None = None, + num_samples: int = 1, + ) -> List[SampleResponseModel]: + """Submit sampling and asynchronously await it without blocking the event loop.""" + import asyncio + task = await asyncio.to_thread( + self.submit_sample, + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_uri=adapter_uri, + policy_version=policy_version, + group_ids=group_ids, + num_samples=num_samples, + ) + result = await task.aresult() + if isinstance(result, dict) and result.get('output_ref'): + output_ref = DataRef(**result['output_ref']) + try: + batch = await self.data_plane.aget_batch(output_ref) + finally: + await self.data_plane.arelease(output_ref) + if batch.tags and all('prompt_index' in tag for tag in batch.tags): + grouped: dict[int, list[tuple[int, dict[str, Any]]]] = {} + for row, tag in zip(batch.rows, batch.tags): + grouped.setdefault(int(tag['prompt_index']), []).append( + (int(tag.get('generation_idx', 0)), row)) + samples = [] + for prompt_index in sorted(grouped): + generation_rows = [row for _, row in sorted(grouped[prompt_index])] + first = generation_rows[0] + samples.append({ + 'sequences': [{ + key: value + for key, value in row.items() + if key not in ('prompt_logprobs', 'topk_prompt_logprobs') + } for row in generation_rows], + 'prompt_logprobs': first.get('prompt_logprobs'), + 'topk_prompt_logprobs': first.get('topk_prompt_logprobs'), + }) + else: + # Compatibility with a server that still stores one nested row per prompt. + samples = batch.rows + else: + samples = result.get('samples', []) if isinstance(result, dict) else [] + return [SampleResponseModel(**item) for item in samples] + + def unload_adapter_paths(self, adapter_paths: list[str]) -> None: + """Evict policy snapshots that are no longer referenced by this client.""" + response = http_post( + url=f'{self.server_url}/unload_adapter_paths', + json_data={'adapter_paths': adapter_paths}, + ) + response.raise_for_status() + def set_template(self, template_cls: str, adapter_name: str = '', **kwargs) -> SetTemplateResponse: """Set the template for encoding trajectories.""" response = http_post( diff --git a/src/twinkle_client/types/__init__.py b/src/twinkle_client/types/__init__.py index 49673b0e9..46cc02235 100644 --- a/src/twinkle_client/types/__init__.py +++ b/src/twinkle_client/types/__init__.py @@ -92,3 +92,18 @@ ) from .checkpoint import ResolvedLoadPath +from .component import ( + AsyncClipGradAndStepRequest, + AsyncForwardBackwardRequest, + AsyncForwardRequest, + AsyncSampleRequest, + AsyncSaveRequest, + ComponentTaskRef, + DataAppendRequest, + DataGetRequest, + DataPutRequest, + DataRef, + DataReleaseRequest, + DataRowsResponse, + UnloadAdapterPathsRequest, +) diff --git a/src/twinkle_client/types/component.py b/src/twinkle_client/types/component.py new file mode 100644 index 000000000..66b153e24 --- /dev/null +++ b/src/twinkle_client/types/component.py @@ -0,0 +1,115 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Protocol types for directly orchestrating asynchronous server components.""" +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field, model_validator + + +class ComponentTaskRef(BaseModel): + request_id: str + model_id: str | None = None + + +class DataRef(BaseModel): + """Opaque reference to rows stored in the server-side TransferQueue.""" + + ref_id: str + size: int + fields: list[str] = Field(default_factory=list) + kind: str = 'data' + num_tokens: int = 0 + + +class DataPutRequest(BaseModel): + rows: list[dict[str, Any]] + kind: str = 'data' + tags: list[dict[str, Any]] | None = None + + +class DataGetRequest(BaseModel): + ref: DataRef + fields: list[str] | None = None + include_tags: bool = False + + +class DataAppendRequest(BaseModel): + ref: DataRef + rows: list[dict[str, Any]] + tags: list[dict[str, Any]] | None = None + + +class DataReleaseRequest(BaseModel): + ref: DataRef + + +class DataRowsResponse(BaseModel): + rows: list[dict[str, Any]] + tags: list[dict[str, Any]] = Field(default_factory=list) + + +class AsyncSampleRequest(BaseModel): + inputs: Any = None + input_ref: DataRef | None = None + sampling_params: dict[str, Any] | None = None + adapter_name: str = '' + adapter_uri: str | None = None + policy_version: int | None = None + group_ids: list[str] | None = None + num_samples: int = 1 + + @model_validator(mode='after') + def validate_input(self) -> 'AsyncSampleRequest': + if (self.inputs is None) == (self.input_ref is None): + raise ValueError('exactly one of inputs and input_ref must be provided') + if self.group_ids is not None and self.inputs is not None: + size = len(self.inputs) if isinstance(self.inputs, list) else 1 + if len(self.group_ids) != size: + raise ValueError('group_ids must contain one value per sampler input') + return self + + +class AsyncForwardRequest(BaseModel): + inputs: Any = None + input_ref: DataRef | None = None + adapter_name: str = '' + forward_only: bool = False + forward_kwargs: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode='after') + def validate_input(self) -> 'AsyncForwardRequest': + if (self.inputs is None) == (self.input_ref is None): + raise ValueError('exactly one of inputs and input_ref must be provided') + return self + + +class AsyncForwardBackwardRequest(BaseModel): + inputs: Any = None + input_ref: DataRef | None = None + adapter_name: str = '' + kwargs: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode='after') + def validate_input(self) -> 'AsyncForwardBackwardRequest': + if (self.inputs is None) == (self.input_ref is None): + raise ValueError('exactly one of inputs and input_ref must be provided') + return self + + +class AsyncClipGradAndStepRequest(BaseModel): + adapter_name: str = '' + max_grad_norm: float = 1.0 + norm_type: int = 2 + kwargs: dict[str, Any] = Field(default_factory=dict) + + +class AsyncSaveRequest(BaseModel): + adapter_name: str = '' + name: str + save_optimizer: bool = False + is_sampler: bool = False + + +class UnloadAdapterPathsRequest(BaseModel): + adapter_paths: list[str] diff --git a/tests/loss/test_grpo_gkd.py b/tests/loss/test_grpo_gkd.py index 541b867a0..fe6d49341 100644 --- a/tests/loss/test_grpo_gkd.py +++ b/tests/loss/test_grpo_gkd.py @@ -82,6 +82,29 @@ def test_grpo_list_advantages(self): result = loss_fn(inputs, outputs, old_logps=old_logps, advantages=adv_list) assert torch.isfinite(result['loss']) + def test_grpo_weights_sequences_equally(self): + labels = torch.tensor([ + [1, -100, -100], + [1, 1, 1], + ]) + logps = torch.zeros_like(labels, dtype=torch.float32) + inputs = {'labels': labels} + outputs = {'logps': logps} + advantages = torch.tensor([[1.0], [3.0]]) + + result = GRPOLoss()( + inputs, + outputs, + old_logps=logps, + advantages=advantages, + ) + + assert result['loss'].item() == pytest.approx(-2.0) + + def test_grpo_does_not_accept_normalization(self): + with pytest.raises(TypeError, match='normalization'): + GRPOLoss(normalization='token_mean') + def test_grpo_entropy_coef(self): loss_fn = GRPOLoss(epsilon=0.2, entropy_coef=0.01) inputs, outputs, old_logps, _, advantages = _make_rl_batch() diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py new file mode 100644 index 000000000..645fdbfdb --- /dev/null +++ b/tests/model/test_micro_batch.py @@ -0,0 +1,284 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from types import SimpleNamespace + +import pytest + +from twinkle.loss import CrossEntropyLoss, GRPOLoss +from twinkle.loss.base import Loss +from twinkle.model.micro_batch import MicroBatchConfig, collect_micro_batch_outputs, plan_micro_batches +from twinkle.model.transformers.transformers import TransformersModel +from twinkle.processor import InputProcessor + + +@pytest.mark.parametrize('packing_algorithm', ['ffd', 'kk']) +def test_dynamic_micro_batch_plan_preserves_samples_and_limits_cost(packing_algorithm): + lengths = [10, 9, 8, 7, 4, 3, 2, 1] + inputs = [{'input_ids': list(range(length))} for length in lengths] + config = MicroBatchConfig( + micro_batch_size=3, + dynamic_batching=True, + max_tokens_per_micro_batch=18, + packing_algorithm=packing_algorithm, + ) + + plan = plan_micro_batches(inputs, config, padding_free=False) + + assert sorted(index for batch in plan for index in batch) == list(range(len(inputs))) + for batch in plan: + assert len(batch) <= config.micro_batch_size + assert max(lengths[index] for index in batch) * len(batch) <= 18 + + +def test_padding_free_dynamic_batching_uses_unpadded_token_cost(): + lengths = [10, 8, 6, 4] + inputs = [{'input_ids': list(range(length))} for length in lengths] + config = MicroBatchConfig( + micro_batch_size=3, + dynamic_batching=True, + max_tokens_per_micro_batch=18, + ) + + plan = plan_micro_batches(inputs, config, padding_free=True) + + assert sorted(index for batch in plan for index in batch) == list(range(len(inputs))) + assert all(sum(lengths[index] for index in batch) <= 18 for batch in plan) + + +def test_token_mean_loss_weight_uses_valid_label_count(): + inputs = [ + {'labels': [1, 2, -100, -100]}, + {'labels': [3, 4, 5, -100]}, + {'labels': [6, -100, -100, -100]}, + ] + + loss = CrossEntropyLoss(reduction='mean') + first_weight = loss.micro_batch_scale(inputs, [0, 2]) + second_weight = loss.micro_batch_scale(inputs, [1]) + + assert first_weight == .5 + assert second_weight == .5 + + +def test_sample_mean_and_token_sum_micro_batch_scales(): + inputs = [ + {'labels': [1, -100]}, + {'labels': [2, 3]}, + {'labels': [4, -100]}, + {'labels': [5, 6]}, + ] + + assert GRPOLoss().micro_batch_scale(inputs, [0]) == .25 + assert CrossEntropyLoss(reduction='sum').micro_batch_scale(inputs, [0]) == 1.0 + + +def test_loss_without_micro_batch_semantics_fails_when_split(): + with pytest.raises(NotImplementedError, match='does not support micro-batching'): + Loss().micro_batch_scale([{}, {}], [0]) + + +def test_transformers_forward_backward_keeps_original_default_path(): + class ModelHarness: + def __init__(self): + self.calls = [] + + def forward(self, *, inputs, **_kwargs): + self.calls.append(('forward', inputs)) + return {} + + def calculate_loss(self, **_kwargs): + self.calls.append(('loss', None)) + return 2.0 + + def backward(self, **_kwargs): + self.calls.append(('backward', None)) + + model = ModelHarness() + outputs = TransformersModel.forward_backward.__wrapped__( + model, + inputs=[{'input_ids': [1, 2]}], + ) + + assert [name for name, _ in model.calls] == ['forward', 'loss', 'backward'] + assert outputs['loss'] == 2.0 + + +def test_transformers_forward_backward_executes_real_micro_batches(): + class OptimizerConfig: + def __init__(self): + self.processor = InputProcessor(padding_free=False) + self.template = None + self.train_status = SimpleNamespace(loss_value=None, num_tokens=0.0) + self.loss_instance = SimpleNamespace( + micro_batch_scale=lambda inputs, indices: len(indices) / len(inputs), + ) + self._dp_group = None + + def _ensure_dp_group(self): + pass + + class ModelHarness: + _build_micro_batch_plan = TransformersModel._build_micro_batch_plan + _forward_backward_micro_batch = TransformersModel._forward_backward_micro_batch + _forward_backward_micro_batches = TransformersModel._forward_backward_micro_batches + + def __init__(self): + self.optimizer_group = {'adapter': OptimizerConfig()} + self.forward_batches = [] + self.backward_calls = [] + + def _get_default_group(self): + return 'adapter' + + @staticmethod + def _not_encoded(_inputs): + return False + + def forward(self, *, inputs, **_kwargs): + self.forward_batches.append([item['sample_id'] for item in inputs]) + return {} + + def calculate_loss(self, **_kwargs): + self.optimizer_group['adapter'].train_status.loss_value = 2.0 + self.optimizer_group['adapter'].train_status.num_tokens += 1.0 + return 2.0 + + def backward(self, *, sync_gradients, **_kwargs): + loss = self.optimizer_group['adapter'].train_status.loss_value + self.backward_calls.append((sync_gradients, loss)) + self.optimizer_group['adapter'].train_status.loss_value = None + + model = ModelHarness() + inputs = [ + { + 'sample_id': index, + 'input_ids': list(range(index + 1)), + } + for index in range(4) + ] + + outputs = TransformersModel.forward_backward.__wrapped__( + model, + inputs=inputs, + adapter_name='adapter', + micro_batch_size=2, + sync_gradients=True, + ) + + assert model.forward_batches == [[0, 1], [2, 3]] + assert model.backward_calls == [(False, 1.0), (True, 1.0)] + assert model.optimizer_group['adapter'].train_status.num_tokens == 1.0 + assert outputs['micro_batch_count'] == 2 + assert outputs['micro_batch_samples_mean'] == 2.0 + + +def test_fixed_micro_batch_plan_can_match_a_larger_dp_micro_batch_count(): + inputs = [{'input_ids': [index]} for index in range(4)] + + plan = plan_micro_batches( + inputs, + MicroBatchConfig(micro_batch_size=2), + padding_free=False, + min_micro_batches=3, + ) + + assert plan == [[0, 1], [2], [3]] + + +def test_dp_micro_batch_planning_propagates_remote_rank_error(monkeypatch): + from twinkle.model.transformers import transformers as module + + class OptimizerConfig: + processor = InputProcessor(padding_free=False) + _dp_group = object() + + @staticmethod + def _ensure_dp_group(): + pass + + def all_gather(states, local_state, *, group): + assert group is OptimizerConfig._dp_group + states[:] = [ + local_state, + { + 'micro_batch_count': None, + 'input_count': 2, + 'error': 'ValueError: sequence length 20 exceeds the limit', + }, + ] + + monkeypatch.setattr(module.dist, 'get_world_size', lambda _group: 2) + monkeypatch.setattr(module.dist, 'all_gather_object', all_gather) + + with pytest.raises(RuntimeError, match='rank 1.*sequence length 20'): + TransformersModel._build_micro_batch_plan( + object(), + [{'input_ids': [1]}, {'input_ids': [2]}], + MicroBatchConfig(micro_batch_size=1), + OptimizerConfig(), + ) + + +def test_dp_micro_batch_planning_rejects_common_count_on_all_ranks(monkeypatch): + from twinkle.model.transformers import transformers as module + + class OptimizerConfig: + processor = InputProcessor(padding_free=False) + _dp_group = object() + + @staticmethod + def _ensure_dp_group(): + pass + + def all_gather(states, local_state, *, group): + assert group is OptimizerConfig._dp_group + states[:] = [ + local_state, + { + 'micro_batch_count': 3, + 'input_count': 3, + 'error': None, + }, + ] + + monkeypatch.setattr(module.dist, 'get_world_size', lambda _group: 2) + monkeypatch.setattr(module.dist, 'all_gather_object', all_gather) + + with pytest.raises(ValueError, match='same number of non-empty micro-batches'): + TransformersModel._build_micro_batch_plan( + object(), + [{'input_ids': [1]}, {'input_ids': [2]}], + MicroBatchConfig(micro_batch_size=1), + OptimizerConfig(), + ) + + +def test_dp_collection_reduces_micro_batch_statistics(): + class Mesh: + @staticmethod + def get_collect_ranks(): + return [0, 1] + + result = collect_micro_batch_outputs( + [ + { + 'micro_batch_count': 3, + 'micro_batch_samples_mean': 2.0, + 'micro_batch_tokens_mean': 100.0, + 'micro_batch_tokens_max': 150, + }, + { + 'micro_batch_count': 3, + 'micro_batch_samples_mean': 3.0, + 'micro_batch_tokens_mean': 120.0, + 'micro_batch_tokens_max': 180, + }, + ], + Mesh(), + ) + + assert result == { + 'micro_batch_count': 3, + 'micro_batch_samples_mean': 2.5, + 'micro_batch_tokens_mean': 110.0, + 'micro_batch_tokens_max': 180, + } diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 93644cbef..bebd63c9f 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -235,6 +235,32 @@ def test_launcher_accepts_typed_config() -> None: assert launcher.config is cfg +def test_data_plane_application_uses_its_own_strict_args_schema() -> None: + app = ApplicationSpec.model_validate({ + 'name': 'data-plane', + 'route_prefix': '/api/v1/data-plane', + 'import_path': 'data_plane', + 'args': { + 'config': { + 'backend': { + 'SimpleStorage': { + 'num_data_storage_units': 2, + }, + }, + }, + }, + }) + assert app.import_path == 'data_plane' + assert app.args.config['backend']['SimpleStorage']['num_data_storage_units'] == 2 + + with pytest.raises(ValidationError): + ApplicationSpec.model_validate({ + 'name': 'data-plane', + 'import_path': 'data_plane', + 'args': {'unknown': True}, + }) + + def test_cookbook_examples_load() -> None: """Migrated cookbook configs all parse with the new field names.""" here = Path(__file__).resolve().parents[3] diff --git a/tests/server/contract/client_api_baseline.json b/tests/server/contract/client_api_baseline.json index dcafe9470..8051c9557 100644 --- a/tests/server/contract/client_api_baseline.json +++ b/tests/server/contract/client_api_baseline.json @@ -1,4 +1,48 @@ { + "data_plane": { + "paths": { + "/twinkle/append": { + "POST": { + "operationId": "append_twinkle_append_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/get": { + "POST": { + "operationId": "get_twinkle_get_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/put": { + "POST": { + "operationId": "put_twinkle_put_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/release": { + "POST": { + "operationId": "release_twinkle_release_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + } + } + }, "gateway": { "paths": { "/asample": { @@ -793,6 +837,16 @@ ] } }, + "/twinkle/remove_adapter": { + "POST": { + "operationId": "remove_adapter_twinkle_remove_adapter_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/resume_from_checkpoint": { "POST": { "operationId": "resume_from_checkpoint_twinkle_resume_from_checkpoint_post", @@ -873,6 +927,46 @@ ] } }, + "/twinkle/submit_clip_grad_and_step": { + "POST": { + "operationId": "submit_clip_grad_and_step_twinkle_submit_clip_grad_and_step_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/submit_forward": { + "POST": { + "operationId": "submit_forward_twinkle_submit_forward_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/submit_forward_backward": { + "POST": { + "operationId": "submit_forward_backward_twinkle_submit_forward_backward_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/submit_save": { + "POST": { + "operationId": "submit_save_twinkle_submit_save_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/upload_status/{request_id}": { "GET": { "operationId": "upload_status_twinkle_upload_status__request_id__get", @@ -1009,6 +1103,26 @@ "422" ] } + }, + "/twinkle/submit_sample": { + "POST": { + "operationId": "submit_sample_twinkle_submit_sample_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/unload_adapter_paths": { + "POST": { + "operationId": "unload_adapter_paths_twinkle_unload_adapter_paths_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } } } } diff --git a/tests/server/data_plane/test_proxy.py b/tests/server/data_plane/test_proxy.py new file mode 100644 index 000000000..1339baaaa --- /dev/null +++ b/tests/server/data_plane/test_proxy.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import pytest + +from twinkle.server.data_plane.proxy import DataPlaneProxy +from twinkle_client.http.headers import H_AUTH, H_AUTH_TWINKLE, H_REQUEST_ID +from twinkle_client.types import DataRef + + +class _Response: + + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self): + return self.payload + + +class _Client: + + def __init__(self): + self.calls = [] + + async def post(self, url, **kwargs): + self.calls.append((url, kwargs)) + if url.endswith('/get'): + return _Response({'rows': [{'value': 1}]}) + return _Response({ + 'ref_id': 'output', + 'size': 1, + 'fields': ['value'], + 'kind': 'model-output', + }) + + +@pytest.mark.asyncio +async def test_proxy_routes_by_data_ref_without_tenant_identity() -> None: + proxy = DataPlaneProxy.__new__(DataPlaneProxy) + proxy.base_url = 'http://data-plane' + proxy.client = _Client() + ref = DataRef(ref_id='input', size=1, fields=['value']) + + assert await proxy.get(ref) == [{'value': 1}] + output = await proxy.put([{'value': 2}], kind='model-output') + + assert output.ref_id == 'output' + get_headers = proxy.client.calls[0][1]['headers'] + put_headers = proxy.client.calls[1][1]['headers'] + assert get_headers[H_REQUEST_ID] == 'data-ref-input' + assert put_headers[H_REQUEST_ID] == 'data-put-model-output' + assert get_headers[H_AUTH] == get_headers[H_AUTH_TWINKLE] == '' + assert put_headers[H_AUTH] == put_headers[H_AUTH_TWINKLE] == '' diff --git a/tests/server/data_plane/test_store.py b/tests/server/data_plane/test_store.py new file mode 100644 index 000000000..ea89ffc75 --- /dev/null +++ b/tests/server/data_plane/test_store.py @@ -0,0 +1,110 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import pytest + +from twinkle.server.data_plane.store import TQDataRefStore, _partition + + +@pytest.mark.asyncio +async def test_data_ref_round_trip_append_release_and_ref_isolation(monkeypatch) -> None: + import transfer_queue as tq + + records = {} + + async def batch_put(*, keys, partition_id, fields, tags=None): + storage_key = (partition_id, tuple(keys)) + if storage_key in records: + records[storage_key].update(fields) + else: + records[storage_key] = fields.clone() + if tags is not None: + current = tag_records.setdefault(storage_key, [{} for _ in tags]) + for existing, update in zip(current, tags): + existing.update(update) + + async def batch_get(*, keys, partition_id, select_fields): + data = records[(partition_id, tuple(keys))] + return data.select(*select_fields) + + async def clear(*, keys, partition_id): + records.pop((partition_id, tuple(keys))) + + async def kv_list(partition_id): + result = {} + for (stored_partition, keys), tags in tag_records.items(): + if stored_partition == partition_id: + result[stored_partition] = dict(zip(keys, tags)) + return result + + monkeypatch.setattr(tq, 'async_kv_batch_put', batch_put) + monkeypatch.setattr(tq, 'async_kv_batch_get', batch_get) + monkeypatch.setattr(tq, 'async_kv_clear', clear) + monkeypatch.setattr(tq, 'async_kv_list', kv_list) + + # Bypass tq.init(): this test exercises only the DataRef mapping layer. + store = TQDataRefStore.__new__(TQDataRefStore) + tag_records = {} + rows = [ + {'input_ids': [1, 2], 'answer': 'a'}, + {'input_ids': [3], 'answer': 'b'}, + ] + tags = [{'group_id': 'g0', 'generation_idx': 0}, {'group_id': 'g0', 'generation_idx': 1}] + ref = await store.put(rows, kind='train', tags=tags) + + assert ref.size == 2 + assert ref.fields == ['input_ids', 'answer'] + assert ref.num_tokens == 3 + assert await store.get(ref) == rows + assert await store.get_tags(ref) == tags + + ref = await store.append( + ref, + [{'reward': 1.0}, {'reward': -1.0}], + tags=[{'status': 'ready'}, {'status': 'ready'}], + ) + assert ref.fields == ['input_ids', 'answer', 'reward'] + assert ref.num_tokens == 3 + assert await store.get( + ref, + fields=['answer', 'reward'], + ) == [ + {'answer': 'a', 'reward': 1.0}, + {'answer': 'b', 'reward': -1.0}, + ] + assert await store.get_tags(ref) == [ + {'group_id': 'g0', 'generation_idx': 0, 'status': 'ready'}, + {'group_id': 'g0', 'generation_idx': 1, 'status': 'ready'}, + ] + + ref = await store.append( + ref, + [{'input_ids': [4, 5, 6]}, {'input_ids': [7, 8]}], + ) + assert ref.num_tokens == 5 + + with pytest.raises(KeyError): + await store.get(ref.model_copy(update={'ref_id': 'another-ref'})) + + await store.release(ref) + assert records == {} + + +@pytest.mark.asyncio +async def test_append_rejects_row_count_mismatch() -> None: + from twinkle_client.types import DataRef + + store = TQDataRefStore.__new__(TQDataRefStore) + ref = DataRef(ref_id='r', size=2, fields=['x']) + with pytest.raises(ValueError, match='row count'): + await store.append(ref, [{'reward': 1.0}]) + + +def test_partition_is_stable_and_scoped_by_data_ref() -> None: + from twinkle_client.types import DataRef + + first = DataRef(ref_id='a', size=1, fields=['x']) + same = DataRef(ref_id='a', size=99, fields=['other']) + other = DataRef(ref_id='b', size=1, fields=['x']) + assert _partition(first) == _partition(same) + assert _partition(first) != _partition(other) diff --git a/tests/server/gateway/test_future_retrieval.py b/tests/server/gateway/test_future_retrieval.py new file mode 100644 index 000000000..1152b7db1 --- /dev/null +++ b/tests/server/gateway/test_future_retrieval.py @@ -0,0 +1,85 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient +from unittest.mock import AsyncMock, MagicMock + +tinker = pytest.importorskip('tinker') + +from twinkle.server.gateway.tinker_handlers import _register_tinker_routes + + +def _make_client(get_future: AsyncMock) -> TestClient: + management = MagicMock() + management.state.get_future = get_future + management.supported_models = [] + + app = FastAPI() + + @app.middleware('http') + async def _set_request_state(request: Request, call_next): + authorization = request.headers.get('Authorization', '') + request.state.token = authorization.removeprefix('Bearer ') + request.state.session_id = request.headers.get('X-Twinkle-Session-Id', '') + return await call_next(request) + + _register_tinker_routes(app, lambda: management) + return TestClient(app) + + +def test_retrieve_future_uses_request_id_as_capability() -> None: + get_future = AsyncMock(return_value={'status': 'completed', 'result': {'value': 7}}) + client = _make_client(get_future) + + response = client.post( + '/retrieve_future', + json={'request_id': 'req-1'}, + headers={ + 'Authorization': 'Bearer tenant-token', + 'X-Twinkle-Session-Id': 'session-1', + }, + ) + + assert response.status_code == 200 + assert response.json() == {'value': 7} + get_future.assert_awaited_once_with('req-1') + + +def test_not_yet_visible_future_returns_try_again(monkeypatch) -> None: + monkeypatch.setenv('TWINKLE_LONG_POLL_TIMEOUT', '0') + get_future = AsyncMock(return_value=None) + client = _make_client(get_future) + + response = client.post( + '/retrieve_future', + json={'request_id': 'unknown'}, + headers={ + 'Authorization': 'Bearer tenant-token', + 'X-Twinkle-Session-Id': 'session-1', + }, + ) + + assert response.status_code == 200 + assert response.json() == {'type': 'try_again'} + + +def test_initial_cross_replica_miss_is_long_polled(monkeypatch) -> None: + monkeypatch.setenv('TWINKLE_LONG_POLL_TIMEOUT', '1') + monkeypatch.setenv('TWINKLE_POLL_INTERVAL', '0') + get_future = AsyncMock(side_effect=[ + None, + {'status': 'completed', 'result': {'value': 9}}, + ]) + client = _make_client(get_future) + + response = client.post( + '/retrieve_future', + json={'request_id': 'req-replicated'}, + headers={'Authorization': 'Bearer tenant-token'}, + ) + + assert response.status_code == 200 + assert response.json() == {'value': 9} + assert get_future.await_count == 2 diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py new file mode 100644 index 000000000..f77e7fecd --- /dev/null +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from starlette.requests import Request + +import twinkle_client.types as types +from twinkle.server.model.twinkle_handlers import ( + _model_result_rows, + _register_twinkle_routes, +) + + +def test_model_result_rows_keeps_one_output_row_per_sample() -> None: + assert _model_result_rows( + {'logps': [[-1.0], [-2.0]], 'loss': 0.25}, + batch_size=2, + ) == [ + {'logps': [-1.0], 'loss': 0.25}, + {'logps': [-2.0], 'loss': 0.25}, + ] + + +class _SchedulingManagement: + + def __init__(self): + self.data_world_size = 2 + self.scheduled = [] + self.model_calls = [] + self.model = self + + async def _on_request_start(self, _request): + return 'token' + + def assert_resource_exists(self, _adapter_name): + return None + + def forward_backward(self, *, inputs, adapter_name, **kwargs): + self.model_calls.append((inputs, adapter_name, kwargs)) + return {'loss': 1.0} + + async def schedule_task(self, task, **kwargs): + self.scheduled.append(kwargs) + await task() + return {'request_id': 'request-1', 'model_id': kwargs.get('model_id')} + + +@pytest.mark.asyncio +async def test_async_forward_backward_schedules_without_algorithm_metadata() -> None: + management = _SchedulingManagement() + app = FastAPI() + _register_twinkle_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/submit_forward_backward') + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.AsyncForwardBackwardRequest( + adapter_name='adapter', + inputs=[{'input_ids': [index]} for index in range(8)], + kwargs={ + 'old_logps': [[-0.1]] * 8, + 'advantages': [1.0] * 8, + }, + ) + + await route.endpoint(request, body, management) + + assert management.scheduled[-1]['batch_size'] == 8 + assert management.scheduled[-1]['data_world_size'] == 2 + assert 'batch_size_multiple' not in management.scheduled[-1] + _, adapter_name, forwarded_kwargs = management.model_calls[-1] + assert adapter_name == 'session-adapter' + assert forwarded_kwargs == body.kwargs diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index afa72ddac..e4564c5df 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -17,7 +17,7 @@ from twinkle.data_format import InputFeature, SamplingParams from twinkle.server.exceptions import ConfigError -from twinkle.server.sampler.app import SAMPLER_SELECTOR +from twinkle.server.sampler.app import SAMPLER_SELECTOR, _construct_sampler_backend from twinkle.server.sampler.backends.mock_sampler import MockSampler _SAMPLER_TYPES = tuple(SAMPLER_SELECTOR.builders) @@ -114,6 +114,41 @@ def test_mock_dispatch_returns_mock_sampler() -> None: assert isinstance(s, MockSampler) +def test_data_plane_vllm_uses_fire_and_forget_sampler(monkeypatch) -> None: + from twinkle_agentic.async_rl import vllm_sampler_tq as module + + captured = {} + + def construct(**kwargs): + captured.update(kwargs) + return 'vllm-tq' + + monkeypatch.setattr(module, 'VLLMSamplerTQ', construct) + + sampler = _construct_sampler_backend( + 'vllm', + {'model_id': 'local-model'}, + 'http://data-plane', + ) + + assert sampler == 'vllm-tq' + assert captured == {'model_id': 'local-model', 'context_manager': None} + + +def test_vllm_without_data_plane_keeps_standard_backend(monkeypatch) -> None: + calls = [] + monkeypatch.setattr( + SAMPLER_SELECTOR, + 'construct', + lambda sampler_type, kwargs: calls.append((sampler_type, kwargs)) or 'standard-vllm', + ) + + sampler = _construct_sampler_backend('vllm', {'model_id': 'local-model'}, None) + + assert sampler == 'standard-vllm' + assert calls == [('vllm', {'model_id': 'local-model'})] + + @settings(max_examples=100) @given(bad=st.text(min_size=1, max_size=10).filter(lambda s: s not in _SAMPLER_TYPES)) def test_invalid_sampler_type_raises_config_error(bad: str) -> None: diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py new file mode 100644 index 000000000..40f59c3e3 --- /dev/null +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import pytest + +import twinkle_client.types as types +from twinkle.server.sampler.twinkle_handlers import _sample_models_to_rows + + +def _response(tokens: list[int]) -> types.SampleResponseModel: + return types.SampleResponseModel( + sequences=[ + types.SampledSequenceModel( + stop_reason='stop', + tokens=[token], + logprobs=[[(token, -0.1)]], + new_input_feature={'input_ids': [token], 'labels': [token]}, + ) + for token in tokens + ], + prompt_logprobs=[-0.2], + ) + + +def test_async_sampler_flattens_generations_to_tagged_tq_rows() -> None: + rows, tags = _sample_models_to_rows( + [_response([10, 11]), _response([20, 21])], + group_ids=['group-a', 'group-b'], + policy_version=7, + adapter_uri='twinkle://policy-7', + ) + + assert [row['tokens'] for row in rows] == [[10], [11], [20], [21]] + assert [(tag['group_id'], tag['generation_idx']) for tag in tags] == [ + ('group-a', 0), + ('group-a', 1), + ('group-b', 0), + ('group-b', 1), + ] + assert {tag['rollout_policy_version'] for tag in tags} == {7} + assert {tag['rollout_adapter_uri'] for tag in tags} == {'twinkle://policy-7'} + + +def test_async_sampler_rejects_group_id_count_mismatch() -> None: + with pytest.raises(ValueError, match='group_ids contains 1 values for 2'): + _sample_models_to_rows( + [_response([10]), _response([20])], + group_ids=['only-one'], + policy_version=0, + adapter_uri=None, + ) diff --git a/tests/server/state/test_managers.py b/tests/server/state/test_managers.py index 1b9016be1..eea300b18 100644 --- a/tests/server/state/test_managers.py +++ b/tests/server/state/test_managers.py @@ -432,7 +432,6 @@ async def test_store_status_queue_state(self, manager): assert result.queue_state == 'paused_rate_limit' assert result.queue_state_reason == 'Rate limit hit' - # ============================================================ # Cascade Cleanup Consistency (merged from test_cleanup_cascade_consistency) # ============================================================ diff --git a/tests/server/test_app_builders_characterization.py b/tests/server/test_app_builders_characterization.py index 31f9e7267..d2a7f5523 100644 --- a/tests/server/test_app_builders_characterization.py +++ b/tests/server/test_app_builders_characterization.py @@ -1,8 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Characterization tests for the four App_Builders. +"""Characterization tests for the component App_Builders. These freeze the externally observable behavior of ``build_gateway_app``, -``build_model_app``, ``build_sampler_app`` and ``build_processor_app`` BEFORE +``build_model_app``, ``build_sampler_app``, ``build_processor_app``, and the +DataPlane builder. The original four were captured before the Shared_App_Scaffold is extracted, so the extraction can be shown to be behavior-preserving. For each builder they assert, as fixed expectations: @@ -15,7 +16,7 @@ stack is wired in front of the routes; 3. the **bound deployment** identified by the name passed to ``serve.deployment`` (``GatewayServer``, ``ModelManagement``, ``SamplerManagement``, - ``ProcessorManagement``). + ``ProcessorManagement``, ``DataPlaneManagement``). They MUST NOT assert internal object identity or internal middleware-stack structure. They are built by capturing the FastAPI app the builder @@ -187,6 +188,29 @@ def test_processor_builder_characterization(monkeypatch) -> None: _assert_middleware_lifo_order(res.app, expect_cleanup=False) +def test_data_plane_builder_characterization(monkeypatch) -> None: + from twinkle.server.data_plane import app as data_plane_mod + + res = _capture_builder( + monkeypatch, + data_plane_mod.build_data_plane_app, + deploy_options={}, + ) + assert res.deployment_name == 'DataPlaneManagement' + assert _route_set(res.app) == _baseline_route_set('data_plane') + _assert_auth_middleware_effect(res.app) + _assert_middleware_lifo_order(res.app, expect_cleanup=False) + + +def test_data_plane_builder_rejects_multiple_replicas() -> None: + from twinkle.server.data_plane import app as data_plane_mod + + with pytest.raises(ValueError, match='exactly one replica'): + data_plane_mod.build_data_plane_app( + deploy_options={'num_replicas': 2}, + ) + + # ----- black-box middleware-effect oracle ---------------------------------- # diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index c579a3b9a..93f6c5e5d 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -1,3 +1,5 @@ +import asyncio + import pytest from twinkle.server.utils.task_queue.config import TaskQueueConfig @@ -46,6 +48,8 @@ async def test_preflight_rejects_batch_without_per_dp_multiple(): assert result == {'request_id': 'req1', 'model_id': 'model1'} _, kwargs = queue.state.records[-1] assert kwargs['result']['category'] == 'User' + assert 'token' not in kwargs + assert 'session_id' not in kwargs assert 'Batch size 2 must be divisible by 4' in kwargs['result']['error'] @@ -65,3 +69,20 @@ async def test_preflight_accepts_batch_with_per_dp_multiple(): assert result is None assert queue.state.records == [] + + +@pytest.mark.asyncio +async def test_background_task_tracks_status_without_owner_or_preflight(): + queue = _DummyQueue() + + async def work(): + return {'ok': True} + + await queue.schedule_background_task( + work, + model_id='model1', + ) + await asyncio.sleep(0) + + assert queue.state.records[0][0][1] == 'running' + assert all('token' not in kwargs and 'session_id' not in kwargs for _, kwargs in queue.state.records) diff --git a/tests/twinkle_agentic/test_async_rl_metrics.py b/tests/twinkle_agentic/test_async_rl_metrics.py new file mode 100644 index 000000000..34f1b60e7 --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_metrics.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +import json +import sys +import threading +import time +import types + +import pytest + +from twinkle.metric import ( + CompletionRewardMetric, + MetricBuffer, + MetricRecord, + create_metrics_reporter, +) +from twinkle_agentic.async_rl.metrics import advantage_signal_metrics, rollout_metrics +from twinkle.metric.reporting import MetricsReporter, _QueuedBackend + + +def test_advantage_signal_metrics_report_zero_and_nonzero_groups(): + metrics = advantage_signal_metrics( + rewards=[1.0, 1.0, 0.0, 1.0], + advantages=[0.0, 0.0, -1.0, 1.0], + num_generations=2, + ) + assert metrics['group_count'] == 2 + assert metrics['group_reward_std_mean'] == pytest.approx(0.25) + assert metrics['zero_advantage_group_ratio'] == pytest.approx(0.5) + assert metrics['positive_advantage_ratio'] == pytest.approx(0.25) + + +def test_advantage_signal_metrics_reject_incomplete_groups(): + with pytest.raises(ValueError, match='complete groups'): + advantage_signal_metrics([1.0, 0.0, 1.0], [1.0, -1.0, 0.0], num_generations=2) + + +def test_rollout_metrics_include_rewards_tokens_and_truncation(): + metrics = rollout_metrics( + rewards={'accuracy': [1.0, 0.0]}, + completion_lengths=[3, 7], + stop_reasons=['stop', 'length'], + rollout_latency_s=2.0, + ) + assert metrics == { + 'sample_count': 2, + 'completion_length_mean': 5.0, + 'completion_length_p95': 7, + 'completion_length_max': 7, + 'completion_truncated_count': 1, + 'completion_truncated_ratio': 0.5, + 'output_tokens': 10, + 'rollout_latency_s': 2.0, + 'output_tokens_per_s': 5.0, + 'accuracy_reward': 0.5, + 'accuracy_reward_std': pytest.approx(2**-0.5), + } + + +def test_completion_reward_metric_preserves_model_metric_contract(): + metric = CompletionRewardMetric() + metric.accumulate( + rewards={'accuracy': [1.0, 0.0]}, + completion_lengths=[3, 7], + generate_time=2.0, + weight_sync_time=0.25, + ) + result = metric.calculate() + assert result == { + 'profiling/Time taken: move_model_to_sampler': 0.25, + 'profiling/Time taken: generate': 2.0, + 'train/accuracy_reward': 0.5, + 'train/accuracy_reward_std': pytest.approx(2**-0.5), + 'train/completion_length': 5.0, + } + + +def test_metric_buffer_drain_is_atomic_and_destructive(): + buffer = MetricBuffer() + + def produce(start): + for value in range(start, start + 50): + buffer.record(MetricRecord(stage='train', values={'loss': value})) + + threads = [threading.Thread(target=produce, args=(index * 50,)) for index in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + records = buffer.drain() + assert len(records) == 200 + assert buffer.drain() == [] + + +def test_reporter_writes_new_jsonl_schema_and_summary(tmp_path): + path = tmp_path / 'metrics.jsonl' + summary_path = tmp_path / 'summary.json' + reporter = create_metrics_reporter({ + 'queue_capacity': 100, + 'jsonl': { + 'path': path, + 'summary_path': summary_path, + 'batch_size': 2, + 'flush_interval_s': 60, + }, + }, run_id='test') + reporter.record(MetricRecord( + stage='rollout', + status='completed', + context_key='tenant/run/adapter', + partition_id='tenant/run/adapter/train_5', + partition_index=5, + policy_version=0, + values={'sample_count': 4, 'reward': 0.5}, + attributes={'scope': 'group', 'group_id': 'group_0'}, + )) + reporter.record(MetricRecord( + stage='rollout', + status='completed', + context_key='tenant/run/adapter', + partition_id='tenant/run/adapter/train_5', + partition_index=5, + policy_version=0, + values={'sample_count': 4, 'reward': 0.25}, + attributes={'scope': 'partition'}, + )) + reporter.record(MetricRecord( + stage='train', + context_key='tenant/run/adapter', + partition_id='tenant/run/adapter/train_5', + partition_index=5, + optimizer_step=1, + policy_version=0, + values={'sample_count': 4, 'loss': '0.25'}, + )) + reporter.record(MetricRecord( + stage='partition', + context_key='tenant/run/adapter', + partition_index=5, + policy_version=1, + values={}, + )) + reporter.record(MetricRecord(stage='run', values={'trained_partitions': 1, 'wall_time_s': 10.0})) + reporter.close() + + records = [json.loads(line) for line in path.read_text().splitlines()] + assert [record['sequence'] for record in records] == [1, 2, 3, 4, 5] + assert records[2]['stage'] == 'train' + assert records[2]['optimizer_step'] == 1 + assert records[2]['values']['loss'] == 0.25 + assert 'event' not in records[2] + summary = json.loads(summary_path.read_text()) + assert summary['status'] == 'completed' + assert summary['trained_partitions'] == 1 + assert summary['per_context']['tenant/run/adapter']['optimizer_step'] == 1 + assert summary['metrics']['train/loss']['mean'] == 0.25 + assert summary['metrics']['rollout/reward']['count'] == 1 + assert summary['metrics']['rollout/reward']['mean'] == 0.5 + + +def test_summary_policy_version_does_not_regress_for_out_of_order_records(): + reporter = MetricsReporter(run_id='test') + reporter.record(MetricRecord( + stage='policy', + context_key='tenant/run/adapter', + policy_version=5, + values={}, + )) + reporter.record(MetricRecord( + stage='rollout', + context_key='tenant/run/adapter', + policy_version=3, + values={'sample_count': 4}, + attributes={'scope': 'group'}, + )) + + assert reporter.summary()['per_context']['tenant/run/adapter']['policy_version'] == 5 + reporter.close() + + +def test_jsonl_backend_waits_for_batch_threshold(tmp_path): + path = tmp_path / 'metrics.jsonl' + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': path, + 'summary_path': tmp_path / 'summary.json', + 'batch_size': 2, + 'flush_interval_s': 60, + }, + }, run_id='test') + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + time.sleep(0.05) + assert path.read_text() == '' + reporter.record(MetricRecord(stage='train', values={'loss': 2.0})) + deadline = time.monotonic() + 1 + while not path.read_text() and time.monotonic() < deadline: + time.sleep(0.01) + assert len(path.read_text().splitlines()) == 2 + reporter.close() + + +def test_reporter_assigns_monotonic_sequence_across_threads(tmp_path): + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': tmp_path / 'metrics.jsonl', + 'summary_path': tmp_path / 'summary.json', + }, + }, run_id='test') + + def produce(): + for _ in range(25): + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + + threads = [threading.Thread(target=produce) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + reporter.close() + records = [json.loads(line) for line in (tmp_path / 'metrics.jsonl').read_text().splitlines()] + assert [record['sequence'] for record in records] == list(range(1, 101)) + + +def test_swanlab_backend_uses_sequence_and_context_metric_names(monkeypatch, tmp_path): + logged = [] + + class Run: + def log(self, values, step): + logged.append((values, step)) + + fake_swanlab = types.SimpleNamespace(init=lambda **_kwargs: Run(), finish=lambda: None) + monkeypatch.setitem(sys.modules, 'swanlab', fake_swanlab) + reporter = create_metrics_reporter({ + 'jsonl': {'enabled': False}, + 'swanlab': { + 'enabled': True, + 'mode': 'local', + 'project': 'test', + 'name': 'test-run', + 'log_dir': tmp_path, + 'batch_size': 1, + }, + }, run_id='test') + reporter.record(MetricRecord( + stage='train', + context_key='tenant/run/adapter', + optimizer_step=3, + policy_version=2, + partition_index=1, + values={'loss': 0.5}, + )) + reporter.close() + values, sequence = logged[0] + assert sequence == 1 + assert values['context/tenant_run_adapter/train/loss'] == 0.5 + assert values['context/tenant_run_adapter/train/optimizer_step'] == 3 + assert values['context/tenant_run_adapter/policy/version'] == 2 + assert values['context/tenant_run_adapter/partition/index'] == 1 + + +def test_backend_queue_drops_oldest_record_without_blocking_reporter(): + release = threading.Event() + + class BlockingBackend(_QueuedBackend): + def _write_batch(self, batch): + release.wait(2) + + backend = BlockingBackend( + 'blocking', + queue_capacity=2, + batch_size=2, + flush_interval_s=60, + ) + reporter = MetricsReporter(run_id='test', backends=[backend]) + for index in range(5): + reporter.record(MetricRecord(stage='train', values={'loss': index})) + assert reporter.health()['backends']['blocking']['dropped_records'] >= 1 + release.set() + reporter.close() + + +def test_backend_failure_is_nonfatal_and_reported(): + class FailedBackend(_QueuedBackend): + def _write_batch(self, batch): + raise OSError('disk unavailable') + + backend = FailedBackend( + 'failed', + queue_capacity=4, + batch_size=1, + flush_interval_s=1, + ) + reporter = MetricsReporter(run_id='test', backends=[backend]) + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + reporter.flush() + reporter.record(MetricRecord(stage='train', values={'loss': 2.0})) + health = reporter.health()['backends']['failed'] + assert health['enabled'] is False + assert health['failure_count'] == 1 + assert 'disk unavailable' in health['last_error'] + reporter.close() + + +def test_swanlab_failure_does_not_prevent_jsonl(monkeypatch, tmp_path): + class FailedRun: + def log(self, values, step): + raise RuntimeError('swanlab unavailable') + + fake_swanlab = types.SimpleNamespace(init=lambda **_kwargs: FailedRun(), finish=lambda: None) + monkeypatch.setitem(sys.modules, 'swanlab', fake_swanlab) + path = tmp_path / 'metrics.jsonl' + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': path, + 'summary_path': tmp_path / 'summary.json', + 'batch_size': 1, + }, + 'swanlab': { + 'enabled': True, + 'mode': 'local', + 'log_dir': tmp_path, + 'batch_size': 1, + }, + }, run_id='test') + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + reporter.close() + assert len(path.read_text().splitlines()) == 1 + assert reporter.health()['backends']['swanlab']['failure_count'] == 1 + + +def test_jsonl_startup_failure_is_nonfatal(tmp_path): + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': tmp_path, + 'summary_path': tmp_path / 'summary.json', + }, + }, run_id='test') + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + health = reporter.health() + assert health['record_count'] == 1 + assert health['backends']['jsonl']['enabled'] is False + assert health['backends']['jsonl']['failure_count'] == 1 + reporter.close() diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py new file mode 100644 index 000000000..596597d8f --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -0,0 +1,1292 @@ +from __future__ import annotations + +import asyncio +import inspect +import json +import time + +import pytest + +from cookbook.rl.sync_barrier_multi_lora_grpo import SyncBarrierMultiLoraGRPO +from twinkle import DeviceMesh +from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams +from twinkle.infra import _dispatch_args +from twinkle.metric import MetricRecord +from twinkle_agentic.async_rl import (AsyncMultiLoraGRPOPipeline, ContextSchedulePolicy, ContextScheduler, + ContextStatus, LoraContext, LoraContextManager, ScheduleCandidate, + SchedulerConfig, TQDataPlane, TrainerWorker) +from twinkle_agentic.async_rl.data_plane import build_rollout_group_sample_write +from twinkle_agentic.async_rl.metrics import training_policy_metrics +from twinkle_agentic.async_rl.native_tq import ContextGRPOGroupNSampler +from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches +from twinkle_agentic.async_rl.pipeline import create_cpu_actor, _reward_for_context, _train_batch +from twinkle_agentic.async_rl.types import (PartitionAdmission, PreparedPartition, PromptGroup, RolloutPolicy) +from twinkle_agentic.async_rl.utils import ( + TrainBatchConfig, + build_native_fsdp_model_kwargs, + configure_lora_lr_scheduler, + resolve_context_learning_rate, + resolve_context_lora_target_modules, + resolve_context_loss_config, + resolve_model_attention_implementation, + resolve_sequence_parallel_size, + sampler_data_parallel_size, + validate_context_batch_config, +) +from twinkle_agentic.async_rl.vllm_sampler_tq import ( + VLLMSamplerTQ, + _GeneratedSample, + _PromptGroupRolloutStats, +) +from twinkle_agentic.async_rl.workers import RolloutWorker + + +class LocalActorHandle: + def __init__(self, target): + self.target = target + + def __getattr__(self, name): + method = getattr(self.target, name) + + class RemoteMethod: + async def remote(_, *args, **kwargs): + result = method(*args, **kwargs) + return await result if inspect.isawaitable(result) else result + + return RemoteMethod() + + +def test_cpu_service_actor_uses_twinkle_ray_mode(monkeypatch): + import ray + + captured = {} + + class ActorClass: + + @staticmethod + def remote(*args, **kwargs): + captured['actor_args'] = args + captured['actor_kwargs'] = kwargs + return 'actor' + + def fake_remote(**options): + captured['options'] = options + return lambda cls: ActorClass + + monkeypatch.setattr(ray, 'remote', fake_remote) + + assert create_cpu_actor(object, 'value', enabled=True) == 'actor' + assert captured['options'] == { + 'num_cpus': 1, + 'runtime_env': { + 'env_vars': { + 'TWINKLE_MODE': 'ray' + } + }, + } + assert captured['actor_args'] == ('value', ) + assert captured['actor_kwargs'] == {'enabled': True} + + +def test_unload_lora_paths_does_not_require_pruned_checkpoint(tmp_path): + removed_paths = [] + + class Engine: + + async def unload_lora_paths(self, paths): + removed_paths.extend(paths) + + class Completed: + + @staticmethod + def result(): + return None + + sampler = object.__new__(VLLMSamplerTQ) + sampler.engine = Engine() + + def submit(coro): + asyncio.run(coro) + return Completed() + + sampler._submit_in_loop = submit + pruned_path = tmp_path / 'already-pruned' + sampler.unload_lora_paths([str(pruned_path)]) + + assert removed_paths == [str(pruned_path.resolve())] + + +def test_sequence_parallel_size_must_divide_model_gpus(): + assert resolve_sequence_parallel_size(2, 1) == 1 + assert resolve_sequence_parallel_size(2, 2) == 2 + + with pytest.raises(ValueError, match='must be divisible'): + resolve_sequence_parallel_size(2, 3) + + +def test_padding_free_sequence_parallel_requires_flash_attention(): + assert resolve_model_attention_implementation( + {'attn_implementation': 'flash_attention_2'}, + padding_free=True, + sequence_parallel_size=2, + ) == 'flash_attention_2' + + with pytest.raises(ValueError, match='model.attn_implementation'): + resolve_model_attention_implementation({}, padding_free=True, sequence_parallel_size=2) + + +def test_train_batch_preserves_position_ids_from_tq(): + class Batch(dict): + batch_size = (1, ) + + class Model: + inputs = None + + def forward_backward(self, *, inputs, **_kwargs): + self.inputs = inputs + + def clip_grad_and_step(self, **_kwargs): + return None + + def calculate_metric(self, **_kwargs): + return {} + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 1, 0) + data = Batch({ + 'input_ids': [[1, 2]], + 'labels': [[-100, 2]], + 'attention_mask': [[1, 1]], + 'position_ids': [[0, 1]], + 'logprobs': [[-.1]], + 'advantages': [1.], + 'rewards': [1.], + }) + model = Model() + + _train_batch(model, {context.key: TrainBatchConfig(1, 1)}, data, admission) + + assert model.inputs == [{ + 'input_ids': [1, 2], + 'labels': [-100, 2], + 'attention_mask': [1, 1], + 'position_ids': [0, 1], + }] + + +def test_train_batch_accumulates_real_micro_batches_before_one_optimizer_step(): + class Batch(dict): + batch_size = (4, ) + + class Model: + def __init__(self): + self.calls = [] + self.optimizer_steps = 0 + + def forward_backward(self, **kwargs): + self.calls.append(kwargs) + return lambda: { + 'micro_batch_count': 4, + 'micro_batch_samples_mean': 1.0, + 'micro_batch_tokens_mean': 1.0, + 'micro_batch_tokens_max': 1, + } + + def clip_grad_and_step(self, **_kwargs): + self.optimizer_steps += 1 + + def calculate_metric(self, **_kwargs): + return {} + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 4, 0) + data = Batch({ + 'input_ids': [[index] for index in range(4)], + 'labels': [[index] for index in range(4)], + 'attention_mask': [[1] for _ in range(4)], + 'position_ids': [[0] for _ in range(4)], + 'logprobs': [[-.1] for _ in range(4)], + 'advantages': [1., 2., 3., 4.], + 'rewards': [1., 1., 1., 1.], + }) + model = Model() + + metrics = _train_batch( + model, + {context.key: TrainBatchConfig(4, 1)}, + data, + admission, + model_data_parallel_size=1, + ) + + assert [len(call['inputs']) for call in model.calls] == [4] + assert model.calls[0]['advantages'] == [1., 2., 3., 4.] + assert model.calls[0]['micro_batch_size'] == 1 + assert model.calls[0]['loss_scale'] == 1.0 + assert model.optimizer_steps == 1 + assert metrics['micro_batch_size_per_rank'] == 1 + assert 'micro_batch_count' not in metrics + + +def test_dynamic_micro_batch_planner_honors_per_rank_sample_and_token_limits(): + lengths = [10, 9, 8, 7, 4, 3, 2, 1] + inputs = [{'input_ids': list(range(length))} for length in lengths] + config = MicroBatchConfig( + micro_batch_size=3, + dynamic_batching=True, + max_tokens_per_micro_batch=18, + ) + + batches = plan_micro_batches(inputs, config, padding_free=False) + + assert sorted(index for batch in batches for index in batch) == list(range(8)) + for batch in batches: + assert len(batch) <= 3 + padded_tokens = max(lengths[index] for index in batch) * len(batch) + assert padded_tokens <= 18 + + +def test_sync_training_batch_preserves_position_ids(): + rows = [{ + 'input_ids': [1, 2], + 'labels': [-100, 2], + 'attention_mask': [1, 1], + 'position_ids': [0, 1], + 'logprobs': [-.1], + }] + + batch = SyncBarrierMultiLoraGRPO._training_batch(rows, rewards=[1.], advantages=[0.]) + + assert 'position_ids' in batch.keys() + assert batch['position_ids'][0] == [0, 1] + + +class PolicyProvider: + + def __init__(self, policies): + self.policies = iter(policies) + self.released = [] + + def get_rollout_policy(self, _context): + return next(self.policies) + + def acquire_rollout_policy(self, context): + return self.get_rollout_policy(context) + + def release_rollout_policy(self, policy): + self.released.append(policy) + + +class GenerationHarness: + _merge_partial_responses = VLLMSamplerTQ._merge_partial_responses + + def __init__(self, policies, responses): + self.context_manager = LocalActorHandle(PolicyProvider(policies)) + self.responses = iter(responses) + self.rollout_max_retries = 1 + self.rollout_retry_delay_s = 0 + self.calls = [] + self.template = type('Template', (), {'decode': staticmethod(lambda tokens: str(tokens))})() + + async def _load_lora_for_policy(self, policy): + return policy.version + + async def _sample_single(self, feat, sampling_params, *, lora_request, multi_modal_data, logprobs_only): + self.calls.append((list(feat['input_ids']), sampling_params.max_tokens, lora_request)) + return next(self.responses) + + +def _context(name: str = 'adapter') -> LoraContext: + return LoraContext('tenant', f'run_{name}', 'model', name) + + +def _sample_response(tokens, stop_reason, input_ids): + return SampleResponse( + prompt_token_ids=[1, 2], + sequences=[ + SampledSequence( + stop_reason=stop_reason, + tokens=tokens, + logprobs=[[(token, -.1)] for token in tokens], + new_input_feature={ + 'input_ids': input_ids, + 'labels': [-100, -100, *tokens], + }, + ) + ], + ) + + +def test_sampler_data_parallel_size_is_derived_from_gpu_and_tp_sizes(): + assert sampler_data_parallel_size(8, 2) == 4 + assert sampler_data_parallel_size(1, 1) == 1 + + +def test_sampler_parallelism_rejects_incomplete_tp_group(): + try: + sampler_data_parallel_size(3, 2) + except ValueError as exc: + assert 'must be divisible' in str(exc) + else: + raise AssertionError('expected invalid sampler GPU/TP layout to fail') + + +def test_lora_lr_scheduler_uses_shared_adapter_config(): + calls = [] + + class Model: + def set_lr_scheduler(self, scheduler_cls, **kwargs): + calls.append((scheduler_cls, kwargs)) + + configure_lora_lr_scheduler( + Model(), + 'tenant_lora', + { + 'lr_scheduler': { + 'cls': 'CosineAnnealingLR', + 'T_max': 2000, + 'eta_min': 0.0, + }, + }, + ) + + assert calls == [('CosineAnnealingLR', { + 'adapter_name': 'tenant_lora', + 'T_max': 2000, + 'eta_min': 0.0, + })] + + +def test_context_learning_rate_overrides_global_default(): + assert resolve_context_learning_rate({'learning_rate': 5e-6}, {'learning_rate': 1e-6}) == pytest.approx(5e-6) + assert resolve_context_learning_rate({}, {'learning_rate': 1e-6}) == pytest.approx(1e-6) + + +def test_context_lora_target_modules_override_global_default(): + defaults = {'target_modules': 'all-linear'} + + assert resolve_context_lora_target_modules({}, defaults) == 'all-linear' + assert resolve_context_lora_target_modules( + {'lora': {'target_modules': ['q_proj', 'v_proj']}}, + defaults, + ) == ['q_proj', 'v_proj'] + + +@pytest.mark.parametrize('value', ['', [], [None], {'q_proj': True}]) +def test_context_lora_target_modules_reject_invalid_values(value): + with pytest.raises(ValueError, match='target_modules'): + resolve_context_lora_target_modules( + {'lora': {'target_modules': value}}, + {'target_modules': 'all-linear'}, + ) + + +def test_context_loss_config_overrides_global_defaults(): + loss_cls, loss_kwargs = resolve_context_loss_config( + { + 'loss': { + 'cls': 'GSPOLoss', + 'normalization': 'token_mean', + } + }, + { + 'cls': 'GRPOLoss', + 'epsilon': 0.2, + 'normalization': 'sequence_mean', + }, + ) + + assert loss_cls == 'GSPOLoss' + assert loss_kwargs == { + 'epsilon': 0.2, + 'normalization': 'token_mean', + } + + +def test_context_loss_config_uses_grpo_defaults(): + assert resolve_context_loss_config({}) == ( + 'GRPOLoss', + { + 'epsilon': 0.2, + 'normalization': 'sequence_mean', + }, + ) + + +def test_context_loss_config_rejects_empty_class_name(): + with pytest.raises(ValueError, match='loss.cls'): + resolve_context_loss_config({'loss': {'cls': ''}}) + + +@pytest.mark.parametrize('value', [0, -1e-6, float('inf')]) +def test_context_learning_rate_rejects_invalid_values(value): + with pytest.raises(ValueError, match='positive finite'): + resolve_context_learning_rate({'learning_rate': value}, {'learning_rate': 1e-6}) + + +def test_rl_model_kwargs_enforce_native_fsdp(): + assert build_native_fsdp_model_kwargs({}) == { + 'strategy': 'native_fsdp', + 'fsdp_config': {}, + } + assert build_native_fsdp_model_kwargs({ + 'strategy': 'native_fsdp', + 'fsdp_config': {'reshard_after_forward': False}, + }) == { + 'strategy': 'native_fsdp', + 'fsdp_config': {'reshard_after_forward': False}, + } + with pytest.raises(ValueError, match='must be native_fsdp'): + build_native_fsdp_model_kwargs({'strategy': 'accelerate'}) + + +def test_reward_factory_loads_class_and_resolved_kwargs(): + reward = _reward_for_context( + { + 'class_path': 'twinkle.reward.DAPOMathReward', + 'kwargs': { + 'max_response_length': 8192, + 'overlong_buffer_length': 4096, + 'overlong_penalty_factor': 1.0, + 'score_tail_chars': 300, + }, + }, + context_key='tenant/run/adapter', + ) + + assert reward.max_response_length == 8192 + assert reward.overlong_buffer_length == 4096 + + +def test_reward_factory_rejects_non_reward_class(): + with pytest.raises(TypeError, match='Reward subclass'): + _reward_for_context( + {'class_path': 'collections.Counter'}, + context_key='tenant/run/adapter', + ) + + +def test_context_batch_config_accepts_group_aligned_dp_batches(): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=8, micro_batch_size=2), + sampler_dp=2, + model_dp=2, + ) + + +def test_context_batch_config_allows_training_group_to_span_model_dp_ranks(): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=16, micro_batch_size=2), + sampler_dp=1, + model_dp=8, + ) + + +def test_context_batch_config_rejects_partition_tail_and_undersized_rank_batch(): + try: + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=6, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=6, micro_batch_size=1), + sampler_dp=2, + model_dp=2, + ) + except ValueError as exc: + assert 'complete prompt groups' in str(exc) + else: + raise AssertionError('expected a split prompt group to fail') + + try: + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=2, + train=TrainBatchConfig(mini_batch_size=2, micro_batch_size=2), + sampler_dp=2, + model_dp=2, + ) + except ValueError as exc: + assert 'per-rank train batch' in str(exc) + else: + raise AssertionError('expected an oversized micro batch to fail') + + +def test_context_batch_config_requires_token_limit_for_dynamic_batching(): + with pytest.raises(ValueError, match='max_tokens_per_micro_batch'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig( + mini_batch_size=8, + micro_batch_size=2, + dynamic_batching=True, + ), + sampler_dp=1, + model_dp=1, + ) + + +def test_sampler_dp_dispatch_slices_complete_groups_without_duplication(): + mesh = DeviceMesh.from_sizes(world_size=4, dp_size=2, tp_size=2) + groups = ['group_0', 'group_1', 'group_2', 'group_3'] + dispatched = _dispatch_args( + workers=['dp_0', 'dp_1'], + dispatch='slice_dp', + execute='all', + device_mesh=mesh, + args=(groups, 'sampling_params', False), + kwargs={}, + ) + + assert [worker for worker, _, _ in dispatched] == ['dp_0', 'dp_1'] + assert [args[0] for _, args, _ in dispatched] == [groups[:2], groups[2:]] + assert [group for _, args, _ in dispatched for group in args[0]] == groups + + +@pytest.mark.parametrize(('dp_size', 'expected_scope'), [(1, 'partition'), (2, 'shard')]) +def test_sampler_reports_submission_throughput_at_partition_or_shard_scope(dp_size, expected_scope): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 2, 2, 0) + groups = [ + PromptGroup(context, admission, f'{admission.partition_id}/group_{index}', {}, object()) + for index in range(2) + ] + + class RolloutMetricsHarness: + def __init__(self): + self.device_mesh = DeviceMesh.from_sizes(world_size=dp_size, dp_size=dp_size) + self.events = [] + + async def _run_prompt_group(self, *, group, **_kwargs): + index = int(group.group_id.rsplit('_', 1)[1]) + lengths = ((10, 20), (30, 40))[index] + reasons = (('stop', 'length'), ('stop', 'stop'))[index] + return _PromptGroupRolloutStats(lengths, reasons, (index + 1, index + 1)) + + def _record_metrics(self, group, values, **kwargs): + self.events.append((group, values, kwargs)) + + sampler = RolloutMetricsHarness() + asyncio.run( + VLLMSamplerTQ._sample_prompt_groups( + sampler, + 'submission', + groups, + SamplingParams(max_tokens=64), + False, + time.perf_counter() - 1, + )) + + recorded_group, metrics, record_options = sampler.events[-1] + assert recorded_group.context == context + assert recorded_group.partition_id == admission.partition_id + assert record_options['attributes']['scope'] == expected_scope + assert metrics['prompt_group_count'] == 2 + assert metrics['sample_count'] == 4 + assert metrics['output_tokens'] == 100 + assert metrics['completion_length_mean'] == 25 + assert metrics['completion_truncated_count'] == 1 + assert metrics['policy_version_min'] == 1 + assert metrics['policy_version_max'] == 2 + assert metrics['sampler_dp_size'] == dp_size + assert metrics['output_tokens_per_s'] == pytest.approx(100 / metrics['rollout_latency_s']) + + +def test_sampler_writes_one_atomic_rollout_file_per_prompt_group(tmp_path): + context = _context() + admission = PartitionAdmission(context, context.partition_id(3), 3, 1, 2, 0) + group = PromptGroup( + context, + admission, + f'{admission.partition_id}/group_0', + {'user_data': [('ground_truth', '"42"')]}, + object(), + ) + policy = RolloutPolicy(context.key, context.adapter_name, 7, '/tmp/adapter-v7') + generated = [ + _GeneratedSample( + SampleResponse( + sequences=[SampledSequence('stop', [20 + index], decoded=f'completion-{index}')], + prompt_token_ids=[10, 11], + ), + (policy,), + attempts=1, + was_aborted=False, + resumed_partial_output=False, + ) + for index in range(2) + ] + rows = [ + { + 'generation_idx': index, + 'rollout_policy_version': 7, + 'initial_policy_version': 7, + 'final_policy_version': 7, + 'rollout_policy_versions': [7], + 'rollout_adapter_path': '/tmp/adapter-v7', + 'stop_reason': 'stop', + 'logprobs': [-0.1], + } + for index in range(2) + ] + + class Template: + @staticmethod + def decode(token_ids, **_kwargs): + return ' '.join(map(str, token_ids)) + + sampler = object.__new__(VLLMSamplerTQ) + sampler.rollout_output_dir = tmp_path + sampler.rollout_output_include_token_ids = False + sampler.template = Template() + + sampler._write_rollout_group('submission-1', group, generated, rows, [1.0, 0.0]) + sampler._write_rollout_group('submission-2', group, generated, rows, [1.0, 0.0]) + + output_path = ( + tmp_path + / context.tenant_id + / context.training_run_id + / context.adapter_name + / 'policy_7' + / 'train_3-group_0.jsonl' + ) + records = [json.loads(line) for line in output_path.read_text().splitlines()] + assert len(records) == 2 + assert records[0]['submission_id'] == 'submission-2' + assert records[0]['prompt'] == '10 11' + assert records[0]['completion'] == '20' + assert records[0]['ground_truth'] == '42' + assert records[0]['reward'] == 1.0 + assert records[0]['head_version'] == 7 + assert records[0]['tail_version'] == 7 + assert 'prompt_token_ids' not in records[0] + + +def test_aborted_generation_restarts_from_original_prompt_when_partial_is_disabled(): + context = _context() + policies = [ + RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), + RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), + ] + sampler = GenerationHarness( + policies, + [ + _sample_response([7], 'abort', [1, 2, 7]), + _sample_response([8], 'stop', [1, 2, 8]), + ], + ) + generated = asyncio.run( + VLLMSamplerTQ._generate_sample( + sampler, + context, + { + 'input_ids': [1, 2], + 'labels': [-100, -100] + }, + SamplingParams(max_tokens=4, logprobs=1), + multi_modal_data=None, + logprobs_only=False, + allow_partial_rollout=False, + )) + + assert sampler.calls == [([1, 2], 4, 3), ([1, 2], 4, 4)] + assert generated.response.sequences[0].tokens == [8] + assert [policy.version for policy in generated.policies] == [4] + assert generated.retry_count == 1 + assert generated.was_aborted + assert not generated.resumed_partial_output + + +def test_aborted_generation_continues_from_partial_tokens_when_enabled(): + context = _context() + policies = [ + RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), + RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), + ] + sampler = GenerationHarness( + policies, + [ + _sample_response([7], 'abort', [1, 2, 7]), + _sample_response([8], 'stop', [1, 2, 7, 8]), + ], + ) + generated = asyncio.run( + VLLMSamplerTQ._generate_sample( + sampler, + context, + { + 'input_ids': [1, 2], + 'labels': [-100, -100] + }, + SamplingParams(max_tokens=4, logprobs=1), + multi_modal_data=None, + logprobs_only=False, + allow_partial_rollout=True, + )) + + assert sampler.calls == [([1, 2], 4, 3), ([1, 2, 7], 3, 4)] + assert generated.response.sequences[0].tokens == [7, 8] + assert [policy.version for policy in generated.policies] == [3, 4] + assert generated.initial_policy.version == 3 + assert generated.final_policy.version == 4 + assert generated.retry_count == 1 + assert generated.was_aborted + assert generated.resumed_partial_output + + +def test_training_policy_metrics_use_final_version_and_partial_span(): + metrics = training_policy_metrics(( + { + 'final_policy_version': 3, + 'policy_version_span': 1 + }, + { + 'final_policy_version': 4, + 'policy_version_span': 0 + }, + ), train_policy_version=5) + + assert metrics == { + 'policy_version_gap_mean': 1.5, + 'policy_version_gap_p95': 2, + 'policy_version_gap_max': 2, + 'rollout_policy_span_mean': 0.5, + 'rollout_policy_span_max': 1, + } + + +def test_training_policy_metrics_reject_future_rollout_version(): + try: + training_policy_metrics(({ + 'final_policy_version': 6, + 'policy_version_span': 0 + }, ), train_policy_version=5) + except ValueError as exc: + assert 'older than rollout versions' in str(exc) + else: + raise AssertionError('expected a future rollout policy version to fail') + + +def test_unified_staleness_admission(): + context = _context() + zero = LoraContextManager(max_staleness=0) + zero.register_context(context) + first = zero.request_rollout_partition(context, target_groups=1, num_generations=2) + assert first is not None + assert zero.request_rollout_partition(context, target_groups=1, num_generations=2) is None + + one = LoraContextManager(max_staleness=1) + one.register_context(context) + assert one.request_rollout_partition(context, target_groups=1, num_generations=2) is not None + assert one.request_rollout_partition(context, target_groups=1, num_generations=2) is not None + assert one.request_rollout_partition(context, target_groups=1, num_generations=2) is None + + +def test_rollout_worker_retains_prefetched_batch_until_admission_succeeds(): + context = _context() + + class AdmissionGate: + def __init__(self): + self.blocked = True + self.attempts = 0 + self.accepted = False + + def is_rollout_admission_closed(self): + return False + + def context_status(self, _context): + return ContextStatus.ACTIVE + + def request_rollout_partition(self, _context, *, target_groups, num_generations): + self.attempts += 1 + if self.blocked or self.accepted: + return None + self.accepted = True + return PartitionAdmission(context, context.partition_id(0), 0, target_groups, num_generations, 0) + + class DataPlane: + async def prepare_rollout_partition(self, admission, _prompts, sampling_params): + return PreparedPartition(admission, (), sampling_params) + + class Sampler: + def __init__(self, loop): + self.submitted = asyncio.Event() + self.loop = loop + + def sample(self, _groups, _sampling_params, _allow_partial_rollout): + self.loop.call_soon_threadsafe(self.submitted.set) + + loaded_batches = [] + + def batches(): + for value in (1, 2): + loaded_batches.append(value) + yield [{'input_ids': [value]}] + + async def run(): + manager = AdmissionGate() + sampler = Sampler(asyncio.get_running_loop()) + worker = RolloutWorker( + context_manager=LocalActorHandle(manager), + data_plane=DataPlane(), + sampler=sampler, + prompt_batches={context.key: batches()}, + rollout_config={ + context.key: { + 'context': context, + 'batch_size': 1, + 'num_generations': 2, + 'sampling_params': {}, + } + }, + scheduler=SchedulerConfig(ContextSchedulePolicy.ROUND_ROBIN, 1), + idle_delay_s=.001, + ) + await worker.start() + while manager.attempts == 0: + await asyncio.sleep(.001) + prefetched_task = worker._next_batch_tasks[context.key] + await asyncio.sleep(.01) + assert worker._next_batch_tasks[context.key] is prefetched_task + assert loaded_batches == [1] + + manager.blocked = False + await asyncio.wait_for(sampler.submitted.wait(), timeout=1) + while loaded_batches == [1]: + await asyncio.sleep(.001) + assert loaded_batches == [1, 2] + await worker.stop() + + asyncio.run(run()) + + +def test_partition_clear_releases_capacity_only_after_publish(): + context = _context() + manager = LoraContextManager(max_staleness=0) + manager.register_context(context, adapter_path='initial') + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + manager.on_partition_training_started(admission) + policy = manager.on_partition_trained(admission, adapter_path='v1') + assert policy.version == 1 + assert manager.request_rollout_partition(context, target_groups=1, num_generations=2) is None + manager.on_partition_cleared(admission) + assert manager.request_rollout_partition(context, target_groups=1, num_generations=2) is not None + + +def test_context_trains_partitions_in_step_order(): + context = _context() + manager = LoraContextManager(max_staleness=1) + manager.register_context(context) + first = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + second = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + + assert manager.list_trainable_partitions() == [first] + manager.on_partition_training_started(first) + assert manager.list_trainable_partitions() == [first] + + try: + manager.on_partition_training_started(second) + except RuntimeError as exc: + assert f'already trains {first.partition_id}' in str(exc) + else: + raise AssertionError('expected the next partition to remain blocked') + + manager.on_partition_trained(first, adapter_path='v1') + manager.on_partition_cleared(first) + assert manager.list_trainable_partitions() == [second] + manager.on_partition_training_started(second) + + +def test_scheduler_supports_round_robin_sticky_and_oldest(): + a, b = _context('a'), _context('b') + candidates = [ScheduleCandidate(a), ScheduleCandidate(b)] + round_robin = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.ROUND_ROBIN, 1)) + assert round_robin.choose(candidates).context == a + round_robin.on_success(candidates[0]) + assert round_robin.choose(candidates).context == b + + sticky = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.STICKY, None)) + sticky.on_success(candidates[1]) + assert sticky.choose(candidates).context == b + sticky.on_blocked(candidates[1]) + assert sticky.choose(candidates).context == a + + manager = LoraContextManager(max_staleness=2) + manager.register_context(a) + manager.register_context(b) + old = manager.request_rollout_partition(a, target_groups=1, num_generations=2) + new = manager.request_rollout_partition(b, target_groups=1, num_generations=2) + oldest = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.OLDEST_PARTITION, 1)) + assert oldest.choose([ScheduleCandidate(b, new), ScheduleCandidate(a, old)]).partition == old + + +def test_sticky_scheduler_switches_context_at_consecutive_cap(): + a, b = _context('a'), _context('b') + candidates = [ScheduleCandidate(a), ScheduleCandidate(b)] + scheduler = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.STICKY, 1)) + + first = scheduler.choose(candidates) + scheduler.on_success(first) + + assert scheduler.choose(candidates).context == b + + +def test_context_group_sampler_uses_request_generation_count(): + sampler = ContextGRPOGroupNSampler() + + selected, consumed = sampler.sample( + [0, 1, 4, 5, 6, 7], + batch_size=4, + partition_id='train_0', + task_name='advantage/context', + n_samples_per_prompt=4, + ) + + assert selected == [4, 5, 6, 7] + assert consumed == selected + + selected, consumed = sampler.sample( + [8, 9, 12], + batch_size=2, + partition_id='train_1', + task_name='advantage/context', + n_samples_per_prompt=2, + ) + + assert selected == [8, 9] + assert consumed == selected + + +def test_context_finishes_after_exhaustion_and_clear(): + context = _context() + manager = LoraContextManager() + manager.register_context(context) + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + manager.on_dataset_exhausted(context) + assert not manager.is_run_finished() + manager.on_partition_training_started(admission) + manager.on_partition_trained(admission, adapter_path='v1') + manager.on_partition_cleared(admission) + assert manager.is_run_finished() + + +def test_pipeline_fails_fast_when_a_worker_service_fails(): + context = _context() + manager = LoraContextManager() + manager.register_context(context) + + class FailedWorker: + async def start(self): + return None + + async def stop(self): + return None + + async def get_service_state(self): + return {'running': False, 'failure': 'CUDA out of memory'} + + def drain_metric_records(self): + return [] + + worker = LocalActorHandle(FailedWorker()) + pipeline = AsyncMultiLoraGRPOPipeline( + context_manager=LocalActorHandle(manager), + rollout_worker=worker, + advantage_worker=worker, + trainer_worker=worker, + ) + + try: + asyncio.run(pipeline.run_async()) + except RuntimeError as exc: + assert 'CUDA out of memory' in str(exc) + else: + raise AssertionError('expected worker failure to fail the pipeline') + + +def test_pipeline_drains_actor_metric_buffers_when_reporting_is_disabled(): + class BufferedWorker: + def __init__(self): + self.drain_count = 0 + + def drain_metric_records(self): + self.drain_count += 1 + return [MetricRecord(stage='train', values={'loss': 1.0})] + + class BufferedSampler: + def __init__(self): + self.drain_count = 0 + + def drain_metric_records(self): + self.drain_count += 1 + return [MetricRecord(stage='rollout', values={'sample_count': 1})] + + workers = [BufferedWorker() for _ in range(3)] + sampler = BufferedSampler() + pipeline = AsyncMultiLoraGRPOPipeline( + context_manager=object(), + rollout_worker=LocalActorHandle(workers[0]), + advantage_worker=LocalActorHandle(workers[1]), + trainer_worker=LocalActorHandle(workers[2]), + sampler=sampler, + metrics=None, + ) + + asyncio.run(pipeline._drain_metrics()) + + assert [worker.drain_count for worker in workers] == [1, 1, 1] + assert sampler.drain_count == 1 + + +def test_global_max_steps_limits_admission_and_closes_after_completion(): + first, second = _context('a'), _context('b') + manager = LoraContextManager(max_staleness=1, max_steps=1) + manager.register_context(first) + manager.register_context(second) + first_admission = manager.request_rollout_partition(first, target_groups=1, num_generations=2) + assert manager.request_rollout_partition(second, target_groups=1, num_generations=2) is None + manager.on_partition_training_started(first_admission) + manager.on_partition_trained(first_admission, adapter_path='v1') + manager.on_partition_cleared(first_admission) + assert manager.is_rollout_admission_closed() + assert manager.is_run_finished() + + +def test_zero_max_steps_finishes_without_admission(): + context = _context() + manager = LoraContextManager(max_steps=0) + manager.register_context(context) + assert manager.request_rollout_partition(context, target_groups=1, num_generations=2) is None + assert manager.is_rollout_admission_closed() + assert manager.is_run_finished() + + +def test_rollout_sample_tags_use_new_context_descriptor_only(): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, batch_meta=None) + fields, tags = build_rollout_group_sample_write( + group, + [ + { + 'generation_idx': 0, + 'labels': [-100, 1], + 'logprobs': [-.1], + 'rollout_policy_version': 3, + 'rollout_adapter_path': 'adapter-v3', + }, + { + 'generation_idx': 1, + 'labels': [-100, 2], + 'logprobs': [-.2], + 'rollout_policy_version': 4, + 'rollout_adapter_path': 'adapter-v4', + }, + ], + rewards=[1., 0.], + expected_num_generations=2, + ) + assert [row['rewards'] for row in fields] == [1., 0.] + assert [tag['generation_idx'] for tag in tags] == [0, 1] + assert all(tag['context_key'] == context.key for tag in tags) + assert [tag['rollout_policy_version'] for tag in tags] == [3, 4] + + +def test_data_plane_completes_rollout_with_full_training_trajectory(): + class Metadata: + def __init__(self): + self.size = 2 + self.custom_meta = [{}, {}] + + def update_custom_meta(self, updates): + for tag, update in zip(self.custom_meta, updates): + tag.update(update) + + class Client: + def __init__(self): + self.written = None + self.calls = [] + + async def async_put(self, data, metadata=None, partition_id=None): + self.calls.append('fields') + self.written = data + return metadata + + async def async_set_custom_meta(self, _metadata): + self.calls.append('tags') + return None + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) + metadata = Metadata() + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) + client = Client() + rows = [{ + 'input_ids': [1, 2, token], + 'labels': [-100, -100, token], + 'attention_mask': [1, 1, 1], + 'position_ids': [0, 1, 2], + 'logprobs': [-.1], + 'generation_idx': generation_idx, + 'rollout_policy_version': 3, + 'rollout_policy_versions': [3], + 'initial_policy_version': 3, + 'final_policy_version': 3, + 'policy_version_span': 0, + 'rollout_adapter_path': 'adapter-v3', + 'completion_length': 1, + } for generation_idx, token in enumerate((7, 8))] + + asyncio.run( + TQDataPlane(client).complete_rollout_group( + group, + rollout_rows=rows, + rewards=[1., 0.], + submission_id='submission', + )) + + assert set(client.written.keys()) == { + 'input_ids', 'labels', 'attention_mask', 'position_ids', 'logprobs', 'rewards' + } + assert client.calls == ['tags', 'fields'] + assert [tag['rollout_status'] for tag in metadata.custom_meta] == ['ROLLOUT_DONE', 'ROLLOUT_DONE'] + assert [tag['submission_id'] for tag in metadata.custom_meta] == ['submission', 'submission'] + + +def test_data_plane_rejects_rollout_without_complete_model_inputs(): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 1, 0) + metadata = type('Metadata', (), {'size': 1})() + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) + row = { + 'input_ids': [1, 2], + 'labels': [-100, 2], + 'logprobs': [-.1], + 'generation_idx': 0, + 'rollout_policy_version': 0, + } + + try: + asyncio.run( + TQDataPlane(object()).complete_rollout_group( + group, + rollout_rows=[row], + rewards=[1.], + submission_id='submission', + )) + except ValueError as exc: + assert 'attention_mask' in str(exc) + assert 'position_ids' in str(exc) + else: + raise AssertionError('expected incomplete rollout model fields to fail') + + +def test_checkpoint_retention_preserves_current_policy_and_history_window(): + context = _context() + manager = LoraContextManager() + manager.register_context(context, adapter_path='initial') + removed = [] + worker = TrainerWorker( + context_manager=LocalActorHandle(manager), + data_plane=TQDataPlane(), + train_fn=lambda _data, _admission: {}, + save_adapter=lambda _admission: 'unused', + mini_batch_sizes={context.key: 2}, + scheduler=SchedulerConfig(ContextSchedulePolicy.STICKY, None), + keep_adapter_versions=1, + initial_adapter_paths={context.key: 'initial'}, + remove_adapter=removed.append, + ) + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + manager.on_partition_training_started(admission) + manager.on_partition_trained(admission, adapter_path='current') + manager.on_partition_cleared(admission) + worker._adapter_history[context.key].append('current') + async def prune(): + await worker._prune_adapter_history(context) + await worker.stop() + + asyncio.run(prune()) + assert removed == ['initial'] + assert worker._adapter_history[context.key] == ['current'] + prune_events = [ + record for record in worker.drain_metric_records() + if record.stage == 'policy' and record.attributes.get('operation') == 'adapter_prune' + ] + assert len(prune_events) == 1 + assert prune_events[0].context_key == context.key + assert prune_events[0].attributes['adapter_path'] == 'initial' + assert prune_events[0].values['adapter_prune_latency_s'] >= 0 + + +def test_policy_retention_keeps_only_current_and_actively_referenced_paths(): + context = _context() + manager = LoraContextManager(max_staleness=1) + manager.register_context(context, adapter_path='initial') + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + + acquired = manager.acquire_rollout_policy(context) + manager.on_partition_training_started(admission) + manager.on_partition_trained(admission, adapter_path='current') + + assert manager.adapter_paths_to_keep() == {'initial', 'current'} + manager.release_rollout_policy(acquired) + assert manager.adapter_paths_to_keep() == {'current'} + + +def test_trainer_periodically_evaluates_published_policy(): + context = _context() + manager = LoraContextManager() + manager.register_context(context, adapter_path='initial') + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + calls = [] + + def evaluate_batch(batch, evaluated_admission, adapter_path, policy_version, sampling_params): + calls.append((list(batch), evaluated_admission, adapter_path, policy_version, sampling_params)) + return { + 'rewards': [1.0] * len(batch), + 'completion_lengths': [10] * len(batch), + } + + worker = TrainerWorker( + context_manager=LocalActorHandle(manager), + data_plane=TQDataPlane(), + train_fn=lambda _data, _admission: {}, + save_adapter=lambda _admission: 'unused', + mini_batch_sizes={context.key: 2}, + scheduler=SchedulerConfig(ContextSchedulePolicy.STICKY, None), + evaluation_config={ + context.key: { + 'interval': 5, + 'dataset_name': 'validation', + 'prompt_batches': lambda: [[{'input_ids': [1]}], [{'input_ids': [2]}]], + 'sampling_params': 'params', + } + }, + evaluate_batch=evaluate_batch, + ) + worker._optimizer_steps[context.key] = 50 + + async def evaluate(): + await worker._evaluate_policy(admission, 'adapter-v4', 4) + await worker._evaluate_policy(admission, 'adapter-v5', 5) + + asyncio.run(evaluate()) + assert len(calls) == 2 + records = [record for record in worker.drain_metric_records() if record.stage == 'evaluation'] + assert len(records) == 1 + assert records[0].policy_version == 5 + assert records[0].optimizer_step == 50 + assert records[0].values['accuracy'] == 1.0 + assert records[0].values['prompt_count'] == 2 + assert records[0].values['sample_count'] == 2 + assert records[0].values['completion_length'] == 10 diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py new file mode 100644 index 000000000..29fa62f0b --- /dev/null +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import asyncio +from concurrent.futures import Future + +import pytest + +from twinkle.data_format import SamplingParams +from twinkle.server.sampler.twinkle_handlers import _await_generation +from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ, _dispatch_generation + + +def _bare_sampler() -> VLLMSamplerTQ: + sampler = object.__new__(VLLMSamplerTQ) + sampler._generation_submissions = {} + return sampler + + +def test_generation_dispatch_allows_one_prompt_with_multiple_dp_workers() -> None: + assert VLLMSamplerTQ.submit_generation._dispatch is _dispatch_generation + assert VLLMSamplerTQ.sample._dispatch == 'slice_dp' + shards = [ + _dispatch_generation( + 3, + worker_index, + ('submission', [{'input_ids': [1]}], 'params'), + {}, + )[0][1] + for worker_index in range(3) + ] + + assert shards == [[{'input_ids': [1]}], [], []] + + +def test_generation_submission_returns_before_generation_finishes() -> None: + sampler = _bare_sampler() + pending = Future() + submitted_coroutines = [] + + def submit(coro): + submitted_coroutines.append(coro) + coro.close() + return pending + + sampler._submit_in_loop = submit + + result = sampler.submit_generation( + 'submission-1', + [{'input_ids': [1]}], + SamplingParams(max_tokens=4), + ) + + assert result == {'submission_id': 'submission-1', 'status': 'running'} + assert not pending.done() + assert len(submitted_coroutines) == 1 + assert sampler.get_generation_status('submission-1')['status'] == 'running' + + responses = [object()] + pending.set_result(responses) + assert sampler.get_generation_status('submission-1')['status'] == 'completed' + assert sampler.collect_generation('submission-1') == responses + assert 'submission-1' not in sampler._generation_submissions + + +def test_generation_keeps_one_response_per_prompt() -> None: + sampler = _bare_sampler() + sampler.template = None + + async def sample_single(feat, _params, **_kwargs): + await asyncio.sleep(0) + return feat['input_ids'][0] + + sampler._sample_single = sample_single + responses = asyncio.run( + sampler._generate_inputs( + [{'input_ids': [10]}, {'input_ids': [20]}], + SamplingParams(max_tokens=4), + adapter_name='', + adapter_path=None, + use_base_model=False, + )) + + assert responses == [10, 20] + + +def test_generation_failure_is_isolated_and_consumed() -> None: + sampler = _bare_sampler() + failed = Future() + failed.set_exception(ValueError('bad prompt')) + sampler._generation_submissions['failed'] = failed + + state = sampler.get_generation_status('failed') + assert state['status'] == 'failed' + assert state['error'] == 'ValueError: bad prompt' + + with pytest.raises(ValueError, match='bad prompt'): + sampler.collect_generation('failed') + assert 'failed' not in sampler._generation_submissions + + +def test_generation_can_be_cancelled_without_waiting() -> None: + sampler = _bare_sampler() + pending = Future() + sampler._generation_submissions['pending'] = pending + + state = sampler.cancel_generation('pending') + + assert state == {'submission_id': 'pending', 'status': 'cancelled'} + assert pending.cancelled() + assert 'pending' not in sampler._generation_submissions + + +def test_all_generations_are_cancelled_on_shutdown() -> None: + sampler = _bare_sampler() + first = Future() + second = Future() + sampler._generation_submissions.update(first=first, second=second) + + state = sampler.cancel_all_generations() + + assert state == {'submissions': 2, 'cancelled': 2} + assert first.cancelled() + assert second.cancelled() + assert sampler._generation_submissions == {} + + +def test_native_prompt_group_sampling_requires_context_manager() -> None: + sampler = _bare_sampler() + sampler.context_manager = None + + with pytest.raises(RuntimeError, match='context_manager is required'): + sampler.sample([], SamplingParams(max_tokens=4)) + + +def test_server_waiter_admits_later_submission_before_first_finishes() -> None: + + class Sampler: + + def __init__(self): + self.futures: dict[str, Future] = {} + self.submission_order = [] + + def submit_generation(self, submission_id, *_args, **_kwargs): + self.submission_order.append(submission_id) + self.futures[submission_id] = Future() + + def get_generation_status(self, submission_id): + future = self.futures[submission_id] + return {'status': 'completed' if future.done() else 'running'} + + def collect_generation(self, submission_id): + return self.futures[submission_id].result() + + def cancel_generation(self, submission_id): + self.futures.pop(submission_id, None) + + sampler = Sampler() + + async def run(): + first = asyncio.create_task( + _await_generation( + sampler, + 'first', + [{'input_ids': [1]}], + SamplingParams(max_tokens=4), + adapter_name='', + adapter_path=None, + )) + second = asyncio.create_task( + _await_generation( + sampler, + 'second', + [{'input_ids': [2]}], + SamplingParams(max_tokens=4), + adapter_name='', + adapter_path=None, + )) + while len(sampler.submission_order) < 2: + await asyncio.sleep(0) + assert not first.done() + sampler.futures['second'].set_result(['short']) + assert await second == ['short'] + assert not first.done() + sampler.futures['first'].set_result(['long']) + assert await first == ['long'] + + asyncio.run(run()) + assert set(sampler.submission_order) == {'first', 'second'} diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py new file mode 100644 index 000000000..9b51e2b11 --- /dev/null +++ b/tests/twinkle_client/test_async_components.py @@ -0,0 +1,165 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle_client.types import ComponentTaskRef, DataRef, DataRowsResponse + + +class _Response: + + def __init__(self, payload, status_code: int = 200): + self._payload = payload + self.status_code = status_code + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise RuntimeError(self.status_code) + + def json(self): + return self._payload + + +def test_remote_task_uses_existing_future_endpoint(monkeypatch) -> None: + from twinkle_client import remote_task as module + + responses = iter([_Response({'type': 'try_again'}), _Response({'value': 7})]) + calls = [] + + def post(url, json_data, **kwargs): + calls.append((url, json_data, kwargs)) + return next(responses) + + monkeypatch.setattr(module, 'http_post', post) + monkeypatch.setattr(module, 'get_base_url', lambda: 'http://server/api/v1') + + task = module.RemoteTask(ComponentTaskRef(request_id='req-1', model_id='adapter')) + assert task.result(timeout=1) == {'value': 7} + assert [call[:2] for call in calls] == [ + ('http://server/api/v1/retrieve_future', {'request_id': 'req-1'}), + ('http://server/api/v1/retrieve_future', {'request_id': 'req-1'}), + ] + assert all(0 < call[2]['timeout'] <= 1 for call in calls) + + +def test_remote_task_async_result_uses_non_blocking_http_polling(monkeypatch) -> None: + import httpx + from twinkle_client import remote_task as module + + responses = iter([_Response({'type': 'try_again'}), _Response({'value': 9})]) + calls = [] + + class AsyncClient: + + def __init__(self, **_kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def post(self, url, *, headers, json, timeout): + calls.append((url, headers, json, timeout)) + return next(responses) + + monkeypatch.setattr(httpx, 'AsyncClient', AsyncClient) + monkeypatch.setattr(module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, '_build_headers', lambda: {'Authorization': 'Bearer token'}) + + task = module.RemoteTask(ComponentTaskRef(request_id='req-async', model_id='adapter')) + assert asyncio.run(task.aresult(timeout=1)) == {'value': 9} + assert [call[2] for call in calls] == [ + {'request_id': 'req-async'}, + {'request_id': 'req-async'}, + ] + assert all(0 < call[3] <= 1 for call in calls) + + +def test_remote_task_sync_timeout_bounds_long_poll(monkeypatch) -> None: + import requests + from twinkle_client import remote_task as module + + observed = [] + + def post(_url, *, json_data, timeout): + observed.append((json_data, timeout)) + raise requests.Timeout('long poll exceeded client deadline') + + monkeypatch.setattr(module, 'http_post', post) + monkeypatch.setattr(module, 'get_base_url', lambda: 'http://server/api/v1') + + task = module.RemoteTask('req-timeout') + with pytest.raises(TimeoutError, match='within 0.1s'): + task.result(timeout=0.1) + assert observed[0][0] == {'request_id': 'req-timeout'} + assert 0 < observed[0][1] <= 0.1 + + +def test_model_component_submits_data_ref_without_control_plane(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + if url.endswith('/create'): + return _Response({}) + return _Response({'request_id': 'req-model', 'model_id': 'session-adapter'}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + monkeypatch.setattr('twinkle_client.remote_task.get_base_url', lambda: 'http://server/api/v1') + + model = module.MultiLoraTransformersModel('ms://base') + model.adapter_name = 'adapter' + ref = DataRef(ref_id='data-1', size=4, fields=['input_ids']) + task = model.submit_forward_backward(ref, advantages=[1, -1, 1, -1]) + + assert task.request_id == 'req-model' + url, body = calls[-1] + assert url.endswith('/model/base/twinkle/submit_forward_backward') + assert body['input_ref'] == ref.model_dump() + assert body['adapter_name'] == 'adapter' + assert body['kwargs']['advantages'] == [1, -1, 1, -1] + + +def test_sampler_component_fetches_and_releases_data_plane_output(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.sampler import vllm_sampler as module + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', lambda **_kwargs: _Response({})) + sampler = module.vLLMSampler('ms://base') + + output_ref = DataRef(ref_id='rollout-1', size=1, fields=['tokens'], kind='rollout') + + class _DoneTask: + + async def aresult(self): + return {'output_ref': output_ref.model_dump()} + + released = [] + monkeypatch.setattr(sampler, 'submit_sample', lambda *_args, **_kwargs: _DoneTask()) + monkeypatch.setattr( + sampler.data_plane, + 'get_batch', + lambda ref: DataRowsResponse( + rows=[{ + 'tokens': [], + 'stop_reason': 'stop', + }], + tags=[{'prompt_index': 0, 'generation_idx': 0}], + ), + ) + monkeypatch.setattr(sampler.data_plane, 'release', lambda ref: released.append(ref)) + + responses = asyncio.run(sampler.asample([{'input_ids': [1]}])) + assert len(responses) == 1 + assert len(responses[0].sequences) == 1 + assert responses[0].sequences[0].tokens == [] + assert released == [output_ref] diff --git a/tests/twinkle_client/test_async_rl_workers.py b/tests/twinkle_client/test_async_rl_workers.py new file mode 100644 index 000000000..7a79cb8a5 --- /dev/null +++ b/tests/twinkle_client/test_async_rl_workers.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle_client.async_rl import Worker, WorkerPipeline + + +class _FunctionWorker(Worker): + + def __init__(self, name, function): + super().__init__(name) + self.function = function + + async def run(self) -> None: + await self.function() + + +def test_worker_pipeline_runs_roles_concurrently() -> None: + producer_started = asyncio.Event() + consumer_started = asyncio.Event() + + async def producer(): + producer_started.set() + await consumer_started.wait() + + async def consumer(): + consumer_started.set() + await producer_started.wait() + + asyncio.run(WorkerPipeline(( + _FunctionWorker('producer', producer), + _FunctionWorker('consumer', consumer), + )).run()) + + +def test_worker_pipeline_cancels_peer_when_one_role_fails() -> None: + waiting = asyncio.Event() + cancelled = asyncio.Event() + + async def peer(): + waiting.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def failure(): + await waiting.wait() + raise RuntimeError('role failed') + + with pytest.raises(RuntimeError, match='role failed'): + asyncio.run(WorkerPipeline(( + _FunctionWorker('peer', peer), + _FunctionWorker('failure', failure), + )).run()) + assert cancelled.is_set() + + +def test_worker_pipeline_rejects_duplicate_role_names() -> None: + async def noop(): + return None + + with pytest.raises(ValueError, match='unique'): + WorkerPipeline(( + _FunctionWorker('same', noop), + _FunctionWorker('same', noop), + )) diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index 1ff69f5f4..31af25757 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -309,6 +309,32 @@ def _build_from_scripts(scripts_spec: List[Dict[str, Any]]): return trajectories, sampler, template +@pytest.mark.asyncio +async def test_async_rollout_passes_explicit_policy_to_sampler() -> None: + trajectories, sampler, template = _build_from_scripts([{ + 'num_tools': 0, + 'terminal': 'stop', + 'logprobs': True, + }]) + calls = [] + + async def asample(inputs, **kwargs): + calls.append(kwargs) + return sampler.sample(inputs, sampling_params=kwargs.get('sampling_params')) + + sampler.asample = asample + rollout = ClientMultiTurnRollout(sampler, template, max_turns=2) + outputs = await rollout.arun( + trajectories, + adapter_name='math-lora', + adapter_uri='twinkle://run/weights/policy-3', + ) + + assert len(outputs) == 1 + assert calls[0]['adapter_name'] == 'math-lora' + assert calls[0]['adapter_uri'] == 'twinkle://run/weights/policy-3' + + # ============================================================================= # Hypothesis strategies # ============================================================================= diff --git a/tests/twinkle_client/test_client_orchestrated_dpo.py b/tests/twinkle_client/test_client_orchestrated_dpo.py new file mode 100644 index 000000000..3d6ac460a --- /dev/null +++ b/tests/twinkle_client/test_client_orchestrated_dpo.py @@ -0,0 +1,116 @@ +import asyncio + +from cookbook.client.async_rl.client_orchestrated_dpo import ( + _extract_ref_outputs, + prepare_dpo_batch, + run_dpo, +) +from twinkle_client.types import DataRef + + +def test_prepare_dpo_batch_interleaves_complete_pairs() -> None: + batch = [ + { + 'pair_id': 'a', + 'positive': {'input_ids': [1, 2], 'labels': [-100, 2]}, + 'negative': {'input_ids': [1, 3], 'labels': [-100, 3]}, + }, + { + 'pair_id': 'b', + 'positive': {'input_ids': [4], 'labels': [4]}, + 'negative': {'input_ids': [5], 'labels': [5]}, + }, + ] + + rows = prepare_dpo_batch(batch) + + assert [row['pair_id'] for row in rows] == ['a', 'a', 'b', 'b'] + assert [row['input_ids'] for row in rows] == [[1, 2], [1, 3], [4], [5]] + + +def test_extract_ref_outputs_unwraps_data_plane_row() -> None: + ref_outputs = _extract_ref_outputs( + {'output_ref': {'ref_id': 'unused'}}, + [{'result': {'logps': [[-1.0, -2.0], [-3.0, -4.0]], 'logits': None}}], + ) + + assert ref_outputs == {'logps': [[-1.0, -2.0], [-3.0, -4.0]]} + + +def test_dpo_roles_overlap_reference_and_training(monkeypatch) -> None: + import cookbook.client.async_rl.client_orchestrated_dpo as module + + monkeypatch.setattr(module, 'MAX_STEPS', 2) + first_train_started = asyncio.Event() + events = [] + + class FakeDataPlane: + + def __init__(self): + self.rows = {} + self.released = [] + + async def aput(self, rows, *, kind, tags=None): + ref = DataRef( + ref_id=f'{kind}-{len(self.rows)}', + size=len(rows), + fields=list(rows[0]), + kind=kind, + ) + self.rows[ref.ref_id] = rows + return ref + + async def aappend(self, ref, updates, *, tags=None): + self.rows[ref.ref_id] = [ + {**row, **update} + for row, update in zip(self.rows[ref.ref_id], updates) + ] + return ref.model_copy(update={'fields': list(self.rows[ref.ref_id][0])}) + + async def arelease(self, ref): + self.released.append(ref.ref_id) + + class FakeModel: + + def __init__(self): + self.references = 0 + self.steps = 0 + self.forward_backward_kwargs = [] + + async def submit_forward_only(self, ref, **_kwargs): + self.references += 1 + name = f'reference-{self.references}' + events.append(f'{name}-start') + if self.references == 2: + await first_train_started.wait() + events.append(f'{name}-done') + return {'logps': [[-0.1]] * ref.size} + + async def submit_forward_backward(self, _ref, **kwargs): + self.forward_backward_kwargs.append(kwargs) + events.append('train-start') + first_train_started.set() + + async def submit_clip_grad_and_step(self, **_kwargs): + self.steps += 1 + + async def submit_save(self, name, **_kwargs): + return {'twinkle_path': name} + + batches = [ + [{'pair_id': 'a', 'positive': {'input_ids': [1]}, 'negative': {'input_ids': [2]}}], + [{'pair_id': 'b', 'positive': {'input_ids': [3]}, 'negative': {'input_ids': [4]}}], + ] + model = FakeModel() + data_plane = FakeDataPlane() + + saved = asyncio.run(run_dpo(batches, model, data_plane)) + + assert events.index('train-start') < events.index('reference-2-done') + assert model.steps == 2 + assert model.forward_backward_kwargs == [ + {'ref_outputs': {'logps': [[-0.1], [-0.1]]}}, + {'ref_outputs': {'logps': [[-0.1], [-0.1]]}}, + ] + assert saved == {'twinkle_path': 'dpo-policy-2'} + assert len(data_plane.released) == 2 diff --git a/tests/twinkle_client/test_client_orchestrated_grpo.py b/tests/twinkle_client/test_client_orchestrated_grpo.py new file mode 100644 index 000000000..6b9aa4dd2 --- /dev/null +++ b/tests/twinkle_client/test_client_orchestrated_grpo.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import sys +from pathlib import Path + +from twinkle_client.types import DataRef, DataRowsResponse + + +MODULE_PATH = ( + Path(__file__).parents[2] / 'cookbook' / 'client' / 'async_rl' / 'client_orchestrated_grpo.py' +) + + +def _load_module(): + spec = importlib.util.spec_from_file_location('client_orchestrated_grpo', MODULE_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_rollout_and_train_overlap_with_fifo_policy_publication(monkeypatch) -> None: + module = _load_module() + monkeypatch.setattr(module, 'BATCH_SIZE', 2) + monkeypatch.setattr(module, 'NUM_GENERATIONS', 2) + monkeypatch.setattr(module, 'TRAIN_MINI_BATCH_SIZE', 2) + monkeypatch.setattr(module, 'MAX_STALENESS', 1) + monkeypatch.setattr(module, 'MAX_PARTITIONS', 3) + + first_train_started = asyncio.Event() + rollout_snapshots = [] + events = [] + + async def fake_rollout(_sampler, prompt, policy, _semaphore, _group_id): + name = prompt['name'] + rollout_snapshots.append((name, policy.version, policy.adapter_uri)) + events.append(f'rollout-start:{name}') + if name == 'p0-g1': + await first_train_started.wait() + events.append(f'rollout-done:{name}') + return DataRef( + ref_id=name, + size=module.NUM_GENERATIONS, + fields=['new_input_feature', 'logprobs'], + kind='rollout', + ) + + monkeypatch.setattr(module, 'rollout_group', fake_rollout) + monkeypatch.setattr(module, 'GSM8KAccuracyReward', lambda: lambda rows: [1.0] * len(rows)) + monkeypatch.setattr( + module, + 'GRPOAdvantage', + lambda: lambda rewards, **_kwargs: [1.0, -1.0], + ) + + class FakeModel: + def __init__(self): + self.saved = [] + self.steps = 0 + self.forward_backward_kwargs = [] + + async def submit_save(self, name): + self.saved.append(name) + return {'twinkle_path': f'/checkpoints/{name}'} + + async def submit_forward_backward(self, _ref, **kwargs): + self.forward_backward_kwargs.append(kwargs) + events.append('train') + first_train_started.set() + + async def submit_clip_grad_and_step(self, **_kwargs): + self.steps += 1 + + class FakeDataPlane: + def __init__(self): + self.released = [] + + async def aget_batch(self, ref): + return DataRowsResponse( + rows=[{ + 'new_input_feature': {'name': f'{ref.ref_id}-{index}'}, + 'logprobs': [[(0, 0.0)]], + } for index in range(ref.size)], + tags=[{'group_id': ref.ref_id, 'generation_idx': index} for index in range(ref.size)], + ) + + async def aappend(self, ref, rows, *, tags): + return ref.model_copy(update={'fields': [*ref.fields, *rows[0]]}) + + async def aput(self, rows, *, kind, tags=None): + return DataRef(ref_id=f'{kind}-{id(rows)}', size=len(rows), fields=list(rows[0]), kind=kind) + + async def arelease(self, ref): + self.released.append(ref) + + async def run(): + model = FakeModel() + data_plane = FakeDataPlane() + batches = [ + [{'name': 'p0-g0'}, {'name': 'p0-g1'}], + [{'name': 'p1-g0'}, {'name': 'p1-g1'}], + [{'name': 'p2-g0'}, {'name': 'p2-g1'}], + ] + await module.run_grpo(batches, model, object(), data_plane) + return model, data_plane + + model, data_plane = asyncio.run(run()) + + assert events.index('train') < events.index('rollout-done:p0-g1') + assert model.saved == ['policy-0', 'policy-1', 'policy-2', 'policy-3'] + assert model.steps == 6 + assert all('old_logps' in kwargs and 'advantages' in kwargs + for kwargs in model.forward_backward_kwargs) + assert len(data_plane.released) == 6 + + snapshots = {name: (version, uri) for name, version, uri in rollout_snapshots} + assert snapshots['p0-g0'] == (0, '/checkpoints/policy-0') + assert snapshots['p1-g0'] == (0, '/checkpoints/policy-0') + assert snapshots['p2-g0'][0] in (1, 2) + assert snapshots['p2-g0'][1] == f'/checkpoints/policy-{snapshots["p2-g0"][0]}' + + +def test_younger_rollout_failure_stops_admission(monkeypatch) -> None: + module = _load_module() + monkeypatch.setattr(module, 'BATCH_SIZE', 1) + monkeypatch.setattr(module, 'NUM_GENERATIONS', 1) + monkeypatch.setattr(module, 'TRAIN_MINI_BATCH_SIZE', 1) + monkeypatch.setattr(module, 'MAX_STALENESS', 1) + monkeypatch.setattr(module, 'MAX_PARTITIONS', 3) + + started = [] + + async def fake_rollout(_sampler, prompt, _policy, _semaphore, _group_id): + name = prompt['name'] + started.append(name) + if name == 'p1': + raise RuntimeError('rollout failed') + return DataRef( + ref_id=name, + size=1, + fields=['new_input_feature', 'logprobs'], + kind='rollout', + ) + + monkeypatch.setattr(module, 'rollout_group', fake_rollout) + monkeypatch.setattr(module, 'GSM8KAccuracyReward', lambda: lambda rows: [1.0]) + monkeypatch.setattr(module, 'GRPOAdvantage', lambda: lambda rewards, **_kwargs: [1.0]) + + class FakeModel: + def __init__(self): + self.saved = [] + + async def submit_save(self, name): + self.saved.append(name) + return {'twinkle_path': name} + + async def submit_forward_backward(self, _ref, **_kwargs): + return None + + async def submit_clip_grad_and_step(self, **_kwargs): + return None + + class FakeDataPlane: + async def aget_batch(self, ref): + return DataRowsResponse( + rows=[{ + 'new_input_feature': {'name': ref.ref_id}, + 'logprobs': [[(0, 0.0)]], + }], + tags=[{'group_id': ref.ref_id, 'generation_idx': 0}], + ) + + async def aappend(self, ref, rows, *, tags): + return ref.model_copy(update={'fields': [*ref.fields, *rows[0]]}) + + async def aput(self, rows, *, kind, tags=None): + return DataRef(ref_id=kind, size=len(rows), fields=list(rows[0]), kind=kind) + + async def arelease(self, _ref): + return None + + async def run(): + model = FakeModel() + try: + await module.run_grpo( + [[{'name': 'p0'}], [{'name': 'p1'}], [{'name': 'p2'}]], + model, + object(), + FakeDataPlane(), + ) + except RuntimeError as error: + assert str(error) == 'rollout failed' + else: + raise AssertionError('expected the younger rollout failure') + return model + + model = asyncio.run(run()) + assert started == ['p0', 'p1'] + assert model.saved == ['policy-0'] diff --git a/tests/twinkle_client/test_data_plane_async.py b/tests/twinkle_client/test_data_plane_async.py new file mode 100644 index 000000000..fd616f347 --- /dev/null +++ b/tests/twinkle_client/test_data_plane_async.py @@ -0,0 +1,102 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from twinkle_client.data_plane import DataPlaneClient +from twinkle_client.types import DataRef, DataRowsResponse + + +def test_async_convenience_methods_delegate_to_sync_operations(monkeypatch) -> None: + client = DataPlaneClient('http://server/data-plane') + original_ref = DataRef(ref_id='data-1', size=1, fields=['value']) + appended_ref = DataRef(ref_id='data-1', size=2, fields=['value']) + calls = [] + caller_thread = threading.get_ident() + + def put(rows, *, kind='data'): + calls.append(('put', rows, kind, threading.get_ident())) + return original_ref + + def get(ref, *, fields=None): + calls.append(('get', ref, fields, threading.get_ident())) + return [{'value': 1}] + + def append(ref, rows): + calls.append(('append', ref, rows, threading.get_ident())) + return appended_ref + + def release(ref): + calls.append(('release', ref, threading.get_ident())) + + monkeypatch.setattr(client, 'put', put) + monkeypatch.setattr(client, 'get', get) + monkeypatch.setattr(client, 'append', append) + monkeypatch.setattr(client, 'release', release) + + async def run(): + assert await client.aput([{'value': 1}], kind='rollout') == original_ref + assert await client.aget(original_ref, fields=['value']) == [{'value': 1}] + assert await client.aappend(original_ref, [{'value': 2}]) == appended_ref + assert await client.arelease(appended_ref) is None + + asyncio.run(run()) + + assert [call[:-1] for call in calls] == [ + ('put', [{'value': 1}], 'rollout'), + ('get', original_ref, ['value']), + ('append', original_ref, [{'value': 2}]), + ('release', appended_ref), + ] + assert all(call[-1] != caller_thread for call in calls) + + +def test_async_convenience_method_propagates_sync_error(monkeypatch) -> None: + client = DataPlaneClient('http://server/data-plane') + + def fail(_rows, *, kind='data'): + raise RuntimeError(f'put failed for {kind}') + + monkeypatch.setattr(client, 'put', fail) + + with pytest.raises(RuntimeError, match='put failed for rollout'): + asyncio.run(client.aput([], kind='rollout')) + + +def test_async_tagged_methods_and_batch_read_delegate_to_sync_operations(monkeypatch) -> None: + client = DataPlaneClient('http://server/data-plane') + ref = DataRef(ref_id='data-1', size=1, fields=['value']) + tags = [{'group_id': 'group-1'}] + calls = [] + + def put(rows, *, kind='data', tags=None): + calls.append(('put', rows, kind, tags)) + return ref + + def get_batch(value, *, fields=None): + calls.append(('get_batch', value, fields)) + return DataRowsResponse(rows=[{'value': 1}], tags=tags) + + def append(value, rows, *, tags=None): + calls.append(('append', value, rows, tags)) + return value + + monkeypatch.setattr(client, 'put', put) + monkeypatch.setattr(client, 'get_batch', get_batch) + monkeypatch.setattr(client, 'append', append) + + async def run(): + assert await client.aput([{'value': 1}], tags=tags) == ref + assert await client.aget_batch(ref) == DataRowsResponse(rows=[{'value': 1}], tags=tags) + assert await client.aappend(ref, [{'reward': 1.0}], tags=tags) == ref + + asyncio.run(run()) + + assert calls == [ + ('put', [{'value': 1}], 'data', tags), + ('get_batch', ref, None), + ('append', ref, [{'reward': 1.0}], tags), + ] From f96b037a052095af83f1f459db36a1f00f8cc58b Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Mon, 10 Aug 2026 11:17:15 +0800 Subject: [PATCH 02/20] feat(server): add local-model client-server config --- .../transformer/server_config_local.yaml | 159 ++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 cookbook/client/server/transformer/server_config_local.yaml diff --git a/cookbook/client/server/transformer/server_config_local.yaml b/cookbook/client/server/transformer/server_config_local.yaml new file mode 100644 index 000000000..32e9d9042 --- /dev/null +++ b/cookbook/client/server/transformer/server_config_local.yaml @@ -0,0 +1,159 @@ +# Twinkle Server Configuration - local Qwen3.5-4B for client-orchestrated async RL + +# Set the absolute Hugging Face-compatible model directory before loading this file: +# export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +# +# The public HTTP model name remains Qwen/Qwen3.5-4B. Both the training Model +# and the vLLM Sampler load the same local directory, so neither component +# downloads the base model independently. + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: file + file_path: /tmp/twinkle_state.json + +applications: + + # Gateway and Future query endpoint. + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 3 + supported_models: + - Qwen/Qwen3.5-4B + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 128 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "0" + + # TransferQueue-backed DataRef service. + - name: data-plane + route_prefix: /api/v1/data-plane + import_path: data_plane + args: + config: + backend: + SimpleStorage: + num_data_storage_units: 2 + deployments: + - name: DataPlaneManagement + num_replicas: 1 + ray_actor_options: + num_cpus: 1 + + # One GPU hosts the training base model and multiple LoRA adapters. + - name: models-Qwen3.5-4B + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + backend: transformers + model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + max_loras: 8 + max_length: 10240 + data_plane_url: http://127.0.0.1:8000/api/v1/data-plane + nproc_per_node: 1 + device_group: + name: model + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + queue_config: + rps_limit: 100 + tps_limit: 100000 + adapter_config: + adapter_timeout: 30 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 16 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "0" + + # A second GPU hosts vLLM and loads the same local base model. + - name: sampler-Qwen3.5-4B + route_prefix: /api/v1/sampler/Qwen/Qwen3.5-4B + import_path: sampler + args: + model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + data_plane_url: http://127.0.0.1:8000/api/v1/data-plane + nproc_per_node: 1 + sampler_type: vllm + engine_args: + max_model_len: 4096 + gpu_memory_utilization: 0.5 + enable_lora: true + max_loras: 8 + logprobs_mode: processed_logprobs + device_group: + name: sampler + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + queue_config: + rps_limit: 100 + tps_limit: 100000 + deployments: + - name: SamplerManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 16 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "0" + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: model + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 128 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "0" From 212ddcf401f0fe269ae646e7f76befee1b88163f Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Mon, 10 Aug 2026 11:27:33 +0800 Subject: [PATCH 03/20] wip --- .../client/async_rl/client_orchestrated_dpo.py | 8 ++++++-- .../client/async_rl/client_orchestrated_grpo.py | 10 +++++++--- .../server/transformer/server_config_local.yaml | 17 +++++++++-------- 3 files changed, 22 insertions(+), 13 deletions(-) diff --git a/cookbook/client/async_rl/client_orchestrated_dpo.py b/cookbook/client/async_rl/client_orchestrated_dpo.py index d9c822188..e94ba60da 100644 --- a/cookbook/client/async_rl/client_orchestrated_dpo.py +++ b/cookbook/client/async_rl/client_orchestrated_dpo.py @@ -19,6 +19,10 @@ BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') MODEL_ID = f'ms://{BASE_MODEL}' +TEMPLATE_CLS = os.environ.get( + 'TWINKLE_TEMPLATE_CLS', + 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', +) DATASET_ID = os.environ.get('TWINKLE_DPO_DATASET_ID', 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji') ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-dpo') MAX_STEPS = int(os.environ.get('TWINKLE_MAX_STEPS', '100')) @@ -29,7 +33,7 @@ def create_dataset() -> Dataset: """Load and encode preference pairs in the client process.""" dataset = Dataset(DatasetMeta(DATASET_ID, data_slice=range(MAX_STEPS * BATCH_SIZE))) - dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=MAX_LENGTH) + dataset.set_template(TEMPLATE_CLS, model_id=MODEL_ID, max_length=MAX_LENGTH) dataset.map(EmojiDPOProcessor, init_args={'system': 'You are a helpful assistant.'}) dataset.encode() return dataset @@ -216,7 +220,7 @@ async def train() -> None: ADAPTER_NAME, LoraConfig(target_modules='all-linear', r=8, lora_alpha=32, lora_dropout=0.05), ) - model.set_template('Qwen3_5Template', model_id=MODEL_ID) + model.set_template(TEMPLATE_CLS, model_id=MODEL_ID) model.set_processor('InputProcessor', padding_side='right') model.set_loss('DPOLoss', beta=0.1, loss_type='sigmoid', reference_free=False) model.add_metric('DPOMetric', beta=0.1) diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py index b5892e193..490ac96b1 100644 --- a/cookbook/client/async_rl/client_orchestrated_grpo.py +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -23,6 +23,10 @@ BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') MODEL_ID = f'ms://{BASE_MODEL}' +TEMPLATE_CLS = os.environ.get( + 'TWINKLE_TEMPLATE_CLS', + 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', +) ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-grpo') MAX_PARTITIONS = int(os.environ.get('TWINKLE_MAX_PARTITIONS', '100')) MAX_STALENESS = int(os.environ.get('TWINKLE_MAX_STALENESS', '2')) @@ -118,7 +122,7 @@ async def publish(self, partition: _RolloutPartition, policy: _Policy) -> None: def create_dataset() -> Dataset: dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) - dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=2048, enable_thinking=False) + dataset.set_template(TEMPLATE_CLS, model_id=MODEL_ID, max_length=2048, enable_thinking=False) dataset.map(GSM8KProcessor(system='Put the final answer within \\boxed{}.')) dataset.encode(add_generation_prompt=True) return dataset @@ -439,8 +443,8 @@ async def train() -> None: model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) model.set_optimizer('AdamW', lr=2e-5) model.set_processor('InputProcessor', padding_free=True) - model.set_template('Qwen3_5Template', model_id=MODEL_ID) - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID) + model.set_template(TEMPLATE_CLS, model_id=MODEL_ID) + sampler.set_template(TEMPLATE_CLS, model_id=MODEL_ID) dataloader = DataLoader(dataset=create_dataset(), batch_size=BATCH_SIZE, num_workers=0) await run_grpo(dataloader, model, sampler, data_plane) diff --git a/cookbook/client/server/transformer/server_config_local.yaml b/cookbook/client/server/transformer/server_config_local.yaml index 32e9d9042..f3aae2203 100644 --- a/cookbook/client/server/transformer/server_config_local.yaml +++ b/cookbook/client/server/transformer/server_config_local.yaml @@ -1,9 +1,10 @@ -# Twinkle Server Configuration - local Qwen3.5-4B for client-orchestrated async RL +# Twinkle Server Configuration - local Qwen3-4B for client-orchestrated async RL # Set the absolute Hugging Face-compatible model directory before loading this file: -# export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +# export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3-4B +# export TWINKLE_MODEL_ID=Qwen/Qwen3-4B # in every client process # -# The public HTTP model name remains Qwen/Qwen3.5-4B. Both the training Model +# The public HTTP model name remains Qwen/Qwen3-4B. Both the training Model # and the vLLM Sampler load the same local directory, so neither component # downloads the base model independently. @@ -31,7 +32,7 @@ applications: server_config: per_token_model_limit: 3 supported_models: - - Qwen/Qwen3.5-4B + - Qwen/Qwen3-4B deployments: - name: TinkerCompatServer max_ongoing_requests: 50 @@ -61,8 +62,8 @@ applications: num_cpus: 1 # One GPU hosts the training base model and multiple LoRA adapters. - - name: models-Qwen3.5-4B - route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + - name: models-Qwen3-4B + route_prefix: /api/v1/model/Qwen/Qwen3-4B import_path: model args: backend: transformers @@ -97,8 +98,8 @@ applications: TWINKLE_FAIL_FAST: "0" # A second GPU hosts vLLM and loads the same local base model. - - name: sampler-Qwen3.5-4B - route_prefix: /api/v1/sampler/Qwen/Qwen3.5-4B + - name: sampler-Qwen3-4B + route_prefix: /api/v1/sampler/Qwen/Qwen3-4B import_path: sampler args: model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} From e7a5a7e10a67d31835024d21f3bf539e82ccbf94 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 11 Aug 2026 11:26:30 +0800 Subject: [PATCH 04/20] docs: update async RL cookbook API names and DPO template handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename low-level component methods to match new API (e.g., `submit_sample` → `sample_to_data_plane`, `submit_forward_only` → `forward_only`) - Clarify that token tensors and log-probabilities remain server-side in async RL flow - Add `TWINKLE_TEMPLATE_MODEL_ID` environment variable for DPO template configuration - Remove `_extract_ref_outputs` helper as it's no longer needed with updated API --- cookbook/client/async_rl/README.md | 24 +- .../async_rl/client_orchestrated_dpo.py | 70 ++--- .../async_rl/client_orchestrated_grpo.py | 145 ++++----- .../server/transformer/server_config.yaml | 2 +- .../transformer/server_config_local.yaml | 3 +- cookbook/client/twinkle/short_math_grpo.py | 14 +- src/twinkle/infra/_ray/ray_helper.py | 18 +- src/twinkle/infra/_ray/resource_manager.py | 23 +- src/twinkle/model/multi_lora.py | 7 +- src/twinkle/server/config/application_spec.py | 2 +- src/twinkle/server/data_plane/proxy.py | 25 +- src/twinkle/server/data_plane/store.py | 15 +- src/twinkle/server/model/twinkle_handlers.py | 282 +++++++----------- src/twinkle/server/sampler/app.py | 18 +- .../server/sampler/twinkle_handlers.py | 191 ++++++------ src/twinkle/utils/nccl_safe.py | 4 + .../async_rl/vllm_sampler_tq.py | 7 +- src/twinkle_client/__init__.py | 3 +- .../model/multi_lora_transformers.py | 163 +++++----- src/twinkle_client/remote_task.py | 94 ------ src/twinkle_client/sampler/vllm_sampler.py | 72 ++--- src/twinkle_client/types/__init__.py | 7 +- src/twinkle_client/types/component.py | 50 +--- src/twinkle_client/types/model.py | 34 ++- tests/model/test_micro_batch.py | 12 + .../test_multi_lora_target_parameters.py | 30 +- .../server/contract/client_api_baseline.json | 52 +--- tests/server/data_plane/test_proxy.py | 8 +- tests/server/data_plane/test_store.py | 7 + .../server/model/test_twinkle_async_inputs.py | 51 +++- tests/server/sampler/test_mock_sampler.py | 10 +- .../server/sampler/test_twinkle_async_rows.py | 97 +++++- .../test_vllm_sampler_tq_generation.py | 54 +++- tests/twinkle_client/test_async_components.py | 172 ++++------- .../test_client_orchestrated_dpo.py | 26 +- .../test_client_orchestrated_grpo.py | 57 ++-- 36 files changed, 868 insertions(+), 981 deletions(-) delete mode 100644 src/twinkle_client/remote_task.py diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md index f44653707..12d7776f6 100644 --- a/cookbook/client/async_rl/README.md +++ b/cookbook/client/async_rl/README.md @@ -32,22 +32,22 @@ python cookbook/client/async_rl/client_orchestrated_grpo.py The training loop composes only the low-level component methods: -- `sampler.submit_sample(...)` / `sampler.asample(...)` -- `model.submit_forward_only(...)` -- `model.submit_forward_backward(...)` -- `model.submit_clip_grad_and_step(...)` -- `model.submit_save(...)` +- `sampler.sample_to_data_plane(...)` / `sampler.asample_to_data_plane(...)` +- `model.forward_only(...)` +- `model.forward_backward(...)` +- `model.clip_grad_and_step(...)` +- `model.save(...)` - `data_plane.put/get/append/release(...)` and `aput/aget/aget_batch/aappend/arelease(...)` There is no additional RL runtime or orchestration protocol. The GRPO example -uses `submit_sample()` so the Sampler's output `DataRef` remains in TQ. Each +uses `asample_to_data_plane()` so the Sampler's output `DataRef` remains in TQ. Each generation is one row tagged with its group, generation index, rollout policy, -and status. The Advantage worker reads those rows and appends reward, advantage, -old log-probability, and updated status to the same keys. The Trainer consumes -the resulting `DataRef` and releases it in a local `finally` block. `asample()` -remains available when an algorithm prefers materialized response objects and -does not need to retain the intermediate TQ rows. +and status. The Advantage worker reads only decoded completions and appends +reward and advantage to the same keys. Token tensors and sampled log-probabilities +remain server-side. The Trainer passes one or more `DataRef` values to +`forward_backward()` and releases them in a local `finally` block. `asample()` +remains the materialized-response convenience API. - `ClientMultiTurnRollout.arun()` keeps tool calls and Reward computation in the client and accepts an explicit `adapter_uri` policy snapshot. @@ -55,7 +55,7 @@ does not need to retain the intermediate TQ rows. `_RolloutPartition` is a private client record, not a server resource or SDK API. The local FIFO limits live DataLoader batches before rollout, ready prompt groups immediately use the Model primitives above, and the client calls -`submit_save()` once after a whole batch has trained. `WorkerPipeline` only +`save()` once after a whole batch has trained. `WorkerPipeline` only starts, joins, and fail-fast cancels concrete roles; queues and algorithm state remain ordinary client Python code. diff --git a/cookbook/client/async_rl/client_orchestrated_dpo.py b/cookbook/client/async_rl/client_orchestrated_dpo.py index e94ba60da..f5118a83a 100644 --- a/cookbook/client/async_rl/client_orchestrated_dpo.py +++ b/cookbook/client/async_rl/client_orchestrated_dpo.py @@ -19,6 +19,7 @@ BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') MODEL_ID = f'ms://{BASE_MODEL}' +TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) TEMPLATE_CLS = os.environ.get( 'TWINKLE_TEMPLATE_CLS', 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', @@ -33,7 +34,7 @@ def create_dataset() -> Dataset: """Load and encode preference pairs in the client process.""" dataset = Dataset(DatasetMeta(DATASET_ID, data_slice=range(MAX_STEPS * BATCH_SIZE))) - dataset.set_template(TEMPLATE_CLS, model_id=MODEL_ID, max_length=MAX_LENGTH) + dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=MAX_LENGTH) dataset.map(EmojiDPOProcessor, init_args={'system': 'You are a helpful assistant.'}) dataset.encode() return dataset @@ -49,21 +50,6 @@ def prepare_dpo_batch(batch: list[dict[str, Any]]) -> list[dict[str, Any]]: return json_safe(rows) -def _extract_ref_outputs(result: Any, rows: list[dict[str, Any]] | None = None) -> dict[str, Any]: - """Normalize an async forward-only result into the DPOLoss input shape.""" - payload: Any = rows - if payload is None and isinstance(result, dict): - payload = result.get('result', result) - if isinstance(payload, list) and payload and all( - isinstance(row, dict) and row.get('logps') is not None for row in payload): - return {'logps': [row['logps'] for row in payload]} - if isinstance(payload, list) and len(payload) == 1 and isinstance(payload[0], dict): - payload = payload[0].get('result', payload[0]) - if not isinstance(payload, dict) or payload.get('logps') is None: - raise RuntimeError('reference forward did not return per-token logps') - return {'logps': payload['logps']} - - async def _put_rows(data_plane, rows, *, kind, tags): try: return await data_plane.aput(rows, kind=kind, tags=tags) @@ -82,6 +68,10 @@ async def _submit(method, *args, **kwargs): return task +def _response_payload(value: Any) -> dict[str, Any]: + return value if isinstance(value, dict) else value.model_dump() + + class _DatasetWorker(Worker): def __init__(self, dataloader, data_plane, output): @@ -122,10 +112,9 @@ async def run(self) -> None: class _ReferenceWorker(Worker): - def __init__(self, model, data_plane, source, output): + def __init__(self, model, source, output): super().__init__('reference') self.model = model - self.data_plane = data_plane self.source = source self.output = output @@ -135,9 +124,8 @@ async def run(self) -> None: if item is None: await self.output.put(None) return - ref = item - ref_outputs = await _reference_forward(self.model, self.data_plane, ref) - await self.output.put((ref, ref_outputs)) + ref = await _reference_forward(self.model, item) + await self.output.put(ref) class _TrainerWorker(Worker): @@ -154,20 +142,20 @@ async def run(self) -> None: while True: item = await self.source.get() if item is None: - self.saved = await _submit( - self.model.submit_save, + self.saved = _response_payload(await _submit( + self.model.save, f'dpo-policy-{self.completed_steps}', save_optimizer=True, - ) + )) return - ref, ref_outputs = item + ref = item try: await _submit( - self.model.submit_forward_backward, + self.model.forward_backward, ref, - ref_outputs=ref_outputs, + kwarg_fields={'ref_outputs.logps': 'ref_logps'}, ) - await _submit(self.model.submit_clip_grad_and_step, max_grad_norm=1.0) + await _submit(self.model.clip_grad_and_step, max_grad_norm=1.0) finally: await self.data_plane.arelease(ref) self.completed_steps += 1 @@ -175,20 +163,16 @@ async def run(self) -> None: async def _reference_forward( model: MultiLoraTransformersModel, - data_plane: DataPlaneClient, batch_ref: DataRef, -) -> dict[str, Any]: - """Run the frozen base model and materialize its DataPlane result.""" - result = await _submit(model.submit_forward_only, batch_ref, disable_lora=True) - if not isinstance(result, dict) or not result.get('output_ref'): - return _extract_ref_outputs(result) - - output_ref = DataRef(**result['output_ref']) - try: - rows = await data_plane.aget(output_ref) - return _extract_ref_outputs(result, rows) - finally: - await data_plane.arelease(output_ref) +) -> DataRef: + """Run the frozen base model and append reference logps to the same rows.""" + return await _submit( + model.forward_only, + batch_ref, + disable_lora=True, + output_ref=batch_ref, + output_fields={'logps': 'ref_logps'}, + ) async def run_dpo( @@ -202,7 +186,7 @@ async def run_dpo( trainer = _TrainerWorker(model, data_plane, reference_ready) await WorkerPipeline(( _DatasetWorker(dataloader, data_plane, preference_ready), - _ReferenceWorker(model, data_plane, preference_ready, reference_ready), + _ReferenceWorker(model, preference_ready, reference_ready), trainer, )).run() return trainer.saved @@ -220,7 +204,7 @@ async def train() -> None: ADAPTER_NAME, LoraConfig(target_modules='all-linear', r=8, lora_alpha=32, lora_dropout=0.05), ) - model.set_template(TEMPLATE_CLS, model_id=MODEL_ID) + model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) model.set_processor('InputProcessor', padding_side='right') model.set_loss('DPOLoss', beta=0.1, loss_type='sigmoid', reference_free=False) model.add_metric('DPOMetric', beta=0.1) diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py index 490ac96b1..e59fa01c3 100644 --- a/cookbook/client/async_rl/client_orchestrated_grpo.py +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -23,6 +23,7 @@ BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') MODEL_ID = f'ms://{BASE_MODEL}' +TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) TEMPLATE_CLS = os.environ.get( 'TWINKLE_TEMPLATE_CLS', 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', @@ -50,6 +51,7 @@ class _RolloutPartition: partition_id: int policy: _Policy + prompts: list[dict[str, Any]] rollouts: list[asyncio.Task[Any]] ready: asyncio.Queue['_ReadyGroup'] = field(default_factory=asyncio.Queue) @@ -57,16 +59,14 @@ class _RolloutPartition: @dataclass class _ReadyGroup: group_index: int - rows: list[dict[str, Any]] - tags: list[dict[str, Any]] ref: Any - forward_kwargs: dict[str, Any] @dataclass class _RolloutResult: partition: _RolloutPartition group_index: int + prompt: dict[str, Any] ref: Any @@ -122,7 +122,7 @@ async def publish(self, partition: _RolloutPartition, policy: _Policy) -> None: def create_dataset() -> Dataset: dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) - dataset.set_template(TEMPLATE_CLS, model_id=MODEL_ID, max_length=2048, enable_thinking=False) + dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=2048, enable_thinking=False) dataset.map(GSM8KProcessor(system='Put the final answer within \\boxed{}.')) dataset.encode(add_generation_prompt=True) return dataset @@ -137,8 +137,8 @@ async def rollout_group( ) -> Any: """Submit one GRPO group and keep its sample-level TQ DataRef alive.""" async with semaphore: - result = await _submit( - sampler.submit_sample, + return await _submit( + sampler.asample_to_data_plane, [prompt], adapter_name=ADAPTER_NAME, adapter_uri=policy.adapter_uri, @@ -152,10 +152,6 @@ async def rollout_group( }, num_samples=NUM_GENERATIONS, ) - if not isinstance(result, dict) or not result.get('output_ref'): - raise RuntimeError('async sampler did not return a DataPlane output_ref') - from twinkle_client.types import DataRef - return DataRef(**result['output_ref']) def start_partition( @@ -166,9 +162,11 @@ def start_partition( semaphore: asyncio.Semaphore, ) -> _RolloutPartition: """Capture the snapshot before submitting any rollout in this partition.""" + prompts = json_safe(batch) return _RolloutPartition( partition_id=partition_id, policy=policy, + prompts=prompts, rollouts=[ asyncio.create_task( rollout_group( @@ -178,21 +176,11 @@ def start_partition( semaphore, f'partition-{partition_id}/group-{group_index}', )) - for group_index, prompt in enumerate(json_safe(batch)) + for group_index, prompt in enumerate(prompts) ], ) -async def _put_rows(data_plane: DataPlaneClient, rows, *, kind: str, tags): - """Pass native sample tags while remaining friendly to small cookbook fakes.""" - try: - return await data_plane.aput(rows, kind=kind, tags=tags) - except TypeError as error: - if 'tags' not in str(error): - raise - return await data_plane.aput(rows, kind=kind) - - async def _submit(method, *args, **kwargs): if inspect.iscoroutinefunction(method): return await method(*args, **kwargs) @@ -202,6 +190,12 @@ async def _submit(method, *args, **kwargs): return task +def _checkpoint_path(saved: Any) -> str: + if isinstance(saved, dict): + return str(saved['twinkle_path']) + return str(saved.twinkle_path) + + class _RolloutWorker(Worker): def __init__(self, dataloader, sampler, state: _GRPOState, output: asyncio.Queue, semaphore): @@ -215,7 +209,8 @@ def __init__(self, dataloader, sampler, state: _GRPOState, output: asyncio.Queue async def _collect(self, partition, group_index, task): try: ref = await task - await self.output.put(_RolloutResult(partition, group_index, ref)) + await self.output.put( + _RolloutResult(partition, group_index, partition.prompts[group_index], ref)) except BaseException as error: await self.state.fail(error) raise @@ -268,52 +263,28 @@ async def run(self) -> None: return group_id = f'partition-{result.partition.partition_id}/group-{result.group_index}' try: - batch = await self.data_plane.aget_batch(result.ref) - if len(batch.rows) != NUM_GENERATIONS: + rows = await self.data_plane.aget(result.ref, fields=['decoded']) + if len(rows) != NUM_GENERATIONS: raise RuntimeError( f'group {group_id} expected {NUM_GENERATIONS} generations, ' - f'got {len(batch.rows)}') - features = [row.get('new_input_feature') for row in batch.rows] - if not all(isinstance(feature, dict) for feature in features): - raise RuntimeError(f'group {group_id} has no trainable new_input_feature') - old_logps = [ - [position[0][1] for position in (row.get('logprobs') or [])] - for row in batch.rows - ] - rewards = await asyncio.to_thread(GSM8KAccuracyReward(), features) + f'got {len(rows)}') + trajectories = [] + for row in rows: + messages = list(result.prompt.get('messages') or []) + messages.append({'role': 'assistant', 'content': row.get('decoded') or ''}) + trajectories.append({**result.prompt, 'messages': messages}) + rewards = await asyncio.to_thread(GSM8KAccuracyReward(), trajectories) advantages = await asyncio.to_thread( GRPOAdvantage(), rewards, num_generations=NUM_GENERATIONS) - train_rows = [dict(feature) for feature in features] - if batch.tags and len(batch.tags) != len(train_rows): - raise RuntimeError( - f'group {group_id} returned {len(batch.tags)} tags for ' - f'{len(train_rows)} rows') - source_tags = batch.tags or [{} for _ in train_rows] - tags = [{ - **tag, - 'record_type': 'sample', - 'group_id': group_id, - 'generation_idx': index, - 'rollout_status': 'ROLLOUT_DONE', - 'advantage_status': 'ADVANTAGE_DONE', - 'rollout_policy_version': result.partition.policy.version, - 'rollout_adapter_uri': result.partition.policy.adapter_uri, - } for index, tag in enumerate(source_tags)] - ref = await self.data_plane.aappend(result.ref, train_rows, tags=tags) - # The physical TQ rows retain rollout fields, while this - # reference selects only the trainable InputFeature. - train_ref = ref.model_copy(update={'fields': list(train_rows[0])}) + ref = await self.data_plane.aappend( + result.ref, + [{ + 'reward': float(reward), + 'advantage': float(advantage), + } for reward, advantage in zip(rewards, advantages)], + ) await result.partition.ready.put( - _ReadyGroup( - result.group_index, - train_rows, - tags, - train_ref, - { - 'old_logps': json_safe(old_logps), - 'advantages': json_safe(advantages), - }, - )) + _ReadyGroup(result.group_index, ref)) except BaseException: await self.data_plane.arelease(result.ref) raise @@ -332,40 +303,24 @@ def __init__(self, model, data_plane, state: _GRPOState): self.optimizer_step = 0 async def _train(self, groups: list[_ReadyGroup]) -> None: - rows = [row for group in groups for row in group.rows] - tags = [tag for group in groups for tag in group.tags] - created_batch = len(groups) > 1 - train_ref = ( - await _put_rows(self.data_plane, rows, kind='grpo-train', tags=tags) - if created_batch else groups[0].ref - ) - old_logps = [ - value - for group in groups - for value in group.forward_kwargs['old_logps'] - ] - advantages = [ - value - for group in groups - for value in group.forward_kwargs['advantages'] - ] + refs = [group.ref for group in groups] try: await _submit( - self.model.submit_forward_backward, - train_ref, - old_logps=old_logps, - advantages=advantages, + self.model.forward_backward, + refs, + input_field='train_input', + kwarg_fields={ + 'old_logps': 'sampled_logprobs', + 'advantages': 'advantage', + }, dynamic_batching=True, micro_batch_size=MICRO_BATCH_SIZE, max_tokens_per_micro_batch=MAX_TOKENS_PER_MICRO_BATCH, ) - await _submit(self.model.submit_clip_grad_and_step, max_grad_norm=1.0) + await _submit(self.model.clip_grad_and_step, max_grad_norm=1.0) self.optimizer_step += 1 finally: - if created_batch: - await self.data_plane.arelease(train_ref) - for group in groups: - await self.data_plane.arelease(group.ref) + await asyncio.gather(*(self.data_plane.arelease(ref) for ref in refs)) async def run(self) -> None: try: @@ -387,8 +342,8 @@ async def run(self) -> None: if ready: raise RuntimeError('partition ended with an incomplete train mini-batch') publish_version = self.state.policy.version + 1 - saved = await _submit(self.model.submit_save, f'policy-{publish_version}') - policy = _Policy(publish_version, saved['twinkle_path']) + saved = await _submit(self.model.save, f'policy-{publish_version}') + policy = _Policy(publish_version, _checkpoint_path(saved)) await self.state.publish(partition, policy) print( f'partition={partition.partition_id} policy={policy.version} ' @@ -415,8 +370,8 @@ async def run_grpo( if BATCH_SIZE % groups_per_step: raise ValueError('BATCH_SIZE * NUM_GENERATIONS must be divisible by TRAIN_MINI_BATCH_SIZE') - initial = await _submit(model.submit_save, 'policy-0') - state = _GRPOState(_Policy(version=0, adapter_uri=initial['twinkle_path'])) + initial = await _submit(model.save, 'policy-0') + state = _GRPOState(_Policy(version=0, adapter_uri=_checkpoint_path(initial))) semaphore = asyncio.Semaphore(ROLLOUT_CONCURRENCY) rollout_results: asyncio.Queue = asyncio.Queue() await WorkerPipeline(( @@ -443,8 +398,8 @@ async def train() -> None: model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) model.set_optimizer('AdamW', lr=2e-5) model.set_processor('InputProcessor', padding_free=True) - model.set_template(TEMPLATE_CLS, model_id=MODEL_ID) - sampler.set_template(TEMPLATE_CLS, model_id=MODEL_ID) + model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) + sampler.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) dataloader = DataLoader(dataset=create_dataset(), batch_size=BATCH_SIZE, num_workers=0) await run_grpo(dataloader, model, sampler, data_plane) diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index b11adfc1c..b7dac3cdf 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -113,7 +113,7 @@ applications: model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier data_plane_url: http://127.0.0.1:8000/api/v1/data-plane nproc_per_node: 1 # Number of GPU processes per node - sampler_type: vllm # Inference engine: 'vllm' (fast) or 'torch' (TorchSampler) + sampler_type: vllm_async # Non-blocking vLLM admission for client-orchestrated RL engine_args: # vLLM engine-specific settings max_model_len: 4096 # Maximum sequence length the engine supports gpu_memory_utilization: 0.5 # Fraction of GPU memory to use (0.0-1.0) diff --git a/cookbook/client/server/transformer/server_config_local.yaml b/cookbook/client/server/transformer/server_config_local.yaml index f3aae2203..833445bf3 100644 --- a/cookbook/client/server/transformer/server_config_local.yaml +++ b/cookbook/client/server/transformer/server_config_local.yaml @@ -3,6 +3,7 @@ # Set the absolute Hugging Face-compatible model directory before loading this file: # export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3-4B # export TWINKLE_MODEL_ID=Qwen/Qwen3-4B # in every client process +# export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3-4B # if clients share that path # # The public HTTP model name remains Qwen/Qwen3-4B. Both the training Model # and the vLLM Sampler load the same local directory, so neither component @@ -105,7 +106,7 @@ applications: model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} data_plane_url: http://127.0.0.1:8000/api/v1/data-plane nproc_per_node: 1 - sampler_type: vllm + sampler_type: vllm_async engine_args: max_model_len: 4096 gpu_memory_utilization: 0.5 diff --git a/cookbook/client/twinkle/short_math_grpo.py b/cookbook/client/twinkle/short_math_grpo.py index 1e90d38ce..a9af2f32d 100644 --- a/cookbook/client/twinkle/short_math_grpo.py +++ b/cookbook/client/twinkle/short_math_grpo.py @@ -82,6 +82,12 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: # ========== Configuration ========== BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') MODEL_ID = f'ms://{BASE_MODEL}' +TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) +TEMPLATE_CLS = os.environ.get( + 'TWINKLE_TEMPLATE_CLS', + 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', +) +DATASET_ID = os.environ.get('TWINKLE_DATASET_ID', 'ms://modelscope/gsm8k') NUM_GENERATIONS = 4 MAX_NEW_TOKENS = 1024 LEARNING_RATE = 2e-5 @@ -101,8 +107,8 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: 'and put your final answer within \\boxed{}.') def create_gsm8k_dataset(): - dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train', data_slice=range(DATA_NUM))) - dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=2048, enable_thinking=False) + dataset = Dataset(DatasetMeta(DATASET_ID, subset_name='main', split='train', data_slice=range(DATA_NUM))) + dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=2048, enable_thinking=False) dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT)) dataset.encode(add_generation_prompt=True) return dataset @@ -179,11 +185,11 @@ def train(): # Set processor and template for encoding inputs model.set_processor('InputProcessor') - model.set_template('Qwen3_5Template', model_id=MODEL_ID) + model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) # Step 4: Configure the sampler sampler = vLLMSampler(model_id=MODEL_ID) - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID) + sampler.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) # Step 5: Setup metrics and advantage function advantage_fn = GRPOAdvantage() diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 29485e052..18d79f039 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -325,10 +325,20 @@ def create_workers(worker_cls: Type[T], 'num_cpus': 0.01, } - if device_type == 'GPU': - worker_options['num_gpus'] = 0.01 - else: - # Use custom resource key for non-GPU accelerators (e.g., NPU). + # The placement group already reserves a GPU worker's full + # allocation and ``CUDA_VISIBLE_DEVICES`` above pins that actor + # to the ranks selected by DeviceGroup. Do not additionally + # request a fractional Ray GPU resource for CUDA workers. + # + # With a fractional request, Ray translates the placement-group + # GPU id through the already narrowed visible-device list. For + # the second GPU that becomes index 1 into ``['1']``, killing + # the worker in ``set_visible_accelerator_ids`` before its + # constructor runs. The CPU claim and placement-group strategy + # are sufficient to place the actor in the GPU-owning bundle. + if device_type != 'GPU': + # Use a custom resource key for non-GPU accelerators + # (for example, NPU). worker_options['resources'] = {device_type: 0.01} worker = worker_cls.options(**worker_options).remote(*args, **kwargs) diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index e9cba0f4d..e54e5bd32 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -97,7 +97,7 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De except IndexError: node = self.nodes[0] node_cpu = int(node['Resources']['CPU']) - bundles.append({device_type: nproc_per_node, 'CPU': max(node_cpu // 2, 1)}) + bundles.append({device_type: nproc_per_node, 'CPU': max(nproc_per_node, 1)}) # CPU placement groups: only create when there are actual CPU processes to allocate. if cpu_proc_count > 0: @@ -150,11 +150,22 @@ def get_visible_devices(): return os.environ.get(Platform.get_platform(device_type).visible_device_env()) if self.placement_groups: - self.visible_devices = ray.get([ - get_visible_devices.options(placement_group=pg, runtime_env={ - 'env_vars': self.noset_env() - }).remote() for pg in self.placement_groups - ]) + visible_device_futures = [] + for pg in self.placement_groups: + probe_options = {'placement_group': pg} + if device_type == 'GPU': + # Ask Ray for the GPUs owned by this placement group. A + # no-GPU probe with NOSET_* only sees the node-global + # CUDA_VISIBLE_DEVICES list, so independently initialized + # Model and Sampler groups both select its first entry. + # The probe is short-lived and releases the bundle before + # component workers are created. + probe_options['num_gpus'] = nproc_per_node + else: + probe_options['resources'] = {device_type: nproc_per_node} + visible_device_futures.append( + get_visible_devices.options(**probe_options).remote()) + self.visible_devices = ray.get(visible_device_futures) visible_devices = [] for visible_device in self.visible_devices: diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 43cd6108d..2af27b574 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -696,7 +696,8 @@ def _load_weights(_module): _load_weights(_module) else: _load_weights(self.module) - self.target_parameter_manager.set_state_dict(tenant_adapter_name, state_dict) + if getattr(_lora.tenant_config, 'target_parameters', None): + self.target_parameter_manager.set_state_dict(tenant_adapter_name, state_dict) def get_state_dict(self, tenant_adapter_name): state_dict = {} @@ -721,7 +722,9 @@ def _get_weights(_module): state_dict.update(_get_weights(_module)) else: state_dict = _get_weights(self.module) - target_state_dict = self.target_parameter_manager.get_state_dict(tenant_adapter_name) + target_state_dict = {} + if getattr(_lora.tenant_config, 'target_parameters', None): + target_state_dict = self.target_parameter_manager.get_state_dict(tenant_adapter_name) overlap = state_dict.keys() & target_state_dict.keys() if overlap: raise ValueError(f'Duplicate LoRA state keys: {sorted(overlap)[:5]}') diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index 539d756ca..581a81e1d 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -71,7 +71,7 @@ class SamplerArgs(_ArgsBase): nproc_per_node: int = 1 device_group: dict[str, Any] device_mesh: dict[str, Any] - sampler_type: Literal['mock', 'vllm', 'torch'] + sampler_type: Literal['mock', 'vllm', 'vllm_async', 'torch'] engine_args: dict[str, Any] | None = None queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) data_plane_url: str | None = None diff --git a/src/twinkle/server/data_plane/proxy.py b/src/twinkle/server/data_plane/proxy.py index 3f5004b9f..e9b148aab 100644 --- a/src/twinkle/server/data_plane/proxy.py +++ b/src/twinkle/server/data_plane/proxy.py @@ -23,12 +23,14 @@ def enabled(self) -> bool: async def get( self, ref: DataRef, + *, + fields: list[str] | None = None, ) -> list[dict[str, Any]]: if self.client is None or self.base_url is None: raise RuntimeError('data_plane_url is required when a component request uses input_ref') response = await self.client.post( f'{self.base_url}/twinkle/get', - json={'ref': ref.model_dump()}, + json={'ref': ref.model_dump(), 'fields': fields}, headers=build_routing_headers(f'data-ref-{ref.ref_id}'), ) response.raise_for_status() @@ -51,6 +53,27 @@ async def put( response.raise_for_status() return DataRef(**response.json()) + async def append( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if self.client is None or self.base_url is None: + raise RuntimeError('data_plane_url is required to append component output') + response = await self.client.post( + f'{self.base_url}/twinkle/append', + json={ + 'ref': ref.model_dump(), + 'rows': rows, + 'tags': tags, + }, + headers=build_routing_headers(f'data-append-{ref.ref_id}'), + ) + response.raise_for_status() + return DataRef(**response.json()) + async def close(self) -> None: if self.client is not None: await self.client.aclose() diff --git a/src/twinkle/server/data_plane/store.py b/src/twinkle/server/data_plane/store.py index e1a367ddf..8ecd4e916 100644 --- a/src/twinkle/server/data_plane/store.py +++ b/src/twinkle/server/data_plane/store.py @@ -20,11 +20,14 @@ def _partition(ref: DataRef) -> str: def _input_token_count(rows: list[dict[str, Any]]) -> int: - return sum( - len(row.get('input_ids', [])) - for row in rows - if isinstance(row.get('input_ids'), (list, tuple)) - ) + total = 0 + for row in rows: + input_ids = row.get('input_ids') + if input_ids is None and isinstance(row.get('train_input'), dict): + input_ids = row['train_input'].get('input_ids') + if isinstance(input_ids, (list, tuple)): + total += len(input_ids) + return total def _rows_from_tensordict(data: Any, size: int) -> list[dict[str, Any]]: @@ -124,7 +127,7 @@ async def append( updates: dict[str, Any] = { 'fields': list(dict.fromkeys([*ref.fields, *rows[0].keys()])), } - if 'input_ids' in rows[0]: + if 'input_ids' in rows[0] or 'train_input' in rows[0]: updates['num_tokens'] = _input_token_count(rows) return ref.model_copy(update=updates) diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 0413c9144..86ee6845c 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -74,6 +74,93 @@ def _model_result_rows(result: Any, batch_size: int) -> list[dict[str, Any]]: return [{'result': result}] +def _value_at_path(value: Any, path: str) -> Any: + for part in path.split('.'): + if not isinstance(value, dict) or part not in value: + raise KeyError(f'data field {path!r} does not exist') + value = value[part] + return value + + +def _set_at_path(target: dict[str, Any], path: str, value: Any) -> None: + parts = path.split('.') + current = target + for part in parts[:-1]: + nested = current.setdefault(part, {}) + if not isinstance(nested, dict): + raise ValueError(f'cannot bind nested model argument {path!r}') + current = nested + current[parts[-1]] = value + + +async def _resolve_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: + """Resolve transport-level references before entering the model backend.""" + if body.input_refs is None: + return body.inputs, {} + + selected_fields = None + if body.input_field is not None: + selected_fields = list(dict.fromkeys([ + body.input_field, + *(source.split('.', 1)[0] for source in body.kwarg_fields.values()), + ])) + batches = await asyncio.gather(*( + data_plane.get(ref, fields=selected_fields) + for ref in body.input_refs + )) + rows = [row for batch in batches for row in batch] + if body.input_field is None: + kwarg_roots = {source.split('.', 1)[0] for source in body.kwarg_fields.values()} + inputs = [ + {key: value for key, value in row.items() if key not in kwarg_roots} + for row in rows + ] + else: + inputs = [_value_at_path(row, body.input_field) for row in rows] + + field_kwargs: dict[str, Any] = {} + for target_path, source_path in body.kwarg_fields.items(): + _set_at_path( + field_kwargs, + target_path, + [_value_at_path(row, source_path) for row in rows], + ) + return inputs, field_kwargs + + +def _merge_forward_kwargs(explicit: dict[str, Any], bound: dict[str, Any]) -> dict[str, Any]: + collisions = set(explicit).intersection(bound) + if collisions: + names = ', '.join(sorted(collisions)) + raise ValueError(f'explicit model kwargs conflict with kwarg_fields: {names}') + return {**explicit, **bound} + + +def _request_shape(body: Any) -> tuple[int, int]: + if body.input_refs is not None: + return ( + sum(ref.num_tokens for ref in body.input_refs), + sum(ref.size for ref in body.input_refs), + ) + inputs = body.inputs if isinstance(body.inputs, list) else [body.inputs] + return ( + sum(len(item.get('input_ids', [])) if isinstance(item, dict) else 0 for item in inputs), + len(inputs), + ) + + +def _select_output_rows( + result: Any, + *, + batch_size: int, + output_fields: dict[str, str], +) -> list[dict[str, Any]]: + rows = _model_result_rows(json_safe(result), batch_size) + if len(rows) != batch_size: + raise ValueError(f'model returned {len(rows)} rows for an output_ref of size {batch_size}') + return [{target: _value_at_path(row, source) for source, target in output_fields.items()} for row in rows] + + def _register_twinkle_routes(app: FastAPI, self_fn: Callable[[], ModelManagement]) -> None: """Register all /twinkle/* routes on the given FastAPI app. @@ -123,13 +210,13 @@ async def forward(request: Request, body: types.ForwardRequest, async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} - inputs = _parse_inputs(body.inputs) - ret = self.model.forward(inputs=inputs, adapter_name=adapter_name, **extra_kwargs) + raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) + inputs = _parse_inputs(raw_inputs) + kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) + ret = self.model.forward(inputs=inputs, adapter_name=adapter_name, **kwargs) return {'result': ret} - inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] - input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) - batch_size = len(inputs_list) + input_tokens, batch_size = _request_shape(body) return await run_task( self.schedule_task_and_wait( _task, @@ -141,162 +228,6 @@ async def _task(): task_type='forward', )) - @app.post('/twinkle/submit_forward') - async def submit_forward( - request: Request, - body: types.AsyncForwardRequest, - self: ModelManagement = Depends(self_fn), - ) -> dict[str, Any]: - """Queue a forward pass and return immediately with a task id.""" - token = await self._on_request_start(request) - adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) - - async def _task(): - self.assert_resource_exists(adapter_name) - raw_inputs = ( - await self.data_plane.get(body.input_ref) - if body.input_ref is not None else body.inputs - ) - inputs = _parse_inputs(raw_inputs) - method = self.model.forward_only if body.forward_only else self.model.forward - ret = method(inputs=inputs, adapter_name=adapter_name, **body.forward_kwargs) - safe = json_safe(ret) - if self.data_plane.enabled: - rows = _model_result_rows(safe, len(inputs)) - output_ref = await self.data_plane.put(rows, kind='model-output') - return {'output_ref': output_ref.model_dump()} - return {'result': safe} - - raw_for_metrics = body.inputs or [] - inputs_list = raw_for_metrics if isinstance(raw_for_metrics, list) else [raw_for_metrics] - input_tokens = ( - body.input_ref.num_tokens - if body.input_ref is not None else - sum(len(item.get('input_ids', [])) if isinstance(item, dict) else 0 for item in inputs_list) - ) - batch_size = body.input_ref.size if body.input_ref is not None else len(inputs_list) - return await self.schedule_task( - _task, - model_id=adapter_name, - token=token, - input_tokens=input_tokens, - batch_size=batch_size, - data_world_size=self.data_world_size, - task_type='async_forward', - ) - - @app.post('/twinkle/submit_forward_backward') - async def submit_forward_backward( - request: Request, - body: types.AsyncForwardBackwardRequest, - self: ModelManagement = Depends(self_fn), - ) -> dict[str, Any]: - """Queue the same forward/backward primitive exposed by the synchronous client.""" - token = await self._on_request_start(request) - adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) - - def first_element(data): - while isinstance(data, list): - if len(data) == 0: - return None - data = data[0] - return data - - async def _task(): - self.assert_resource_exists(adapter_name) - raw_inputs = ( - await self.data_plane.get(body.input_ref) - if body.input_ref is not None else body.inputs - ) - inputs = _parse_inputs(raw_inputs) - for model_input in inputs: - for key in model_input: - if (isinstance(model_input[key], list) - and isinstance(first_element(model_input[key]), (int, float))): - model_input[key] = torch.tensor(model_input[key]) - ret = self.model.forward_backward(inputs=inputs, adapter_name=adapter_name, **body.kwargs) - return {'result': json_safe(ret)} - - raw_for_metrics = body.inputs or [] - inputs_list = raw_for_metrics if isinstance(raw_for_metrics, list) else [raw_for_metrics] - input_tokens = ( - body.input_ref.num_tokens - if body.input_ref is not None else - sum(len(item.get('input_ids', [])) for item in inputs_list if isinstance(item, dict)) - ) - batch_size = body.input_ref.size if body.input_ref is not None else len(inputs_list) - return await self.schedule_task( - _task, - model_id=adapter_name, - token=token, - input_tokens=input_tokens, - batch_size=batch_size, - data_world_size=self.data_world_size, - task_type='async_forward_backward', - ) - - @app.post('/twinkle/submit_clip_grad_and_step') - async def submit_clip_grad_and_step( - request: Request, - body: types.AsyncClipGradAndStepRequest, - self: ModelManagement = Depends(self_fn), - ) -> dict[str, Any]: - token = await self._on_request_start(request) - adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) - - async def _task(): - self.assert_resource_exists(adapter_name) - self.model.clip_grad_and_step( - max_grad_norm=body.max_grad_norm, - norm_type=body.norm_type, - adapter_name=adapter_name, - **body.kwargs, - ) - return {'status': 'ok'} - - return await self.schedule_task( - _task, - model_id=adapter_name, - token=token, - task_type='async_clip_grad_and_step', - ) - - @app.post('/twinkle/submit_save') - async def submit_save( - request: Request, - body: types.AsyncSaveRequest, - self: ModelManagement = Depends(self_fn), - ) -> dict[str, Any]: - """Queue an adapter snapshot; used for explicit policy publication.""" - token = await self._on_request_start(request) - adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) - - async def _task(): - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') - checkpoint_name = checkpoint_manager.get_ckpt_name(body.name) - save_dir = checkpoint_manager.get_save_dir(model_id=adapter_name, is_sampler=body.is_sampler) - twinkle_path = checkpoint_manager.save( - model_id=adapter_name, - name=checkpoint_name, - is_sampler=body.is_sampler, - ) - model_save_name = 'latest' if body.is_sampler else checkpoint_name - checkpoint_dir = self.model.save( - name=model_save_name, - output_dir=save_dir, - adapter_name=adapter_name, - save_optimizer=body.save_optimizer, - ) - return {'twinkle_path': twinkle_path, 'checkpoint_dir': checkpoint_dir} - - return await self.schedule_task( - _task, - model_id=adapter_name, - token=token, - task_type='async_save', - ) - @app.post('/twinkle/remove_adapter') async def remove_adapter( request: Request, @@ -331,18 +262,29 @@ async def forward_only( async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} - inputs = _parse_inputs(body.inputs) - ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **extra_kwargs) + raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) + inputs = _parse_inputs(raw_inputs) + kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) + ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **kwargs) + if body.output_ref is not None: + rows = _select_output_rows( + ret, + batch_size=len(inputs), + output_fields=body.output_fields, + ) + output_ref = await self.data_plane.append(body.output_ref, rows) + return {'result': output_ref.model_dump()} return {'result': ret} - inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] - input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) + input_tokens, batch_size = _request_shape(body) return await run_task( self.schedule_task_and_wait( _task, model_id=adapter_name, token=token, input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, task_type='forward_only', )) @@ -395,17 +337,17 @@ def first_element(data): async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} - all_inputs = _parse_inputs(body.inputs) + raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) + all_inputs = _parse_inputs(raw_inputs) for inputs in all_inputs: for key in inputs: if isinstance(inputs[key], list) and isinstance(first_element(inputs[key]), (int, float)): inputs[key] = torch.tensor(inputs[key]) - ret = self.model.forward_backward(inputs=all_inputs, adapter_name=adapter_name, **extra_kwargs) + kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) + ret = self.model.forward_backward(inputs=all_inputs, adapter_name=adapter_name, **kwargs) return {'result': ret} - inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] - input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) - batch_size = len(inputs_list) + input_tokens, batch_size = _request_shape(body) return await run_task( self.schedule_task_and_wait( _task, diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index 9e8eaac4e..df39d1033 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -40,6 +40,13 @@ def _make_vllm_sampler(kw: dict[str, Any]) -> Any: return vLLMSampler(**kw) +def _make_vllm_async_sampler(kw: dict[str, Any]) -> Any: + """Construct the vLLM backend with non-blocking generation admission.""" + from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ + + return VLLMSamplerTQ(**kw, context_manager=None) + + def _make_torch_sampler(kw: dict[str, Any]) -> Any: from twinkle.sampler import TorchSampler # type: ignore[attr-defined] @@ -52,6 +59,7 @@ def _make_torch_sampler(kw: dict[str, Any]) -> Any: { 'mock': _make_mock_sampler, 'vllm': _make_vllm_sampler, + 'vllm_async': _make_vllm_async_sampler, 'torch': _make_torch_sampler, }, ) @@ -62,12 +70,8 @@ def _construct_sampler_backend( sampler_kwargs: dict[str, Any], data_plane_url: str | None, ) -> Any: - if sampler_type == 'vllm' and data_plane_url: - # Client-orchestrated async RL needs sampler admission to release the - # Ray actor immediately. VLLMSamplerTQ retains the inherited - # synchronous API for ordinary sample calls. - from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ - return VLLMSamplerTQ(**sampler_kwargs, context_manager=None) + # Backend selection is explicit. DataPlane config controls where results + # are stored, not which sampler implementation is instantiated. return SAMPLER_SELECTOR.construct(sampler_type, sampler_kwargs) @@ -170,7 +174,7 @@ def build_sampler_app(model_id: str, device_group: Device group configuration dict device_mesh: Device mesh configuration dict for parallelism deploy_options: Ray Serve deployment options - sampler_type: Sampler selector — ``mock`` | ``vllm`` | ``torch``. + sampler_type: Sampler selector — ``mock`` | ``vllm`` | ``vllm_async`` | ``torch``. Validated up front; bad values raise :class:`ConfigError` before any side effect. engine_args: Additional engine arguments for the sampler diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index bec3c74a2..522b8d235 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -76,8 +76,16 @@ def _sample_models_to_rows( tags = [] for prompt_index, (response, group_id) in enumerate(zip(sample_models, resolved_group_ids)): for generation_idx, sequence in enumerate(response.sequences): + sampled_logprobs = [ + 0.0 if not position else float(position[0][1]) + for position in (sequence.logprobs or []) + ] rows.append({ - **sequence.model_dump(), + 'train_input': sequence.new_input_feature, + 'sampled_logprobs': sampled_logprobs, + 'tokens': sequence.tokens, + 'decoded': sequence.decoded, + 'stop_reason': sequence.stop_reason, 'prompt_logprobs': response.prompt_logprobs, 'topk_prompt_logprobs': response.topk_prompt_logprobs, }) @@ -127,31 +135,32 @@ def _submission_states(value) -> list[dict]: async def _await_generation( sampler, submission_id: str, - inputs, - params: SamplingParams, - *, - adapter_name: str, - adapter_path: str | None, ): - """Submit a generation and poll without occupying a worker thread.""" - submitted = False + """Poll an admitted generation without occupying the sampler admission queue.""" collected = False try: - # Mark before dispatch so a partial multi-DP admission is still rolled - # back if one actor rejects while another has already registered it. - submitted = True - await asyncio.to_thread( - sampler.submit_generation, - submission_id, - inputs, - params, - adapter_name=adapter_name, - adapter_path=adapter_path, - ) poll_interval = 0.01 while True: - states = _submission_states( - await asyncio.to_thread(sampler.get_generation_status, submission_id)) + try: + states = _submission_states( + await asyncio.to_thread(sampler.get_generation_status, submission_id)) + except Exception as error: + # A pending read-only actor call can be cancelled by Ray while + # the generation submitted just above remains alive. Treating + # that as a generation failure makes the finally block discard + # otherwise valid rollout work. Retry only Ray's explicit task + # cancellation; actor death and application errors must still + # propagate immediately. + from ray.exceptions import TaskCancelledError + if not isinstance(error, TaskCancelledError): + raise + logger.warning( + 'Generation status poll was cancelled; retrying submission %s', + submission_id, + ) + await asyncio.sleep(poll_interval) + poll_interval = min(poll_interval * 1.5, 0.25) + continue failed = next( (state for state in states if state.get('status') not in ('running', 'completed')), None, @@ -166,7 +175,7 @@ async def _await_generation( await asyncio.sleep(poll_interval) poll_interval = min(poll_interval * 1.5, 0.25) finally: - if submitted and not collected: + if not collected: try: await asyncio.to_thread(sampler.cancel_generation, submission_id) except Exception: @@ -263,81 +272,81 @@ async def _task(): task_type='sample', )) - @app.post('/twinkle/submit_sample') - async def submit_sample( + @app.post('/twinkle/sample_to_data_plane', response_model=types.DataRef) + async def sample_to_data_plane( request: Request, - body: types.AsyncSampleRequest, + body: types.DataPlaneSampleRequest, self: SamplerManagement = Depends(self_fn), - ) -> dict: - """Queue sampling and return immediately for client-side orchestration.""" + ) -> types.DataRef: + """Generate a complete group, store it server-side, and return its DataRef.""" token = await self._on_request_start(request) + if not self.data_plane.enabled: + raise HTTPException(status_code=503, detail='sample_to_data_plane requires data_plane_url') + if not callable(getattr(self.sampler, 'submit_generation', None)): + raise HTTPException(status_code=503, detail='sampler_type must be vllm_async') - async def _task(): - adapter_path = None - full_adapter_name = _get_twinkle_sampler_adapter_name(request, body.adapter_name) or '' - if body.adapter_uri: - from twinkle.server.checkpoint import create_checkpoint_manager - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') - _, adapter_path = checkpoint_manager.parse_adapter_uri(body.adapter_uri) + adapter_path = None + full_adapter_name = _get_twinkle_sampler_adapter_name(request, body.adapter_name) or '' + if body.adapter_uri: + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + _, adapter_path = checkpoint_manager.parse_adapter_uri(body.adapter_uri) - inputs = ( - await self.data_plane.get(body.input_ref) - if body.input_ref is not None else body.inputs - ) - if isinstance(inputs, list) and inputs: - first = inputs[0] - if isinstance(first, dict) and 'input_ids' in first: - inputs = [InputFeature(**item) for item in inputs] - else: - inputs = [Trajectory(**item) for item in inputs] - elif isinstance(inputs, dict): - inputs = [InputFeature(**inputs)] if 'input_ids' in inputs else [Trajectory(**inputs)] - - params_dict = dict(body.sampling_params or {}) - params_dict['num_samples'] = body.num_samples - params = SamplingParams.from_dict(params_dict) - if callable(getattr(self.sampler, 'submit_generation', None)): - responses = await _await_generation( - self.sampler, - uuid.uuid4().hex, - inputs, - params, - adapter_name=full_adapter_name, - adapter_path=adapter_path, - ) + inputs = ( + await self.data_plane.get(body.input_ref) + if body.input_ref is not None else body.inputs + ) + if isinstance(inputs, list) and inputs: + first = inputs[0] + if isinstance(first, dict) and 'input_ids' in first: + inputs = [InputFeature(**item) for item in inputs] else: - # Mock/Torch and vLLM deployments without a DataPlane retain - # the compatibility path. DataPlane-enabled vLLM deployments - # are constructed with VLLMSamplerTQ and never wait here. - responses = await asyncio.to_thread( - self.sampler.sample, - inputs, - params, - adapter_name=full_adapter_name, - adapter_path=adapter_path, - ) - sample_models = _responses_to_models(responses) - payload = types.SampleResponseModelList(samples=sample_models).model_dump() - if self.data_plane.enabled: - rows, tags = _sample_models_to_rows( - sample_models, - group_ids=body.group_ids, - policy_version=body.policy_version, - adapter_uri=body.adapter_uri, - ) - output_ref = await self.data_plane.put( - [json_safe(item) for item in rows], - kind='rollout', - tags=tags, - ) - return {'output_ref': output_ref.model_dump()} - return payload - - return await self.schedule_background_task( - _task, - model_id=full_adapter_name if (full_adapter_name := _get_twinkle_sampler_adapter_name( - request, body.adapter_name)) else None, - task_type='async_sample', + inputs = [Trajectory(**item) for item in inputs] + elif isinstance(inputs, dict): + inputs = [InputFeature(**inputs)] if 'input_ids' in inputs else [Trajectory(**inputs)] + + params_dict = dict(body.sampling_params or {}) + params_dict['num_samples'] = body.num_samples + params = SamplingParams.from_dict(params_dict) + submission_id = uuid.uuid4().hex + + async def _admit(): + await asyncio.to_thread( + self.sampler.submit_generation, + submission_id, + inputs, + params, + adapter_name=full_adapter_name, + adapter_path=adapter_path, + ) + return submission_id + + inline_inputs = body.inputs if isinstance(body.inputs, list) else [body.inputs] + input_tokens = ( + body.input_ref.num_tokens + if body.input_ref is not None else + sum(len(item.get('input_ids', [])) for item in inline_inputs if isinstance(item, dict)) + ) + await run_task( + self.schedule_task_and_wait( + _admit, + model_id=full_adapter_name or None, + token=token, + input_tokens=input_tokens, + task_type='sample_admission', + )) + + responses = await _await_generation(self.sampler, submission_id) + rows, tags = _sample_models_to_rows( + _responses_to_models(responses), + group_ids=body.group_ids, + policy_version=body.policy_version, + adapter_uri=body.adapter_uri, + ) + return await self.data_plane.put( + [json_safe(item) for item in rows], + kind='rollout', + tags=tags, ) @app.post('/twinkle/unload_adapter_paths') diff --git a/src/twinkle/utils/nccl_safe.py b/src/twinkle/utils/nccl_safe.py index b22b10137..acb368034 100644 --- a/src/twinkle/utils/nccl_safe.py +++ b/src/twinkle/utils/nccl_safe.py @@ -93,6 +93,10 @@ def __call__(self, inputs, outputs, **kwargs): type(e).__name__, e, traceback.format_exc()) return _zero_loss(outputs) + def micro_batch_scale(self, inputs, indices): + """Preserve the wrapped loss's micro-batch reduction semantics.""" + return self._loss_instance.micro_batch_scale(inputs, indices) + def _zero_loss(outputs) -> 'LossOutput': """Create a graph-connected zero loss for FSDP compatibility. diff --git a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py index 141fbe821..33c7dd30d 100644 --- a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py +++ b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py @@ -139,7 +139,10 @@ def __init__( ): self.context_manager = context_manager super().__init__(model_id=model_id, engine_args=engine_args, device_mesh=device_mesh, **kwargs) - self.data_plane = TQDataPlane() + # Native YAML async-RL writes rollout groups to TransferQueue. The C/S + # component mode only uses submit_generation/collect_generation and + # stores results through the server DataPlane deployment. + self.data_plane = TQDataPlane() if context_manager is not None else None self.reward_registry = dict(reward_registry or {}) self.rollout_max_retries = int(rollout_max_retries) self.rollout_retry_delay_s = float(rollout_retry_delay_s) @@ -541,6 +544,8 @@ async def _run_prompt_group( if rewards is None: raise ValueError(f'no reward function registered for context {group.context.key}') reward_metrics = _compute_reward_metrics(self.reward_registry, group.context, rows, rewards) + if self.data_plane is None: + raise RuntimeError('native TQ data plane is required for prompt-group sampling') await self.data_plane.complete_rollout_group( group, rollout_rows=rows, diff --git a/src/twinkle_client/__init__.py b/src/twinkle_client/__init__.py index a3d281af1..bb13f19ad 100644 --- a/src/twinkle_client/__init__.py +++ b/src/twinkle_client/__init__.py @@ -72,7 +72,6 @@ def init_twinkle_client( ) -from .remote_task import RemoteTask, RemoteTaskError from .data_plane import DataPlaneClient -__all__ = ['DataPlaneClient', 'RemoteTask', 'RemoteTaskError', 'init_tinker_client', 'init_twinkle_client'] +__all__ = ['DataPlaneClient', 'init_tinker_client', 'init_twinkle_client'] diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index c620269ff..7eda3258e 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -2,9 +2,8 @@ from pathlib import Path import time from twinkle_client.http import http_get, http_post -from twinkle_client.remote_task import RemoteTask from twinkle_client.common.json_utils import json_safe -from twinkle_client.types.component import ComponentTaskRef, DataRef +from twinkle_client.types.component import DataRef from twinkle_client.types.model import ( CalculateLossResponse, CalculateMetricResponse, @@ -18,6 +17,17 @@ ) +def _component_input_payload(inputs: Any) -> dict[str, Any]: + """Encode inline rows or one or more opaque server-side data references.""" + if isinstance(inputs, DataRef): + return {'input_refs': [inputs.model_dump()]} + if isinstance(inputs, list) and inputs and any(isinstance(item, DataRef) for item in inputs): + if not all(isinstance(item, DataRef) for item in inputs): + raise TypeError('model inputs cannot mix DataRef values with inline rows') + return {'input_refs': [item.model_dump() for item in inputs]} + return {'inputs': json_safe(inputs)} + + class MultiLoraTransformersModel: """Client wrapper for TwinkleModel that calls server HTTP endpoints. @@ -29,8 +39,7 @@ def __init__(self, model_id: str, **kwargs): """Initialize model client.""" from twinkle_client.http import get_base_url self.server_url = get_base_url() - from twinkle_client.data_plane import DataPlaneClient - self.data_plane = DataPlaneClient(kwargs.pop('data_plane_url', None)) + kwargs.pop('data_plane_url', None) if '://' in model_id: model_id = model_id.split('://')[1] @@ -42,85 +51,6 @@ def __init__(self, model_id: str, **kwargs): ) response.raise_for_status() - def submit_forward( - self, - inputs: Any | DataRef, - *, - forward_only: bool = False, - **forward_kwargs, - ) -> RemoteTask: - """Submit the Model component's forward primitive directly.""" - body = { - 'adapter_name': self.adapter_name or '', - 'forward_only': forward_only, - 'forward_kwargs': forward_kwargs, - } - body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( - inputs.model_dump() if isinstance(inputs, DataRef) else json_safe(inputs)) - response = http_post( - url=f'{self.server_url}/submit_forward', - json_data=body, - ) - response.raise_for_status() - return RemoteTask(ComponentTaskRef(**response.json())) - - def submit_forward_only( - self, - inputs: Any | DataRef, - **forward_kwargs, - ) -> RemoteTask: - return self.submit_forward( - inputs, - forward_only=True, - **forward_kwargs, - ) - - def submit_forward_backward( - self, - inputs: Any | DataRef, - **kwargs, - ) -> RemoteTask: - """Submit forward/backward without prescribing the surrounding train loop.""" - body = {'adapter_name': self.adapter_name or '', 'kwargs': json_safe(kwargs)} - body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( - inputs.model_dump() if isinstance(inputs, DataRef) else json_safe(inputs)) - response = http_post( - url=f'{self.server_url}/submit_forward_backward', - json_data=body, - ) - response.raise_for_status() - return RemoteTask(ComponentTaskRef(**response.json())) - - def submit_clip_grad_and_step( - self, - max_grad_norm: float = 1.0, - norm_type: int = 2, - **kwargs, - ) -> RemoteTask: - response = http_post( - url=f'{self.server_url}/submit_clip_grad_and_step', - json_data={ - 'adapter_name': self.adapter_name or '', - 'max_grad_norm': max_grad_norm, - 'norm_type': norm_type, - 'kwargs': kwargs, - }, - ) - response.raise_for_status() - return RemoteTask(ComponentTaskRef(**response.json())) - - def submit_save(self, name: str, *, save_optimizer: bool = False) -> RemoteTask: - response = http_post( - url=f'{self.server_url}/submit_save', - json_data={ - 'adapter_name': self.adapter_name or '', - 'name': name, - 'save_optimizer': save_optimizer, - }, - ) - response.raise_for_status() - return RemoteTask(ComponentTaskRef(**response.json())) - def add_adapter_to_model(self, adapter_name: str, config: Dict[str, Any], **kwargs) -> None: """Add a new adapter to the model.""" save_dir = kwargs.get('save_dir') @@ -144,23 +74,57 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: if name == self.adapter_name: self.adapter_name = None - def forward(self, inputs: Any, **kwargs) -> ForwardResponse: - """Execute forward pass on the model.""" + def forward( + self, + inputs: Any | DataRef | list[DataRef], + *, + input_field: str | None = None, + kwarg_fields: dict[str, str] | None = None, + **kwargs, + ) -> ForwardResponse: + """Execute forward over inline rows or server-side references.""" response = http_post( url=f'{self.server_url}/forward', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} + json_data={ + **_component_input_payload(inputs), + 'adapter_name': self.adapter_name, + 'input_field': input_field, + 'kwarg_fields': kwarg_fields or {}, + **json_safe(kwargs), + }, ) response.raise_for_status() return ForwardResponse(**response.json()) - def forward_only(self, inputs: Any, **kwargs) -> ForwardResponse: - """Execute forward pass without gradient computation.""" + def forward_only( + self, + inputs: Any | DataRef | list[DataRef], + *, + input_field: str | None = None, + kwarg_fields: dict[str, str] | None = None, + output_ref: DataRef | None = None, + output_fields: dict[str, str] | None = None, + **kwargs, + ) -> ForwardResponse | DataRef: + """Execute forward-only, optionally reading and updating server-side rows.""" + body = { + **_component_input_payload(inputs), + 'adapter_name': self.adapter_name, + 'input_field': input_field, + 'kwarg_fields': kwarg_fields or {}, + 'output_ref': output_ref.model_dump() if output_ref is not None else None, + 'output_fields': output_fields or {}, + **json_safe(kwargs), + } response = http_post( url=f'{self.server_url}/forward_only', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} + json_data=body, ) response.raise_for_status() - return ForwardResponse(**response.json()) + result = ForwardResponse(**response.json()) + if output_ref is not None: + return DataRef(**result.result) + return result def calculate_loss(self, **kwargs) -> CalculateLossResponse: """Calculate loss from model outputs.""" @@ -188,11 +152,24 @@ def backward(self, **kwargs) -> None: ) response.raise_for_status() - def forward_backward(self, inputs: Any, **kwargs) -> ForwardBackwardResponse: - """Execute combined forward and backward pass.""" + def forward_backward( + self, + inputs: Any | DataRef | list[DataRef], + *, + input_field: str | None = None, + kwarg_fields: dict[str, str] | None = None, + **kwargs, + ) -> ForwardBackwardResponse: + """Execute forward/backward over inline rows or server-side references.""" response = http_post( url=f'{self.server_url}/forward_backward', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} + json_data={ + **_component_input_payload(inputs), + 'adapter_name': self.adapter_name, + 'input_field': input_field, + 'kwarg_fields': kwarg_fields or {}, + **json_safe(kwargs), + }, ) response.raise_for_status() return ForwardBackwardResponse(**response.json()) diff --git a/src/twinkle_client/remote_task.py b/src/twinkle_client/remote_task.py deleted file mode 100644 index d3ff36cec..000000000 --- a/src/twinkle_client/remote_task.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Future handle returned by an individual Model or Sampler component.""" -from __future__ import annotations - -import asyncio -import time -from typing import Any - -from twinkle_client.http import get_base_url, http_post -from twinkle_client.http.http_utils import _build_headers -from twinkle_client.types.component import ComponentTaskRef - - -class RemoteTaskError(RuntimeError): - pass - - -class RemoteTask: - """A wrapper over the server's existing component future registry.""" - - def __init__(self, task: ComponentTaskRef | str): - self.request_id = task if isinstance(task, str) else task.request_id - self.model_id = None if isinstance(task, str) else task.model_id - self._url = f'{get_base_url()}/retrieve_future' - - def poll(self, timeout: float | None = None) -> Any | None: - request_kwargs = {} if timeout is None else {'timeout': timeout} - response = http_post( - self._url, - json_data={'request_id': self.request_id}, - **request_kwargs, - ) - response.raise_for_status() - return self._resolve_payload(response.json()) - - @staticmethod - def _resolve_payload(payload: Any) -> Any | None: - if isinstance(payload, dict) and payload.get('type') == 'try_again': - return None - if isinstance(payload, dict) and 'error' in payload: - raise RemoteTaskError(payload['error']) - return payload - - def result(self, timeout: float | None = None) -> Any: - import requests - - deadline = None if timeout is None else time.monotonic() + timeout - while True: - remaining = None if deadline is None else deadline - time.monotonic() - if remaining is not None and remaining <= 0: - raise TimeoutError(f'component task {self.request_id} did not finish within {timeout}s') - try: - result = self.poll(timeout=remaining) - except requests.Timeout as exc: - raise TimeoutError( - f'component task {self.request_id} did not finish within {timeout}s') from exc - if result is not None: - return result - if deadline is not None and time.monotonic() >= deadline: - raise TimeoutError(f'component task {self.request_id} did not finish within {timeout}s') - - async def aresult(self, timeout: float | None = None) -> Any: - import httpx - deadline = None if timeout is None else time.monotonic() + timeout - poll_interval = 0.05 - async with httpx.AsyncClient(timeout=600) as client: - while True: - remaining = None if deadline is None else deadline - time.monotonic() - if remaining is not None and remaining <= 0: - raise TimeoutError( - f'component task {self.request_id} did not finish within {timeout}s') - try: - response = await client.post( - self._url, - headers=_build_headers(), - json={'request_id': self.request_id}, - timeout=remaining if remaining is not None else 600, - ) - except httpx.TimeoutException as exc: - raise TimeoutError( - f'component task {self.request_id} did not finish within {timeout}s') from exc - response.raise_for_status() - result = self._resolve_payload(response.json()) - if result is not None: - return result - if deadline is not None and time.monotonic() >= deadline: - raise TimeoutError(f'component task {self.request_id} did not finish within {timeout}s') - remaining = None if deadline is None else deadline - time.monotonic() - await asyncio.sleep( - poll_interval if remaining is None else min(poll_interval, max(remaining, 0.0))) - poll_interval = min(poll_interval * 1.5, 1.0) - - def __await__(self): - return self.aresult().__await__() diff --git a/src/twinkle_client/sampler/vllm_sampler.py b/src/twinkle_client/sampler/vllm_sampler.py index c0c865ba8..25ea639c0 100644 --- a/src/twinkle_client/sampler/vllm_sampler.py +++ b/src/twinkle_client/sampler/vllm_sampler.py @@ -1,11 +1,11 @@ +import asyncio from typing import Any, Dict, List, Optional, Union from twinkle_client.http import http_post from twinkle_client.types.sampler import AddAdapterResponse, SampleResponseModel, SetTemplateResponse from peft import PeftConfig from twinkle.data_format import Trajectory, InputFeature from twinkle_client.common.json_utils import json_safe -from twinkle_client.remote_task import RemoteTask -from twinkle_client.types.component import ComponentTaskRef, DataRef +from twinkle_client.types.component import DataRef # Intentionally does NOT subclass ``twinkle.sampler.base.Sampler``: importing @@ -98,7 +98,7 @@ def sample( response.raise_for_status() return [SampleResponseModel(**r) for r in response.json()['samples']] - def submit_sample( + def sample_to_data_plane( self, inputs: Union[List[Trajectory], List[InputFeature], DataRef], sampling_params: Optional[Dict[str, Any]] = None, @@ -108,8 +108,8 @@ def submit_sample( policy_version: int | None = None, group_ids: list[str] | None = None, num_samples: int = 1, - ) -> RemoteTask: - """Submit directly to the Sampler component and return immediately.""" + ) -> DataRef: + """Generate complete prompt groups and keep their rows in the server DataPlane.""" body = { 'sampling_params': sampling_params, 'adapter_name': adapter_name, @@ -121,13 +121,31 @@ def submit_sample( body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( inputs.model_dump() if isinstance(inputs, DataRef) else _json_safe(inputs)) response = http_post( - url=f'{self.server_url}/submit_sample', + url=f'{self.server_url}/sample_to_data_plane', json_data=json_safe(body), ) response.raise_for_status() - return RemoteTask(ComponentTaskRef(**response.json())) + return DataRef(**response.json()) async def asample( + self, + inputs: Union[List[Trajectory], List[InputFeature]], + sampling_params: Optional[Dict[str, Any]] = None, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + num_samples: int = 1, + ) -> List[SampleResponseModel]: + """Asynchronous convenience wrapper for the materialized sample API.""" + return await asyncio.to_thread( + self.sample, + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_uri=adapter_uri, + num_samples=num_samples, + ) + + async def asample_to_data_plane( self, inputs: Union[List[Trajectory], List[InputFeature], DataRef], sampling_params: Optional[Dict[str, Any]] = None, @@ -137,11 +155,10 @@ async def asample( policy_version: int | None = None, group_ids: list[str] | None = None, num_samples: int = 1, - ) -> List[SampleResponseModel]: - """Submit sampling and asynchronously await it without blocking the event loop.""" - import asyncio - task = await asyncio.to_thread( - self.submit_sample, + ) -> DataRef: + """Asynchronously sample and return the opaque server-side result reference.""" + return await asyncio.to_thread( + self.sample_to_data_plane, inputs, sampling_params, adapter_name=adapter_name, @@ -150,37 +167,6 @@ async def asample( group_ids=group_ids, num_samples=num_samples, ) - result = await task.aresult() - if isinstance(result, dict) and result.get('output_ref'): - output_ref = DataRef(**result['output_ref']) - try: - batch = await self.data_plane.aget_batch(output_ref) - finally: - await self.data_plane.arelease(output_ref) - if batch.tags and all('prompt_index' in tag for tag in batch.tags): - grouped: dict[int, list[tuple[int, dict[str, Any]]]] = {} - for row, tag in zip(batch.rows, batch.tags): - grouped.setdefault(int(tag['prompt_index']), []).append( - (int(tag.get('generation_idx', 0)), row)) - samples = [] - for prompt_index in sorted(grouped): - generation_rows = [row for _, row in sorted(grouped[prompt_index])] - first = generation_rows[0] - samples.append({ - 'sequences': [{ - key: value - for key, value in row.items() - if key not in ('prompt_logprobs', 'topk_prompt_logprobs') - } for row in generation_rows], - 'prompt_logprobs': first.get('prompt_logprobs'), - 'topk_prompt_logprobs': first.get('topk_prompt_logprobs'), - }) - else: - # Compatibility with a server that still stores one nested row per prompt. - samples = batch.rows - else: - samples = result.get('samples', []) if isinstance(result, dict) else [] - return [SampleResponseModel(**item) for item in samples] def unload_adapter_paths(self, adapter_paths: list[str]) -> None: """Evict policy snapshots that are no longer referenced by this client.""" diff --git a/src/twinkle_client/types/__init__.py b/src/twinkle_client/types/__init__.py index 46cc02235..a0952717f 100644 --- a/src/twinkle_client/types/__init__.py +++ b/src/twinkle_client/types/__init__.py @@ -93,14 +93,9 @@ from .checkpoint import ResolvedLoadPath from .component import ( - AsyncClipGradAndStepRequest, - AsyncForwardBackwardRequest, - AsyncForwardRequest, - AsyncSampleRequest, - AsyncSaveRequest, - ComponentTaskRef, DataAppendRequest, DataGetRequest, + DataPlaneSampleRequest, DataPutRequest, DataRef, DataReleaseRequest, diff --git a/src/twinkle_client/types/component.py b/src/twinkle_client/types/component.py index 66b153e24..d7e9ec2a0 100644 --- a/src/twinkle_client/types/component.py +++ b/src/twinkle_client/types/component.py @@ -7,11 +7,6 @@ from pydantic import BaseModel, Field, model_validator -class ComponentTaskRef(BaseModel): - request_id: str - model_id: str | None = None - - class DataRef(BaseModel): """Opaque reference to rows stored in the server-side TransferQueue.""" @@ -49,7 +44,7 @@ class DataRowsResponse(BaseModel): tags: list[dict[str, Any]] = Field(default_factory=list) -class AsyncSampleRequest(BaseModel): +class DataPlaneSampleRequest(BaseModel): inputs: Any = None input_ref: DataRef | None = None sampling_params: dict[str, Any] | None = None @@ -60,7 +55,7 @@ class AsyncSampleRequest(BaseModel): num_samples: int = 1 @model_validator(mode='after') - def validate_input(self) -> 'AsyncSampleRequest': + def validate_input(self) -> 'DataPlaneSampleRequest': if (self.inputs is None) == (self.input_ref is None): raise ValueError('exactly one of inputs and input_ref must be provided') if self.group_ids is not None and self.inputs is not None: @@ -70,46 +65,5 @@ def validate_input(self) -> 'AsyncSampleRequest': return self -class AsyncForwardRequest(BaseModel): - inputs: Any = None - input_ref: DataRef | None = None - adapter_name: str = '' - forward_only: bool = False - forward_kwargs: dict[str, Any] = Field(default_factory=dict) - - @model_validator(mode='after') - def validate_input(self) -> 'AsyncForwardRequest': - if (self.inputs is None) == (self.input_ref is None): - raise ValueError('exactly one of inputs and input_ref must be provided') - return self - - -class AsyncForwardBackwardRequest(BaseModel): - inputs: Any = None - input_ref: DataRef | None = None - adapter_name: str = '' - kwargs: dict[str, Any] = Field(default_factory=dict) - - @model_validator(mode='after') - def validate_input(self) -> 'AsyncForwardBackwardRequest': - if (self.inputs is None) == (self.input_ref is None): - raise ValueError('exactly one of inputs and input_ref must be provided') - return self - - -class AsyncClipGradAndStepRequest(BaseModel): - adapter_name: str = '' - max_grad_norm: float = 1.0 - norm_type: int = 2 - kwargs: dict[str, Any] = Field(default_factory=dict) - - -class AsyncSaveRequest(BaseModel): - adapter_name: str = '' - name: str - save_optimizer: bool = False - is_sampler: bool = False - - class UnloadAdapterPathsRequest(BaseModel): adapter_paths: list[str] diff --git a/src/twinkle_client/types/model.py b/src/twinkle_client/types/model.py index 10a60b947..830ed6010 100644 --- a/src/twinkle_client/types/model.py +++ b/src/twinkle_client/types/model.py @@ -4,9 +4,11 @@ These models are used by both the server-side handler and the twinkle client. """ -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator from typing import Any, Dict, List, Optional, Union +from .component import DataRef + class CreateRequest(BaseModel): @@ -15,17 +17,43 @@ class Config: class ForwardRequest(BaseModel): - inputs: Any + inputs: Any = None + input_refs: List[DataRef] | None = None + input_field: str | None = None + kwarg_fields: Dict[str, str] = Field(default_factory=dict) adapter_name: str + @model_validator(mode='after') + def validate_input(self) -> 'ForwardRequest': + if (self.inputs is None) == (self.input_refs is None): + raise ValueError('exactly one of inputs and input_refs must be provided') + if self.input_refs is not None and not self.input_refs: + raise ValueError('input_refs must not be empty') + return self + class Config: extra = 'allow' class ForwardOnlyRequest(BaseModel): - inputs: Any + inputs: Any = None + input_refs: List[DataRef] | None = None + input_field: str | None = None + kwarg_fields: Dict[str, str] = Field(default_factory=dict) + output_ref: DataRef | None = None + output_fields: Dict[str, str] = Field(default_factory=dict) adapter_name: Optional[str] = None + @model_validator(mode='after') + def validate_input(self) -> 'ForwardOnlyRequest': + if (self.inputs is None) == (self.input_refs is None): + raise ValueError('exactly one of inputs and input_refs must be provided') + if self.input_refs is not None and not self.input_refs: + raise ValueError('input_refs must not be empty') + if (self.output_ref is None) != (len(self.output_fields) == 0): + raise ValueError('output_ref and output_fields must be configured together') + return self + class Config: extra = 'allow' diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py index 645fdbfdb..8a07529f4 100644 --- a/tests/model/test_micro_batch.py +++ b/tests/model/test_micro_batch.py @@ -8,6 +8,7 @@ from twinkle.model.micro_batch import MicroBatchConfig, collect_micro_batch_outputs, plan_micro_batches from twinkle.model.transformers.transformers import TransformersModel from twinkle.processor import InputProcessor +from twinkle.utils.nccl_safe import safe_loss @pytest.mark.parametrize('packing_algorithm', ['ffd', 'kk']) @@ -71,6 +72,17 @@ def test_sample_mean_and_token_sum_micro_batch_scales(): assert CrossEntropyLoss(reduction='sum').micro_batch_scale(inputs, [0]) == 1.0 +def test_safe_loss_preserves_wrapped_micro_batch_scale(): + inputs = [ + {'labels': [1, -100]}, + {'labels': [2, 3]}, + {'labels': [4, -100]}, + {'labels': [5, 6]}, + ] + + assert safe_loss(GRPOLoss()).micro_batch_scale(inputs, [0, 2]) == .5 + + def test_loss_without_micro_batch_semantics_fails_when_split(): with pytest.raises(NotImplementedError, match='does not support micro-batching'): Loss().micro_batch_scale([{}, {}], [0]) diff --git a/tests/model/test_multi_lora_target_parameters.py b/tests/model/test_multi_lora_target_parameters.py index b28ef6b36..d26b32147 100644 --- a/tests/model/test_multi_lora_target_parameters.py +++ b/tests/model/test_multi_lora_target_parameters.py @@ -45,6 +45,16 @@ def forward(self, x, expert_idx=0): return self.mlp.experts(x, expert_idx=expert_idx) +class FakeLinearModel(nn.Module): + + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + + def forward(self, x): + return self.proj(x) + + def test_peft_target_parameter_key_shapes_for_3d_experts(): model = FakeModel() cfg = LoraConfig( @@ -213,6 +223,24 @@ def test_multilora_state_dict_round_trips_target_parameters(): assert torch.allclose(actual, expected, atol=1e-6) +def test_multilora_state_dict_without_target_parameters_does_not_require_slot(): + from twinkle.model.multi_lora import LoraTenant, MultiLora + + slot_cfg = LoraConfig(r=4, lora_alpha=8, target_modules=["proj"]) + model = get_peft_model(FakeLinearModel(), slot_cfg, adapter_name="lora_0") + multi_lora = MultiLora(max_loras=1, max_r=4) + multi_lora.module = model + multi_lora.loras = [LoraTenant(index=0, adapter_name="lora_0", config=slot_cfg)] + + tenant_cfg = LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]) + multi_lora.acquire_lora("adapter_a", tenant_cfg) + + assert "adapter_a" not in multi_lora.target_parameter_manager.tenant_to_slot + state = multi_lora.get_state_dict("adapter_a") + assert state + multi_lora.set_state_dict("adapter_a", state) + + def test_multilora_transformers_installs_target_parameters_once(): from twinkle.model.multi_lora import LoraTenant, MultiLora @@ -266,4 +294,4 @@ def test_multilora_transformers_installs_target_parameters_once(): assert test_target_parameter_multi_lora_updates_only_active_adapter() == True assert test_multilora_releases_target_parameter_slot_to_initial_weights() == True assert test_multilora_state_dict_round_trips_target_parameters() == True - assert test_multilora_transformers_installs_target_parameters_once() == True \ No newline at end of file + assert test_multilora_transformers_installs_target_parameters_once() == True diff --git a/tests/server/contract/client_api_baseline.json b/tests/server/contract/client_api_baseline.json index 8051c9557..d87b0f89c 100644 --- a/tests/server/contract/client_api_baseline.json +++ b/tests/server/contract/client_api_baseline.json @@ -927,46 +927,6 @@ ] } }, - "/twinkle/submit_clip_grad_and_step": { - "POST": { - "operationId": "submit_clip_grad_and_step_twinkle_submit_clip_grad_and_step_post", - "parameters": [], - "responses": [ - "200", - "422" - ] - } - }, - "/twinkle/submit_forward": { - "POST": { - "operationId": "submit_forward_twinkle_submit_forward_post", - "parameters": [], - "responses": [ - "200", - "422" - ] - } - }, - "/twinkle/submit_forward_backward": { - "POST": { - "operationId": "submit_forward_backward_twinkle_submit_forward_backward_post", - "parameters": [], - "responses": [ - "200", - "422" - ] - } - }, - "/twinkle/submit_save": { - "POST": { - "operationId": "submit_save_twinkle_submit_save_post", - "parameters": [], - "responses": [ - "200", - "422" - ] - } - }, "/twinkle/upload_status/{request_id}": { "GET": { "operationId": "upload_status_twinkle_upload_status__request_id__get", @@ -1084,9 +1044,9 @@ ] } }, - "/twinkle/sample_stream": { + "/twinkle/sample_to_data_plane": { "POST": { - "operationId": "sample_stream_twinkle_sample_stream_post", + "operationId": "sample_to_data_plane_twinkle_sample_to_data_plane_post", "parameters": [], "responses": [ "200", @@ -1094,9 +1054,9 @@ ] } }, - "/twinkle/set_template": { + "/twinkle/sample_stream": { "POST": { - "operationId": "set_template_twinkle_set_template_post", + "operationId": "sample_stream_twinkle_sample_stream_post", "parameters": [], "responses": [ "200", @@ -1104,9 +1064,9 @@ ] } }, - "/twinkle/submit_sample": { + "/twinkle/set_template": { "POST": { - "operationId": "submit_sample_twinkle_submit_sample_post", + "operationId": "set_template_twinkle_set_template_post", "parameters": [], "responses": [ "200", diff --git a/tests/server/data_plane/test_proxy.py b/tests/server/data_plane/test_proxy.py index 1339baaaa..550998db2 100644 --- a/tests/server/data_plane/test_proxy.py +++ b/tests/server/data_plane/test_proxy.py @@ -43,13 +43,19 @@ async def test_proxy_routes_by_data_ref_without_tenant_identity() -> None: proxy.client = _Client() ref = DataRef(ref_id='input', size=1, fields=['value']) - assert await proxy.get(ref) == [{'value': 1}] + assert await proxy.get(ref, fields=['value']) == [{'value': 1}] output = await proxy.put([{'value': 2}], kind='model-output') + appended = await proxy.append(ref, [{'value': 3}]) assert output.ref_id == 'output' + assert appended.ref_id == 'output' get_headers = proxy.client.calls[0][1]['headers'] put_headers = proxy.client.calls[1][1]['headers'] + append_headers = proxy.client.calls[2][1]['headers'] + assert proxy.client.calls[0][1]['json']['fields'] == ['value'] assert get_headers[H_REQUEST_ID] == 'data-ref-input' assert put_headers[H_REQUEST_ID] == 'data-put-model-output' + assert append_headers[H_REQUEST_ID] == 'data-append-input' assert get_headers[H_AUTH] == get_headers[H_AUTH_TWINKLE] == '' assert put_headers[H_AUTH] == put_headers[H_AUTH_TWINKLE] == '' + assert append_headers[H_AUTH] == append_headers[H_AUTH_TWINKLE] == '' diff --git a/tests/server/data_plane/test_store.py b/tests/server/data_plane/test_store.py index ea89ffc75..3c7ec1633 100644 --- a/tests/server/data_plane/test_store.py +++ b/tests/server/data_plane/test_store.py @@ -83,10 +83,17 @@ async def kv_list(partition_id): ) assert ref.num_tokens == 5 + nested_ref = await store.put([ + {'train_input': {'input_ids': [1, 2, 3]}}, + {'train_input': {'input_ids': [4]}}, + ], kind='rollout') + assert nested_ref.num_tokens == 4 + with pytest.raises(KeyError): await store.get(ref.model_copy(update={'ref_id': 'another-ref'})) await store.release(ref) + await store.release(nested_ref) assert records == {} diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index f77e7fecd..bc0790565 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -28,6 +28,19 @@ def __init__(self): self.scheduled = [] self.model_calls = [] self.model = self + self.data_plane = self + self.rows = { + 'data-a': [{ + 'train_input': {'input_ids': [index]}, + 'sampled_logprobs': [-0.1], + 'advantage': 1.0, + } for index in range(4)], + 'data-b': [{ + 'train_input': {'input_ids': [index]}, + 'sampled_logprobs': [-0.2], + 'advantage': -1.0, + } for index in range(4, 8)], + } async def _on_request_start(self, _request): return 'token' @@ -39,26 +52,35 @@ def forward_backward(self, *, inputs, adapter_name, **kwargs): self.model_calls.append((inputs, adapter_name, kwargs)) return {'loss': 1.0} - async def schedule_task(self, task, **kwargs): + async def get(self, ref, *, fields=None): + rows = self.rows[ref.ref_id] + if fields is None: + return rows + return [{field: row[field] for field in fields} for row in rows] + + async def schedule_task_and_wait(self, task, **kwargs): self.scheduled.append(kwargs) - await task() - return {'request_id': 'request-1', 'model_id': kwargs.get('model_id')} + return await task() @pytest.mark.asyncio -async def test_async_forward_backward_schedules_without_algorithm_metadata() -> None: +async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() -> None: management = _SchedulingManagement() app = FastAPI() _register_twinkle_routes(app, lambda: management) - route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/submit_forward_backward') + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') request = Request({'type': 'http', 'headers': []}) request.state.session_id = 'session' - body = types.AsyncForwardBackwardRequest( + body = types.ForwardRequest( adapter_name='adapter', - inputs=[{'input_ids': [index]} for index in range(8)], - kwargs={ - 'old_logps': [[-0.1]] * 8, - 'advantages': [1.0] * 8, + input_refs=[ + types.DataRef(ref_id='data-a', size=4, num_tokens=4), + types.DataRef(ref_id='data-b', size=4, num_tokens=4), + ], + input_field='train_input', + kwarg_fields={ + 'old_logps': 'sampled_logprobs', + 'advantages': 'advantage', }, ) @@ -66,7 +88,10 @@ async def test_async_forward_backward_schedules_without_algorithm_metadata() -> assert management.scheduled[-1]['batch_size'] == 8 assert management.scheduled[-1]['data_world_size'] == 2 - assert 'batch_size_multiple' not in management.scheduled[-1] - _, adapter_name, forwarded_kwargs = management.model_calls[-1] + inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] assert adapter_name == 'session-adapter' - assert forwarded_kwargs == body.kwargs + assert [row['input_ids'].tolist() for row in inputs] == [[index] for index in range(8)] + assert forwarded_kwargs == { + 'old_logps': [[-0.1]] * 4 + [[-0.2]] * 4, + 'advantages': [1.0] * 4 + [-1.0] * 4, + } diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index e4564c5df..8efae14b5 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -114,7 +114,7 @@ def test_mock_dispatch_returns_mock_sampler() -> None: assert isinstance(s, MockSampler) -def test_data_plane_vllm_uses_fire_and_forget_sampler(monkeypatch) -> None: +def test_explicit_async_vllm_uses_non_blocking_sampler(monkeypatch) -> None: from twinkle_agentic.async_rl import vllm_sampler_tq as module captured = {} @@ -126,16 +126,16 @@ def construct(**kwargs): monkeypatch.setattr(module, 'VLLMSamplerTQ', construct) sampler = _construct_sampler_backend( - 'vllm', + 'vllm_async', {'model_id': 'local-model'}, - 'http://data-plane', + None, ) assert sampler == 'vllm-tq' assert captured == {'model_id': 'local-model', 'context_manager': None} -def test_vllm_without_data_plane_keeps_standard_backend(monkeypatch) -> None: +def test_standard_vllm_is_independent_of_data_plane(monkeypatch) -> None: calls = [] monkeypatch.setattr( SAMPLER_SELECTOR, @@ -143,7 +143,7 @@ def test_vllm_without_data_plane_keeps_standard_backend(monkeypatch) -> None: lambda sampler_type, kwargs: calls.append((sampler_type, kwargs)) or 'standard-vllm', ) - sampler = _construct_sampler_backend('vllm', {'model_id': 'local-model'}, None) + sampler = _construct_sampler_backend('vllm', {'model_id': 'local-model'}, 'http://data-plane') assert sampler == 'standard-vllm' assert calls == [('vllm', {'model_id': 'local-model'})] diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index 40f59c3e3..9a9e628f0 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -1,9 +1,15 @@ from __future__ import annotations import pytest +from fastapi import FastAPI +from starlette.requests import Request import twinkle_client.types as types -from twinkle.server.sampler.twinkle_handlers import _sample_models_to_rows +from twinkle.data_format import SampledSequence, SampleResponse +from twinkle.server.sampler.twinkle_handlers import ( + _register_twinkle_sampler_routes, + _sample_models_to_rows, +) def _response(tokens: list[int]) -> types.SampleResponseModel: @@ -30,6 +36,9 @@ def test_async_sampler_flattens_generations_to_tagged_tq_rows() -> None: ) assert [row['tokens'] for row in rows] == [[10], [11], [20], [21]] + assert [row['sampled_logprobs'] for row in rows] == [[-0.1]] * 4 + assert [row['train_input']['input_ids'] for row in rows] == [[10], [11], [20], [21]] + assert all('new_input_feature' not in row for row in rows) assert [(tag['group_id'], tag['generation_idx']) for tag in tags] == [ ('group-a', 0), ('group-a', 1), @@ -48,3 +57,89 @@ def test_async_sampler_rejects_group_id_count_mismatch() -> None: policy_version=0, adapter_uri=None, ) + + +class _SamplerManagement: + + def __init__(self): + self.sampler = self + self.data_plane = self + self.enabled = True + self.scheduled = [] + self.put_rows = None + + async def _on_request_start(self, _request): + return 'token' + + async def schedule_task_and_wait(self, task, **kwargs): + self.scheduled.append(kwargs) + return await task() + + def submit_generation(self, submission_id, inputs, params, **kwargs): + self.submission_id = submission_id + self.inputs = inputs + self.params = params + self.generation_kwargs = kwargs + + def get_generation_status(self, submission_id): + assert submission_id == self.submission_id + return {'status': 'completed'} + + def collect_generation(self, submission_id): + assert submission_id == self.submission_id + return [SampleResponse(sequences=[SampledSequence( + stop_reason='stop', + tokens=[7], + logprobs=[[(7, -0.25)]], + decoded='answer', + new_input_feature={'input_ids': [1, 7], 'labels': [-100, 7]}, + )])] + + def cancel_generation(self, _submission_id): + raise AssertionError('completed generation must not be cancelled') + + async def put(self, rows, *, kind, tags): + self.put_rows = rows + self.put_tags = tags + return types.DataRef(ref_id='rollout-ref', size=len(rows), fields=list(rows[0]), kind=kind) + + +@pytest.mark.asyncio +async def test_sample_to_data_plane_returns_ref_after_short_admission() -> None: + management = _SamplerManagement() + app = FastAPI() + _register_twinkle_sampler_routes(app, lambda: management) + route = next( + route for route in app.routes + if getattr(route, 'path', None) == '/twinkle/sample_to_data_plane' + ) + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.DataPlaneSampleRequest( + inputs=[{'input_ids': [1]}], + adapter_name='adapter', + group_ids=['group-1'], + policy_version=3, + num_samples=1, + sampling_params={'max_tokens': 4}, + ) + + ref = await route.endpoint(request, body, management) + + assert ref.ref_id == 'rollout-ref' + assert management.scheduled == [{ + 'model_id': 'session-adapter', + 'token': 'token', + 'input_tokens': 1, + 'task_type': 'sample_admission', + }] + assert management.put_rows == [{ + 'train_input': {'input_ids': [1, 7], 'labels': [-100, 7]}, + 'sampled_logprobs': [-0.25], + 'tokens': [7], + 'decoded': 'answer', + 'stop_reason': 'stop', + 'prompt_logprobs': None, + 'topk_prompt_logprobs': None, + }] + assert management.put_tags[0]['group_id'] == 'group-1' diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index 29fa62f0b..7e5749261 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -157,24 +157,12 @@ def cancel_generation(self, submission_id): sampler = Sampler() async def run(): + sampler.submit_generation('first') + sampler.submit_generation('second') first = asyncio.create_task( - _await_generation( - sampler, - 'first', - [{'input_ids': [1]}], - SamplingParams(max_tokens=4), - adapter_name='', - adapter_path=None, - )) + _await_generation(sampler, 'first')) second = asyncio.create_task( - _await_generation( - sampler, - 'second', - [{'input_ids': [2]}], - SamplingParams(max_tokens=4), - adapter_name='', - adapter_path=None, - )) + _await_generation(sampler, 'second')) while len(sampler.submission_order) < 2: await asyncio.sleep(0) assert not first.done() @@ -186,3 +174,37 @@ async def run(): asyncio.run(run()) assert set(sampler.submission_order) == {'first', 'second'} + + +def test_server_waiter_retries_cancelled_status_poll() -> None: + from ray.exceptions import TaskCancelledError + + class Sampler: + + def __init__(self): + self.status_calls = 0 + self.cancelled = False + + def submit_generation(self, *_args, **_kwargs): + return None + + def get_generation_status(self, _submission_id): + self.status_calls += 1 + if self.status_calls == 1: + raise TaskCancelledError() + return {'status': 'completed'} + + def collect_generation(self, _submission_id): + return ['completed'] + + def cancel_generation(self, _submission_id): + self.cancelled = True + + sampler = Sampler() + sampler.submit_generation('submission') + result = asyncio.run( + _await_generation(sampler, 'submission')) + + assert result == ['completed'] + assert sampler.status_calls == 2 + assert not sampler.cancelled diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 9b51e2b11..c4a1ea6a8 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -3,9 +3,7 @@ import asyncio -import pytest - -from twinkle_client.types import ComponentTaskRef, DataRef, DataRowsResponse +from twinkle_client.types import DataRef class _Response: @@ -22,144 +20,94 @@ def json(self): return self._payload -def test_remote_task_uses_existing_future_endpoint(monkeypatch) -> None: - from twinkle_client import remote_task as module +def test_model_forward_backward_sends_multiple_data_refs(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module - responses = iter([_Response({'type': 'try_again'}), _Response({'value': 7})]) calls = [] - def post(url, json_data, **kwargs): - calls.append((url, json_data, kwargs)) - return next(responses) + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + if url.endswith('/create'): + return _Response({}) + return _Response({'result': {'loss': 1.0}}) + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') monkeypatch.setattr(module, 'http_post', post) - monkeypatch.setattr(module, 'get_base_url', lambda: 'http://server/api/v1') - - task = module.RemoteTask(ComponentTaskRef(request_id='req-1', model_id='adapter')) - assert task.result(timeout=1) == {'value': 7} - assert [call[:2] for call in calls] == [ - ('http://server/api/v1/retrieve_future', {'request_id': 'req-1'}), - ('http://server/api/v1/retrieve_future', {'request_id': 'req-1'}), - ] - assert all(0 < call[2]['timeout'] <= 1 for call in calls) - - -def test_remote_task_async_result_uses_non_blocking_http_polling(monkeypatch) -> None: - import httpx - from twinkle_client import remote_task as module - responses = iter([_Response({'type': 'try_again'}), _Response({'value': 9})]) - calls = [] - - class AsyncClient: - - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def post(self, url, *, headers, json, timeout): - calls.append((url, headers, json, timeout)) - return next(responses) - - monkeypatch.setattr(httpx, 'AsyncClient', AsyncClient) - monkeypatch.setattr(module, 'get_base_url', lambda: 'http://server/api/v1') - monkeypatch.setattr(module, '_build_headers', lambda: {'Authorization': 'Bearer token'}) - - task = module.RemoteTask(ComponentTaskRef(request_id='req-async', model_id='adapter')) - assert asyncio.run(task.aresult(timeout=1)) == {'value': 9} - assert [call[2] for call in calls] == [ - {'request_id': 'req-async'}, - {'request_id': 'req-async'}, + model = module.MultiLoraTransformersModel('ms://base') + model.adapter_name = 'adapter' + refs = [ + DataRef(ref_id='data-1', size=2, fields=['train_input']), + DataRef(ref_id='data-2', size=2, fields=['train_input']), ] - assert all(0 < call[3] <= 1 for call in calls) - - -def test_remote_task_sync_timeout_bounds_long_poll(monkeypatch) -> None: - import requests - from twinkle_client import remote_task as module - - observed = [] - - def post(_url, *, json_data, timeout): - observed.append((json_data, timeout)) - raise requests.Timeout('long poll exceeded client deadline') - - monkeypatch.setattr(module, 'http_post', post) - monkeypatch.setattr(module, 'get_base_url', lambda: 'http://server/api/v1') + model.forward_backward( + refs, + input_field='train_input', + kwarg_fields={'advantages': 'advantage'}, + ) - task = module.RemoteTask('req-timeout') - with pytest.raises(TimeoutError, match='within 0.1s'): - task.result(timeout=0.1) - assert observed[0][0] == {'request_id': 'req-timeout'} - assert 0 < observed[0][1] <= 0.1 + url, body = calls[-1] + assert url.endswith('/model/base/twinkle/forward_backward') + assert body['input_refs'] == [ref.model_dump() for ref in refs] + assert body['input_field'] == 'train_input' + assert body['kwarg_fields'] == {'advantages': 'advantage'} + assert body['adapter_name'] == 'adapter' -def test_model_component_submits_data_ref_without_control_plane(monkeypatch) -> None: +def test_model_forward_accepts_data_ref_without_a_separate_submit_api(monkeypatch) -> None: import twinkle_client.http as http_module from twinkle_client.model import multi_lora_transformers as module calls = [] - def post(*, url, json_data=None, **_kwargs): + def post(url, json_data=None, **_kwargs): calls.append((url, json_data)) - if url.endswith('/create'): - return _Response({}) - return _Response({'request_id': 'req-model', 'model_id': 'session-adapter'}) + return _Response({} if url.endswith('/create') else {'result': {'value': 1}}) monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') monkeypatch.setattr(module, 'http_post', post) - monkeypatch.setattr('twinkle_client.remote_task.get_base_url', lambda: 'http://server/api/v1') model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' - ref = DataRef(ref_id='data-1', size=4, fields=['input_ids']) - task = model.submit_forward_backward(ref, advantages=[1, -1, 1, -1]) + ref = DataRef(ref_id='data-1', size=2, fields=['train_input']) + model.forward(ref, input_field='train_input') - assert task.request_id == 'req-model' url, body = calls[-1] - assert url.endswith('/model/base/twinkle/submit_forward_backward') - assert body['input_ref'] == ref.model_dump() - assert body['adapter_name'] == 'adapter' - assert body['kwargs']['advantages'] == [1, -1, 1, -1] + assert url.endswith('/model/base/twinkle/forward') + assert body['input_refs'] == [ref.model_dump()] + assert body['input_field'] == 'train_input' -def test_sampler_component_fetches_and_releases_data_plane_output(monkeypatch) -> None: +def test_sampler_async_data_plane_path_returns_reference_without_materializing(monkeypatch) -> None: import twinkle_client.http as http_module from twinkle_client.sampler import vllm_sampler as module + output_ref = DataRef( + ref_id='rollout-1', + size=4, + fields=['train_input', 'sampled_logprobs', 'decoded'], + kind='rollout', + ) + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + if url.endswith('/create'): + return _Response({}) + return _Response(output_ref.model_dump()) + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') - monkeypatch.setattr(module, 'http_post', lambda **_kwargs: _Response({})) + monkeypatch.setattr(module, 'http_post', post) sampler = module.vLLMSampler('ms://base') - output_ref = DataRef(ref_id='rollout-1', size=1, fields=['tokens'], kind='rollout') - - class _DoneTask: - - async def aresult(self): - return {'output_ref': output_ref.model_dump()} - - released = [] - monkeypatch.setattr(sampler, 'submit_sample', lambda *_args, **_kwargs: _DoneTask()) - monkeypatch.setattr( - sampler.data_plane, - 'get_batch', - lambda ref: DataRowsResponse( - rows=[{ - 'tokens': [], - 'stop_reason': 'stop', - }], - tags=[{'prompt_index': 0, 'generation_idx': 0}], - ), - ) - monkeypatch.setattr(sampler.data_plane, 'release', lambda ref: released.append(ref)) + result = asyncio.run(sampler.asample_to_data_plane( + [{'input_ids': [1]}], + num_samples=4, + group_ids=['group-1'], + )) - responses = asyncio.run(sampler.asample([{'input_ids': [1]}])) - assert len(responses) == 1 - assert len(responses[0].sequences) == 1 - assert responses[0].sequences[0].tokens == [] - assert released == [output_ref] + assert result == output_ref + url, body = calls[-1] + assert url.endswith('/sampler/base/twinkle/sample_to_data_plane') + assert body['num_samples'] == 4 + assert body['group_ids'] == ['group-1'] diff --git a/tests/twinkle_client/test_client_orchestrated_dpo.py b/tests/twinkle_client/test_client_orchestrated_dpo.py index 3d6ac460a..1e7426fd0 100644 --- a/tests/twinkle_client/test_client_orchestrated_dpo.py +++ b/tests/twinkle_client/test_client_orchestrated_dpo.py @@ -1,7 +1,6 @@ import asyncio from cookbook.client.async_rl.client_orchestrated_dpo import ( - _extract_ref_outputs, prepare_dpo_batch, run_dpo, ) @@ -28,15 +27,6 @@ def test_prepare_dpo_batch_interleaves_complete_pairs() -> None: assert [row['input_ids'] for row in rows] == [[1, 2], [1, 3], [4], [5]] -def test_extract_ref_outputs_unwraps_data_plane_row() -> None: - ref_outputs = _extract_ref_outputs( - {'output_ref': {'ref_id': 'unused'}}, - [{'result': {'logps': [[-1.0, -2.0], [-3.0, -4.0]], 'logits': None}}], - ) - - assert ref_outputs == {'logps': [[-1.0, -2.0], [-3.0, -4.0]]} - - def test_dpo_roles_overlap_reference_and_training(monkeypatch) -> None: import cookbook.client.async_rl.client_orchestrated_dpo as module @@ -77,24 +67,26 @@ def __init__(self): self.steps = 0 self.forward_backward_kwargs = [] - async def submit_forward_only(self, ref, **_kwargs): + async def forward_only(self, ref, **kwargs): self.references += 1 name = f'reference-{self.references}' events.append(f'{name}-start') if self.references == 2: await first_train_started.wait() events.append(f'{name}-done') - return {'logps': [[-0.1]] * ref.size} + assert kwargs['output_ref'] == ref + assert kwargs['output_fields'] == {'logps': 'ref_logps'} + return ref.model_copy(update={'fields': [*ref.fields, 'ref_logps']}) - async def submit_forward_backward(self, _ref, **kwargs): + async def forward_backward(self, _ref, **kwargs): self.forward_backward_kwargs.append(kwargs) events.append('train-start') first_train_started.set() - async def submit_clip_grad_and_step(self, **_kwargs): + async def clip_grad_and_step(self, **_kwargs): self.steps += 1 - async def submit_save(self, name, **_kwargs): + async def save(self, name, **_kwargs): return {'twinkle_path': name} batches = [ @@ -109,8 +101,8 @@ async def submit_save(self, name, **_kwargs): assert events.index('train-start') < events.index('reference-2-done') assert model.steps == 2 assert model.forward_backward_kwargs == [ - {'ref_outputs': {'logps': [[-0.1], [-0.1]]}}, - {'ref_outputs': {'logps': [[-0.1], [-0.1]]}}, + {'kwarg_fields': {'ref_outputs.logps': 'ref_logps'}}, + {'kwarg_fields': {'ref_outputs.logps': 'ref_logps'}}, ] assert saved == {'twinkle_path': 'dpo-policy-2'} assert len(data_plane.released) == 2 diff --git a/tests/twinkle_client/test_client_orchestrated_grpo.py b/tests/twinkle_client/test_client_orchestrated_grpo.py index 6b9aa4dd2..90210e02d 100644 --- a/tests/twinkle_client/test_client_orchestrated_grpo.py +++ b/tests/twinkle_client/test_client_orchestrated_grpo.py @@ -5,7 +5,7 @@ import sys from pathlib import Path -from twinkle_client.types import DataRef, DataRowsResponse +from twinkle_client.types import DataRef MODULE_PATH = ( @@ -44,7 +44,7 @@ async def fake_rollout(_sampler, prompt, policy, _semaphore, _group_id): return DataRef( ref_id=name, size=module.NUM_GENERATIONS, - fields=['new_input_feature', 'logprobs'], + fields=['train_input', 'sampled_logprobs', 'decoded'], kind='rollout', ) @@ -62,37 +62,29 @@ def __init__(self): self.steps = 0 self.forward_backward_kwargs = [] - async def submit_save(self, name): + async def save(self, name): self.saved.append(name) return {'twinkle_path': f'/checkpoints/{name}'} - async def submit_forward_backward(self, _ref, **kwargs): + async def forward_backward(self, _refs, **kwargs): self.forward_backward_kwargs.append(kwargs) events.append('train') first_train_started.set() - async def submit_clip_grad_and_step(self, **_kwargs): + async def clip_grad_and_step(self, **_kwargs): self.steps += 1 class FakeDataPlane: def __init__(self): self.released = [] - async def aget_batch(self, ref): - return DataRowsResponse( - rows=[{ - 'new_input_feature': {'name': f'{ref.ref_id}-{index}'}, - 'logprobs': [[(0, 0.0)]], - } for index in range(ref.size)], - tags=[{'group_id': ref.ref_id, 'generation_idx': index} for index in range(ref.size)], - ) + async def aget(self, ref, *, fields=None): + assert fields == ['decoded'] + return [{'decoded': f'{ref.ref_id}-{index}'} for index in range(ref.size)] - async def aappend(self, ref, rows, *, tags): + async def aappend(self, ref, rows, **_kwargs): return ref.model_copy(update={'fields': [*ref.fields, *rows[0]]}) - async def aput(self, rows, *, kind, tags=None): - return DataRef(ref_id=f'{kind}-{id(rows)}', size=len(rows), fields=list(rows[0]), kind=kind) - async def arelease(self, ref): self.released.append(ref) @@ -112,8 +104,11 @@ async def run(): assert events.index('train') < events.index('rollout-done:p0-g1') assert model.saved == ['policy-0', 'policy-1', 'policy-2', 'policy-3'] assert model.steps == 6 - assert all('old_logps' in kwargs and 'advantages' in kwargs - for kwargs in model.forward_backward_kwargs) + assert all(kwargs['input_field'] == 'train_input' for kwargs in model.forward_backward_kwargs) + assert all(kwargs['kwarg_fields'] == { + 'old_logps': 'sampled_logprobs', + 'advantages': 'advantage', + } for kwargs in model.forward_backward_kwargs) assert len(data_plane.released) == 6 snapshots = {name: (version, uri) for name, version, uri in rollout_snapshots} @@ -141,7 +136,7 @@ async def fake_rollout(_sampler, prompt, _policy, _semaphore, _group_id): return DataRef( ref_id=name, size=1, - fields=['new_input_feature', 'logprobs'], + fields=['train_input', 'sampled_logprobs', 'decoded'], kind='rollout', ) @@ -153,32 +148,24 @@ class FakeModel: def __init__(self): self.saved = [] - async def submit_save(self, name): + async def save(self, name): self.saved.append(name) return {'twinkle_path': name} - async def submit_forward_backward(self, _ref, **_kwargs): + async def forward_backward(self, _ref, **_kwargs): return None - async def submit_clip_grad_and_step(self, **_kwargs): + async def clip_grad_and_step(self, **_kwargs): return None class FakeDataPlane: - async def aget_batch(self, ref): - return DataRowsResponse( - rows=[{ - 'new_input_feature': {'name': ref.ref_id}, - 'logprobs': [[(0, 0.0)]], - }], - tags=[{'group_id': ref.ref_id, 'generation_idx': 0}], - ) + async def aget(self, ref, *, fields=None): + assert fields == ['decoded'] + return [{'decoded': ref.ref_id}] - async def aappend(self, ref, rows, *, tags): + async def aappend(self, ref, rows, **_kwargs): return ref.model_copy(update={'fields': [*ref.fields, *rows[0]]}) - async def aput(self, rows, *, kind, tags=None): - return DataRef(ref_id=kind, size=len(rows), fields=list(rows[0]), kind=kind) - async def arelease(self, _ref): return None From 645c72e0f70d394125c073c636a7495f1651b9c1 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 11 Aug 2026 16:02:57 +0800 Subject: [PATCH 05/20] wip --- src/twinkle/server/model/twinkle_handlers.py | 30 +++++++- .../server/model/test_twinkle_async_inputs.py | 74 ++++++++++++++++++- 2 files changed, 99 insertions(+), 5 deletions(-) diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 86ee6845c..89605aa93 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -13,6 +13,7 @@ import traceback from collections.abc import Callable from fastapi import Depends, FastAPI, HTTPException, Request +from numbers import Number from pathlib import Path from peft import LoraConfig from typing import TYPE_CHECKING, Any @@ -93,6 +94,30 @@ def _set_at_path(target: dict[str, Any], path: str, value: Any) -> None: current[parts[-1]] = value +def _restore_dataref_value(value: Any) -> Any: + """Restore numeric DataPlane fields to tensors before model dispatch. + + DataPlane rows cross an HTTP/JSON boundary, so tensor-valued fields arrive + as Python lists. Fields selected through ``kwarg_fields`` are model data, + not request configuration: rebuild rectangular numeric arrays as tensors + and keep ragged arrays as lists of tensors for losses that align each sample + independently. + """ + if isinstance(value, dict): + return {key: _restore_dataref_value(item) for key, item in value.items()} + if not isinstance(value, list) or not value: + return value + if all(isinstance(item, Number) for item in value): + return torch.tensor(value) + + restored = [_restore_dataref_value(item) for item in value] + if all(torch.is_tensor(item) for item in restored): + shapes = {tuple(item.shape) for item in restored} + if len(shapes) == 1: + return torch.stack(restored) + return restored + + async def _resolve_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: """Resolve transport-level references before entering the model backend.""" if body.input_refs is None: @@ -120,10 +145,13 @@ async def _resolve_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[s field_kwargs: dict[str, Any] = {} for target_path, source_path in body.kwarg_fields.items(): + field_value = _restore_dataref_value([ + _value_at_path(row, source_path) for row in rows + ]) _set_at_path( field_kwargs, target_path, - [_value_at_path(row, source_path) for row in rows], + field_value, ) return inputs, field_kwargs diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index bc0790565..fd9917e76 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -1,6 +1,7 @@ from __future__ import annotations import pytest +import torch from fastapi import FastAPI from starlette.requests import Request @@ -8,6 +9,7 @@ from twinkle.server.model.twinkle_handlers import ( _model_result_rows, _register_twinkle_routes, + _restore_dataref_value, ) @@ -21,6 +23,26 @@ def test_model_result_rows_keeps_one_output_row_per_sample() -> None: ] +def test_restore_dataref_value_rebuilds_rectangular_and_ragged_numeric_arrays() -> None: + rectangular = _restore_dataref_value([ + [-0.1, -0.2], + [-0.3, -0.4], + ]) + torch.testing.assert_close( + rectangular, + torch.tensor([[-0.1, -0.2], [-0.3, -0.4]]), + ) + + ragged = _restore_dataref_value([ + [-0.1], + [-0.2, -0.3], + ]) + assert isinstance(ragged, list) + assert all(torch.is_tensor(item) for item in ragged) + torch.testing.assert_close(ragged[0], torch.tensor([-0.1])) + torch.testing.assert_close(ragged[1], torch.tensor([-0.2, -0.3])) + + class _SchedulingManagement: def __init__(self): @@ -91,7 +113,51 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] assert adapter_name == 'session-adapter' assert [row['input_ids'].tolist() for row in inputs] == [[index] for index in range(8)] - assert forwarded_kwargs == { - 'old_logps': [[-0.1]] * 4 + [[-0.2]] * 4, - 'advantages': [1.0] * 4 + [-1.0] * 4, - } + torch.testing.assert_close( + forwarded_kwargs['old_logps'], + torch.tensor([[-0.1]] * 4 + [[-0.2]] * 4), + ) + torch.testing.assert_close( + forwarded_kwargs['advantages'], + torch.tensor([1.0] * 4 + [-1.0] * 4), + ) + + +@pytest.mark.asyncio +async def test_forward_backward_restores_nested_dpo_ref_logps_as_tensor() -> None: + management = _SchedulingManagement() + management.rows['dpo'] = [ + { + 'input_ids': [1, 2, 3], + 'labels': [-100, 2, 3], + 'ref_logps': [-0.1, -0.2, -0.3], + }, + { + 'input_ids': [1, 4, 5], + 'labels': [-100, 4, 5], + 'ref_logps': [-0.4, -0.5, -0.6], + }, + ] + app = FastAPI() + _register_twinkle_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.ForwardRequest( + adapter_name='adapter', + input_refs=[types.DataRef(ref_id='dpo', size=2, num_tokens=6)], + kwarg_fields={'ref_outputs.logps': 'ref_logps'}, + ) + + await route.endpoint(request, body, management) + + inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] + assert adapter_name == 'session-adapter' + assert [row['input_ids'].tolist() for row in inputs] == [[1, 2, 3], [1, 4, 5]] + torch.testing.assert_close( + forwarded_kwargs['ref_outputs']['logps'], + torch.tensor([ + [-0.1, -0.2, -0.3], + [-0.4, -0.5, -0.6], + ]), + ) From 61e57b83682f2006d9ab5dccd09fa9d6f7c11e95 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 11 Aug 2026 19:25:38 +0800 Subject: [PATCH 06/20] wip --- src/twinkle/server/utils/task_queue/mixin.py | 184 +++++++++--------- src/twinkle/server/utils/task_queue/types.py | 8 + src/twinkle/server/utils/task_queue/worker.py | 58 ++++-- tests/server/utils/test_task_queue_mixin.py | 103 ++++++++++ 4 files changed, 244 insertions(+), 109 deletions(-) diff --git a/src/twinkle/server/utils/task_queue/mixin.py b/src/twinkle/server/utils/task_queue/mixin.py index 05aa45fab..84da48cb3 100644 --- a/src/twinkle/server/utils/task_queue/mixin.py +++ b/src/twinkle/server/utils/task_queue/mixin.py @@ -17,6 +17,7 @@ from twinkle.server.telemetry.middleware import get_task_metrics from twinkle.server.utils.task_errors import task_error_payload from twinkle.utils.logger import get_logger + from .config import TaskQueueConfig from .rate_limiter import RateLimiter from .types import QueuedTask, QueueState, TaskStatus @@ -98,6 +99,7 @@ async def _perform_preflight_checks( batch_size: int | None = None, data_world_size: int | None = None, batch_size_multiple: int | None = None, + persist_failure: bool = True, ) -> dict[str, Any] | None: """Run rate-limit and validation checks before queuing a task. @@ -106,69 +108,53 @@ async def _perform_preflight_checks( if not token or not self._task_queue_config.enabled: return None - if input_tokens > self._task_queue_config.max_input_tokens: - error_msg = (f'Input tokens ({input_tokens}) exceed maximum allowed ' - f'({self._task_queue_config.max_input_tokens})') + async def reject(error_msg: str, queue_state: str) -> dict[str, Any]: error_payload = {'error': error_msg, 'category': 'User'} - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - result=error_payload, - queue_state=QueueState.UNKNOWN.value, - queue_state_reason=error_msg, - ) - return {'request_id': request_id, 'model_id': model_id} - - if batch_size is not None and data_world_size is not None: - if batch_size < data_world_size: - error_msg = (f'Batch size {batch_size} must be >= data world size {data_world_size}') - error_payload = {'error': error_msg, 'category': 'User'} + if persist_failure: await self.state.store_future_status( request_id, TaskStatus.FAILED.value, model_id, result=error_payload, - queue_state=QueueState.UNKNOWN.value, + queue_state=queue_state, queue_state_reason=error_msg, ) return {'request_id': request_id, 'model_id': model_id} + # Private marker consumed by schedule_task_and_wait(). It is not + # returned by the public polling-style schedule_task() API. + return { + 'request_id': request_id, + 'model_id': model_id, + '_error': error_msg, + } + + if input_tokens > self._task_queue_config.max_input_tokens: + error_msg = (f'Input tokens ({input_tokens}) exceed maximum allowed ' + f'({self._task_queue_config.max_input_tokens})') + return await reject(error_msg, QueueState.UNKNOWN.value) + + if batch_size is not None and data_world_size is not None: + if batch_size < data_world_size: + error_msg = (f'Batch size {batch_size} must be >= data world size {data_world_size}') + return await reject(error_msg, QueueState.UNKNOWN.value) if batch_size_multiple is not None: required_multiple = data_world_size * batch_size_multiple if batch_size % required_multiple != 0: error_msg = (f'Batch size {batch_size} must be divisible by {required_multiple} ' f'so each data-parallel shard gets a multiple of ' f'{batch_size_multiple} examples') - error_payload = {'error': error_msg, 'category': 'User'} - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - result=error_payload, - queue_state=QueueState.UNKNOWN.value, - queue_state_reason=error_msg, - ) - return {'request_id': request_id, 'model_id': model_id} + return await reject(error_msg, QueueState.UNKNOWN.value) allowed, reason = await self._rate_limiter.check_and_record(token, input_tokens) if not allowed: if self._task_metrics: self._task_metrics.rate_limit_rejections.inc(tags={'deployment': self._deployment_name}) error_msg = f'Rate limit exceeded: {reason}' - error_payload = {'error': error_msg, 'category': 'User'} - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - result=error_payload, - queue_state=QueueState.PAUSED_RATE_LIMIT.value, - queue_state_reason=error_msg, - ) - return {'request_id': request_id, 'model_id': model_id} + return await reject(error_msg, QueueState.PAUSED_RATE_LIMIT.value) return None - async def schedule_task( + async def _schedule_task( self, coro_factory: Callable[[], Coroutine], model_id: str | None = None, @@ -178,25 +164,11 @@ async def schedule_task( data_world_size: int | None = None, batch_size_multiple: int | None = None, task_type: str | None = None, + *, + completion: asyncio.Future[Any] | None = None, + persist_status: bool, ) -> dict[str, Any]: - """Schedule a GPU compute task through the serial compute queue. - - Tasks are processed one at a time in round-robin order across all - per-adapter/per-token queues. Use for any operation that touches GPU - state: forward, backward, step, save, load, add_adapter, etc. - - Args: - coro_factory: Zero-argument callable that creates the coroutine. - model_id: Adapter/model id for queue routing and result association. - token: User token for rate limiting. - input_tokens: Token count for TPS rate limiting. - batch_size: Optional batch size, validated against data_world_size. - data_world_size: Optional data world size for batch validation. - task_type: Label for logging and metrics. - - Returns: - {'request_id': str, 'model_id': str | None} - """ + """Common enqueue path for polling and in-process wait callers.""" request_id = f'req_{uuid.uuid4().hex}' preflight_result = await self._perform_preflight_checks( @@ -207,6 +179,7 @@ async def schedule_task( batch_size=batch_size, data_world_size=data_world_size, batch_size_multiple=batch_size_multiple, + persist_failure=persist_status, ) if preflight_result is not None: return preflight_result @@ -214,12 +187,13 @@ async def schedule_task( if self._event_loop is None: self._event_loop = asyncio.get_running_loop() - await self.state.store_future_status( - request_id, - TaskStatus.PENDING.value, - model_id, - queue_state=QueueState.ACTIVE.value, - ) + if persist_status: + await self.state.store_future_status( + request_id, + TaskStatus.PENDING.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) queue_key = self._queue_key(model_id=model_id, token=token) self._compute_worker.ensure_queue_registered(queue_key) @@ -235,13 +209,16 @@ async def schedule_task( input_tokens=input_tokens, task_type=task_type, created_at=time.monotonic(), + completion=completion, + persist_status=persist_status, )) - await self.state.store_future_status( - request_id, - TaskStatus.QUEUED.value, - model_id, - queue_state=QueueState.ACTIVE.value, - ) + if persist_status: + await self.state.store_future_status( + request_id, + TaskStatus.QUEUED.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) logger.info(f'[TaskQueue] Task {request_id} queued, type={task_type or "unknown"}, ' f'model_id={model_id}, queue_key={queue_key}, ' f'queue_depth={q.qsize()}, input_tokens={input_tokens}') @@ -254,6 +231,47 @@ async def schedule_task( return {'request_id': request_id, 'model_id': model_id} + async def schedule_task( + self, + coro_factory: Callable[[], Coroutine], + model_id: str | None = None, + token: str | None = None, + input_tokens: int = 0, + batch_size: int | None = None, + data_world_size: int | None = None, + batch_size_multiple: int | None = None, + task_type: str | None = None, + ) -> dict[str, Any]: + """Schedule a GPU compute task through the serial compute queue. + + Tasks are processed one at a time in round-robin order across all + per-adapter/per-token queues. Use for any operation that touches GPU + state: forward, backward, step, save, load, add_adapter, etc. + + Args: + coro_factory: Zero-argument callable that creates the coroutine. + model_id: Adapter/model id for queue routing and result association. + token: User token for rate limiting. + input_tokens: Token count for TPS rate limiting. + batch_size: Optional batch size, validated against data_world_size. + data_world_size: Optional data world size for batch validation. + task_type: Label for logging and metrics. + + Returns: + {'request_id': str, 'model_id': str | None} + """ + return await self._schedule_task( + coro_factory, + model_id=model_id, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=data_world_size, + batch_size_multiple=batch_size_multiple, + task_type=task_type, + persist_status=True, + ) + async def schedule_task_and_wait( self, coro_factory: Callable[[], Coroutine], @@ -268,12 +286,14 @@ async def schedule_task_and_wait( """Schedule a compute task and block until it completes. Twinkle-side counterpart to schedule_task(). Enqueues the task through - the serial worker, polls until a terminal state, and returns the result. + the same serial worker but delivers the result through an in-process + Future. Large model outputs therefore never enter ServerState. Raises: RuntimeError: If the task fails or scheduling is rejected. """ - future_ref = await self.schedule_task( + completion = asyncio.get_running_loop().create_future() + task_ref = await self._schedule_task( coro_factory, model_id=model_id, token=token, @@ -282,25 +302,13 @@ async def schedule_task_and_wait( data_world_size=data_world_size, batch_size_multiple=batch_size_multiple, task_type=task_type, + completion=completion, + persist_status=False, ) - request_id = future_ref.get('request_id') - if request_id is None: - raise RuntimeError(f'Task scheduling failed: {future_ref}') - - poll_interval = 0.05 - max_poll_interval = 1.0 - while True: - record = await self.state.get_future(request_id) - if record and record.get('status') not in ('pending', 'queued', 'running'): - break - await asyncio.sleep(poll_interval) - poll_interval = min(poll_interval * 2, max_poll_interval) - - if record['status'] == 'failed': - error = record.get('result', {}).get('error', 'Unknown error') + if error := task_ref.get('_error'): + completion.cancel() raise RuntimeError(error) - - return record['result'] + return await completion async def schedule_background_task( self, diff --git a/src/twinkle/server/utils/task_queue/types.py b/src/twinkle/server/utils/task_queue/types.py index 4d9a7aa5d..daf8d2bb2 100644 --- a/src/twinkle/server/utils/task_queue/types.py +++ b/src/twinkle/server/utils/task_queue/types.py @@ -9,9 +9,11 @@ """ from __future__ import annotations +import asyncio from collections.abc import Callable, Coroutine from dataclasses import dataclass from enum import Enum +from typing import Any class TaskStatus(Enum): @@ -47,3 +49,9 @@ class QueuedTask: task_type: str | None created_at: float first_rate_limited_at: float | None = None + # ``schedule_task_and_wait`` is an in-process request/response path. Its + # potentially large result is delivered through this Future instead of + # being persisted in ServerState merely for the same process to read it + # back. Polling-style ``schedule_task`` leaves this as ``None``. + completion: asyncio.Future[Any] | None = None + persist_status: bool = True diff --git a/src/twinkle/server/utils/task_queue/worker.py b/src/twinkle/server/utils/task_queue/worker.py index 4084febe2..58d06ce05 100644 --- a/src/twinkle/server/utils/task_queue/worker.py +++ b/src/twinkle/server/utils/task_queue/worker.py @@ -18,6 +18,7 @@ from twinkle.server.telemetry.tracing import traced_operation from twinkle.server.utils.task_errors import task_error_payload from twinkle.utils.logger import get_logger + from .config import TaskQueueConfig from .types import QueuedTask, QueueState, TaskStatus @@ -125,6 +126,16 @@ def _record_queue_metrics(self, task_type: str, queue_wait: float) -> None: # ------------------------------------------------------------------ + @staticmethod + def _complete_result(task: QueuedTask, result: Any) -> None: + if task.completion is not None and not task.completion.done(): + task.completion.set_result(result) + + @staticmethod + def _complete_error(task: QueuedTask, error: str) -> None: + if task.completion is not None and not task.completion.done(): + task.completion.set_exception(RuntimeError(error)) + async def _store_task_failed( self, task: QueuedTask, @@ -133,14 +144,16 @@ async def _store_task_failed( queue_state_reason: str | None = None, ) -> None: """Store FAILED status with a standardised error payload.""" - await self._state.store_future_status( - task.request_id, - TaskStatus.FAILED.value, - task.model_id, - result=task_error_payload(error), - queue_state=queue_state, - queue_state_reason=queue_state_reason, - ) + if task.persist_status: + await self._state.store_future_status( + task.request_id, + TaskStatus.FAILED.value, + task.model_id, + result=task_error_payload(error), + queue_state=queue_state, + queue_state_reason=queue_state_reason, + ) + self._complete_error(task, error) async def fail_queue_tasks(self, queue_key: str, reason: str) -> None: """Drain a queue and mark all pending tasks as FAILED.""" @@ -191,12 +204,13 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue Handles execution timeout, general exceptions, and always calls q.task_done() in the finally block. """ - await self._state.store_future_status( - task.request_id, - TaskStatus.RUNNING.value, - task.model_id, - queue_state=QueueState.ACTIVE.value, - ) + if task.persist_status: + await self._state.store_future_status( + task.request_id, + TaskStatus.RUNNING.value, + task.model_id, + queue_state=QueueState.ACTIVE.value, + ) task_type = task.task_type or 'unknown' exec_start = time.monotonic() @@ -225,13 +239,15 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue result = await coro exec_time = time.monotonic() - exec_start logger.info(f'[ComputeWorker] Task {task.request_id} completed in {exec_time:.2f}s, type={task_type}') - await self._state.store_future_status( - task.request_id, - TaskStatus.COMPLETED.value, - task.model_id, - result=result, - queue_state=QueueState.ACTIVE.value, - ) + if task.persist_status: + await self._state.store_future_status( + task.request_id, + TaskStatus.COMPLETED.value, + task.model_id, + result=result, + queue_state=QueueState.ACTIVE.value, + ) + self._complete_result(task, result) except asyncio.TimeoutError: task_status = 'timeout' exec_time = time.monotonic() - exec_start diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index 93f6c5e5d..2aa133a7e 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -4,6 +4,7 @@ from twinkle.server.utils.task_queue.config import TaskQueueConfig from twinkle.server.utils.task_queue.mixin import TaskQueueMixin +from twinkle.server.utils.task_queue.worker import ComputeWorker class _DummyState: @@ -30,6 +31,15 @@ def __init__(self): self._task_metrics = None self._deployment_name = 'test' + def enable_compute_worker(self): + self._compute_worker = ComputeWorker( + state=self.state, + config=self._task_queue_config, + task_metrics=None, + deployment_name=self._deployment_name, + ) + self._event_loop = None + @pytest.mark.asyncio async def test_preflight_rejects_batch_without_per_dp_multiple(): @@ -86,3 +96,96 @@ async def work(): assert queue.state.records[0][0][1] == 'running' assert all('token' not in kwargs and 'session_id' not in kwargs for _, kwargs in queue.state.records) + + +@pytest.mark.asyncio +async def test_schedule_task_and_wait_returns_large_result_without_persisting_it(): + queue = _DummyQueue() + queue.enable_compute_worker() + result = {'logps': [[float(index) for index in range(128)]]} + + async def work(): + return result + + try: + actual = await queue.schedule_task_and_wait( + work, + model_id='model1', + token='token1', + task_type='forward_backward', + ) + finally: + await queue._compute_worker.stop() + + assert actual is result + assert queue.state.records == [] + + +@pytest.mark.asyncio +async def test_polling_schedule_task_still_persists_its_result(): + queue = _DummyQueue() + queue.enable_compute_worker() + result = {'value': 42} + + async def work(): + return result + + try: + await queue.schedule_task(work, model_id='model1', token='token1') + for _ in range(100): + completed = [ + kwargs + for args, kwargs in queue.state.records + if args[1] == 'completed' + ] + if completed: + break + await asyncio.sleep(0) + finally: + await queue._compute_worker.stop() + + assert completed[-1]['result'] is result + + +@pytest.mark.asyncio +async def test_schedule_task_and_wait_propagates_failure_without_persisting_it(): + queue = _DummyQueue() + queue.enable_compute_worker() + + async def work(): + raise ValueError('model failed') + + try: + with pytest.raises(RuntimeError, match='ValueError: model failed'): + await queue.schedule_task_and_wait( + work, + model_id='model1', + token='token1', + task_type='forward_backward', + ) + finally: + await queue._compute_worker.stop() + + assert queue.state.records == [] + + +@pytest.mark.asyncio +async def test_schedule_task_and_wait_reports_preflight_failure_without_persisting_it(): + queue = _DummyQueue() + queue.enable_compute_worker() + + async def work(): + raise AssertionError('preflight rejection must not execute the task') + + with pytest.raises(RuntimeError, match='Batch size 2 must be divisible by 4'): + await queue.schedule_task_and_wait( + work, + model_id='model1', + token='token1', + batch_size=2, + data_world_size=2, + batch_size_multiple=2, + ) + + assert queue.state.records == [] + assert queue._compute_worker._worker_task is None From a658458522eb710e934e5a26cb2ba9bcc018b5f1 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 12 Aug 2026 11:30:21 +0800 Subject: [PATCH 07/20] wip --- src/twinkle/loss/dpo.py | 59 ++-------- src/twinkle/metric/dpo.py | 49 +++------ src/twinkle/server/model/twinkle_handlers.py | 43 +------- src/twinkle/utils/rl_tensor_utils.py | 101 ++++++++++++++++++ tests/loss/test_dpo.py | 9 ++ tests/metric/test_metrics.py | 20 ++++ .../server/model/test_twinkle_async_inputs.py | 49 ++------- tests/utils/test_rl_tensor_utils.py | 58 ++++++++++ 8 files changed, 224 insertions(+), 164 deletions(-) create mode 100644 src/twinkle/utils/rl_tensor_utils.py create mode 100644 tests/utils/test_rl_tensor_utils.py diff --git a/src/twinkle/loss/dpo.py b/src/twinkle/loss/dpo.py index 7a58acdf9..9099225a2 100644 --- a/src/twinkle/loss/dpo.py +++ b/src/twinkle/loss/dpo.py @@ -11,6 +11,7 @@ from twinkle.data_format import LossOutput from twinkle.loss.base import Loss +from twinkle.utils.rl_tensor_utils import align_per_token_values from twinkle.utils.torch_utils import selective_log_softmax if TYPE_CHECKING: @@ -144,49 +145,6 @@ def __init__( self.reference_free = reference_free self.sft_weight = sft_weight - def _align_logps( - self, - logps: 'torch.Tensor', - target_shape: tuple, - device: 'torch.device', - dtype: 'torch.dtype', - ) -> 'torch.Tensor': - """Align log probabilities to target shape. - - Args: - logps: Input log probabilities tensor - target_shape: Target (batch, seq_len) shape - device: Target device - dtype: Target dtype - - Returns: - Aligned tensor of shape target_shape - """ - import torch - - if not torch.is_tensor(logps): - raise TypeError(f'Expected torch.Tensor, got {type(logps)}') - - if logps.dim() == 1: - logps = logps.unsqueeze(0) - - if logps.shape == target_shape: - return logps.to(device=device, dtype=dtype) - - # Handle tensor with different sequence length - if logps.dim() == 2 and logps.shape[0] == target_shape[0]: - batch_size, target_seq_len = target_shape - src_seq_len = logps.shape[1] - logps = logps.to(device=device, dtype=dtype) - if src_seq_len > target_seq_len: - # Truncate right (keep left part) - may happen in Ray result merging - return logps[:, :target_seq_len] - else: - raise ValueError(f'ref_logps seq_len ({src_seq_len}) < target seq_len ({target_seq_len}). ' - f'This should not happen when both models process the same batch.') - - raise ValueError(f'Cannot align ref_logps shape {logps.shape} to target shape {target_shape}') - def _compute_dpo_loss( self, policy_chosen_logps: 'torch.Tensor', @@ -311,13 +269,18 @@ def __call__( # Handle reference log probs if ref_chosen_logps is not None and ref_rejected_logps is not None: # Pre-computed sequence-level reference log probs provided - reference_chosen_logps = ref_chosen_logps.to(device=device, dtype=dtype) - reference_rejected_logps = ref_rejected_logps.to(device=device, dtype=dtype) + reference_chosen_logps = torch.as_tensor(ref_chosen_logps, device=device, dtype=dtype) + reference_rejected_logps = torch.as_tensor(ref_rejected_logps, device=device, dtype=dtype) elif ref_logps is not None: # Per-token reference log probs provided, need to align and sum - if not torch.is_tensor(ref_logps): - ref_logps = torch.as_tensor(ref_logps) - ref_logps_aligned = self._align_logps(ref_logps, labels.shape, device, dtype) + ref_logps_aligned = align_per_token_values( + ref_logps, + tuple(labels.shape), + device=device, + dtype=dtype, + name='ref_logps', + valid_mask=labels != self.ignore_index, + ) ref_chosen, ref_rejected = self._split_chosen_rejected(ref_logps_aligned) reference_chosen_logps = self._compute_sequence_logps(ref_chosen, chosen_labels) reference_rejected_logps = self._compute_sequence_logps(ref_rejected, rejected_labels) diff --git a/src/twinkle/metric/dpo.py b/src/twinkle/metric/dpo.py index 024cb0473..3d16993e5 100644 --- a/src/twinkle/metric/dpo.py +++ b/src/twinkle/metric/dpo.py @@ -4,6 +4,8 @@ from twinkle.data_format import InputFeature, ModelOutput from twinkle.utils import pad_and_stack_tensors +from twinkle.utils.rl_tensor_utils import align_per_token_values + from .base import Metric @@ -33,38 +35,9 @@ def __init__(self, device_mesh, process_group, ignore_index: int = -100, beta: f def _compute_sequence_logps(self, per_token_logps, labels): """Compute sequence-level log probs by summing valid token logps.""" - import torch loss_mask = (labels != self.ignore_index).float() return (per_token_logps * loss_mask).sum(dim=-1) - def _align_logps(self, logps, target_shape, device, dtype): - """Align per-token logps to target shape by padding or truncating. - - Args: - logps: [batch, seq_len] tensor to align - target_shape: Target shape (batch, target_seq_len) - device: Target device - dtype: Target dtype - - Returns: - Aligned tensor with shape matching target_shape - """ - import torch - - if not torch.is_tensor(logps): - logps = torch.as_tensor(logps) - logps = logps.to(device=device, dtype=dtype) - batch_size, src_len = logps.shape - _, target_len = target_shape - - if src_len == target_len: - return logps - elif src_len < target_len: - raise ValueError(f'ref_logps seq_len ({src_len}) < target seq_len ({target_len}). ' - f'This should not happen when both models process the same batch.') - else: - return logps[:, :target_len] - def _split_chosen_rejected(self, tensor): """Split interleaved tensor into chosen and rejected. @@ -121,7 +94,6 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M # Split into chosen and rejected (interleaved format) chosen_logps, rejected_logps = self._split_chosen_rejected(seq_logps) - chosen_labels, rejected_labels = self._split_chosen_rejected(labels) # Accumulate policy logps self.total_chosen_logps += chosen_logps.sum().item() @@ -131,15 +103,18 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M ref_outputs = kwargs.get('ref_outputs') if ref_outputs is not None: ref_logps = ref_outputs.get('logps') - if ref_logps is not None: - if isinstance(ref_logps, list): - if len(ref_logps) == 0: - ref_logps = None - else: - ref_logps = pad_and_stack_tensors(ref_logps) + if isinstance(ref_logps, (list, tuple)) and not ref_logps: + ref_logps = None if ref_logps is not None: # Align ref_logps to match labels shape (handles different seq lengths) - ref_logps = self._align_logps(ref_logps, labels.shape, labels.device, logps.dtype) + ref_logps = align_per_token_values( + ref_logps, + tuple(labels.shape), + device=labels.device, + dtype=logps.dtype, + name='ref_logps', + valid_mask=labels != self.ignore_index, + ) ref_seq_logps = self._compute_sequence_logps(ref_logps, labels) ref_chosen_logps, ref_rejected_logps = self._split_chosen_rejected(ref_seq_logps) diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 89605aa93..ec485a9ee 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -9,11 +9,9 @@ from __future__ import annotations import asyncio -import torch import traceback from collections.abc import Callable from fastapi import Depends, FastAPI, HTTPException, Request -from numbers import Number from pathlib import Path from peft import LoraConfig from typing import TYPE_CHECKING, Any @@ -94,30 +92,6 @@ def _set_at_path(target: dict[str, Any], path: str, value: Any) -> None: current[parts[-1]] = value -def _restore_dataref_value(value: Any) -> Any: - """Restore numeric DataPlane fields to tensors before model dispatch. - - DataPlane rows cross an HTTP/JSON boundary, so tensor-valued fields arrive - as Python lists. Fields selected through ``kwarg_fields`` are model data, - not request configuration: rebuild rectangular numeric arrays as tensors - and keep ragged arrays as lists of tensors for losses that align each sample - independently. - """ - if isinstance(value, dict): - return {key: _restore_dataref_value(item) for key, item in value.items()} - if not isinstance(value, list) or not value: - return value - if all(isinstance(item, Number) for item in value): - return torch.tensor(value) - - restored = [_restore_dataref_value(item) for item in value] - if all(torch.is_tensor(item) for item in restored): - shapes = {tuple(item.shape) for item in restored} - if len(shapes) == 1: - return torch.stack(restored) - return restored - - async def _resolve_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: """Resolve transport-level references before entering the model backend.""" if body.input_refs is None: @@ -145,9 +119,11 @@ async def _resolve_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[s field_kwargs: dict[str, Any] = {} for target_path, source_path in body.kwarg_fields.items(): - field_value = _restore_dataref_value([ + # Keep the transport contract JSON-native. The processor/loss/metric + # that understands this field owns any tensor conversion. + field_value = [ _value_at_path(row, source_path) for row in rows - ]) + ] _set_at_path( field_kwargs, target_path, @@ -355,22 +331,11 @@ async def forward_backward( token = await self._on_request_start(request) adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) - def first_element(data): - while isinstance(data, list): - if len(data) == 0: - return None - data = data[0] - return data - async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) all_inputs = _parse_inputs(raw_inputs) - for inputs in all_inputs: - for key in inputs: - if isinstance(inputs[key], list) and isinstance(first_element(inputs[key]), (int, float)): - inputs[key] = torch.tensor(inputs[key]) kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) ret = self.model.forward_backward(inputs=all_inputs, adapter_name=adapter_name, **kwargs) return {'result': ret} diff --git a/src/twinkle/utils/rl_tensor_utils.py b/src/twinkle/utils/rl_tensor_utils.py new file mode 100644 index 000000000..5e262e5cc --- /dev/null +++ b/src/twinkle/utils/rl_tensor_utils.py @@ -0,0 +1,101 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tensor normalization helpers for JSON-compatible RL inputs.""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import torch + + +def align_per_token_values( + values: Any, + target_shape: tuple[int, int], + *, + device: torch.device, + dtype: torch.dtype, + name: str = 'values', + padding_value: float = 0.0, + valid_mask: Any | None = None, +) -> torch.Tensor: + """Convert tensor-like per-token values and align them to a model batch. + + Component HTTP APIs intentionally use JSON-compatible values, so tensors + sent by a client or read through DataPlane arrive as Python lists. The + computation layer calls this helper only for fields it knows are per-token + tensors. Rectangular values are converted directly; ragged rows are padded + before validating and aligning them to ``target_shape``. Missing suffixes + are accepted only when ``valid_mask`` marks those positions as padding. + """ + import torch + + row_lengths: list[int] | None = None + if torch.is_tensor(values): + tensor = values + else: + try: + tensor = torch.as_tensor(values) + except (TypeError, ValueError): + if not isinstance(values, (list, tuple)) or not values: + raise TypeError(f'{name} must be a tensor or a non-empty sequence') from None + + rows = [] + for index, value in enumerate(values): + try: + row = torch.as_tensor(value) + except (TypeError, ValueError) as exc: + raise TypeError(f'{name}[{index}] cannot be converted to a tensor') from exc + if row.ndim == 2 and row.shape[0] == 1: + row = row.squeeze(0) + if row.ndim != 1: + raise ValueError( + f'{name}[{index}] must be one-dimensional, got shape {tuple(row.shape)}') + rows.append(row) + row_lengths = [row.numel() for row in rows] + tensor = torch.nn.utils.rnn.pad_sequence( + rows, + batch_first=True, + padding_value=padding_value, + ) + + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + if tensor.ndim != 2: + raise ValueError(f'{name} must be two-dimensional, got shape {tuple(tensor.shape)}') + + target_batch_size, target_seq_len = target_shape + batch_size, seq_len = tensor.shape + if batch_size != target_batch_size: + raise ValueError( + f'{name} batch size ({batch_size}) does not match target batch size ' + f'({target_batch_size})') + mask = None + if valid_mask is not None: + mask = torch.as_tensor(valid_mask, dtype=torch.bool) + if tuple(mask.shape) != target_shape: + raise ValueError( + f'valid_mask shape {tuple(mask.shape)} does not match target shape ' + f'{target_shape}') + + if row_lengths is not None: + for index, row_len in enumerate(row_lengths): + if row_len >= target_seq_len: + continue + if mask is None or bool(mask[index, row_len:].any().item()): + raise ValueError( + f'{name}[{index}] has {row_len} tokens but target sequence length ' + f'is {target_seq_len}') + if seq_len < target_seq_len: + if mask is None or bool(mask[:, seq_len:].any().item()): + raise ValueError( + f'{name} seq_len ({seq_len}) is smaller than target seq_len ' + f'({target_seq_len})') + tensor = torch.nn.functional.pad( + tensor, + (0, target_seq_len - seq_len), + value=padding_value, + ) + if seq_len > target_seq_len: + tensor = tensor[:, :target_seq_len] + + return tensor.to(device=device, dtype=dtype) diff --git a/tests/loss/test_dpo.py b/tests/loss/test_dpo.py index f12d5ea3a..89b9ddac4 100644 --- a/tests/loss/test_dpo.py +++ b/tests/loss/test_dpo.py @@ -42,6 +42,15 @@ def test_basic_dpo_sigmoid(self): assert isinstance(result, dict) and 'loss' in result assert result['loss'].dim() == 0 + def test_json_ref_outputs_match_tensor_ref_outputs(self): + loss_fn = DPOLoss(beta=0.1, loss_type='sigmoid') + inputs, outputs, ref_logps = _make_preference_batch() + + tensor_result = loss_fn(inputs, outputs, ref_outputs={'logps': ref_logps}) + json_result = loss_fn(inputs, outputs, ref_outputs={'logps': ref_logps.tolist()}) + + torch.testing.assert_close(json_result['loss'], tensor_result['loss']) + def test_dpo_hinge(self): loss_fn = DPOLoss(beta=0.1, loss_type='hinge') inputs, outputs, ref_logps = _make_preference_batch() diff --git a/tests/metric/test_metrics.py b/tests/metric/test_metrics.py index e4651ee86..691da2c06 100644 --- a/tests/metric/test_metrics.py +++ b/tests/metric/test_metrics.py @@ -271,6 +271,26 @@ def test_dpo_metric_with_ref(self): assert 'rewards/chosen' in result assert 'rewards/accuracies' in result + def test_dpo_metric_accepts_json_ref_logps(self): + labels = torch.tensor([[1, 2, -100], [3, 4, -100]]) + logps = torch.randn(2, 3) + ref_logps = torch.randn(2, 3) + tensor_metric = _no_dist_metric(DPOMetric, beta=0.1) + json_metric = _no_dist_metric(DPOMetric, beta=0.1) + + tensor_metric.accumulate( + {'labels': labels}, + {'logps': logps}, + ref_outputs={'logps': ref_logps}, + ) + json_metric.accumulate( + {'labels': labels}, + {'logps': logps}, + ref_outputs={'logps': ref_logps.tolist()}, + ) + + assert json_metric.calculate() == pytest.approx(tensor_metric.calculate()) + def test_dpo_metric_no_logps_skips(self): m = _no_dist_metric(DPOMetric, beta=0.1) m.accumulate({'labels': torch.tensor([[1, 2]])}, {}) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index fd9917e76..51cc2de56 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -1,7 +1,6 @@ from __future__ import annotations import pytest -import torch from fastapi import FastAPI from starlette.requests import Request @@ -9,7 +8,6 @@ from twinkle.server.model.twinkle_handlers import ( _model_result_rows, _register_twinkle_routes, - _restore_dataref_value, ) @@ -23,26 +21,6 @@ def test_model_result_rows_keeps_one_output_row_per_sample() -> None: ] -def test_restore_dataref_value_rebuilds_rectangular_and_ragged_numeric_arrays() -> None: - rectangular = _restore_dataref_value([ - [-0.1, -0.2], - [-0.3, -0.4], - ]) - torch.testing.assert_close( - rectangular, - torch.tensor([[-0.1, -0.2], [-0.3, -0.4]]), - ) - - ragged = _restore_dataref_value([ - [-0.1], - [-0.2, -0.3], - ]) - assert isinstance(ragged, list) - assert all(torch.is_tensor(item) for item in ragged) - torch.testing.assert_close(ragged[0], torch.tensor([-0.1])) - torch.testing.assert_close(ragged[1], torch.tensor([-0.2, -0.3])) - - class _SchedulingManagement: def __init__(self): @@ -112,19 +90,13 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - assert management.scheduled[-1]['data_world_size'] == 2 inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] assert adapter_name == 'session-adapter' - assert [row['input_ids'].tolist() for row in inputs] == [[index] for index in range(8)] - torch.testing.assert_close( - forwarded_kwargs['old_logps'], - torch.tensor([[-0.1]] * 4 + [[-0.2]] * 4), - ) - torch.testing.assert_close( - forwarded_kwargs['advantages'], - torch.tensor([1.0] * 4 + [-1.0] * 4), - ) + assert [row['input_ids'] for row in inputs] == [[index] for index in range(8)] + assert forwarded_kwargs['old_logps'] == [[-0.1]] * 4 + [[-0.2]] * 4 + assert forwarded_kwargs['advantages'] == [1.0] * 4 + [-1.0] * 4 @pytest.mark.asyncio -async def test_forward_backward_restores_nested_dpo_ref_logps_as_tensor() -> None: +async def test_forward_backward_binds_nested_dpo_ref_logps_without_coercion() -> None: management = _SchedulingManagement() management.rows['dpo'] = [ { @@ -153,11 +125,8 @@ async def test_forward_backward_restores_nested_dpo_ref_logps_as_tensor() -> Non inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] assert adapter_name == 'session-adapter' - assert [row['input_ids'].tolist() for row in inputs] == [[1, 2, 3], [1, 4, 5]] - torch.testing.assert_close( - forwarded_kwargs['ref_outputs']['logps'], - torch.tensor([ - [-0.1, -0.2, -0.3], - [-0.4, -0.5, -0.6], - ]), - ) + assert [row['input_ids'] for row in inputs] == [[1, 2, 3], [1, 4, 5]] + assert forwarded_kwargs['ref_outputs']['logps'] == [ + [-0.1, -0.2, -0.3], + [-0.4, -0.5, -0.6], + ] diff --git a/tests/utils/test_rl_tensor_utils.py b/tests/utils/test_rl_tensor_utils.py new file mode 100644 index 000000000..d5d4486d1 --- /dev/null +++ b/tests/utils/test_rl_tensor_utils.py @@ -0,0 +1,58 @@ +import pytest +import torch + +from twinkle.utils.rl_tensor_utils import align_per_token_values + + +def test_align_per_token_values_accepts_json_rows() -> None: + actual = align_per_token_values( + [[-0.1, -0.2], [-0.3, -0.4]], + (2, 2), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + ) + + torch.testing.assert_close( + actual, + torch.tensor([[-0.1, -0.2], [-0.3, -0.4]]), + ) + + +def test_align_per_token_values_pads_ragged_json_rows() -> None: + actual = align_per_token_values( + [[-0.1, -0.2, -0.3], [-0.4, -0.5]], + (2, 3), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + valid_mask=torch.tensor([[True, True, True], [True, True, False]]), + ) + + torch.testing.assert_close( + actual, + torch.tensor([[-0.1, -0.2, -0.3], [-0.4, -0.5, 0.0]]), + ) + + +def test_align_per_token_values_rejects_missing_valid_tokens() -> None: + with pytest.raises(ValueError, match=r'ref_logps\[1\] has 2 tokens'): + align_per_token_values( + [[-0.1, -0.2, -0.3], [-0.4, -0.5]], + (2, 3), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + valid_mask=torch.ones(2, 3, dtype=torch.bool), + ) + + +def test_align_per_token_values_rejects_short_batches() -> None: + with pytest.raises(ValueError, match='batch size'): + align_per_token_values( + [[-0.1, -0.2]], + (2, 2), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + ) From f6f2fdc068ee3e7b06db2273e12690eba8ff2afc Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 13 Aug 2026 15:53:25 +0800 Subject: [PATCH 08/20] wip --- cookbook/rl/async_multi_lora_dapo_grpo.yaml | 1 - .../async_multi_lora_dapo_hparam_sweep.yaml | 1 - cookbook/rl/async_multi_lora_grpo.yaml | 1 - cookbook/rl/async_single_lora_dapo_grpo.yaml | 1 - .../rl/async_single_lora_gsm8k_areal.yaml | 1 - cookbook/rl/async_single_lora_gsm8k_verl.yaml | 126 ++++++++++++++++++ src/twinkle/infra/__init__.py | 2 +- .../transformers/multi_lora_transformers.py | 7 +- src/twinkle_agentic/async_rl/pipeline.py | 31 +++-- src/twinkle_agentic/async_rl/utils.py | 1 - .../test_async_rl_native_tq.py | 73 +++++++++- 11 files changed, 224 insertions(+), 21 deletions(-) create mode 100644 cookbook/rl/async_single_lora_gsm8k_verl.yaml diff --git a/cookbook/rl/async_multi_lora_dapo_grpo.yaml b/cookbook/rl/async_multi_lora_dapo_grpo.yaml index d54641f59..f49bf1ccb 100644 --- a/cookbook/rl/async_multi_lora_dapo_grpo.yaml +++ b/cookbook/rl/async_multi_lora_dapo_grpo.yaml @@ -83,7 +83,6 @@ lora: loss: cls: GRPOLoss epsilon: 0.2 - normalization: sequence_mean lora_contexts: # Set both environment variables to distinct local parquet splits for a disjoint-tenant experiment. diff --git a/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml b/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml index 02d9168c6..166f96332 100644 --- a/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml +++ b/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml @@ -74,7 +74,6 @@ lora: loss: cls: GRPOLoss epsilon: 0.2 - normalization: sequence_mean lora_contexts: # A/B/C isolate learning rate while keeping mini_batch_size fixed at 32 samples. diff --git a/cookbook/rl/async_multi_lora_grpo.yaml b/cookbook/rl/async_multi_lora_grpo.yaml index 4467abb9d..cca7775c1 100644 --- a/cookbook/rl/async_multi_lora_grpo.yaml +++ b/cookbook/rl/async_multi_lora_grpo.yaml @@ -80,7 +80,6 @@ lora: loss: cls: GRPOLoss epsilon: 0.2 - normalization: sequence_mean lora_contexts: - tenant_id: tenant_a diff --git a/cookbook/rl/async_single_lora_dapo_grpo.yaml b/cookbook/rl/async_single_lora_dapo_grpo.yaml index 004639886..701b750c6 100644 --- a/cookbook/rl/async_single_lora_dapo_grpo.yaml +++ b/cookbook/rl/async_single_lora_dapo_grpo.yaml @@ -82,7 +82,6 @@ lora: loss: cls: GRPOLoss epsilon: 0.2 - normalization: sequence_mean lora_contexts: - tenant_id: tenant_single_dapo diff --git a/cookbook/rl/async_single_lora_gsm8k_areal.yaml b/cookbook/rl/async_single_lora_gsm8k_areal.yaml index 81ac1141d..d6becbbd5 100644 --- a/cookbook/rl/async_single_lora_gsm8k_areal.yaml +++ b/cookbook/rl/async_single_lora_gsm8k_areal.yaml @@ -85,7 +85,6 @@ lora: loss: cls: GRPOLoss epsilon: 0.2 - normalization: token_mean lora_contexts: - tenant_id: tenant_single diff --git a/cookbook/rl/async_single_lora_gsm8k_verl.yaml b/cookbook/rl/async_single_lora_gsm8k_verl.yaml new file mode 100644 index 000000000..7dae0f1ee --- /dev/null +++ b/cookbook/rl/async_single_lora_gsm8k_verl.yaml @@ -0,0 +1,126 @@ +runtime: + run_id: async_single_lora_gsm8k_accuracy + model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} + mode: ray + model_gpus: 1 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 1 + seed: 1 + max_staleness: 0 + max_steps: 125 + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_single_lora_gsm8k_accuracy + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/single_lora_gsm8k_accuracy_metrics.jsonl + summary_path: outputs/async_rl/single_lora_gsm8k_accuracy_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: true + +model: + strategy: native_fsdp + attn_implementation: flash_attention_2 + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 1 + padding_free: false + max_length: 2048 + +sampler: + max_model_len: 2048 + max_num_batched_tokens: 4096 + gpu_memory_utilization: 0.7 + max_num_seqs: 64 + enforce_eager: false + +rollout_output: + enabled: true + output_dir: ${runtime.output_dir}/rollouts + include_token_ids: false + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +evaluation: + enabled: true + interval: 10 + batch_size: 16 + sampling_params: + max_tokens: 1024 + temperature: 0.6 + top_p: 1.0 + +lora: + target_modules: all-linear + r: 16 + alpha: 16 + dropout: 0.0 + learning_rate: 1.7e-5 + +loss: + cls: GRPOLoss + epsilon: 0.2 + +lora_contexts: + - tenant_id: tenant_single + training_run_id: gsm8k_accuracy + adapter_name: gsm8k_accuracy_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + dataset: + dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID} + subset_name: main + split: train + data_num: 2000 + max_length: 1024 + processor: GSM8KProcessor + system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' + eval_dataset: + name: gsm8k/test + dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID} + subset_name: main + split: test + data_num: null + max_length: 1024 + processor: GSM8KProcessor + system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 1024 + temperature: 1.0 + top_p: 1.0 + train: + mini_batch_size: 64 + micro_batch_size: 64 + dynamic_batching: true + max_tokens_per_micro_batch: 4096 + packing_algorithm: ffd diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index a2760c900..3b1ddfbe3 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -826,7 +826,7 @@ def _notifying_result_func(*rargs, **rkwargs): wrapper._execute = execute wrapper._collect = collect wrapper._dispatch = dispatch - wrapper._lazy_collect = _lazy_collect + wrapper._lazy_collect = _lazy_collect if lazy_collect is None else lazy_collect wrapper._sync = sync return wrapper diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index ea53930de..50df7a05f 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -267,7 +267,12 @@ def _get_adapter_state_dict_for_save(self, adapter_name: str) -> dict: adapter_state = self.multi_adapter.get_state_dict(adapter_name) return {key: torch_util.to_local_tensor(value).cpu() for key, value in adapter_state.items()} - @remote_function(collect='first') + # Saving publishes an immutable adapter checkpoint to callers, so the + # checkpoint path must be collected before returning. In particular, + # TrainerWorker calls this model handle from another Ray actor; inheriting + # that actor's default lazy-collect mode would otherwise return a callable + # instead of the path string. + @remote_function(collect='first', lazy_collect=False) def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): self._check_adapter_valid(kwargs.get('adapter_name')) with self.multi_adapter.save_context(kwargs.get('adapter_name')): diff --git a/src/twinkle_agentic/async_rl/pipeline.py b/src/twinkle_agentic/async_rl/pipeline.py index 012372bea..8acd2cea0 100644 --- a/src/twinkle_agentic/async_rl/pipeline.py +++ b/src/twinkle_agentic/async_rl/pipeline.py @@ -288,10 +288,13 @@ def from_config( eval_dataset.get('reward'), context_key=f'{context.key} evaluation', ) - initial_paths[context.key] = model.save( - f'async-{context.adapter_name}-initial', - output_dir=runtime['output_dir'], - adapter_name=context.adapter_name, + initial_paths[context.key] = _require_adapter_path( + model.save( + f'async-{context.adapter_name}-initial', + output_dir=runtime['output_dir'], + adapter_name=context.adapter_name, + ), + operation=f'initial adapter save for {context.key}', ) manager = create_cpu_actor( @@ -672,13 +675,25 @@ def _train_batch_with_config( def _save_adapter(model: Any, output_dir: str, admission: PartitionAdmission) -> str: - return model.save( - f'async-{admission.context.adapter_name}-v{admission.step + 1}', - output_dir=output_dir, - adapter_name=admission.context.adapter_name, + return _require_adapter_path( + model.save( + f'async-{admission.context.adapter_name}-v{admission.step + 1}', + output_dir=output_dir, + adapter_name=admission.context.adapter_name, + ), + operation=f'adapter save for {admission.partition_id}', ) +def _require_adapter_path(value: Any, *, operation: str) -> str: + """Fail at the save boundary instead of publishing an invalid policy.""" + if not isinstance(value, str) or not value: + raise TypeError( + f'{operation} must return a non-empty checkpoint path string, ' + f'got {type(value).__name__}: {value!r}') + return value + + def _remove_adapter_snapshot(sampler: Any, adapter_path: str) -> None: """Unload an unreferenced policy from vLLM before deleting its checkpoint.""" from .workers import _remove_local_adapter diff --git a/src/twinkle_agentic/async_rl/utils.py b/src/twinkle_agentic/async_rl/utils.py index bd2a176de..764c0a98e 100644 --- a/src/twinkle_agentic/async_rl/utils.py +++ b/src/twinkle_agentic/async_rl/utils.py @@ -207,7 +207,6 @@ def resolve_context_loss_config( loss_config: dict[str, Any] = { 'cls': 'GRPOLoss', 'epsilon': 0.2, - 'normalization': 'sequence_mean', } loss_config.update(dict(loss_defaults or {})) loss_config.update(dict(context_config.get('loss') or {})) diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py index 596597d8f..921f1f085 100644 --- a/tests/twinkle_agentic/test_async_rl_native_tq.py +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -19,7 +19,12 @@ from twinkle_agentic.async_rl.metrics import training_policy_metrics from twinkle_agentic.async_rl.native_tq import ContextGRPOGroupNSampler from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches -from twinkle_agentic.async_rl.pipeline import create_cpu_actor, _reward_for_context, _train_batch +from twinkle_agentic.async_rl.pipeline import ( + _require_adapter_path, + _reward_for_context, + _train_batch, + create_cpu_actor, +) from twinkle_agentic.async_rl.types import (PartitionAdmission, PreparedPartition, PromptGroup, RolloutPolicy) from twinkle_agentic.async_rl.utils import ( TrainBatchConfig, @@ -386,20 +391,19 @@ def test_context_loss_config_overrides_global_defaults(): { 'loss': { 'cls': 'GSPOLoss', - 'normalization': 'token_mean', + 'epsilon_high': 0.3, } }, { 'cls': 'GRPOLoss', 'epsilon': 0.2, - 'normalization': 'sequence_mean', }, ) assert loss_cls == 'GSPOLoss' assert loss_kwargs == { 'epsilon': 0.2, - 'normalization': 'token_mean', + 'epsilon_high': 0.3, } @@ -408,11 +412,70 @@ def test_context_loss_config_uses_grpo_defaults(): 'GRPOLoss', { 'epsilon': 0.2, - 'normalization': 'sequence_mean', }, ) +def test_async_single_lora_gsm8k_verl_config_matches_current_grpo_api(): + from omegaconf import OmegaConf + from twinkle.loss import GRPOLoss + + config_path = 'cookbook/rl/async_single_lora_gsm8k_verl.yaml' + config = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False) + context = config['lora_contexts'][0] + loss_cls, loss_kwargs = resolve_context_loss_config(context, config['loss']) + + assert loss_cls == 'GRPOLoss' + GRPOLoss(**loss_kwargs) + assert context['reward']['class_path'] == 'twinkle.reward.GSM8KAccuracyReward' + assert context['eval_dataset']['reward']['class_path'] == 'twinkle.reward.GSM8KAccuracyReward' + validate_context_batch_config( + f"{context['tenant_id']}/{context['training_run_id']}/{context['adapter_name']}", + rollout_groups=context['rollout']['batch_size'], + num_generations=context['rollout']['num_generations'], + train=TrainBatchConfig( + mini_batch_size=context['train']['mini_batch_size'], + micro_batch_size=context['train']['micro_batch_size'], + dynamic_batching=context['train']['dynamic_batching'], + max_tokens_per_micro_batch=context['train']['max_tokens_per_micro_batch'], + packing_algorithm=context['train']['packing_algorithm'], + ), + sampler_dp=config['runtime']['sampler_gpus'] // config['runtime']['sampler_tp'], + model_dp=config['runtime']['model_gpus'], + ) + + +def test_adapter_path_rejects_uncollected_remote_result(): + def lazy_result(): + return '/tmp/policy' + + with pytest.raises(TypeError, match='must return a non-empty checkpoint path string'): + _require_adapter_path(lazy_result, operation='test save') + + +def test_multi_lora_transformers_save_disables_lazy_collect(): + from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + assert MultiLoraTransformersModel.save._lazy_collect is False + + +def test_remote_function_metadata_uses_explicit_lazy_collect_value(): + from twinkle import remote_function + + class Component: + + @remote_function(lazy_collect=False) + def eager(self): + return None + + @remote_function(lazy_collect=True) + def lazy(self): + return None + + assert Component.eager._lazy_collect is False + assert Component.lazy._lazy_collect is True + + def test_context_loss_config_rejects_empty_class_name(): with pytest.raises(ValueError, match='loss.cls'): resolve_context_loss_config({'loss': {'cls': ''}}) From 1d8b3312e5db18023d7692b67dc83eb0dd454f21 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 13 Aug 2026 16:23:26 +0800 Subject: [PATCH 09/20] wip --- src/twinkle/model/multi_lora.py | 16 ++++++++++++---- tests/model/test_multi_lora.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) create mode 100644 tests/model/test_multi_lora.py diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 2af27b574..4a9e592f6 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -201,10 +201,18 @@ def _after(_module): _after(self.module) # self.deactivate_adapter() - def check_length(self, inputs: InputFeature): - total_length = sum(len(_input['input_ids']) for _input in inputs) - if total_length > self.max_length: - raise ValueError(f'Max length exceeds {self.max_length}') + def check_length( + self, + inputs: Union[InputFeature, List[InputFeature]], + ): + if isinstance(inputs, dict): + inputs = [inputs] + for index, item in enumerate(inputs): + if 'input_ids' not in item: + continue + length = len(item['input_ids']) + if length > self.max_length: + raise ValueError(f'Input length {length} exceeds max_length {self.max_length} at sample {index}') def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str: if self.has_lora(tenant_adapter_name): diff --git a/tests/model/test_multi_lora.py b/tests/model/test_multi_lora.py new file mode 100644 index 000000000..f2795a5ad --- /dev/null +++ b/tests/model/test_multi_lora.py @@ -0,0 +1,28 @@ +import pytest + +from twinkle.model.multi_lora import MultiLora + + +def test_check_length_checks_each_sample_independently(): + multi_lora = MultiLora(max_length=4) + + multi_lora.check_length([ + {'input_ids': [1, 2, 3]}, + {'input_ids': [4, 5, 6]}, + ]) + + +def test_check_length_accepts_a_single_input_feature(): + multi_lora = MultiLora(max_length=4) + + multi_lora.check_length({'input_ids': [1, 2, 3, 4]}) + + +def test_check_length_reports_the_oversized_sample(): + multi_lora = MultiLora(max_length=4) + + with pytest.raises(ValueError, match=r'Input length 5 exceeds max_length 4 at sample 1'): + multi_lora.check_length([ + {'input_ids': [1, 2]}, + {'input_ids': [1, 2, 3, 4, 5]}, + ]) From 2165d2e07112499afd401203573247d681cd50c3 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 13 Aug 2026 17:57:18 +0800 Subject: [PATCH 10/20] wip --- src/twinkle/model/multi_lora.py | 3 + .../transformers/strategy/native_fsdp.py | 5 ++ .../sampler/vllm_sampler/vllm_engine.py | 40 ++++++++++-- .../test_multi_lora_target_parameters.py | 18 +++++ tests/sampler/test_vllm_lora_loading.py | 65 +++++++++++++++++++ .../transformers/test_native_fsdp_strategy.py | 25 +++++++ 6 files changed, 150 insertions(+), 6 deletions(-) create mode 100644 tests/sampler/test_vllm_lora_loading.py create mode 100644 tests/transformers/test_native_fsdp_strategy.py diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 4a9e592f6..353589761 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -848,6 +848,9 @@ def _get_parameters(_module): return trainable_param_names def get_target_parameter_trainable_parameters(self, tenant_adapter_name): + lora = self.find_lora_by_tenant(tenant_adapter_name) + if not getattr(lora.tenant_config, 'target_parameters', None): + return {} return { name: parameter for name, parameter in self.target_parameter_manager.named_slot_parameters(tenant_adapter_name) diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index f57877bb0..83d05c024 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -116,6 +116,11 @@ def wrap_model(self, model, optimizer=None): if self.device_mesh is None: return model, optimizer fsdp_mesh = _build_fsdp_mesh(self.device_mesh) + if fsdp_mesh is None: + # FSDP has nothing to shard with a single rank, but callers still + # expect the native strategy to place the model on the worker's + # local device before CUDA inputs reach it. + model = model.to(torch.device(Platform.get_local_device())) if fsdp_mesh is not None: ep_enabled = (self.enable_ep and self.ep_fsdp_device_mesh is not None) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index a38dd6959..4487a22f1 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio import contextlib import inspect import os @@ -101,6 +102,7 @@ def __init__( self.engine_kwargs = kwargs or {} self._lora_request_cache: Dict[str, Any] = {} + self._lora_load_tasks: Dict[str, asyncio.Task] = {} self._next_lora_id = 1 # Cached LoRARequest for the RL-training synced LoRA. @@ -462,13 +464,26 @@ async def _get_or_load_lora( Returns: ``LoRARequest`` or ``None`` if loading fails. """ - from vllm.lora.request import LoRARequest - - # Fast path: return cached request for this path. if lora_path in self._lora_request_cache: - logger.info(f'Using cached LoRA request for {lora_path}') + logger.debug(f'Using cached LoRA request for {lora_path}') return self._lora_request_cache[lora_path] + load_task = self._lora_load_tasks.get(lora_path) + if load_task is None: + load_task = asyncio.create_task(self._load_lora(lora_path)) + self._lora_load_tasks[lora_path] = load_task + try: + lora_request = await load_task + finally: + if self._lora_load_tasks.get(lora_path) is load_task: + self._lora_load_tasks.pop(lora_path) + if lora_request is not None: + self._lora_request_cache[lora_path] = lora_request + return lora_request + + async def _load_lora(self, lora_path: str): + from vllm.lora.request import LoRARequest + if not os.path.exists(lora_path): logger.error(f'LoRA path does not exist: {lora_path}') return None @@ -487,9 +502,9 @@ async def _get_or_load_lora( lora_path=lora_path, ) + logger.info(f'Loading LoRA from {lora_path}') try: await self.engine.add_lora(lora_request) - self._lora_request_cache[lora_path] = lora_request return lora_request except Exception as e: logger.error(f'Failed to load LoRA from {lora_path}: {e}') @@ -502,10 +517,23 @@ async def unload_lora_paths(self, adapter_paths: list[str]) -> None: request = self._lora_request_cache.pop(normalized, None) if request is None: request = self._lora_request_cache.pop(adapter_path, None) + load_task = self._lora_load_tasks.pop(normalized, None) + if load_task is None: + load_task = self._lora_load_tasks.pop(adapter_path, None) + if load_task is not None and not load_task.done(): + load_task.cancel() + await asyncio.gather(load_task, return_exceptions=True) + elif request is None and load_task is not None: + try: + request = load_task.result() + except (asyncio.CancelledError, Exception): + request = None if request is None: continue try: - await self.engine.remove_lora(request.lora_int_id) + result = self.engine.remove_lora(request.lora_int_id) + if inspect.isawaitable(result): + await result except Exception as exc: logger.warning('Failed to unload LoRA %s: %s', adapter_path, exc) diff --git a/tests/model/test_multi_lora_target_parameters.py b/tests/model/test_multi_lora_target_parameters.py index d26b32147..7edc4c6d8 100644 --- a/tests/model/test_multi_lora_target_parameters.py +++ b/tests/model/test_multi_lora_target_parameters.py @@ -85,6 +85,24 @@ def _make_target_cfg(r=2): ) +def test_standard_lora_has_no_target_parameter_trainable_parameters(): + from twinkle.model.multi_lora import LoraTenant, MultiLora + + config = LoraConfig(r=2, lora_alpha=4, target_modules=['linear']) + multi_lora = MultiLora(max_loras=1, max_r=4) + multi_lora.loras = [ + LoraTenant( + index=0, + adapter_name='lora_0', + config=config, + tenant_adapter_name='adapter_a', + tenant_config=config, + ) + ] + + assert multi_lora.get_target_parameter_trainable_parameters('adapter_a') == {} + + def test_target_parameter_multi_lora_updates_only_active_adapter(): from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager diff --git a/tests/sampler/test_vllm_lora_loading.py b/tests/sampler/test_vllm_lora_loading.py new file mode 100644 index 000000000..f63c9c8ad --- /dev/null +++ b/tests/sampler/test_vllm_lora_loading.py @@ -0,0 +1,65 @@ +import asyncio +from unittest.mock import MagicMock + +from twinkle.sampler.vllm_sampler.vllm_engine import VLLMEngine + + +def test_concurrent_lora_requests_share_one_load_task(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + engine._lora_request_cache = {} + engine._lora_load_tasks = {} + request = object() + load_count = 0 + + async def load_lora(_path): + nonlocal load_count + load_count += 1 + await asyncio.sleep(.01) + return request + + engine._load_lora = load_lora + results = await asyncio.gather(*(engine._get_or_load_lora('/adapter') for _ in range(8))) + + assert load_count == 1 + assert results == [request] * 8 + assert engine._lora_request_cache == {'/adapter': request} + assert engine._lora_load_tasks == {} + + asyncio.run(run()) + + +def test_unload_lora_accepts_synchronous_engine_api(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + request = MagicMock(lora_int_id=7) + engine._lora_request_cache = {'/adapter': request} + engine._lora_load_tasks = {} + engine.engine = MagicMock() + engine.engine.remove_lora.return_value = True + + await engine.unload_lora_paths(['/adapter']) + + engine.engine.remove_lora.assert_called_once_with(7) + assert engine._lora_request_cache == {} + + asyncio.run(run()) + + +def test_unload_lora_removes_a_just_completed_load(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + request = MagicMock(lora_int_id=9) + load_task = asyncio.create_task(asyncio.sleep(0, result=request)) + await load_task + engine._lora_request_cache = {} + engine._lora_load_tasks = {'/adapter': load_task} + engine.engine = MagicMock() + engine.engine.remove_lora.return_value = None + + await engine.unload_lora_paths(['/adapter']) + + engine.engine.remove_lora.assert_called_once_with(9) + assert engine._lora_load_tasks == {} + + asyncio.run(run()) diff --git a/tests/transformers/test_native_fsdp_strategy.py b/tests/transformers/test_native_fsdp_strategy.py new file mode 100644 index 000000000..79402c785 --- /dev/null +++ b/tests/transformers/test_native_fsdp_strategy.py @@ -0,0 +1,25 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from unittest.mock import MagicMock, patch + +import torch + +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + + +def test_single_rank_native_fsdp_places_model_on_local_device(): + device_mesh = MagicMock() + device_mesh.mesh_dim_names = None + model = MagicMock() + moved_model = MagicMock() + model.to.return_value = moved_model + strategy = NativeFSDPStrategy(device_mesh=device_mesh, enable_ep=False) + + with patch( + 'twinkle.model.transformers.strategy.native_fsdp.Platform.get_local_device', + return_value='cuda:0', + ): + wrapped_model, wrapped_optimizer = strategy.wrap_model(model) + + model.to.assert_called_once_with(torch.device('cuda:0')) + assert wrapped_model is moved_model + assert wrapped_optimizer is None From 1f96e188ae199eb597fabfb6f95a5be56474bf58 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Sun, 16 Aug 2026 23:36:38 +0800 Subject: [PATCH 11/20] wip --- cookbook/client/async_rl/README.md | 6 +- .../async_rl/client_orchestrated_dpo.py | 4 +- .../async_rl/client_orchestrated_grpo.py | 2 +- src/twinkle/server/model/twinkle_handlers.py | 123 +++++++++++++++--- .../model/multi_lora_transformers.py | 74 +++++++---- src/twinkle_client/types/__init__.py | 2 + src/twinkle_client/types/model.py | 41 +++--- .../server/contract/client_api_baseline.json | 30 +++++ .../server/model/test_twinkle_async_inputs.py | 34 ++++- tests/twinkle_client/test_async_components.py | 69 +++++++++- .../test_client_orchestrated_dpo.py | 4 +- .../test_client_orchestrated_grpo.py | 2 +- 12 files changed, 308 insertions(+), 83 deletions(-) diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md index 12d7776f6..80a3a19a3 100644 --- a/cookbook/client/async_rl/README.md +++ b/cookbook/client/async_rl/README.md @@ -33,8 +33,8 @@ python cookbook/client/async_rl/client_orchestrated_grpo.py The training loop composes only the low-level component methods: - `sampler.sample_to_data_plane(...)` / `sampler.asample_to_data_plane(...)` -- `model.forward_only(...)` -- `model.forward_backward(...)` +- `model.forward_only_from_data_plane(...)` +- `model.forward_backward_from_data_plane(...)` - `model.clip_grad_and_step(...)` - `model.save(...)` - `data_plane.put/get/append/release(...)` and @@ -46,7 +46,7 @@ generation is one row tagged with its group, generation index, rollout policy, and status. The Advantage worker reads only decoded completions and appends reward and advantage to the same keys. Token tensors and sampled log-probabilities remain server-side. The Trainer passes one or more `DataRef` values to -`forward_backward()` and releases them in a local `finally` block. `asample()` +`forward_backward_from_data_plane()` and releases them in a local `finally` block. `asample()` remains the materialized-response convenience API. - `ClientMultiTurnRollout.arun()` keeps tool calls and Reward computation in diff --git a/cookbook/client/async_rl/client_orchestrated_dpo.py b/cookbook/client/async_rl/client_orchestrated_dpo.py index f5118a83a..625d75b5c 100644 --- a/cookbook/client/async_rl/client_orchestrated_dpo.py +++ b/cookbook/client/async_rl/client_orchestrated_dpo.py @@ -151,7 +151,7 @@ async def run(self) -> None: ref = item try: await _submit( - self.model.forward_backward, + self.model.forward_backward_from_data_plane, ref, kwarg_fields={'ref_outputs.logps': 'ref_logps'}, ) @@ -167,7 +167,7 @@ async def _reference_forward( ) -> DataRef: """Run the frozen base model and append reference logps to the same rows.""" return await _submit( - model.forward_only, + model.forward_only_from_data_plane, batch_ref, disable_lora=True, output_ref=batch_ref, diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py index e59fa01c3..f30d31673 100644 --- a/cookbook/client/async_rl/client_orchestrated_grpo.py +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -306,7 +306,7 @@ async def _train(self, groups: list[_ReadyGroup]) -> None: refs = [group.ref for group in groups] try: await _submit( - self.model.forward_backward, + self.model.forward_backward_from_data_plane, refs, input_field='train_input', kwarg_fields={ diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index ec485a9ee..904637422 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -92,11 +92,8 @@ def _set_at_path(target: dict[str, Any], path: str, value: Any) -> None: current[parts[-1]] = value -async def _resolve_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: - """Resolve transport-level references before entering the model backend.""" - if body.input_refs is None: - return body.inputs, {} - +async def _resolve_data_plane_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: + """Resolve DataPlane references before entering the model backend.""" selected_fields = None if body.input_field is not None: selected_fields = list(dict.fromkeys([ @@ -141,10 +138,11 @@ def _merge_forward_kwargs(explicit: dict[str, Any], bound: dict[str, Any]) -> di def _request_shape(body: Any) -> tuple[int, int]: - if body.input_refs is not None: + input_refs = getattr(body, 'input_refs', None) + if input_refs is not None: return ( - sum(ref.num_tokens for ref in body.input_refs), - sum(ref.size for ref in body.input_refs), + sum(ref.num_tokens for ref in input_refs), + sum(ref.size for ref in input_refs), ) inputs = body.inputs if isinstance(body.inputs, list) else [body.inputs] return ( @@ -214,10 +212,8 @@ async def forward(request: Request, body: types.ForwardRequest, async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} - raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) - inputs = _parse_inputs(raw_inputs) - kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) - ret = self.model.forward(inputs=inputs, adapter_name=adapter_name, **kwargs) + inputs = _parse_inputs(body.inputs) + ret = self.model.forward(inputs=inputs, adapter_name=adapter_name, **extra_kwargs) return {'result': ret} input_tokens, batch_size = _request_shape(body) @@ -232,6 +228,38 @@ async def _task(): task_type='forward', )) + @app.post('/twinkle/forward_from_data_plane', response_model=types.ForwardResponse) + async def forward_from_data_plane( + request: Request, + body: types.DataPlaneForwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.ForwardResponse: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs, field_kwargs = await _resolve_data_plane_model_inputs(body, self.data_plane) + kwargs = _merge_forward_kwargs(body.model_extra or {}, field_kwargs) + ret = self.model.forward( + inputs=_parse_inputs(raw_inputs), + adapter_name=adapter_name, + **kwargs, + ) + return {'result': ret} + + input_tokens, batch_size = _request_shape(body) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward_from_data_plane', + )) + @app.post('/twinkle/remove_adapter') async def remove_adapter( request: Request, @@ -266,9 +294,36 @@ async def forward_only( async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} - raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) + inputs = _parse_inputs(body.inputs) + ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **extra_kwargs) + return {'result': ret} + + input_tokens, batch_size = _request_shape(body) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward_only', + )) + + @app.post('/twinkle/forward_only_from_data_plane', response_model=types.ForwardResponse) + async def forward_only_from_data_plane( + request: Request, + body: types.DataPlaneForwardOnlyRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.ForwardResponse: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs, field_kwargs = await _resolve_data_plane_model_inputs(body, self.data_plane) inputs = _parse_inputs(raw_inputs) - kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) + kwargs = _merge_forward_kwargs(body.model_extra or {}, field_kwargs) ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **kwargs) if body.output_ref is not None: rows = _select_output_rows( @@ -289,7 +344,7 @@ async def _task(): input_tokens=input_tokens, batch_size=batch_size, data_world_size=self.data_world_size, - task_type='forward_only', + task_type='forward_only_from_data_plane', )) @app.post('/twinkle/calculate_loss', response_model=types.CalculateLossResponse) @@ -334,10 +389,8 @@ async def forward_backward( async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} - raw_inputs, field_kwargs = await _resolve_model_inputs(body, self.data_plane) - all_inputs = _parse_inputs(raw_inputs) - kwargs = _merge_forward_kwargs(extra_kwargs, field_kwargs) - ret = self.model.forward_backward(inputs=all_inputs, adapter_name=adapter_name, **kwargs) + all_inputs = _parse_inputs(body.inputs) + ret = self.model.forward_backward(inputs=all_inputs, adapter_name=adapter_name, **extra_kwargs) return {'result': ret} input_tokens, batch_size = _request_shape(body) @@ -352,6 +405,38 @@ async def _task(): task_type='forward_backward', )) + @app.post('/twinkle/forward_backward_from_data_plane', response_model=types.ForwardBackwardResponse) + async def forward_backward_from_data_plane( + request: Request, + body: types.DataPlaneForwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.ForwardBackwardResponse: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs, field_kwargs = await _resolve_data_plane_model_inputs(body, self.data_plane) + kwargs = _merge_forward_kwargs(body.model_extra or {}, field_kwargs) + ret = self.model.forward_backward( + inputs=_parse_inputs(raw_inputs), + adapter_name=adapter_name, + **kwargs, + ) + return {'result': ret} + + input_tokens, batch_size = _request_shape(body) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward_backward_from_data_plane', + )) + @app.post('/twinkle/clip_grad_norm', response_model=types.ClipGradNormResponse) async def clip_grad_norm( request: Request, diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index 7eda3258e..2dad63125 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -17,15 +17,14 @@ ) -def _component_input_payload(inputs: Any) -> dict[str, Any]: - """Encode inline rows or one or more opaque server-side data references.""" - if isinstance(inputs, DataRef): - return {'input_refs': [inputs.model_dump()]} - if isinstance(inputs, list) and inputs and any(isinstance(item, DataRef) for item in inputs): - if not all(isinstance(item, DataRef) for item in inputs): - raise TypeError('model inputs cannot mix DataRef values with inline rows') - return {'input_refs': [item.model_dump() for item in inputs]} - return {'inputs': json_safe(inputs)} +def _data_ref_payload(inputs: DataRef | list[DataRef]) -> dict[str, Any]: + """Encode one or more opaque references for a DataPlane model endpoint.""" + refs = [inputs] if isinstance(inputs, DataRef) else list(inputs) + if not refs: + raise ValueError('at least one DataRef is required') + if not all(isinstance(item, DataRef) for item in refs): + raise TypeError('data-plane model inputs must contain only DataRef values') + return {'input_refs': [item.model_dump() for item in refs]} class MultiLoraTransformersModel: @@ -74,19 +73,37 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: if name == self.adapter_name: self.adapter_name = None - def forward( + def forward(self, inputs: Any, **kwargs) -> ForwardResponse: + """Execute forward pass on inline model inputs.""" + response = http_post( + url=f'{self.server_url}/forward', + json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs}, + ) + response.raise_for_status() + return ForwardResponse(**response.json()) + + def forward_only(self, inputs: Any, **kwargs) -> ForwardResponse: + """Execute forward pass without gradient computation on inline inputs.""" + response = http_post( + url=f'{self.server_url}/forward_only', + json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs}, + ) + response.raise_for_status() + return ForwardResponse(**response.json()) + + def forward_from_data_plane( self, - inputs: Any | DataRef | list[DataRef], + inputs: DataRef | list[DataRef], *, input_field: str | None = None, kwarg_fields: dict[str, str] | None = None, **kwargs, ) -> ForwardResponse: - """Execute forward over inline rows or server-side references.""" + """Execute forward using rows referenced from the server DataPlane.""" response = http_post( - url=f'{self.server_url}/forward', + url=f'{self.server_url}/forward_from_data_plane', json_data={ - **_component_input_payload(inputs), + **_data_ref_payload(inputs), 'adapter_name': self.adapter_name, 'input_field': input_field, 'kwarg_fields': kwarg_fields or {}, @@ -96,9 +113,9 @@ def forward( response.raise_for_status() return ForwardResponse(**response.json()) - def forward_only( + def forward_only_from_data_plane( self, - inputs: Any | DataRef | list[DataRef], + inputs: DataRef | list[DataRef], *, input_field: str | None = None, kwarg_fields: dict[str, str] | None = None, @@ -106,9 +123,9 @@ def forward_only( output_fields: dict[str, str] | None = None, **kwargs, ) -> ForwardResponse | DataRef: - """Execute forward-only, optionally reading and updating server-side rows.""" + """Execute forward-only using DataPlane rows and optionally append outputs.""" body = { - **_component_input_payload(inputs), + **_data_ref_payload(inputs), 'adapter_name': self.adapter_name, 'input_field': input_field, 'kwarg_fields': kwarg_fields or {}, @@ -117,7 +134,7 @@ def forward_only( **json_safe(kwargs), } response = http_post( - url=f'{self.server_url}/forward_only', + url=f'{self.server_url}/forward_only_from_data_plane', json_data=body, ) response.raise_for_status() @@ -152,19 +169,28 @@ def backward(self, **kwargs) -> None: ) response.raise_for_status() - def forward_backward( + def forward_backward(self, inputs: Any, **kwargs) -> ForwardBackwardResponse: + """Execute combined forward and backward pass on inline inputs.""" + response = http_post( + url=f'{self.server_url}/forward_backward', + json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs}, + ) + response.raise_for_status() + return ForwardBackwardResponse(**response.json()) + + def forward_backward_from_data_plane( self, - inputs: Any | DataRef | list[DataRef], + inputs: DataRef | list[DataRef], *, input_field: str | None = None, kwarg_fields: dict[str, str] | None = None, **kwargs, ) -> ForwardBackwardResponse: - """Execute forward/backward over inline rows or server-side references.""" + """Execute forward/backward using rows referenced from the server DataPlane.""" response = http_post( - url=f'{self.server_url}/forward_backward', + url=f'{self.server_url}/forward_backward_from_data_plane', json_data={ - **_component_input_payload(inputs), + **_data_ref_payload(inputs), 'adapter_name': self.adapter_name, 'input_field': input_field, 'kwarg_fields': kwarg_fields or {}, diff --git a/src/twinkle_client/types/__init__.py b/src/twinkle_client/types/__init__.py index a0952717f..1c25324a1 100644 --- a/src/twinkle_client/types/__init__.py +++ b/src/twinkle_client/types/__init__.py @@ -16,6 +16,8 @@ ClipGradNormResponse, CreateRequest, CreateResponse, + DataPlaneForwardOnlyRequest, + DataPlaneForwardRequest, ForwardBackwardResponse, ForwardOnlyRequest, ForwardRequest, diff --git a/src/twinkle_client/types/model.py b/src/twinkle_client/types/model.py index 830ed6010..7dcf4af0d 100644 --- a/src/twinkle_client/types/model.py +++ b/src/twinkle_client/types/model.py @@ -17,46 +17,41 @@ class Config: class ForwardRequest(BaseModel): - inputs: Any = None - input_refs: List[DataRef] | None = None - input_field: str | None = None - kwarg_fields: Dict[str, str] = Field(default_factory=dict) + inputs: Any adapter_name: str - @model_validator(mode='after') - def validate_input(self) -> 'ForwardRequest': - if (self.inputs is None) == (self.input_refs is None): - raise ValueError('exactly one of inputs and input_refs must be provided') - if self.input_refs is not None and not self.input_refs: - raise ValueError('input_refs must not be empty') - return self - class Config: extra = 'allow' class ForwardOnlyRequest(BaseModel): - inputs: Any = None - input_refs: List[DataRef] | None = None + inputs: Any + adapter_name: Optional[str] = None + + class Config: + extra = 'allow' + + +class DataPlaneForwardRequest(BaseModel): + input_refs: List[DataRef] = Field(min_length=1) input_field: str | None = None kwarg_fields: Dict[str, str] = Field(default_factory=dict) + adapter_name: str + + class Config: + extra = 'allow' + + +class DataPlaneForwardOnlyRequest(DataPlaneForwardRequest): output_ref: DataRef | None = None output_fields: Dict[str, str] = Field(default_factory=dict) - adapter_name: Optional[str] = None @model_validator(mode='after') - def validate_input(self) -> 'ForwardOnlyRequest': - if (self.inputs is None) == (self.input_refs is None): - raise ValueError('exactly one of inputs and input_refs must be provided') - if self.input_refs is not None and not self.input_refs: - raise ValueError('input_refs must not be empty') + def validate_output(self) -> 'DataPlaneForwardOnlyRequest': if (self.output_ref is None) != (len(self.output_fields) == 0): raise ValueError('output_ref and output_fields must be configured together') return self - class Config: - extra = 'allow' - class AdapterRequest(BaseModel): adapter_name: str diff --git a/tests/server/contract/client_api_baseline.json b/tests/server/contract/client_api_baseline.json index d87b0f89c..65f9db147 100644 --- a/tests/server/contract/client_api_baseline.json +++ b/tests/server/contract/client_api_baseline.json @@ -777,6 +777,16 @@ ] } }, + "/twinkle/forward_from_data_plane": { + "POST": { + "operationId": "forward_from_data_plane_twinkle_forward_from_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/forward_backward": { "POST": { "operationId": "forward_backward_twinkle_forward_backward_post", @@ -787,6 +797,16 @@ ] } }, + "/twinkle/forward_backward_from_data_plane": { + "POST": { + "operationId": "forward_backward_from_data_plane_twinkle_forward_backward_from_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/forward_only": { "POST": { "operationId": "forward_only_twinkle_forward_only_post", @@ -797,6 +817,16 @@ ] } }, + "/twinkle/forward_only_from_data_plane": { + "POST": { + "operationId": "forward_only_from_data_plane_twinkle_forward_only_from_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/get_state_dict": { "POST": { "operationId": "get_state_dict_twinkle_get_state_dict_post", diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index 51cc2de56..e404b9706 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -64,7 +64,7 @@ async def schedule_task_and_wait(self, task, **kwargs): @pytest.mark.asyncio -async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() -> None: +async def test_forward_backward_inline_route_keeps_original_request_shape() -> None: management = _SchedulingManagement() app = FastAPI() _register_twinkle_routes(app, lambda: management) @@ -72,6 +72,31 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - request = Request({'type': 'http', 'headers': []}) request.state.session_id = 'session' body = types.ForwardRequest( + adapter_name='adapter', + inputs=[{'input_ids': [1, 2, 3]}], + micro_batch_size=1, + ) + + await route.endpoint(request, body, management) + + inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] + assert adapter_name == 'session-adapter' + assert [row['input_ids'] for row in inputs] == [[1, 2, 3]] + assert forwarded_kwargs == {'micro_batch_size': 1} + + +@pytest.mark.asyncio +async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() -> None: + management = _SchedulingManagement() + app = FastAPI() + _register_twinkle_routes(app, lambda: management) + route = next( + route for route in app.routes + if getattr(route, 'path', None) == '/twinkle/forward_backward_from_data_plane' + ) + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.DataPlaneForwardRequest( adapter_name='adapter', input_refs=[ types.DataRef(ref_id='data-a', size=4, num_tokens=4), @@ -112,10 +137,13 @@ async def test_forward_backward_binds_nested_dpo_ref_logps_without_coercion() -> ] app = FastAPI() _register_twinkle_routes(app, lambda: management) - route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') + route = next( + route for route in app.routes + if getattr(route, 'path', None) == '/twinkle/forward_backward_from_data_plane' + ) request = Request({'type': 'http', 'headers': []}) request.state.session_id = 'session' - body = types.ForwardRequest( + body = types.DataPlaneForwardRequest( adapter_name='adapter', input_refs=[types.DataRef(ref_id='dpo', size=2, num_tokens=6)], kwarg_fields={'ref_outputs.logps': 'ref_logps'}, diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index c4a1ea6a8..99b91fcc9 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -41,21 +41,50 @@ def post(*, url, json_data=None, **_kwargs): DataRef(ref_id='data-1', size=2, fields=['train_input']), DataRef(ref_id='data-2', size=2, fields=['train_input']), ] - model.forward_backward( + model.forward_backward_from_data_plane( refs, input_field='train_input', kwarg_fields={'advantages': 'advantage'}, ) url, body = calls[-1] - assert url.endswith('/model/base/twinkle/forward_backward') + assert url.endswith('/model/base/twinkle/forward_backward_from_data_plane') assert body['input_refs'] == [ref.model_dump() for ref in refs] assert body['input_field'] == 'train_input' assert body['kwarg_fields'] == {'advantages': 'advantage'} assert body['adapter_name'] == 'adapter' -def test_model_forward_accepts_data_ref_without_a_separate_submit_api(monkeypatch) -> None: +def test_model_inline_forward_methods_keep_the_original_endpoints(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + return _Response({} if url.endswith('/create') else {'result': {}}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + + model = module.MultiLoraTransformersModel('ms://base') + model.adapter_name = 'adapter' + inputs = [{'input_ids': [1, 2]}] + model.forward(inputs, return_logits=True) + model.forward_only(inputs, disable_lora=True) + model.forward_backward(inputs, micro_batch_size=1) + + assert [url.rsplit('/', 1)[-1] for url, _ in calls[-3:]] == [ + 'forward', + 'forward_only', + 'forward_backward', + ] + assert all(body['inputs'] == inputs for _, body in calls[-3:]) + assert all('input_refs' not in body for _, body in calls[-3:]) + + +def test_model_data_plane_forward_uses_a_separate_api(monkeypatch) -> None: import twinkle_client.http as http_module from twinkle_client.model import multi_lora_transformers as module @@ -70,14 +99,44 @@ def post(url, json_data=None, **_kwargs): model = module.MultiLoraTransformersModel('ms://base') ref = DataRef(ref_id='data-1', size=2, fields=['train_input']) - model.forward(ref, input_field='train_input') + model.forward_from_data_plane(ref, input_field='train_input') url, body = calls[-1] - assert url.endswith('/model/base/twinkle/forward') + assert url.endswith('/model/base/twinkle/forward_from_data_plane') assert body['input_refs'] == [ref.model_dump()] assert body['input_field'] == 'train_input' +def test_model_data_plane_forward_only_can_append_selected_outputs(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + ref = DataRef(ref_id='data-1', size=2, fields=['input_ids']) + updated_ref = ref.model_copy(update={'fields': ['input_ids', 'ref_logps']}) + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + return _Response({} if url.endswith('/create') else {'result': updated_ref.model_dump()}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + + model = module.MultiLoraTransformersModel('ms://base') + result = model.forward_only_from_data_plane( + ref, + output_ref=ref, + output_fields={'logps': 'ref_logps'}, + disable_lora=True, + ) + + url, body = calls[-1] + assert url.endswith('/model/base/twinkle/forward_only_from_data_plane') + assert body['input_refs'] == [ref.model_dump()] + assert body['output_ref'] == ref.model_dump() + assert result == updated_ref + + def test_sampler_async_data_plane_path_returns_reference_without_materializing(monkeypatch) -> None: import twinkle_client.http as http_module from twinkle_client.sampler import vllm_sampler as module diff --git a/tests/twinkle_client/test_client_orchestrated_dpo.py b/tests/twinkle_client/test_client_orchestrated_dpo.py index 1e7426fd0..fa1bd0823 100644 --- a/tests/twinkle_client/test_client_orchestrated_dpo.py +++ b/tests/twinkle_client/test_client_orchestrated_dpo.py @@ -67,7 +67,7 @@ def __init__(self): self.steps = 0 self.forward_backward_kwargs = [] - async def forward_only(self, ref, **kwargs): + async def forward_only_from_data_plane(self, ref, **kwargs): self.references += 1 name = f'reference-{self.references}' events.append(f'{name}-start') @@ -78,7 +78,7 @@ async def forward_only(self, ref, **kwargs): assert kwargs['output_fields'] == {'logps': 'ref_logps'} return ref.model_copy(update={'fields': [*ref.fields, 'ref_logps']}) - async def forward_backward(self, _ref, **kwargs): + async def forward_backward_from_data_plane(self, _ref, **kwargs): self.forward_backward_kwargs.append(kwargs) events.append('train-start') first_train_started.set() diff --git a/tests/twinkle_client/test_client_orchestrated_grpo.py b/tests/twinkle_client/test_client_orchestrated_grpo.py index 90210e02d..2ba86ef6f 100644 --- a/tests/twinkle_client/test_client_orchestrated_grpo.py +++ b/tests/twinkle_client/test_client_orchestrated_grpo.py @@ -66,7 +66,7 @@ async def save(self, name): self.saved.append(name) return {'twinkle_path': f'/checkpoints/{name}'} - async def forward_backward(self, _refs, **kwargs): + async def forward_backward_from_data_plane(self, _refs, **kwargs): self.forward_backward_kwargs.append(kwargs) events.append('train') first_train_started.set() From 29e325bb70d84f824a740f69a4f0643ccba851cd Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 18 Aug 2026 14:13:55 +0800 Subject: [PATCH 12/20] chore: scope branch changes to async RL --- .../async_rl/client_orchestrated_dpo.py | 2 +- .../async_rl/client_orchestrated_grpo.py | 2 +- .../server/transformer/server_config.yaml | 21 +- cookbook/client/twinkle/short_math_grpo.py | 14 +- cookbook/rl/sync_barrier_multi_lora_grpo.py | 608 ------------------ src/twinkle_client/common/serialize.py | 1 - tests/server/gateway/test_future_retrieval.py | 85 --- .../server/model/test_twinkle_async_inputs.py | 3 + tests/server/state/test_managers.py | 1 + .../test_async_rl_native_tq.py | 16 - 10 files changed, 11 insertions(+), 742 deletions(-) delete mode 100644 cookbook/rl/sync_barrier_multi_lora_grpo.py delete mode 100644 tests/server/gateway/test_future_retrieval.py diff --git a/cookbook/client/async_rl/client_orchestrated_dpo.py b/cookbook/client/async_rl/client_orchestrated_dpo.py index 625d75b5c..eb9a4bf1a 100644 --- a/cookbook/client/async_rl/client_orchestrated_dpo.py +++ b/cookbook/client/async_rl/client_orchestrated_dpo.py @@ -13,7 +13,7 @@ from twinkle.preprocessor import EmojiDPOProcessor from twinkle_client import DataPlaneClient, init_twinkle_client from twinkle_client.async_rl import Worker, WorkerPipeline -from twinkle_client.common.serialize import json_safe +from twinkle_client.common.json_utils import json_safe from twinkle_client.model import MultiLoraTransformersModel from twinkle_client.types import DataRef diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py index f30d31673..2a1ed0458 100644 --- a/cookbook/client/async_rl/client_orchestrated_grpo.py +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -17,7 +17,7 @@ from twinkle.reward import GSM8KAccuracyReward from twinkle_client import DataPlaneClient, init_twinkle_client from twinkle_client.async_rl import Worker, WorkerPipeline -from twinkle_client.common.serialize import json_safe +from twinkle_client.common.json_utils import json_safe from twinkle_client.model import MultiLoraTransformersModel from twinkle_client.sampler import vLLMSampler diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index b7dac3cdf..d3ddb2adb 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -53,21 +53,6 @@ applications: env_vars: TWINKLE_FAIL_FAST: "0" - # TransferQueue-backed data references used by client-orchestrated RL. - - name: data-plane - route_prefix: /api/v1/data-plane - import_path: data_plane - args: - config: - backend: - SimpleStorage: - num_data_storage_units: 2 - deployments: - - name: DataPlaneManagement - num_replicas: 1 - ray_actor_options: - num_cpus: 1 - # 2. Model Service - Hosts the base model for training. - name: models-Qwen3.5-4B route_prefix: /api/v1/model/Qwen/Qwen3.5-4B @@ -75,9 +60,7 @@ applications: args: backend: transformers # Model backend: transformers | megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier - max_loras: 8 # Concurrent client-owned training adapters max_length: 10240 - data_plane_url: http://127.0.0.1:8000/api/v1/data-plane nproc_per_node: 1 # Number of GPU processes per node device_group: name: model @@ -111,14 +94,12 @@ applications: import_path: sampler args: model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier - data_plane_url: http://127.0.0.1:8000/api/v1/data-plane nproc_per_node: 1 # Number of GPU processes per node - sampler_type: vllm_async # Non-blocking vLLM admission for client-orchestrated RL + sampler_type: vllm # Inference engine: 'vllm' (fast) or 'torch' (TorchSampler) engine_args: # vLLM engine-specific settings max_model_len: 4096 # Maximum sequence length the engine supports gpu_memory_utilization: 0.5 # Fraction of GPU memory to use (0.0-1.0) enable_lora: true # Allow loading LoRA adapters during inference - max_loras: 8 # Published policies cached across tenants logprobs_mode: processed_logprobs # Logprobs mode for sampling results device_group: # Logical device group for the sampler name: sampler diff --git a/cookbook/client/twinkle/short_math_grpo.py b/cookbook/client/twinkle/short_math_grpo.py index a9af2f32d..1e90d38ce 100644 --- a/cookbook/client/twinkle/short_math_grpo.py +++ b/cookbook/client/twinkle/short_math_grpo.py @@ -82,12 +82,6 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: # ========== Configuration ========== BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') MODEL_ID = f'ms://{BASE_MODEL}' -TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) -TEMPLATE_CLS = os.environ.get( - 'TWINKLE_TEMPLATE_CLS', - 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', -) -DATASET_ID = os.environ.get('TWINKLE_DATASET_ID', 'ms://modelscope/gsm8k') NUM_GENERATIONS = 4 MAX_NEW_TOKENS = 1024 LEARNING_RATE = 2e-5 @@ -107,8 +101,8 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: 'and put your final answer within \\boxed{}.') def create_gsm8k_dataset(): - dataset = Dataset(DatasetMeta(DATASET_ID, subset_name='main', split='train', data_slice=range(DATA_NUM))) - dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=2048, enable_thinking=False) + dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train', data_slice=range(DATA_NUM))) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=2048, enable_thinking=False) dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT)) dataset.encode(add_generation_prompt=True) return dataset @@ -185,11 +179,11 @@ def train(): # Set processor and template for encoding inputs model.set_processor('InputProcessor') - model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) + model.set_template('Qwen3_5Template', model_id=MODEL_ID) # Step 4: Configure the sampler sampler = vLLMSampler(model_id=MODEL_ID) - sampler.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID) # Step 5: Setup metrics and advantage function advantage_fn = GRPOAdvantage() diff --git a/cookbook/rl/sync_barrier_multi_lora_grpo.py b/cookbook/rl/sync_barrier_multi_lora_grpo.py deleted file mode 100644 index 2914d0ec3..000000000 --- a/cookbook/rl/sync_barrier_multi_lora_grpo.py +++ /dev/null @@ -1,608 +0,0 @@ -"""Synchronous barrier baseline for native async multi-LoRA GRPO. - -The model, sampler, datasets, rewards, batch semantics, and checkpoint cadence -match ``async_multi_lora_grpo.py``. The only intentional difference is the -execution schedule: every round finishes rollout for all active contexts -before any context starts training. -""" - -from __future__ import annotations - -import argparse -import os -import shutil -import time -from dataclasses import dataclass -from typing import Any, Iterator, Sequence - -from omegaconf import OmegaConf - -from twinkle.metric import MetricRecord, create_metrics_reporter -from twinkle_agentic.async_rl.metrics import advantage_signal_metrics, rollout_metrics -from twinkle_agentic.async_rl.pipeline import (_prompt_batches, _reward_for_context, _train_batch) -from twinkle_agentic.async_rl.tq_utils import REQUIRED_MODEL_INPUT_FIELDS, columns_to_tq_fields -from twinkle_agentic.async_rl.types import LoraContext, PartitionAdmission -from twinkle_agentic.async_rl.utils import ( - TrainBatchConfig, - build_native_fsdp_model_kwargs, - configure_lora_lr_scheduler, - resolve_context_learning_rate, - resolve_context_lora_target_modules, - resolve_context_loss_config, - resolve_model_attention_implementation, - resolve_sequence_parallel_size, - sample_responses_to_rollout_rows, - sampler_data_parallel_size, - validate_context_batch_config, -) -from twinkle_agentic.async_rl.vllm_sampler_tq import _compute_reward_metrics - - -@dataclass -class SyncContextState: - context: LoraContext - prompt_batches: Iterator[Sequence[dict[str, Any]]] - rollout_batch_size: int - num_generations: int - sampling_params: Any - mini_batch_size: int - reward_fn: Any - adapter_path: str - adapter_history: list[str] - partition_step: int = 0 - optimizer_steps: int = 0 - policy_version: int = 0 - exhausted: bool = False - - -@dataclass -class SyncPartition: - admission: PartitionAdmission - state: SyncContextState - rows: list[dict[str, Any]] - rewards: list[float] - advantages: list[float] | None = None - - -class SyncBarrierMultiLoraGRPO: - - def __init__(self, raw_config: dict[str, Any]): - import twinkle - from peft import LoraConfig - from twinkle import DeviceGroup, DeviceMesh - from twinkle.data_format import SamplingParams - from twinkle.model import MultiLoraTransformersModel - from twinkle.processor import InputProcessor - from twinkle.sampler import vLLMSampler - - raw_config = OmegaConf.to_container(OmegaConf.create(raw_config), resolve=True) - if not isinstance(raw_config, dict): - raise TypeError('sync RL config must resolve to a mapping') - - runtime = raw_config['runtime'] - model_config = raw_config['model'] - lora_data = raw_config['lora'] - loss_data = raw_config.get('loss') - template_data = raw_config.get('template', {}) - template_cls = template_data.get('cls', 'Qwen3_5Template') - enable_thinking = bool(template_data.get('enable_thinking', False)) - model_gpus = int(runtime['model_gpus']) - sampler_gpus = int(runtime['sampler_gpus']) - sampler_tp = int(runtime['sampler_tp']) - sampler_dp = sampler_data_parallel_size(sampler_gpus, sampler_tp) - sequence_parallel_size = resolve_sequence_parallel_size( - model_gpus, - int(model_config['sequence_parallel_size']), - ) - padding_free = bool(model_config['padding_free']) - attn_implementation = resolve_model_attention_implementation( - model_config, - padding_free=padding_free, - sequence_parallel_size=sequence_parallel_size, - ) - model_max_length = int(model_config['max_length']) - sampler_config = raw_config['sampler'] - total_gpus = model_gpus + sampler_gpus - - twinkle.initialize( - mode='ray', - nproc_per_node=total_gpus, - groups=[ - DeviceGroup('model', list(range(model_gpus)), device_type='GPU'), - DeviceGroup( - 'sampler', - list(range(model_gpus, total_gpus)), - device_type='GPU', - gpus_per_worker=sampler_tp, - ), - ], - lazy_collect=False, - ) - model_mesh = DeviceMesh.from_sizes( - world_size=model_gpus, - dp_size=model_gpus, - ulysses_size=sequence_parallel_size, - ) - model_data_parallel_size = model_mesh.data_world_size - self.model_data_parallel_size = model_data_parallel_size - sampler_mesh = DeviceMesh.from_sizes( - world_size=sampler_gpus, - dp_size=sampler_dp, - tp_size=sampler_tp, - ) - model_kwargs = build_native_fsdp_model_kwargs(model_config) - if attn_implementation is not None: - model_kwargs['attn_implementation'] = attn_implementation - self.model = MultiLoraTransformersModel( - model_id=runtime['model_id'], - device_mesh=model_mesh, - remote_group='model', - max_length=model_max_length, - **model_kwargs, - ) - self.train_batch_configs: dict[str, TrainBatchConfig] = {} - self.states: list[SyncContextState] = [] - for item in raw_config['lora_contexts']: - context = LoraContext( - item['tenant_id'], - item['training_run_id'], - runtime['model_id'], - item['adapter_name'], - ) - rollout = item['rollout'] - train = item['train'] - rollout_batch_size = int(rollout['batch_size']) - num_generations = int(rollout['num_generations']) - train_batch_config = TrainBatchConfig( - mini_batch_size=int(train['mini_batch_size']), - micro_batch_size=int(train['micro_batch_size']), - dynamic_batching=bool(train.get('dynamic_batching', False)), - max_tokens_per_micro_batch=( - int(train['max_tokens_per_micro_batch']) - if train.get('max_tokens_per_micro_batch') is not None else None - ), - packing_algorithm=str(train.get('packing_algorithm', 'ffd')), - ) - validate_context_batch_config( - context.key, - rollout_groups=rollout_batch_size, - num_generations=num_generations, - train=train_batch_config, - sampler_dp=sampler_dp, - model_dp=model_data_parallel_size, - ) - adapter_lora_config = LoraConfig( - target_modules=resolve_context_lora_target_modules(item, lora_data), - r=lora_data['r'], - lora_alpha=lora_data['alpha'], - lora_dropout=lora_data['dropout'], - ) - self.model.add_adapter_to_model( - context.adapter_name, - adapter_lora_config, - gradient_accumulation_steps=1, - ) - self.model.set_optimizer( - 'AdamW', - lr=resolve_context_learning_rate(train, lora_data), - adapter_name=context.adapter_name, - ) - configure_lora_lr_scheduler(self.model, context.adapter_name, lora_data) - loss_cls, loss_kwargs = resolve_context_loss_config(item, loss_data) - self.model.set_loss( - loss_cls, - adapter_name=context.adapter_name, - **loss_kwargs, - ) - self.model.set_processor( - InputProcessor, - adapter_name=context.adapter_name, - padding_free=padding_free, - ) - self.model.set_template( - template_cls, - model_id=runtime['model_id'], - adapter_name=context.adapter_name, - enable_thinking=enable_thinking, - max_length=model_max_length, - ) - initial_path = self.model.save( - f'sync-{context.adapter_name}-initial', - output_dir=runtime['output_dir'], - adapter_name=context.adapter_name, - ) - state = SyncContextState( - context=context, - prompt_batches=iter( - _prompt_batches( - item['dataset'], - model_id=runtime['model_id'], - batch_size=rollout_batch_size, - template_cls=template_cls, - enable_thinking=enable_thinking, - )), - rollout_batch_size=rollout_batch_size, - num_generations=num_generations, - sampling_params=SamplingParams( - max_tokens=rollout['max_tokens'], - temperature=rollout['temperature'], - top_p=rollout['top_p'], - repetition_penalty=float(rollout.get('repetition_penalty', 1.0)), - logprobs=1, - num_samples=1, - ), - mini_batch_size=train_batch_config.mini_batch_size, - reward_fn=_reward_for_context( - item.get('reward'), - context_key=context.key, - ), - adapter_path=initial_path, - adapter_history=[initial_path], - ) - self.states.append(state) - self.train_batch_configs[context.key] = train_batch_config - - sampler_engine_args = { - 'tensor_parallel_size': sampler_tp, - 'enable_lora': True, - 'max_loras': int(runtime['sampler_max_loras']), - 'max_lora_rank': lora_data['r'], - 'max_model_len': int(sampler_config['max_model_len']), - 'gpu_memory_utilization': float(sampler_config['gpu_memory_utilization']), - 'max_num_seqs': int(sampler_config['max_num_seqs']), - 'enforce_eager': bool(sampler_config['enforce_eager']), - } - if sampler_config.get('max_num_batched_tokens') is not None: - sampler_engine_args['max_num_batched_tokens'] = int(sampler_config['max_num_batched_tokens']) - self.sampler = vLLMSampler( - model_id=runtime['model_id'], - remote_group='sampler', - device_mesh=sampler_mesh, - engine_args=sampler_engine_args, - ) - self.sampler.set_template( - template_cls, - model_id=runtime['model_id'], - enable_thinking=enable_thinking, - max_length=model_max_length, - ) - self.output_dir = runtime['output_dir'] - self.max_steps = runtime.get('max_steps') - self.max_steps = None if self.max_steps is None else int(self.max_steps) - self.keep_adapter_versions = max(0, int(runtime.get('keep_adapter_versions', 0))) - self.metrics = create_metrics_reporter( - raw_config.get('metrics'), - run_id=str(runtime.get('run_id', 'sync_barrier_multi_lora_grpo')), - ) - self.completed_partitions = 0 - self._creation_order = 0 - - def _record_metric( - self, - stage: str, - *, - admission: PartitionAdmission | None = None, - context: LoraContext | None = None, - values: dict[str, Any] | None = None, - status: str = 'completed', - attributes: dict[str, Any] | None = None, - optimizer_step: int | None = None, - policy_version: int | None = None, - ) -> None: - if self.metrics is None: - return - self.metrics.record(MetricRecord( - stage=stage, - values=dict(values or {}), - context_key=( - admission.context.key if admission is not None - else context.key if context is not None else None - ), - partition_id=admission.partition_id if admission is not None else None, - partition_index=admission.step if admission is not None else None, - optimizer_step=optimizer_step, - policy_version=policy_version, - status=status, - attributes=dict(attributes or {}), - )) - - def run(self) -> dict[str, Any]: - started = time.perf_counter() - try: - while self.max_steps is None or self.completed_partitions < self.max_steps: - partitions = self._rollout_round() - if not partitions: - break - self._advantage_round(partitions) - self._train_round(partitions) - except Exception as exc: - self._record_metric( - 'run', - status='failed', - values={'wall_time_s': time.perf_counter() - started}, - attributes={'error': f'{type(exc).__name__}: {exc}'}, - ) - if self.metrics is not None: - self.metrics.close() - raise - result = { - 'trained_partitions': self.completed_partitions, - 'wall_time_s': time.perf_counter() - started, - 'per_context': { - state.context.key: { - 'optimizer_steps': state.optimizer_steps, - 'policy_version': state.policy_version, - 'adapter_path': state.adapter_path, - } - for state in self.states - }, - } - self._record_metric( - 'run', - values={ - 'trained_partitions': result['trained_partitions'], - 'wall_time_s': result['wall_time_s'], - }, - ) - if self.metrics is not None: - self.metrics.flush() - result['metrics_health'] = self.metrics.health() - self.metrics.close() - return result - - def _rollout_round(self) -> list[SyncPartition]: - partitions = [] - for state in self.states: - if state.exhausted: - continue - if self.max_steps is not None and self.completed_partitions + len(partitions) >= self.max_steps: - break - prompts = next(state.prompt_batches, None) - if prompts is None or len(prompts) != state.rollout_batch_size: - state.exhausted = True - continue - admission = PartitionAdmission( - context=state.context, - partition_id=state.context.partition_id(state.partition_step), - step=state.partition_step, - target_groups=state.rollout_batch_size, - num_generations=state.num_generations, - created_order=self._creation_order, - ) - self._creation_order += 1 - self._record_metric( - 'rollout', - admission=admission, - status='submitted', - policy_version=state.policy_version, - values={ - 'prompt_count': admission.target_groups, - 'sample_count': admission.sample_count, - 'num_generations': admission.num_generations, - }, - attributes={'scope': 'partition'}, - ) - rollout_started = time.perf_counter() - sources = [{ - **dict(prompt), - 'group_id': f'{admission.partition_id}/group_{group_index}', - 'generation_idx': generation_index, - } for group_index, prompt in enumerate(prompts) - for generation_index in range(state.num_generations)] - responses = self.sampler.sample( - [dict(prompt) for prompt in prompts for _ in range(state.num_generations)], - state.sampling_params, - adapter_name=state.context.adapter_name, - adapter_path=state.adapter_path, - ) - rows = sample_responses_to_rollout_rows( - sources, - responses, - policy_version=state.policy_version, - ) - if len(rows) != admission.sample_count: - raise ValueError( - f'{admission.partition_id} expected {admission.sample_count} samples, got {len(rows)}') - for row in rows: - row.update({ - 'rollout_adapter_path': state.adapter_path, - 'rollout_policy_versions': [state.policy_version], - 'initial_policy_version': state.policy_version, - 'final_policy_version': state.policy_version, - 'policy_version_span': 0, - }) - rewards = [float(value) for value in state.reward_fn(rows, context=state.context)] - if len(rewards) != len(rows): - raise ValueError(f'{admission.partition_id} reward count does not match sample count') - rollout_latency_s = time.perf_counter() - rollout_started - self._record_rollout_groups(state, admission, rows, rewards) - self._record_metric( - 'rollout', - admission=admission, - policy_version=state.policy_version, - values=rollout_metrics( - completion_lengths=[int(row['completion_length']) for row in rows], - stop_reasons=[row.get('stop_reason') for row in rows], - rollout_latency_s=rollout_latency_s, - ), - attributes={'scope': 'partition'}, - ) - partitions.append(SyncPartition(admission, state, rows, rewards)) - state.partition_step += 1 - return partitions - - def _record_rollout_groups( - self, - state: SyncContextState, - admission: PartitionAdmission, - rows: list[dict[str, Any]], - rewards: list[float], - ) -> None: - for group_index in range(admission.target_groups): - start = group_index * admission.num_generations - end = start + admission.num_generations - group_rows = rows[start:end] - group_rewards = rewards[start:end] - metrics = { - **_compute_reward_metrics( - {state.context.key: state.reward_fn}, - state.context, - group_rows, - group_rewards, - ), - **rollout_metrics( - rewards={'reward': group_rewards}, - completion_lengths=[int(row['completion_length']) for row in group_rows], - stop_reasons=[row.get('stop_reason') for row in group_rows], - ), - } - self._record_metric( - 'rollout', - admission=admission, - policy_version=state.policy_version, - values=metrics, - attributes={ - 'scope': 'group', - 'group_id': f'{admission.partition_id}/group_{group_index}', - }, - ) - - def _advantage_round(self, partitions: list[SyncPartition]) -> None: - from twinkle.advantage import GRPOAdvantage - - advantage_fn = GRPOAdvantage() - for partition in partitions: - admission = partition.admission - partition.advantages = advantage_fn( - partition.rewards, - num_generations=admission.num_generations, - scale='group', - ).tolist() - samples_per_batch = admission.num_generations - for start in range(0, len(partition.rows), samples_per_batch): - end = min(start + samples_per_batch, len(partition.rows)) - self._record_metric( - 'advantage', - admission=admission, - policy_version=partition.state.policy_version, - values={ - 'sample_count': end - start, - **advantage_signal_metrics( - partition.rewards[start:end], - partition.advantages[start:end], - num_generations=admission.num_generations, - ), - }, - ) - - def _train_round(self, partitions: list[SyncPartition]) -> None: - for partition in partitions: - admission = partition.admission - state = partition.state - assert partition.advantages is not None - samples_per_batch = state.mini_batch_size - for start in range(0, len(partition.rows), samples_per_batch): - end = start + samples_per_batch - batch = self._training_batch( - partition.rows[start:end], - partition.rewards[start:end], - partition.advantages[start:end], - ) - train_started = time.perf_counter() - metrics = _train_batch( - self.model, - self.train_batch_configs, - batch, - admission, - model_data_parallel_size=self.model_data_parallel_size, - ) - state.optimizer_steps += 1 - metrics.update({ - 'sample_count': end - start, - 'train_latency_s': time.perf_counter() - train_started, - 'policy_version_gap_mean': 0.0, - 'policy_version_gap_p95': 0.0, - 'policy_version_gap_max': 0, - 'rollout_policy_span_mean': 0.0, - 'rollout_policy_span_max': 0, - }) - self._record_metric( - 'train', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values=metrics, - ) - finalize_started = time.perf_counter() - next_policy_version = state.policy_version + 1 - save_started = time.perf_counter() - state.adapter_path = self.model.save( - f'sync-{state.context.adapter_name}-v{next_policy_version}', - output_dir=self.output_dir, - adapter_name=state.context.adapter_name, - ) - adapter_save_latency_s = time.perf_counter() - save_started - publish_started = time.perf_counter() - state.policy_version = next_policy_version - policy_publish_latency_s = time.perf_counter() - publish_started - state.adapter_history.append(state.adapter_path) - self._record_metric( - 'policy', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values={ - 'adapter_save_latency_s': adapter_save_latency_s, - 'policy_publish_latency_s': policy_publish_latency_s, - }, - attributes={'operation': 'publish', 'adapter_path': state.adapter_path}, - ) - prune_started = time.perf_counter() - self._prune_adapter_history(state) - adapter_prune_latency_s = time.perf_counter() - prune_started - self.completed_partitions += 1 - self._record_metric( - 'partition', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values={ - 'adapter_save_latency_s': adapter_save_latency_s, - 'policy_publish_latency_s': policy_publish_latency_s, - 'adapter_prune_latency_s': adapter_prune_latency_s, - 'partition_finalize_latency_s': time.perf_counter() - finalize_started, - }, - ) - - @staticmethod - def _training_batch(rows: list[dict[str, Any]], rewards: list[float], advantages: list[float]): - fields = { - name: [row[name] for row in rows] - for name in (*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs') - } - fields.update({'rewards': rewards, 'advantages': advantages}) - return columns_to_tq_fields(fields, len(rows)) - - def _prune_adapter_history(self, state: SyncContextState) -> None: - retained_count = max(1, self.keep_adapter_versions) - stale = state.adapter_history[:-retained_count] - state.adapter_history = state.adapter_history[-retained_count:] - for path in stale: - if os.path.isdir(path): - shutil.rmtree(path) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument('--config', default='cookbook/rl/sync_barrier_multi_lora_grpo.yaml') - args = parser.parse_args() - config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) - print(SyncBarrierMultiLoraGRPO(config).run()) - - -if __name__ == '__main__': - main() - -# MODEL_ID=/path/to/model \ -# DATASET_ID=/path/to/gsm8k \ -# python cookbook/rl/async_multi_lora_grpo.py diff --git a/src/twinkle_client/common/serialize.py b/src/twinkle_client/common/serialize.py index 0f0a49096..42a27beb3 100644 --- a/src/twinkle_client/common/serialize.py +++ b/src/twinkle_client/common/serialize.py @@ -7,7 +7,6 @@ from typing import Any, Mapping from twinkle.dataset import DatasetMeta -from .json_utils import json_safe supported_types = { DatasetMeta, diff --git a/tests/server/gateway/test_future_retrieval.py b/tests/server/gateway/test_future_retrieval.py deleted file mode 100644 index 1152b7db1..000000000 --- a/tests/server/gateway/test_future_retrieval.py +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from __future__ import annotations - -import pytest -from fastapi import FastAPI, Request -from fastapi.testclient import TestClient -from unittest.mock import AsyncMock, MagicMock - -tinker = pytest.importorskip('tinker') - -from twinkle.server.gateway.tinker_handlers import _register_tinker_routes - - -def _make_client(get_future: AsyncMock) -> TestClient: - management = MagicMock() - management.state.get_future = get_future - management.supported_models = [] - - app = FastAPI() - - @app.middleware('http') - async def _set_request_state(request: Request, call_next): - authorization = request.headers.get('Authorization', '') - request.state.token = authorization.removeprefix('Bearer ') - request.state.session_id = request.headers.get('X-Twinkle-Session-Id', '') - return await call_next(request) - - _register_tinker_routes(app, lambda: management) - return TestClient(app) - - -def test_retrieve_future_uses_request_id_as_capability() -> None: - get_future = AsyncMock(return_value={'status': 'completed', 'result': {'value': 7}}) - client = _make_client(get_future) - - response = client.post( - '/retrieve_future', - json={'request_id': 'req-1'}, - headers={ - 'Authorization': 'Bearer tenant-token', - 'X-Twinkle-Session-Id': 'session-1', - }, - ) - - assert response.status_code == 200 - assert response.json() == {'value': 7} - get_future.assert_awaited_once_with('req-1') - - -def test_not_yet_visible_future_returns_try_again(monkeypatch) -> None: - monkeypatch.setenv('TWINKLE_LONG_POLL_TIMEOUT', '0') - get_future = AsyncMock(return_value=None) - client = _make_client(get_future) - - response = client.post( - '/retrieve_future', - json={'request_id': 'unknown'}, - headers={ - 'Authorization': 'Bearer tenant-token', - 'X-Twinkle-Session-Id': 'session-1', - }, - ) - - assert response.status_code == 200 - assert response.json() == {'type': 'try_again'} - - -def test_initial_cross_replica_miss_is_long_polled(monkeypatch) -> None: - monkeypatch.setenv('TWINKLE_LONG_POLL_TIMEOUT', '1') - monkeypatch.setenv('TWINKLE_POLL_INTERVAL', '0') - get_future = AsyncMock(side_effect=[ - None, - {'status': 'completed', 'result': {'value': 9}}, - ]) - client = _make_client(get_future) - - response = client.post( - '/retrieve_future', - json={'request_id': 'req-replicated'}, - headers={'Authorization': 'Bearer tenant-token'}, - ) - - assert response.status_code == 200 - assert response.json() == {'value': 9} - assert get_future.await_count == 2 diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index e404b9706..dc49b82a8 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -48,6 +48,9 @@ async def _on_request_start(self, _request): def assert_resource_exists(self, _adapter_name): return None + def resolve_model_adapter_name(self, adapter_name): + return adapter_name + def forward_backward(self, *, inputs, adapter_name, **kwargs): self.model_calls.append((inputs, adapter_name, kwargs)) return {'loss': 1.0} diff --git a/tests/server/state/test_managers.py b/tests/server/state/test_managers.py index eea300b18..1b9016be1 100644 --- a/tests/server/state/test_managers.py +++ b/tests/server/state/test_managers.py @@ -432,6 +432,7 @@ async def test_store_status_queue_state(self, manager): assert result.queue_state == 'paused_rate_limit' assert result.queue_state_reason == 'Rate limit hit' + # ============================================================ # Cascade Cleanup Consistency (merged from test_cleanup_cascade_consistency) # ============================================================ diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py index 921f1f085..28b6f06a9 100644 --- a/tests/twinkle_agentic/test_async_rl_native_tq.py +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -7,7 +7,6 @@ import pytest -from cookbook.rl.sync_barrier_multi_lora_grpo import SyncBarrierMultiLoraGRPO from twinkle import DeviceMesh from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams from twinkle.infra import _dispatch_args @@ -251,21 +250,6 @@ def test_dynamic_micro_batch_planner_honors_per_rank_sample_and_token_limits(): assert padded_tokens <= 18 -def test_sync_training_batch_preserves_position_ids(): - rows = [{ - 'input_ids': [1, 2], - 'labels': [-100, 2], - 'attention_mask': [1, 1], - 'position_ids': [0, 1], - 'logprobs': [-.1], - }] - - batch = SyncBarrierMultiLoraGRPO._training_batch(rows, rewards=[1.], advantages=[0.]) - - assert 'position_ids' in batch.keys() - assert batch['position_ids'][0] == [0, 1] - - class PolicyProvider: def __init__(self, policies): From ea4e9b9cf7fe8d48a082f7aa5d78062de23691ab Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 18 Aug 2026 14:14:39 +0800 Subject: [PATCH 13/20] style: fix async worker file ending --- cookbook/client/async_rl/README.md | 3 - src/twinkle/infra/_ray/resource_manager.py | 3 +- src/twinkle/loss/grpo.py | 3 + src/twinkle/metric/__init__.py | 2 +- src/twinkle/metric/dpo.py | 1 - src/twinkle/metric/reporting.py | 83 +++------ src/twinkle/model/micro_batch.py | 38 +---- src/twinkle/model/multi_lora.py | 10 +- .../model/transformers/transformers.py | 47 ++--- src/twinkle/reward/dapo_math.py | 7 +- src/twinkle/server/data_plane/app.py | 3 +- src/twinkle/server/data_plane/handlers.py | 22 +-- src/twinkle/server/data_plane/proxy.py | 14 +- src/twinkle/server/data_plane/store.py | 5 +- src/twinkle/server/model/twinkle_handlers.py | 161 ++++-------------- src/twinkle/server/model/utils.py | 92 ++++++++++ src/twinkle/server/sampler/app.py | 4 +- src/twinkle/server/sampler/tinker_handlers.py | 3 +- .../server/sampler/twinkle_handlers.py | 46 ++--- src/twinkle/server/utils/task_queue/mixin.py | 1 - src/twinkle/server/utils/task_queue/worker.py | 1 - src/twinkle/utils/rl_tensor_utils.py | 23 +-- .../async_rl/context_manager.py | 4 +- src/twinkle_agentic/async_rl/data_plane.py | 15 +- src/twinkle_agentic/async_rl/metrics.py | 30 ++-- src/twinkle_agentic/async_rl/native_tq.py | 3 +- src/twinkle_agentic/async_rl/pipeline.py | 66 ++++--- src/twinkle_agentic/async_rl/types.py | 1 + src/twinkle_agentic/async_rl/utils.py | 31 ++-- .../async_rl/vllm_sampler_tq.py | 113 +++++------- src/twinkle_agentic/async_rl/workers.py | 52 +++--- src/twinkle_client/async_rl/workers.py | 1 - src/twinkle_client/rollout/multi_turn.py | 42 +---- tests/loss/test_grpo_gkd.py | 4 - tests/model/test_micro_batch.py | 37 +--- .../test_multi_lora_target_parameters.py | 18 -- .../server/model/test_twinkle_async_inputs.py | 30 +--- .../server/sampler/test_twinkle_async_rows.py | 6 +- .../test_async_rl_native_tq.py | 10 +- .../test_vllm_sampler_tq_generation.py | 4 +- .../test_client_multi_turn_rollout.py | 26 --- 41 files changed, 386 insertions(+), 679 deletions(-) create mode 100644 src/twinkle/server/model/utils.py diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md index 80a3a19a3..2e1fcb2f0 100644 --- a/cookbook/client/async_rl/README.md +++ b/cookbook/client/async_rl/README.md @@ -49,9 +49,6 @@ remain server-side. The Trainer passes one or more `DataRef` values to `forward_backward_from_data_plane()` and releases them in a local `finally` block. `asample()` remains the materialized-response convenience API. -- `ClientMultiTurnRollout.arun()` keeps tool calls and Reward computation in - the client and accepts an explicit `adapter_uri` policy snapshot. - `_RolloutPartition` is a private client record, not a server resource or SDK API. The local FIFO limits live DataLoader batches before rollout, ready prompt groups immediately use the Model primitives above, and the client calls diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index e54e5bd32..2d01f84af 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -163,8 +163,7 @@ def get_visible_devices(): probe_options['num_gpus'] = nproc_per_node else: probe_options['resources'] = {device_type: nproc_per_node} - visible_device_futures.append( - get_visible_devices.options(**probe_options).remote()) + visible_device_futures.append(get_visible_devices.options(**probe_options).remote()) self.visible_devices = ray.get(visible_device_futures) visible_devices = [] diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 2568dcdd7..900a91c09 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -32,6 +32,7 @@ def __init__( beta: float = 0.0, entropy_coef: float = 0.0, ignore_index: int = -100, + **kwargs, ): self.epsilon = epsilon self.epsilon_high = epsilon_high if epsilon_high is not None else epsilon @@ -399,6 +400,7 @@ class CISPOLoss(GRPOLoss): Clamps the IS weight and uses policy gradient. """ + def micro_batch_scale(self, inputs, indices): token_counts = [] for model_input in inputs: @@ -442,6 +444,7 @@ class BNPOLoss(GRPOLoss): Normalizes by total completion tokens across batch. """ + def micro_batch_scale(self, inputs, indices): token_counts = [] for model_input in inputs: diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index a8b8db798..b80bee74a 100644 --- a/src/twinkle/metric/__init__.py +++ b/src/twinkle/metric/__init__.py @@ -7,7 +7,7 @@ from .embedding import EmbeddingMetric from .grpo import CISPOMetric, GRPOMetric, GSPOMetric, PPOMetric from .loss import LossMetric -from .reporting import MetricsReporter, create_metrics_reporter from .ppo import PPOValueMetric +from .reporting import MetricsReporter, create_metrics_reporter from .train_metric import TrainMetric from .types import MetricRecord diff --git a/src/twinkle/metric/dpo.py b/src/twinkle/metric/dpo.py index 3d16993e5..be7ea8c43 100644 --- a/src/twinkle/metric/dpo.py +++ b/src/twinkle/metric/dpo.py @@ -5,7 +5,6 @@ from twinkle.data_format import InputFeature, ModelOutput from twinkle.utils import pad_and_stack_tensors from twinkle.utils.rl_tensor_utils import align_per_token_values - from .base import Metric diff --git a/src/twinkle/metric/reporting.py b/src/twinkle/metric/reporting.py index ce7d20104..de7f9265e 100644 --- a/src/twinkle/metric/reporting.py +++ b/src/twinkle/metric/reporting.py @@ -100,10 +100,7 @@ def add(self, record: MetricRecord) -> None: if record.policy_version is not None: previous_version = self.context_policy_versions.get(record.context_key) self.context_policy_versions[record.context_key] = ( - record.policy_version - if previous_version is None - else max(previous_version, record.policy_version) - ) + record.policy_version if previous_version is None else max(previous_version, record.policy_version)) sample_count = _finite_number(record.values.get('sample_count')) if sample_count is not None and record.status == 'completed': if record.stage == 'rollout' and record.attributes.get('scope', 'group') == 'group': @@ -113,11 +110,7 @@ def add(self, record: MetricRecord) -> None: self.context_trained_samples[record.context_key] += int(sample_count) summarize_values = ( record.status == 'completed' - and ( - record.stage != 'rollout' - or record.attributes.get('scope', 'group') == 'group' - ) - ) + and (record.stage != 'rollout' or record.attributes.get('scope', 'group') == 'group')) if summarize_values: for name, value in record.values.items(): number = _finite_number(value) @@ -133,23 +126,14 @@ def as_dict(self, backend_health: Mapping[str, Any]) -> dict[str, Any]: wall_time = time.time() - self.started_at rollout_groups = sum(self.context_rollout_groups.values()) train_steps = sum(counts['train:completed'] for counts in self.context_counts.values()) - trained_partitions = sum( - counts['partition:completed'] - for counts in self.context_counts.values() - ) + trained_partitions = sum(counts['partition:completed'] for counts in self.context_counts.values()) terminal_partitions = _finite_number(self.result.get('trained_partitions')) if terminal_partitions is not None: trained_partitions = int(terminal_partitions) rollout_samples = sum(self.context_rollout_samples.values()) trained_samples = sum(self.context_trained_samples.values()) - dropped_records = sum( - int(item.get('dropped_records', 0)) - for item in backend_health.values() - ) - backend_write_latency_s = sum( - float(item.get('write_latency_s', 0.0)) - for item in backend_health.values() - ) + dropped_records = sum(int(item.get('dropped_records', 0)) for item in backend_health.values()) + backend_write_latency_s = sum(float(item.get('write_latency_s', 0.0)) for item in backend_health.values()) contexts = {} for context_key, counts in self.context_counts.items(): contexts[context_key] = { @@ -286,22 +270,12 @@ def _run(self) -> None: batch = [] break elapsed = time.monotonic() - last_write - should_write = bool(self._queue) and ( - self._closing - or self._flush_requested - or len(self._queue) >= self.batch_size - or elapsed >= self.flush_interval_s - ) + should_write = bool(self._queue) and (self._closing or self._flush_requested or len(self._queue) + >= self.batch_size or elapsed >= self.flush_interval_s) if should_write: - batch = [ - self._queue.popleft() - for _ in range(min(len(self._queue), self.batch_size)) - ] + batch = [self._queue.popleft() for _ in range(min(len(self._queue), self.batch_size))] break - wait_s = ( - max(0.0, self.flush_interval_s - elapsed) - if self._queue else self.flush_interval_s - ) + wait_s = (max(0.0, self.flush_interval_s - elapsed) if self._queue else self.flush_interval_s) self._condition.wait(wait_s) if not batch: break @@ -358,10 +332,7 @@ def __init__(self, path: str | Path, **kwargs: Any): super().__init__('jsonl', **kwargs) def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: - self._stream.writelines( - json.dumps(payload, ensure_ascii=True, default=str) + '\n' - for payload in batch - ) + self._stream.writelines(json.dumps(payload, ensure_ascii=True, default=str) + '\n' for payload in batch) self._stream.flush() def _close_sink(self) -> None: @@ -392,10 +363,7 @@ def __init__( def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: for payload in batch: - prefix = ( - f'context/{_safe_name(payload["context_key"])}' - if payload.get('context_key') else 'global' - ) + prefix = (f'context/{_safe_name(payload["context_key"])}' if payload.get('context_key') else 'global') stage = _safe_name(payload['stage']) values = {} for name, value in payload['values'].items(): @@ -474,10 +442,7 @@ def close(self, timeout_s: float | None = None) -> None: self._write_summary() def health(self) -> dict[str, Any]: - backend_health = { - backend.name: backend.health() - for backend in self._backends - } + backend_health = {backend.name: backend.health() for backend in self._backends} for name, error in self._initial_errors.items(): backend_health[name] = { 'enabled': False, @@ -487,10 +452,7 @@ def health(self) -> dict[str, Any]: } return { 'record_count': self._sequence, - 'dropped_records': sum( - int(item.get('dropped_records', 0)) - for item in backend_health.values() - ), + 'dropped_records': sum(int(item.get('dropped_records', 0)) for item in backend_health.values()), 'backends': backend_health, } @@ -567,15 +529,16 @@ def create_metrics_reporter(config: Mapping[str, Any] | None, *, run_id: str) -> logger.warning('JSONL metrics backend could not start: %s', exc) if bool(swanlab_config.get('enabled', False)) and swanlab_config.get('mode') != 'disabled': try: - backends.append(_SwanLabBackend( - project=str(swanlab_config.get('project', 'twinkle-rl')), - experiment_name=str(swanlab_config.get('name', run_id)), - log_dir=swanlab_config.get('log_dir', 'outputs/swanlab'), - mode=str(swanlab_config.get('mode', 'local')), - queue_capacity=queue_capacity, - batch_size=int(swanlab_config.get('batch_size', 16)), - flush_interval_s=float(swanlab_config.get('flush_interval_s', 1.0)), - )) + backends.append( + _SwanLabBackend( + project=str(swanlab_config.get('project', 'twinkle-rl')), + experiment_name=str(swanlab_config.get('name', run_id)), + log_dir=swanlab_config.get('log_dir', 'outputs/swanlab'), + mode=str(swanlab_config.get('mode', 'local')), + queue_capacity=queue_capacity, + batch_size=int(swanlab_config.get('batch_size', 16)), + flush_interval_s=float(swanlab_config.get('flush_interval_s', 1.0)), + )) except Exception as exc: backend_errors['swanlab'] = exc logger.warning('SwanLab metrics backend could not start: %s', exc) diff --git a/src/twinkle/model/micro_batch.py b/src/twinkle/model/micro_batch.py index 561013745..f0b47e050 100644 --- a/src/twinkle/model/micro_batch.py +++ b/src/twinkle/model/micro_batch.py @@ -20,12 +20,11 @@ def __post_init__(self): raise ValueError(f'micro_batch_size must be positive, got {self.micro_batch_size}') if self.packing_algorithm not in ('ffd', 'kk'): raise ValueError(f'packing_algorithm must be ffd or kk, got {self.packing_algorithm!r}') - if self.dynamic_batching and ( - self.max_tokens_per_micro_batch is None or self.max_tokens_per_micro_batch <= 0): + if self.dynamic_batching and (self.max_tokens_per_micro_batch is None or self.max_tokens_per_micro_batch <= 0): raise ValueError('max_tokens_per_micro_batch must be positive when dynamic_batching=true') @classmethod - def from_kwargs(cls, kwargs: dict[str, Any]) -> 'MicroBatchConfig | None': + def from_kwargs(cls, kwargs: dict[str, Any]) -> MicroBatchConfig | None: option_names = ( 'micro_batch_size', 'dynamic_batching', @@ -56,8 +55,7 @@ def _batch_cost(group: list[int], lengths: list[int], padding_free: bool) -> int return sum(values) if padding_free else max(values) * len(values) -def _fits(group: list[int], index: int, lengths: list[int], config: MicroBatchConfig, - padding_free: bool) -> bool: +def _fits(group: list[int], index: int, lengths: list[int], config: MicroBatchConfig, padding_free: bool) -> bool: if len(group) >= config.micro_batch_size: return False candidate = [*group, index] @@ -97,11 +95,11 @@ def add(self, index: int, value: int) -> None: self.items.append(index) self.total += value - def merge(self, other: '_KKSet') -> None: + def merge(self, other: _KKSet) -> None: self.items.extend(other.items) self.total += other.total - def __lt__(self, other: '_KKSet') -> bool: + def __lt__(self, other: _KKSet) -> bool: return (self.total, len(self.items), self.items) < (other.total, len(other.items), other.items) @@ -118,12 +116,12 @@ def __init__(self, items: list[tuple[int, int]], group_count: int): def spread(self) -> int: return self.sets[0].total - self.sets[-1].total - def merge(self, other: '_KKState') -> None: + def merge(self, other: _KKState) -> None: for index in range(len(self.sets)): self.sets[index].merge(other.sets[-1 - index]) self.sets.sort(reverse=True) - def __lt__(self, other: '_KKState') -> bool: + def __lt__(self, other: _KKState) -> bool: return self.spread > other.spread @@ -146,8 +144,7 @@ def _kk_allocate(lengths: list[int], config: MicroBatchConfig, padding_free: boo while group_count <= len(lengths): groups = _kk_partition(lengths, group_count) if all( - len(group) <= config.micro_batch_size - and _batch_cost(group, lengths, padding_free) <= capacity + len(group) <= config.micro_batch_size and _batch_cost(group, lengths, padding_free) <= capacity for group in groups): return groups group_count += 1 @@ -179,8 +176,7 @@ def plan_micro_batches( capacity = config.max_tokens_per_micro_batch oversized = [length for length in lengths if length > capacity] if oversized: - raise ValueError( - f'sequence length {max(oversized)} exceeds max_tokens_per_micro_batch={capacity}') + raise ValueError(f'sequence length {max(oversized)} exceeds max_tokens_per_micro_batch={capacity}') if config.packing_algorithm == 'ffd': return _ffd_allocate(lengths, config, padding_free, min_micro_batches) return _kk_allocate(lengths, config, padding_free, min_micro_batches) @@ -194,19 +190,3 @@ def select_batch(value: Any, indices: list[int], batch_size: int) -> Any: if hasattr(value, 'shape') and len(value.shape) > 0 and value.shape[0] == batch_size: return value[indices] return value - - -def collect_micro_batch_outputs(outputs: list[dict[str, Any]], device_mesh: Any) -> dict[str, Any]: - from twinkle.infra.collectors import collect_tensor_dict - - result = collect_tensor_dict(outputs, device_mesh) - if len(outputs) <= 1 or 'micro_batch_count' not in outputs[0]: - return result - collected = [output for index, output in enumerate(outputs) if index in device_mesh.get_collect_ranks()] - result['micro_batch_count'] = collected[0]['micro_batch_count'] - result['micro_batch_samples_mean'] = ( - sum(output['micro_batch_samples_mean'] for output in collected) / len(collected)) - result['micro_batch_tokens_mean'] = ( - sum(output['micro_batch_tokens_mean'] for output in collected) / len(collected)) - result['micro_batch_tokens_max'] = max(output['micro_batch_tokens_max'] for output in collected) - return result diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index b382ac094..ff7766102 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -707,8 +707,7 @@ def _load_weights(_module): _load_weights(_module) else: _load_weights(self.module) - if getattr(_lora.tenant_config, 'target_parameters', None): - self.target_parameter_manager.set_state_dict(tenant_adapter_name, state_dict) + self.target_parameter_manager.set_state_dict(tenant_adapter_name, state_dict) def get_state_dict(self, tenant_adapter_name): state_dict = {} @@ -733,9 +732,7 @@ def _get_weights(_module): state_dict.update(_get_weights(_module)) else: state_dict = _get_weights(self.module) - target_state_dict = {} - if getattr(_lora.tenant_config, 'target_parameters', None): - target_state_dict = self.target_parameter_manager.get_state_dict(tenant_adapter_name) + target_state_dict = self.target_parameter_manager.get_state_dict(tenant_adapter_name) overlap = state_dict.keys() & target_state_dict.keys() if overlap: raise ValueError(f'Duplicate LoRA state keys: {sorted(overlap)[:5]}') @@ -851,9 +848,6 @@ def _get_parameters(_module): return trainable_param_names def get_target_parameter_trainable_parameters(self, tenant_adapter_name): - lora = self.find_lora_by_tenant(tenant_adapter_name) - if not getattr(lora.tenant_config, 'target_parameters', None): - return {} return { name: parameter for name, parameter in self.target_parameter_manager.named_slot_parameters(tenant_adapter_name) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index c9a8ec3d0..d9fe73598 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -33,8 +33,7 @@ from twinkle.loss import CrossEntropyLoss, Loss from twinkle.metric import Accuracy, LossMetric, Metric, TrainMetric from twinkle.model.base import TwinkleModel -from twinkle.model.micro_batch import (MicroBatchConfig, collect_micro_batch_outputs, plan_micro_batches, - select_batch) +from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches, select_batch from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus from twinkle.model.transformers.moe import apply_expert_parallel from twinkle.model.transformers.strategy import AccelerateStrategy, NativeFSDPStrategy @@ -738,8 +737,7 @@ def backward(self, **kwargs): sync_gradients = kwargs.pop('sync_gradients', None) should_sync = ( optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')) - if sync_gradients is None else bool(sync_gradients) - ) + if sync_gradients is None else bool(sync_gradients)) import contextlib no_sync_ctx = contextlib.nullcontext() @@ -795,23 +793,18 @@ def _build_micro_batch_plan(self, inputs, config, optimizer_config): states = [None] * dist.get_world_size(dp_group) dist.all_gather_object(states, local_state, group=dp_group) errors = [ - f'rank {rank}: {state["error"]}' - for rank, state in enumerate(states) - if state['error'] is not None + f'rank {rank}: {state["error"]}' for rank, state in enumerate(states) if state['error'] is not None ] if errors: - raise RuntimeError( - 'micro-batch planning failed on one or more model DP ranks: ' - + '; '.join(errors)) + raise RuntimeError('micro-batch planning failed on one or more model DP ranks: ' + '; '.join(errors)) counts = [state['micro_batch_count'] for state in states] if all(count == len(plan) for count in counts): return plan min_micro_batches = max(counts) if any(min_micro_batches > state['input_count'] for state in states): - raise ValueError( - 'model DP ranks cannot execute the same number of non-empty micro-batches; ' - 'make the input batch divisible by the model data-parallel size') + raise ValueError('model DP ranks cannot execute the same number of non-empty micro-batches; ' + 'make the input batch divisible by the model data-parallel size') def _forward_backward_micro_batch( self, @@ -831,10 +824,8 @@ def _forward_backward_micro_batch( previous_normalizer = optimizer_config.train_status.num_tokens loss = self.calculate_loss(**kwargs) normalizer_delta = optimizer_config.train_status.num_tokens - previous_normalizer - optimizer_config.train_status.loss_value = ( - optimizer_config.train_status.loss_value * loss_scale) - optimizer_config.train_status.num_tokens = ( - previous_normalizer + normalizer_delta * loss_scale) + optimizer_config.train_status.loss_value = (optimizer_config.train_status.loss_value * loss_scale) + optimizer_config.train_status.num_tokens = (previous_normalizer + normalizer_delta * loss_scale) outputs['loss'] = loss * loss_scale self.backward( sync_gradients=sync_gradients, @@ -864,17 +855,11 @@ def _forward_backward_micro_batches( inputs = optimizer_config.template.batch_encode(inputs) local_batch_size = len(inputs) - processor: InputProcessor = optimizer_config.processor plan = self._build_micro_batch_plan(inputs, config, optimizer_config) outputs = {} - micro_batch_samples = [] - micro_batch_tokens = [] loss_instance = optimizer_config.loss_instance for micro_batch_index, indices in enumerate(plan): - micro_kwargs = { - key: select_batch(value, indices, local_batch_size) - for key, value in kwargs.items() - } + micro_kwargs = {key: select_batch(value, indices, local_batch_size) for key, value in kwargs.items()} micro_loss_scale = loss_scale * loss_instance.micro_batch_scale(inputs, indices) is_last_micro_batch = micro_batch_index == len(plan) - 1 outputs = self._forward_backward_micro_batch( @@ -885,21 +870,9 @@ def _forward_backward_micro_batches( increment_step=is_last_micro_batch, **micro_kwargs, ) - lengths = [ - int(inputs[index]['input_ids'].shape[-1]) - if hasattr(inputs[index]['input_ids'], 'shape') else len(inputs[index]['input_ids']) - for index in indices - ] - micro_batch_samples.append(len(indices)) - micro_batch_tokens.append( - sum(lengths) if processor.padding_free else max(lengths) * len(lengths)) - outputs['micro_batch_count'] = len(plan) - outputs['micro_batch_samples_mean'] = sum(micro_batch_samples) / len(plan) - outputs['micro_batch_tokens_mean'] = sum(micro_batch_tokens) / len(plan) - outputs['micro_batch_tokens_max'] = max(micro_batch_tokens) return outputs - @remote_function(dispatch='slice_dp', collect=collect_micro_batch_outputs) + @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): """Do forward, calculate loss, and backward. diff --git a/src/twinkle/reward/dapo_math.py b/src/twinkle/reward/dapo_math.py index a9374fd51..59ecedcfd 100644 --- a/src/twinkle/reward/dapo_math.py +++ b/src/twinkle/reward/dapo_math.py @@ -11,7 +11,6 @@ from twinkle.reward.base import Reward from twinkle.reward.math_reward import MathReward - _ANSWER_LINE = re.compile(r'^\s*Answer\s*:\s*(.+?)\s*$', re.IGNORECASE | re.MULTILINE) @@ -102,10 +101,8 @@ def components(self, trajectories: list[Trajectory]) -> tuple[list[float], list[ def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: accuracy_rewards, overlong_rewards = self.components(trajectories) - return [ - (1.0 if accuracy else -1.0) + overlong - for accuracy, overlong in zip(accuracy_rewards, overlong_rewards) - ] + return [(1.0 if accuracy else -1.0) + overlong + for accuracy, overlong in zip(accuracy_rewards, overlong_rewards)] def metric_payload( self, diff --git a/src/twinkle/server/data_plane/app.py b/src/twinkle/server/data_plane/app.py index f37feb705..271127c4c 100644 --- a/src/twinkle/server/data_plane/app.py +++ b/src/twinkle/server/data_plane/app.py @@ -1,9 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations -from typing import Any - from fastapi import FastAPI +from typing import Any from twinkle.server.deployment import bind_deployment, build_deployment_app from .store import TQDataRefStore diff --git a/src/twinkle/server/data_plane/handlers.py b/src/twinkle/server/data_plane/handlers.py index ddef7284d..aaf214ed7 100644 --- a/src/twinkle/server/data_plane/handlers.py +++ b/src/twinkle/server/data_plane/handlers.py @@ -2,9 +2,8 @@ from __future__ import annotations from collections.abc import Callable -from typing import TYPE_CHECKING - from fastapi import Depends, FastAPI +from typing import TYPE_CHECKING import twinkle_client.types as types @@ -12,11 +11,10 @@ from .app import DataPlaneManagement -def register_data_plane_routes(app: FastAPI, self_fn: Callable[[], 'DataPlaneManagement']) -> None: +def register_data_plane_routes(app: FastAPI, self_fn: Callable[[], DataPlaneManagement]) -> None: @app.post('/twinkle/put', response_model=types.DataRef) - async def put(body: types.DataPutRequest, - self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: + async def put(body: types.DataPutRequest, self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: return await self.store.put( body.rows, kind=body.kind, @@ -24,21 +22,16 @@ async def put(body: types.DataPutRequest, ) @app.post('/twinkle/get', response_model=types.DataRowsResponse) - async def get(body: types.DataGetRequest, - self: DataPlaneManagement = Depends(self_fn)) -> types.DataRowsResponse: + async def get(body: types.DataGetRequest, self: DataPlaneManagement = Depends(self_fn)) -> types.DataRowsResponse: rows = await self.store.get( body.ref, fields=body.fields, ) - tags = ( - await self.store.get_tags(body.ref) - if body.include_tags else [] - ) + tags = (await self.store.get_tags(body.ref) if body.include_tags else []) return types.DataRowsResponse(rows=rows, tags=tags) @app.post('/twinkle/append', response_model=types.DataRef) - async def append(body: types.DataAppendRequest, - self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: + async def append(body: types.DataAppendRequest, self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: return await self.store.append( body.ref, body.rows, @@ -46,7 +39,6 @@ async def append(body: types.DataAppendRequest, ) @app.post('/twinkle/release') - async def release(body: types.DataReleaseRequest, - self: DataPlaneManagement = Depends(self_fn)) -> dict[str, str]: + async def release(body: types.DataReleaseRequest, self: DataPlaneManagement = Depends(self_fn)) -> dict[str, str]: await self.store.release(body.ref) return {'status': 'ok'} diff --git a/src/twinkle/server/data_plane/proxy.py b/src/twinkle/server/data_plane/proxy.py index e9b148aab..0aa2ae7f2 100644 --- a/src/twinkle/server/data_plane/proxy.py +++ b/src/twinkle/server/data_plane/proxy.py @@ -2,9 +2,8 @@ """Internal HTTP adapter used by Model and Sampler component deployments.""" from __future__ import annotations -from typing import Any - import httpx +from typing import Any from twinkle_client.http.headers import build_routing_headers from twinkle_client.types.component import DataRef @@ -30,7 +29,10 @@ async def get( raise RuntimeError('data_plane_url is required when a component request uses input_ref') response = await self.client.post( f'{self.base_url}/twinkle/get', - json={'ref': ref.model_dump(), 'fields': fields}, + json={ + 'ref': ref.model_dump(), + 'fields': fields + }, headers=build_routing_headers(f'data-ref-{ref.ref_id}'), ) response.raise_for_status() @@ -47,7 +49,11 @@ async def put( raise RuntimeError('data_plane_url is required to store component output') response = await self.client.post( f'{self.base_url}/twinkle/put', - json={'rows': rows, 'kind': kind, 'tags': tags}, + json={ + 'rows': rows, + 'kind': kind, + 'tags': tags + }, headers=build_routing_headers(f'data-put-{kind}'), ) response.raise_for_status() diff --git a/src/twinkle/server/data_plane/store.py b/src/twinkle/server/data_plane/store.py index 8ecd4e916..81fe7fa60 100644 --- a/src/twinkle/server/data_plane/store.py +++ b/src/twinkle/server/data_plane/store.py @@ -34,10 +34,7 @@ def _rows_from_tensordict(data: Any, size: int) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] fields = list(data.keys()) for index in range(size): - rows.append({ - field: json_safe(data[field][index]) - for field in fields - }) + rows.append({field: json_safe(data[field][index]) for field in fields}) return rows diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index fc0b03fdb..582aa5c00 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import torch import traceback from collections.abc import Callable from fastapi import Depends, FastAPI, HTTPException, Request @@ -24,9 +25,10 @@ from twinkle.server.checkpoint import (_resolve_client_save_dir, create_checkpoint_manager, create_training_run_manager, validate_user_path) from twinkle.server.exceptions import FullModeBusyError +from twinkle.server.model.utils import (data_plane_request_shape, merge_forward_kwargs, resolve_data_plane_model_inputs, + select_output_rows) from twinkle.server.utils.validation import get_session_id_from_request from twinkle.utils.logger import get_logger -from twinkle_client.common.json_utils import json_safe from twinkle_client.common.serialize import deserialize_object logger = get_logger() @@ -56,114 +58,6 @@ def _get_twinkle_adapter_name(request: Request, adapter_name: str | None) -> str return owner_id + '-' + adapter_name -def _model_result_rows(result: Any, batch_size: int) -> list[dict[str, Any]]: - """Keep per-sample model outputs at TQ sample granularity.""" - if isinstance(result, list) and len(result) == batch_size and all(isinstance(item, dict) for item in result): - return result - if isinstance(result, dict): - batched = { - name - for name, value in result.items() - if isinstance(value, list) and len(value) == batch_size - } - if batched: - return [{ - name: value[index] if name in batched else value - for name, value in result.items() - } for index in range(batch_size)] - return [{'result': result}] - - -def _value_at_path(value: Any, path: str) -> Any: - for part in path.split('.'): - if not isinstance(value, dict) or part not in value: - raise KeyError(f'data field {path!r} does not exist') - value = value[part] - return value - - -def _set_at_path(target: dict[str, Any], path: str, value: Any) -> None: - parts = path.split('.') - current = target - for part in parts[:-1]: - nested = current.setdefault(part, {}) - if not isinstance(nested, dict): - raise ValueError(f'cannot bind nested model argument {path!r}') - current = nested - current[parts[-1]] = value - - -async def _resolve_data_plane_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: - """Resolve DataPlane references before entering the model backend.""" - selected_fields = None - if body.input_field is not None: - selected_fields = list(dict.fromkeys([ - body.input_field, - *(source.split('.', 1)[0] for source in body.kwarg_fields.values()), - ])) - batches = await asyncio.gather(*( - data_plane.get(ref, fields=selected_fields) - for ref in body.input_refs - )) - rows = [row for batch in batches for row in batch] - if body.input_field is None: - kwarg_roots = {source.split('.', 1)[0] for source in body.kwarg_fields.values()} - inputs = [ - {key: value for key, value in row.items() if key not in kwarg_roots} - for row in rows - ] - else: - inputs = [_value_at_path(row, body.input_field) for row in rows] - - field_kwargs: dict[str, Any] = {} - for target_path, source_path in body.kwarg_fields.items(): - # Keep the transport contract JSON-native. The processor/loss/metric - # that understands this field owns any tensor conversion. - field_value = [ - _value_at_path(row, source_path) for row in rows - ] - _set_at_path( - field_kwargs, - target_path, - field_value, - ) - return inputs, field_kwargs - - -def _merge_forward_kwargs(explicit: dict[str, Any], bound: dict[str, Any]) -> dict[str, Any]: - collisions = set(explicit).intersection(bound) - if collisions: - names = ', '.join(sorted(collisions)) - raise ValueError(f'explicit model kwargs conflict with kwarg_fields: {names}') - return {**explicit, **bound} - - -def _request_shape(body: Any) -> tuple[int, int]: - input_refs = getattr(body, 'input_refs', None) - if input_refs is not None: - return ( - sum(ref.num_tokens for ref in input_refs), - sum(ref.size for ref in input_refs), - ) - inputs = body.inputs if isinstance(body.inputs, list) else [body.inputs] - return ( - sum(len(item.get('input_ids', [])) if isinstance(item, dict) else 0 for item in inputs), - len(inputs), - ) - - -def _select_output_rows( - result: Any, - *, - batch_size: int, - output_fields: dict[str, str], -) -> list[dict[str, Any]]: - rows = _model_result_rows(json_safe(result), batch_size) - if len(rows) != batch_size: - raise ValueError(f'model returned {len(rows)} rows for an output_ref of size {batch_size}') - return [{target: _value_at_path(row, source) for source, target in output_fields.items()} for row in rows] - - def _register_twinkle_routes(app: FastAPI, self_fn: Callable[[], ModelManagement]) -> None: """Register all /twinkle/* routes on the given FastAPI app. @@ -218,7 +112,9 @@ async def _task(): inputs=inputs, adapter_name=self.resolve_model_adapter_name(adapter_name), **extra_kwargs) return {'result': ret} - input_tokens, batch_size = _request_shape(body) + inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] + input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) + batch_size = len(inputs_list) return await run_task( self.schedule_task_and_wait( _task, @@ -241,8 +137,8 @@ async def forward_from_data_plane( async def _task(): self.assert_resource_exists(adapter_name) - raw_inputs, field_kwargs = await _resolve_data_plane_model_inputs(body, self.data_plane) - kwargs = _merge_forward_kwargs(body.model_extra or {}, field_kwargs) + raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) + kwargs = merge_forward_kwargs(body.model_extra or {}, field_kwargs) ret = self.model.forward( inputs=_parse_inputs(raw_inputs), adapter_name=adapter_name, @@ -250,7 +146,7 @@ async def _task(): ) return {'result': ret} - input_tokens, batch_size = _request_shape(body) + input_tokens, batch_size = data_plane_request_shape(body) return await run_task( self.schedule_task_and_wait( _task, @@ -301,15 +197,14 @@ async def _task(): inputs=inputs, adapter_name=self.resolve_model_adapter_name(adapter_name), **extra_kwargs) return {'result': ret} - input_tokens, batch_size = _request_shape(body) + inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] + input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) return await run_task( self.schedule_task_and_wait( _task, model_id=adapter_name, token=token, input_tokens=input_tokens, - batch_size=batch_size, - data_world_size=self.data_world_size, task_type='forward_only', )) @@ -324,12 +219,12 @@ async def forward_only_from_data_plane( async def _task(): self.assert_resource_exists(adapter_name) - raw_inputs, field_kwargs = await _resolve_data_plane_model_inputs(body, self.data_plane) + raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) inputs = _parse_inputs(raw_inputs) - kwargs = _merge_forward_kwargs(body.model_extra or {}, field_kwargs) + kwargs = merge_forward_kwargs(body.model_extra or {}, field_kwargs) ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **kwargs) if body.output_ref is not None: - rows = _select_output_rows( + rows = select_output_rows( ret, batch_size=len(inputs), output_fields=body.output_fields, @@ -338,7 +233,7 @@ async def _task(): return {'result': output_ref.model_dump()} return {'result': ret} - input_tokens, batch_size = _request_shape(body) + input_tokens, batch_size = data_plane_request_shape(body) return await run_task( self.schedule_task_and_wait( _task, @@ -389,15 +284,28 @@ async def forward_backward( token = await self._on_request_start(request) adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + def first_element(data): + while isinstance(data, list): + if len(data) == 0: + return None + data = data[0] + return data + async def _task(): self.assert_resource_exists(adapter_name) extra_kwargs = body.model_extra or {} all_inputs = _parse_inputs(body.inputs) + for inputs in all_inputs: + for key in inputs: + if isinstance(inputs[key], list) and isinstance(first_element(inputs[key]), (int, float)): + inputs[key] = torch.tensor(inputs[key]) ret = self.model.forward_backward( inputs=all_inputs, adapter_name=self.resolve_model_adapter_name(adapter_name), **extra_kwargs) return {'result': ret} - input_tokens, batch_size = _request_shape(body) + inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] + input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) + batch_size = len(inputs_list) return await run_task( self.schedule_task_and_wait( _task, @@ -420,8 +328,8 @@ async def forward_backward_from_data_plane( async def _task(): self.assert_resource_exists(adapter_name) - raw_inputs, field_kwargs = await _resolve_data_plane_model_inputs(body, self.data_plane) - kwargs = _merge_forward_kwargs(body.model_extra or {}, field_kwargs) + raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) + kwargs = merge_forward_kwargs(body.model_extra or {}, field_kwargs) ret = self.model.forward_backward( inputs=_parse_inputs(raw_inputs), adapter_name=adapter_name, @@ -429,7 +337,7 @@ async def _task(): ) return {'result': ret} - input_tokens, batch_size = _request_shape(body) + input_tokens, batch_size = data_plane_request_shape(body) return await run_task( self.schedule_task_and_wait( _task, @@ -692,10 +600,7 @@ async def _task(): async_upload=False, ) - future_ref = await self.schedule_background_task( - _task, - task_type='upload_to_hub', - ) + future_ref = await self.schedule_background_task(_task, task_type='upload_to_hub') request_id = future_ref.get('request_id') if request_id is None: raise HTTPException(status_code=500, detail=f'Upload task scheduling failed: {future_ref}') diff --git a/src/twinkle/server/model/utils.py b/src/twinkle/server/model/utils.py new file mode 100644 index 000000000..aed4b3c20 --- /dev/null +++ b/src/twinkle/server/model/utils.py @@ -0,0 +1,92 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Utilities shared by DataPlane-backed model routes.""" +from __future__ import annotations + +import asyncio +from typing import Any + +from twinkle_client.common.json_utils import json_safe + + +def model_result_rows(result: Any, batch_size: int) -> list[dict[str, Any]]: + """Keep per-sample model outputs at DataPlane row granularity.""" + if isinstance(result, list) and len(result) == batch_size and all(isinstance(item, dict) for item in result): + return result + if isinstance(result, dict): + batched = {name for name, value in result.items() if isinstance(value, list) and len(value) == batch_size} + if batched: + return [{ + name: value[index] if name in batched else value + for name, value in result.items() + } for index in range(batch_size)] + return [{'result': result}] + + +def value_at_path(value: Any, path: str) -> Any: + for part in path.split('.'): + if not isinstance(value, dict) or part not in value: + raise KeyError(f'data field {path!r} does not exist') + value = value[part] + return value + + +def set_at_path(target: dict[str, Any], path: str, value: Any) -> None: + parts = path.split('.') + current = target + for part in parts[:-1]: + nested = current.setdefault(part, {}) + if not isinstance(nested, dict): + raise ValueError(f'cannot bind nested model argument {path!r}') + current = nested + current[parts[-1]] = value + + +async def resolve_data_plane_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: + """Resolve DataPlane references into model inputs and bound keyword arguments.""" + selected_fields = None + if body.input_field is not None: + selected_fields = list( + dict.fromkeys([ + body.input_field, + *(source.split('.', 1)[0] for source in body.kwarg_fields.values()), + ])) + batches = await asyncio.gather(*(data_plane.get(ref, fields=selected_fields) for ref in body.input_refs)) + rows = [row for batch in batches for row in batch] + if body.input_field is None: + kwarg_roots = {source.split('.', 1)[0] for source in body.kwarg_fields.values()} + inputs = [{key: value for key, value in row.items() if key not in kwarg_roots} for row in rows] + else: + inputs = [value_at_path(row, body.input_field) for row in rows] + + field_kwargs: dict[str, Any] = {} + for target_path, source_path in body.kwarg_fields.items(): + field_value = [value_at_path(row, source_path) for row in rows] + set_at_path(field_kwargs, target_path, field_value) + return inputs, field_kwargs + + +def merge_forward_kwargs(explicit: dict[str, Any], bound: dict[str, Any]) -> dict[str, Any]: + collisions = set(explicit).intersection(bound) + if collisions: + names = ', '.join(sorted(collisions)) + raise ValueError(f'explicit model kwargs conflict with kwarg_fields: {names}') + return {**explicit, **bound} + + +def data_plane_request_shape(body: Any) -> tuple[int, int]: + return ( + sum(ref.num_tokens for ref in body.input_refs), + sum(ref.size for ref in body.input_refs), + ) + + +def select_output_rows( + result: Any, + *, + batch_size: int, + output_fields: dict[str, str], +) -> list[dict[str, Any]]: + rows = model_result_rows(json_safe(result), batch_size) + if len(rows) != batch_size: + raise ValueError(f'model returned {len(rows)} rows for an output_ref of size {batch_size}') + return [{target: value_at_path(row, source) for source, target in output_fields.items()} for row in rows] diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index df39d1033..8941a40bf 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -8,10 +8,9 @@ from __future__ import annotations import asyncio -from typing import Any - from fastapi import FastAPI, Request from ray import serve +from typing import Any from twinkle import DeviceGroup from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app, init_twinkle_runtime @@ -21,7 +20,6 @@ from twinkle.server.utils.task_queue import TaskQueueConfig, TaskQueueMixin from twinkle.server.utils.validation import get_token_from_request from twinkle.utils.logger import get_logger - from .tinker_handlers import _register_tinker_sampler_routes from .twinkle_handlers import _register_twinkle_sampler_routes diff --git a/src/twinkle/server/sampler/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py index bb3a0429f..de43b4b93 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -98,8 +98,7 @@ async def _do_sample(): else: self.sampler.load_full_weights_from_path(adapter_uri) - sample_fn = getattr(self.sampler, 'sample_sync', self.sampler.sample) - responses = sample_fn( + responses = self.sampler.sample( inputs=[prompt_inputs] * body.num_samples, sampling_params=sampling_params, adapter_path=lora_path, diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index 7bcf36380..708f827f8 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -27,8 +27,8 @@ from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.server.utils.validation import get_session_id_from_request -from twinkle_client.common.json_utils import json_safe from twinkle.utils.logger import get_logger +from twinkle_client.common.json_utils import json_safe logger = get_logger() @@ -59,7 +59,7 @@ def _get_twinkle_sampler_adapter_name(request: Request, adapter_name: str | None return owner_id + '-' + adapter_name -def _sample_models_to_rows( +def _build_rollout_rows_and_tags( sample_models: list[types.SampleResponseModel], *, group_ids: list[str] | None, @@ -69,16 +69,14 @@ def _sample_models_to_rows( """Flatten sampler output to one TQ row per generated sequence.""" resolved_group_ids = group_ids or [uuid.uuid4().hex for _ in sample_models] if len(resolved_group_ids) != len(sample_models): - raise ValueError( - f'group_ids contains {len(resolved_group_ids)} values for ' - f'{len(sample_models)} sampler inputs') + raise ValueError(f'group_ids contains {len(resolved_group_ids)} values for ' + f'{len(sample_models)} sampler inputs') rows = [] tags = [] for prompt_index, (response, group_id) in enumerate(zip(sample_models, resolved_group_ids)): for generation_idx, sequence in enumerate(response.sequences): sampled_logprobs = [ - 0.0 if not position else float(position[0][1]) - for position in (sequence.logprobs or []) + 0.0 if not position else float(position[0][1]) for position in (sequence.logprobs or []) ] rows.append({ 'train_input': sequence.new_input_feature, @@ -101,7 +99,7 @@ def _sample_models_to_rows( return rows, tags -def _responses_to_models(responses) -> list[types.SampleResponseModel]: +def _to_sample_response_models(responses) -> list[types.SampleResponseModel]: """Convert internal sampler responses to the HTTP response schema.""" sample_models = [] for response in responses: @@ -111,12 +109,9 @@ def _responses_to_models(responses) -> list[types.SampleResponseModel]: tokens=list(sequence.tokens), logprobs=list(sequence.logprobs) if sequence.logprobs is not None else None, decoded=sequence.decoded, - new_input_feature=( - _serialize_input_feature(sequence.new_input_feature) - if sequence.new_input_feature is not None else None - ), - ) - for sequence in response.sequences + new_input_feature=(_serialize_input_feature(sequence.new_input_feature) + if sequence.new_input_feature is not None else None), + ) for sequence in response.sequences ] sample_models.append( types.SampleResponseModel( @@ -142,8 +137,7 @@ async def _await_generation( poll_interval = 0.01 while True: try: - states = _submission_states( - await asyncio.to_thread(sampler.get_generation_status, submission_id)) + states = _submission_states(await asyncio.to_thread(sampler.get_generation_status, submission_id)) except Exception as error: # A pending read-only actor call can be cancelled by Ray while # the generation submitted just above remains alive. Treating @@ -260,14 +254,13 @@ async def _task(): params = SamplingParams.from_dict(body.sampling_params) # Sample - sample_fn = getattr(self.sampler, 'sample_sync', self.sampler.sample) - responses = sample_fn( + responses = self.sampler.sample( inputs, params, adapter_name=full_adapter_name, adapter_path=adapter_path, ) - return types.SampleResponseModelList(samples=_responses_to_models(responses)) + return types.SampleResponseModelList(samples=_to_sample_response_models(responses)) # Calculate metrics for queue scheduling inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] @@ -300,10 +293,7 @@ async def sample_to_data_plane( checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') _, adapter_path = checkpoint_manager.parse_adapter_uri(body.adapter_uri) - inputs = ( - await self.data_plane.get(body.input_ref) - if body.input_ref is not None else body.inputs - ) + inputs = (await self.data_plane.get(body.input_ref) if body.input_ref is not None else body.inputs) if isinstance(inputs, list) and inputs: first = inputs[0] if isinstance(first, dict) and 'input_ids' in first: @@ -331,10 +321,8 @@ async def _admit(): inline_inputs = body.inputs if isinstance(body.inputs, list) else [body.inputs] input_tokens = ( - body.input_ref.num_tokens - if body.input_ref is not None else - sum(len(item.get('input_ids', [])) for item in inline_inputs if isinstance(item, dict)) - ) + body.input_ref.num_tokens if body.input_ref is not None else sum( + len(item.get('input_ids', [])) for item in inline_inputs if isinstance(item, dict))) await run_task( self.schedule_task_and_wait( _admit, @@ -345,8 +333,8 @@ async def _admit(): )) responses = await _await_generation(self.sampler, submission_id) - rows, tags = _sample_models_to_rows( - _responses_to_models(responses), + rows, tags = _build_rollout_rows_and_tags( + _to_sample_response_models(responses), group_ids=body.group_ids, policy_version=body.policy_version, adapter_uri=body.adapter_uri, diff --git a/src/twinkle/server/utils/task_queue/mixin.py b/src/twinkle/server/utils/task_queue/mixin.py index 84da48cb3..dcdf805ce 100644 --- a/src/twinkle/server/utils/task_queue/mixin.py +++ b/src/twinkle/server/utils/task_queue/mixin.py @@ -17,7 +17,6 @@ from twinkle.server.telemetry.middleware import get_task_metrics from twinkle.server.utils.task_errors import task_error_payload from twinkle.utils.logger import get_logger - from .config import TaskQueueConfig from .rate_limiter import RateLimiter from .types import QueuedTask, QueueState, TaskStatus diff --git a/src/twinkle/server/utils/task_queue/worker.py b/src/twinkle/server/utils/task_queue/worker.py index 58d06ce05..fdbb36d16 100644 --- a/src/twinkle/server/utils/task_queue/worker.py +++ b/src/twinkle/server/utils/task_queue/worker.py @@ -18,7 +18,6 @@ from twinkle.server.telemetry.tracing import traced_operation from twinkle.server.utils.task_errors import task_error_payload from twinkle.utils.logger import get_logger - from .config import TaskQueueConfig from .types import QueuedTask, QueueState, TaskStatus diff --git a/src/twinkle/utils/rl_tensor_utils.py b/src/twinkle/utils/rl_tensor_utils.py index 5e262e5cc..5f5eba36f 100644 --- a/src/twinkle/utils/rl_tensor_utils.py +++ b/src/twinkle/utils/rl_tensor_utils.py @@ -48,8 +48,7 @@ def align_per_token_values( if row.ndim == 2 and row.shape[0] == 1: row = row.squeeze(0) if row.ndim != 1: - raise ValueError( - f'{name}[{index}] must be one-dimensional, got shape {tuple(row.shape)}') + raise ValueError(f'{name}[{index}] must be one-dimensional, got shape {tuple(row.shape)}') rows.append(row) row_lengths = [row.numel() for row in rows] tensor = torch.nn.utils.rnn.pad_sequence( @@ -66,30 +65,26 @@ def align_per_token_values( target_batch_size, target_seq_len = target_shape batch_size, seq_len = tensor.shape if batch_size != target_batch_size: - raise ValueError( - f'{name} batch size ({batch_size}) does not match target batch size ' - f'({target_batch_size})') + raise ValueError(f'{name} batch size ({batch_size}) does not match target batch size ' + f'({target_batch_size})') mask = None if valid_mask is not None: mask = torch.as_tensor(valid_mask, dtype=torch.bool) if tuple(mask.shape) != target_shape: - raise ValueError( - f'valid_mask shape {tuple(mask.shape)} does not match target shape ' - f'{target_shape}') + raise ValueError(f'valid_mask shape {tuple(mask.shape)} does not match target shape ' + f'{target_shape}') if row_lengths is not None: for index, row_len in enumerate(row_lengths): if row_len >= target_seq_len: continue if mask is None or bool(mask[index, row_len:].any().item()): - raise ValueError( - f'{name}[{index}] has {row_len} tokens but target sequence length ' - f'is {target_seq_len}') + raise ValueError(f'{name}[{index}] has {row_len} tokens but target sequence length ' + f'is {target_seq_len}') if seq_len < target_seq_len: if mask is None or bool(mask[:, seq_len:].any().item()): - raise ValueError( - f'{name} seq_len ({seq_len}) is smaller than target seq_len ' - f'({target_seq_len})') + raise ValueError(f'{name} seq_len ({seq_len}) is smaller than target seq_len ' + f'({target_seq_len})') tensor = torch.nn.functional.pad( tensor, (0, target_seq_len - seq_len), diff --git a/src/twinkle_agentic/async_rl/context_manager.py b/src/twinkle_agentic/async_rl/context_manager.py index 5933b6dc6..f97e91323 100644 --- a/src/twinkle_agentic/async_rl/context_manager.py +++ b/src/twinkle_agentic/async_rl/context_manager.py @@ -130,9 +130,7 @@ def list_context_snapshots(self) -> list[dict[str, object]]: def context_adapter_paths(self, context: LoraContext | str) -> list[str]: return [ - policy.adapter_path - for policy in self._state(context).policy_history - if policy.adapter_path is not None + policy.adapter_path for policy in self._state(context).policy_history if policy.adapter_path is not None ] def get_rollout_policy(self, context: LoraContext | str) -> RolloutPolicy: diff --git a/src/twinkle_agentic/async_rl/data_plane.py b/src/twinkle_agentic/async_rl/data_plane.py index 0bc38f6ca..641947366 100644 --- a/src/twinkle_agentic/async_rl/data_plane.py +++ b/src/twinkle_agentic/async_rl/data_plane.py @@ -177,13 +177,14 @@ async def prepare_rollout_partition( 'generation_idx': generation_idx, 'rollout_status': 'PENDING', } for generation_idx in range(admission.num_generations)]) - groups.append(PromptGroup( - context=admission.context, - partition=admission, - group_id=group_id, - prompt=dict(prompt), - batch_meta=batch_meta, - )) + groups.append( + PromptGroup( + context=admission.context, + partition=admission, + group_id=group_id, + prompt=dict(prompt), + batch_meta=batch_meta, + )) return PreparedPartition(admission, tuple(groups), sampling_params) async def complete_rollout_group( diff --git a/src/twinkle_agentic/async_rl/metrics.py b/src/twinkle_agentic/async_rl/metrics.py index abb4696d1..d78ac06c9 100644 --- a/src/twinkle_agentic/async_rl/metrics.py +++ b/src/twinkle_agentic/async_rl/metrics.py @@ -15,11 +15,11 @@ def _p95(values: list[float]) -> float: def rollout_metrics( - *, - rewards: Mapping[str, Sequence[float]] | None = None, - completion_lengths: Sequence[int] = (), - stop_reasons: Sequence[str | None] = (), - rollout_latency_s: float | None = None, + *, + rewards: Mapping[str, Sequence[float]] | None = None, + completion_lengths: Sequence[int] = (), + stop_reasons: Sequence[str | None] = (), + rollout_latency_s: float | None = None, ) -> dict[str, float | int]: """Summarize one RL rollout collection without retaining state.""" metrics: dict[str, float | int] = {} @@ -28,12 +28,10 @@ def rollout_metrics( raise ValueError(f'reward metric lengths must match, got {reward_counts}') sample_count = len(completion_lengths) or (reward_counts[0] if reward_counts else 0) if completion_lengths and reward_counts and any(count != sample_count for count in reward_counts): - raise ValueError( - f'reward and completion metric lengths must match: {reward_counts} != {sample_count}') + raise ValueError(f'reward and completion metric lengths must match: {reward_counts} != {sample_count}') if stop_reasons and len(stop_reasons) != len(completion_lengths): - raise ValueError( - 'stop reason and completion metric lengths must match: ' - f'{len(stop_reasons)} != {len(completion_lengths)}') + raise ValueError('stop reason and completion metric lengths must match: ' + f'{len(stop_reasons)} != {len(completion_lengths)}') if sample_count: metrics['sample_count'] = sample_count if completion_lengths: @@ -99,9 +97,8 @@ def advantage_signal_metrics( if len(rewards) != len(advantages): raise ValueError(f'rewards and advantages must have equal length: {len(rewards)} != {len(advantages)}') if len(rewards) == 0 or len(rewards) % num_generations: - raise ValueError( - f'advantage metrics require complete groups: sample_count={len(rewards)}, ' - f'num_generations={num_generations}') + raise ValueError(f'advantage metrics require complete groups: sample_count={len(rewards)}, ' + f'num_generations={num_generations}') reward_values = [float(value) for value in rewards] advantage_values = [float(value) for value in advantages] @@ -111,8 +108,7 @@ def advantage_signal_metrics( group_rewards = reward_values[start:start + num_generations] group_advantages = advantage_values[start:start + num_generations] reward_mean = sum(group_rewards) / num_generations - group_reward_stds.append( - math.sqrt(sum((value - reward_mean)**2 for value in group_rewards) / num_generations)) + group_reward_stds.append(math.sqrt(sum((value - reward_mean)**2 for value in group_rewards) / num_generations)) if max(abs(value) for value in group_advantages) <= zero_tolerance: zero_advantage_groups += 1 @@ -124,6 +120,6 @@ def advantage_signal_metrics( 'zero_advantage_group_ratio': zero_advantage_groups / group_count, 'positive_advantage_ratio': sum(value > zero_tolerance for value in advantage_values) / len(advantage_values), 'advantage_mean': advantage_mean, - 'advantage_std': math.sqrt( - sum((value - advantage_mean)**2 for value in advantage_values) / len(advantage_values)), + 'advantage_std': + math.sqrt(sum((value - advantage_mean)**2 for value in advantage_values) / len(advantage_values)), } diff --git a/src/twinkle_agentic/async_rl/native_tq.py b/src/twinkle_agentic/async_rl/native_tq.py index e42f68739..9bc137e67 100644 --- a/src/twinkle_agentic/async_rl/native_tq.py +++ b/src/twinkle_agentic/async_rl/native_tq.py @@ -8,9 +8,8 @@ from __future__ import annotations -from typing import Any, Protocol, Sequence - from transfer_queue import GRPOGroupNSampler +from typing import Any, Protocol, Sequence class AsyncTQClient(Protocol): diff --git a/src/twinkle_agentic/async_rl/pipeline.py b/src/twinkle_agentic/async_rl/pipeline.py index 8acd2cea0..1bce5fbf4 100644 --- a/src/twinkle_agentic/async_rl/pipeline.py +++ b/src/twinkle_agentic/async_rl/pipeline.py @@ -15,11 +15,10 @@ from .data_plane import TQDataPlane from .scheduler import ContextSchedulePolicy, SchedulerConfig from .types import LoraContext, PartitionAdmission -from .utils import (TrainBatchConfig, build_native_fsdp_model_kwargs, - configure_lora_lr_scheduler, resolve_context_learning_rate, - resolve_context_lora_target_modules, resolve_context_loss_config, - resolve_model_attention_implementation, sampler_data_parallel_size, - resolve_sequence_parallel_size, validate_context_batch_config) +from .utils import (TrainBatchConfig, build_native_fsdp_model_kwargs, configure_lora_lr_scheduler, + resolve_context_learning_rate, resolve_context_lora_target_modules, resolve_context_loss_config, + resolve_model_attention_implementation, resolve_sequence_parallel_size, sampler_data_parallel_size, + validate_context_batch_config) from .workers import AdvantageWorker, RolloutWorker, TrainerWorker @@ -210,10 +209,8 @@ def from_config( mini_batch_size=int(train['mini_batch_size']), micro_batch_size=int(train['micro_batch_size']), dynamic_batching=bool(train.get('dynamic_batching', False)), - max_tokens_per_micro_batch=( - int(train['max_tokens_per_micro_batch']) - if train.get('max_tokens_per_micro_batch') is not None else None - ), + max_tokens_per_micro_batch=(int(train['max_tokens_per_micro_batch']) + if train.get('max_tokens_per_micro_batch') is not None else None), packing_algorithm=str(train.get('packing_algorithm', 'ffd')), ) validate_context_batch_config( @@ -264,9 +261,12 @@ def from_config( raise ValueError('evaluation.batch_size and evaluation.interval must be positive') eval_sampling = dict(global_evaluation.get('sampling_params') or {}) evaluation_config[context.key] = { - 'interval': eval_interval, - 'dataset_name': eval_dataset.get('name', eval_dataset['dataset_id']), - 'prompt_batches': partial( + 'interval': + eval_interval, + 'dataset_name': + eval_dataset.get('name', eval_dataset['dataset_id']), + 'prompt_batches': + partial( _prompt_batches, eval_dataset, model_id=runtime['model_id'], @@ -275,7 +275,8 @@ def from_config( enable_thinking=enable_thinking, full_batches_only=False, ), - 'sampling_params': SamplingParams( + 'sampling_params': + SamplingParams( max_tokens=int(eval_sampling.get('max_tokens', rollout['max_tokens'])), temperature=float(eval_sampling.get('temperature', 0.0)), top_p=float(eval_sampling.get('top_p', 1.0)), @@ -328,13 +329,9 @@ def from_config( context_manager=manager, rollout_max_retries=int(runtime.get('rollout_max_retries', 2)), rollout_retry_delay_s=float(runtime.get('rollout_retry_delay_s', 0.5)), - rollout_output_dir=( - rollout_output_config.get('output_dir') - if bool(rollout_output_config.get('enabled', False)) else None - ), - rollout_output_include_token_ids=bool( - rollout_output_config.get('include_token_ids', False) - ), + rollout_output_dir=(rollout_output_config.get('output_dir') if bool( + rollout_output_config.get('enabled', False)) else None), + rollout_output_include_token_ids=bool(rollout_output_config.get('include_token_ids', False)), ) sampler.set_template( template_cls, @@ -380,7 +377,8 @@ def from_config( train_batch_configs=train_batch_configs, save_adapter=partial(_save_adapter, model, runtime['output_dir']), mini_batch_sizes={ - key: config.mini_batch_size for key, config in train_batch_configs.items() + key: config.mini_batch_size + for key, config in train_batch_configs.items() }, scheduler=_scheduler(raw_config['scheduler']['train']), keep_adapter_versions=runtime['keep_adapter_versions'], @@ -406,8 +404,7 @@ def from_config( sampler=sampler, metrics=metrics, config=AsyncMultiLoraGRPOConfig( - metrics_drain_interval_s=float(metrics_config.get('drain_interval_s', 1.0)), - ), + metrics_drain_interval_s=float(metrics_config.get('drain_interval_s', 1.0)), ), model=model, contexts=contexts, ) @@ -433,12 +430,13 @@ async def run_async(self) -> dict[str, Any]: await asyncio.sleep(self.config.metrics_drain_interval_s) except Exception as exc: if self.metrics is not None: - self.metrics.record(MetricRecord( - stage='run', - status='failed', - values={'wall_time_s': time.perf_counter() - started}, - attributes={'error': f'{type(exc).__name__}: {exc}'}, - )) + self.metrics.record( + MetricRecord( + stage='run', + status='failed', + values={'wall_time_s': time.perf_counter() - started}, + attributes={'error': f'{type(exc).__name__}: {exc}'}, + )) self.metrics.flush() raise finally: @@ -646,9 +644,8 @@ def _train_batch_with_config( old_logps = list(data['logprobs']) advantages = list(data['advantages']) if size != config.mini_batch_size: - raise ValueError( - f'train batch for {admission.context.key} has {size} samples; ' - f'expected mini_batch_size={config.mini_batch_size}') + raise ValueError(f'train batch for {admission.context.key} has {size} samples; ' + f'expected mini_batch_size={config.mini_batch_size}') if size % model_data_parallel_size: raise ValueError(f'train batch size {size} must be divisible by model DP size ' @@ -688,9 +685,8 @@ def _save_adapter(model: Any, output_dir: str, admission: PartitionAdmission) -> def _require_adapter_path(value: Any, *, operation: str) -> str: """Fail at the save boundary instead of publishing an invalid policy.""" if not isinstance(value, str) or not value: - raise TypeError( - f'{operation} must return a non-empty checkpoint path string, ' - f'got {type(value).__name__}: {value!r}') + raise TypeError(f'{operation} must return a non-empty checkpoint path string, ' + f'got {type(value).__name__}: {value!r}') return value diff --git a/src/twinkle_agentic/async_rl/types.py b/src/twinkle_agentic/async_rl/types.py index e31352f6c..69545858b 100644 --- a/src/twinkle_agentic/async_rl/types.py +++ b/src/twinkle_agentic/async_rl/types.py @@ -71,6 +71,7 @@ def partition_id(self) -> str: def num_samples(self) -> int: return self.partition.num_generations + @dataclass class PreparedPartition: """Partition prepared in TQ and ready for sampler submission.""" diff --git a/src/twinkle_agentic/async_rl/utils.py b/src/twinkle_agentic/async_rl/utils.py index 764c0a98e..1a6f98bdc 100644 --- a/src/twinkle_agentic/async_rl/utils.py +++ b/src/twinkle_agentic/async_rl/utils.py @@ -10,7 +10,6 @@ from typing import Any, Literal from twinkle.data_format import SampleResponse - from .types import RolloutOutput @@ -84,9 +83,8 @@ def resolve_model_attention_implementation( if implementation is not None: implementation = str(implementation) if padding_free and sequence_parallel_size > 1 and implementation != 'flash_attention_2': - raise ValueError( - 'model.attn_implementation must be flash_attention_2 when ' - 'model.padding_free=true and model.sequence_parallel_size>1') + raise ValueError('model.attn_implementation must be flash_attention_2 when ' + 'model.padding_free=true and model.sequence_parallel_size>1') return implementation @@ -123,13 +121,11 @@ def validate_context_batch_config( f'({sampler_dp}), got {rollout_groups}') partition_samples = rollout_groups * num_generations if partition_samples % train.mini_batch_size: - raise ValueError( - f'partition for {context_key} has {partition_samples} samples and must be divisible by ' - f'train.mini_batch_size={train.mini_batch_size}') + raise ValueError(f'partition for {context_key} has {partition_samples} samples and must be divisible by ' + f'train.mini_batch_size={train.mini_batch_size}') if train.mini_batch_size % num_generations: - raise ValueError( - f'train.mini_batch_size for {context_key} must preserve complete prompt groups: ' - f'{train.mini_batch_size} % {num_generations} != 0') + raise ValueError(f'train.mini_batch_size for {context_key} must preserve complete prompt groups: ' + f'{train.mini_batch_size} % {num_generations} != 0') if train.mini_batch_size % model_dp: raise ValueError(f'train.mini_batch_size for {context_key} must be divisible by ' f'model DP size {model_dp}') @@ -139,13 +135,11 @@ def validate_context_batch_config( f'({samples_per_rank}), got {train.micro_batch_size}') if train.dynamic_batching: if train.max_tokens_per_micro_batch is None or train.max_tokens_per_micro_batch <= 0: - raise ValueError( - f'train.max_tokens_per_micro_batch for {context_key} must be positive when ' - 'train.dynamic_batching=true') + raise ValueError(f'train.max_tokens_per_micro_batch for {context_key} must be positive when ' + 'train.dynamic_batching=true') if train.packing_algorithm not in ('ffd', 'kk'): - raise ValueError( - f'train.packing_algorithm for {context_key} must be ffd or kk, ' - f'got {train.packing_algorithm!r}') + raise ValueError(f'train.packing_algorithm for {context_key} must be ffd or kk, ' + f'got {train.packing_algorithm!r}') def configure_lora_lr_scheduler( @@ -195,9 +189,8 @@ def resolve_context_lora_target_modules( modules = list(target_modules) if all(isinstance(module, str) and module for module in modules): return modules - raise ValueError( - 'lora.target_modules must be a non-empty string or sequence of module names, ' - f'got {target_modules!r}') + raise ValueError('lora.target_modules must be a non-empty string or sequence of module names, ' + f'got {target_modules!r}') def resolve_context_loss_config( diff --git a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py index 33c7dd30d..c77a2e56e 100644 --- a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py +++ b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py @@ -18,7 +18,6 @@ from twinkle.hub import HubOperation from twinkle.metric import MetricBuffer, MetricRecord from twinkle.sampler.vllm_sampler import vLLMSampler - from .data_plane import TQDataPlane from .metrics import rollout_metrics from .types import LoraContext, PromptGroup, RolloutOutput, RolloutPolicy @@ -147,9 +146,7 @@ def __init__( self.rollout_max_retries = int(rollout_max_retries) self.rollout_retry_delay_s = float(rollout_retry_delay_s) self.rollout_output_dir = ( - Path(rollout_output_dir).expanduser().resolve() - if rollout_output_dir is not None else None - ) + Path(rollout_output_dir).expanduser().resolve() if rollout_output_dir is not None else None) self.rollout_output_include_token_ids = bool(rollout_output_include_token_ids) if self.rollout_max_retries < 0: raise ValueError(f'rollout_max_retries must be non-negative, got {self.rollout_max_retries}') @@ -173,16 +170,17 @@ def _record_metrics( attributes: dict[str, Any] | None = None, policy_version: int | None = None, ) -> None: - self.metric_buffer.record(MetricRecord( - stage='rollout', - values=dict(values), - context_key=group.context.key, - partition_id=group.partition_id, - partition_index=group.partition.step, - policy_version=policy_version, - status=status, - attributes=dict(attributes or {}), - )) + self.metric_buffer.record( + MetricRecord( + stage='rollout', + values=dict(values), + context_key=group.context.key, + partition_id=group.partition_id, + partition_index=group.partition.step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) @remote_function(dispatch='all', collect='flatten', lazy_collect=False) def drain_metric_records(self) -> list[MetricRecord]: @@ -208,14 +206,11 @@ def unload_lora_paths(self, adapter_paths: list[str]) -> None: # Unloading is keyed by the normalized path stored in VLLMEngine's # request cache. The checkpoint itself may already have been pruned, # so unlike loading this must not require the path to still exist. - local_paths = [ - os.path.abspath(os.path.expanduser(str(path))) - for path in adapter_paths - ] + local_paths = [os.path.abspath(os.path.expanduser(str(path))) for path in adapter_paths] self._submit_in_loop(self.engine.unload_lora_paths(local_paths)).result() @remote_function(dispatch='slice_dp', collect='none', lazy_collect=False) - def sample( + def submit_prompt_groups( self, groups: list[PromptGroup], sampling_params: SamplingParams, @@ -242,27 +237,6 @@ def sample( 'submitted_samples': sum(group.num_samples for group in groups), } - @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) - def sample_sync( - self, - inputs: Any, - sampling_params: SamplingParams | dict[str, Any] | None = None, - adapter_name: str = '', - adapter_path: str | None = None, - *, - return_encoded: bool = False, - use_base_model: bool = False, - ) -> list[SampleResponse]: - """Run the inherited blocking sampler API for synchronous CS calls.""" - return super().sample( - inputs, - sampling_params, - adapter_name=adapter_name, - adapter_path=adapter_path, - return_encoded=return_encoded, - use_base_model=use_base_model, - ) - @remote_function(dispatch=_dispatch_generation, collect='none', lazy_collect=False) def submit_generation( self, @@ -422,17 +396,14 @@ async def _generate_inputs( logger.warning(f'Failed to pre-load LoRA from {local_adapter_path}, ' 'sampling will proceed without LoRA') - return await asyncio.gather(*( - self._sample_single( - feat, - sampling_params, - lora_request=lora_request, - multi_modal_data=multi_modal_data, - logprobs_only=logprobs_only, - disable_lora=use_base_model, - ) - for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list) - )) + return await asyncio.gather(*(self._sample_single( + feat, + sampling_params, + lora_request=lora_request, + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + disable_lora=use_base_model, + ) for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list))) def _on_submission_done(self, submission_id: str): @@ -470,34 +441,39 @@ async def _sample_prompt_groups( group, {}, status='failed', - attributes={'scope': 'group', 'group_id': group.group_id, 'error': str(error)}, + attributes={ + 'scope': 'group', + 'group_id': group.group_id, + 'error': str(error) + }, ) raise RuntimeError(f'rollout failed for {group.group_id}: {error}') from error rollout_stats = [result for result in results if isinstance(result, _PromptGroupRolloutStats)] - metric_rows = [ - { - 'completion_length': completion_length, - 'stop_reason': stop_reason, - } - for stats in rollout_stats - for completion_length, stop_reason in zip(stats.completion_lengths, stats.stop_reasons) - ] + metric_rows = [{ + 'completion_length': completion_length, + 'stop_reason': stop_reason, + } for stats in rollout_stats + for completion_length, stop_reason in zip(stats.completion_lengths, stats.stop_reasons)] policy_versions = [version for stats in rollout_stats for version in stats.policy_versions] first_group = groups[0] dp_size = self.device_mesh.dp_world_size or 1 self._record_metrics( first_group, { - 'prompt_group_count': len(groups), + 'prompt_group_count': + len(groups), **rollout_metrics( completion_lengths=[row['completion_length'] for row in metric_rows], stop_reasons=[row['stop_reason'] for row in metric_rows], rollout_latency_s=time.perf_counter() - submitted_at, ), - 'policy_version_min': min(policy_versions), - 'policy_version_max': max(policy_versions), - 'sampler_dp_size': dp_size, + 'policy_version_min': + min(policy_versions), + 'policy_version_max': + max(policy_versions), + 'sampler_dp_size': + dp_size, }, attributes={'scope': 'partition' if dp_size == 1 else 'shard'}, policy_version=max(policy_versions), @@ -589,7 +565,10 @@ async def _run_prompt_group( max(policy_versions), **reward_metrics, }, - attributes={'scope': 'group', 'group_id': group.group_id}, + attributes={ + 'scope': 'group', + 'group_id': group.group_id + }, policy_version=max(policy_versions), ) return _PromptGroupRolloutStats( @@ -748,8 +727,8 @@ async def _generate_sample( else: if sequence.stop_reason not in {'abort', 'error'}: if not allow_partial_rollout or not partial_responses: - return _GeneratedSample( - response, (policy, ), attempt + 1, was_aborted, resumed_partial_output) + return _GeneratedSample(response, (policy, ), attempt + 1, was_aborted, + resumed_partial_output) partial_responses.append(response) partial_policies.append(policy) return _GeneratedSample( diff --git a/src/twinkle_agentic/async_rl/workers.py b/src/twinkle_agentic/async_rl/workers.py index 905b63533..984190c8f 100644 --- a/src/twinkle_agentic/async_rl/workers.py +++ b/src/twinkle_agentic/async_rl/workers.py @@ -67,17 +67,18 @@ def _record_metric( optimizer_step: int | None = None, policy_version: int | None = None, ) -> None: - self.metric_buffer.record(MetricRecord( - stage=stage, - values=dict(values or {}), - context_key=context.key if context is not None else None, - partition_id=admission.partition_id if admission is not None else partition_id, - partition_index=admission.step if admission is not None else None, - optimizer_step=optimizer_step, - policy_version=policy_version, - status=status, - attributes=dict(attributes or {}), - )) + self.metric_buffer.record( + MetricRecord( + stage=stage, + values=dict(values or {}), + context_key=context.key if context is not None else None, + partition_id=admission.partition_id if admission is not None else partition_id, + partition_index=admission.step if admission is not None else None, + optimizer_step=optimizer_step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) async def _run_service(self) -> None: try: @@ -237,7 +238,7 @@ async def _serve(self) -> None: config['sampling_params'], ) await asyncio.to_thread( - self.sampler.sample, + self.sampler.submit_prompt_groups, list(prepared.groups), prepared.sampling_params, self.allow_partial_rollout, @@ -360,10 +361,10 @@ def __init__(self, initial_adapter_paths: dict[str, str] | None = None, remove_adapter: Callable[[str], None] | None = None, evaluation_config: dict[str, dict[str, Any]] | None = None, - evaluate_batch: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any], - dict[str, Any]] | None = None, - evaluate_with_reward_fn: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any, - Any], dict[str, Any]] | None = None, + evaluate_batch: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any], dict[str, Any]] + | None = None, + evaluate_with_reward_fn: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any, Any], + dict[str, Any]] | None = None, evaluation_rewards: dict[str, Any] | None = None, persistent: bool = False, idle_delay_s: float = 0.05): @@ -498,8 +499,7 @@ async def _serve(self) -> None: raise RuntimeError(f'training failed for {admission.partition_id}: {exc}') from exc self.scheduler.on_success(candidate) metrics['sample_count'] = sample_count - metrics['reward'] = ( - sum(float(value) for value in batch.data['rewards']) / sample_count) + metrics['reward'] = (sum(float(value) for value in batch.data['rewards']) / sample_count) metrics['train_latency_s'] = time.perf_counter() - started metrics.update(training_policy_metrics(batch.sample_tags, policy.version)) context_key = admission.context.key @@ -534,7 +534,10 @@ async def _finish_partition(self, admission: PartitionAdmission) -> None: 'adapter_save_latency_s': adapter_save_latency_s, 'policy_publish_latency_s': policy_publish_latency_s, }, - attributes={'operation': 'publish', 'adapter_path': adapter_path}, + attributes={ + 'operation': 'publish', + 'adapter_path': adapter_path + }, optimizer_step=self._optimizer_steps[admission.context.key], policy_version=policy.version, ) @@ -648,7 +651,11 @@ async def _remove_adapter(self, context: LoraContext, path: str) -> None: values={ 'adapter_prune_latency_s': time.perf_counter() - started, }, - attributes={'operation': 'adapter_prune', 'adapter_path': path, 'error': str(exc)}, + attributes={ + 'operation': 'adapter_prune', + 'adapter_path': path, + 'error': str(exc) + }, ) return self._record_metric( @@ -657,7 +664,10 @@ async def _remove_adapter(self, context: LoraContext, path: str) -> None: values={ 'adapter_prune_latency_s': time.perf_counter() - started, }, - attributes={'operation': 'adapter_prune', 'adapter_path': path}, + attributes={ + 'operation': 'adapter_prune', + 'adapter_path': path + }, ) diff --git a/src/twinkle_client/async_rl/workers.py b/src/twinkle_client/async_rl/workers.py index b98a9fb1f..8d105bcf9 100644 --- a/src/twinkle_client/async_rl/workers.py +++ b/src/twinkle_client/async_rl/workers.py @@ -58,4 +58,3 @@ async def run(self) -> None: if not task.done(): task.cancel() await asyncio.gather(*tasks, return_exceptions=True) - diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index 54c6e450b..55c5800b7 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -4,7 +4,8 @@ This module hosts :class:`ClientMultiTurnRollout`, a hand-maintained multi-turn rollout orchestrator whose algorithmic structure mirrors ``twinkle_agentic.rollout.multi_turn.MultiTurnRollout`` but issues sampling over -HTTP via the client sampler instead of holding a Ray actor handle. +HTTP via ``twinkle_client.sampler.vLLMSampler.sample()`` instead of holding a +Ray actor handle. Design notes: * It deliberately does NOT subclass ``MultiTurnRollout``. That class is @@ -13,12 +14,9 @@ not match the HTTP-client semantics here. * The ``tool_manager`` type is reused directly from ``twinkle_agentic.tools.tool_manager.ToolManager`` (imported, not copied). - * ``arun`` uses the Sampler component's asynchronous future API; ``__call__`` - remains a synchronous compatibility wrapper. * Bridge-token stitching is reused from ``twinkle_agentic.rollout.bridge.extend_with_bridge``. """ -import asyncio import dataclasses from typing import Any, Dict, List, Optional @@ -35,7 +33,8 @@ class ClientMultiTurnRollout: Mirrors the per-trajectory state machine of ``twinkle_agentic.rollout.multi_turn.MultiTurnRollout`` but issues sampling - through the HTTP Sampler component rather than a Ray actor call. + via ``vLLMSampler.sample()`` (an HTTP call to ``/twinkle/sample``) rather + than a Ray actor call. """ def __init__( @@ -68,24 +67,11 @@ def __init__( f'got {self.sampling_params.num_samples}') def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - """Synchronous wrapper for :meth:`arun`. - - Async client orchestrators should call ``await rollout.arun(...)`` so - independent rollout pipelines can overlap with Model training. - """ - try: - asyncio.get_running_loop() - except RuntimeError: - return asyncio.run(self.arun(trajectories, **kwargs)) - raise RuntimeError('ClientMultiTurnRollout.__call__ cannot run inside an event loop; ' - 'use `await rollout.arun(...)` instead') - - async def arun(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - """Asynchronously run the batched multi-turn state machine over HTTP. + """Run the batched multi-turn rollout state machine over HTTP. Structurally mirrors ``MultiTurnRollout.__call__`` but issues each - round's sampling through ``vLLMSampler.asample()`` rather than a Ray - actor call. Every round makes a + round's sampling through ``vLLMSampler.sample()`` (an HTTP POST to + ``/twinkle/sample``) rather than a Ray actor call. Every round makes a SINGLE batched HTTP call for all currently-live trajectories so the sampler can run them in parallel; finished trajectories are parked and excluded from later batches. @@ -152,19 +138,7 @@ async def arun(self, trajectories: List[Trajectory], **kwargs) -> List[Trajector # upstream concern (retry/backoff) and failures are never # silently swallowed. batch_pifs = [pifs[i] for i in active] - sample_kwargs = {'sampling_params': sampling_params} - for name in ('adapter_name', 'adapter_uri'): - if name in kwargs: - sample_kwargs[name] = kwargs[name] - async_sample = getattr(self.sampler, 'asample', None) - if callable(async_sample): - resps = await async_sample(batch_pifs, **sample_kwargs) - else: - resps = await asyncio.to_thread( - self.sampler.sample, - batch_pifs, - **sample_kwargs, - ) + resps = self.sampler.sample(batch_pifs, sampling_params=sampling_params) pending_bridges: List[tuple] = [] # (global_idx, tool_messages) for local_idx, global_idx in enumerate(active): diff --git a/tests/loss/test_grpo_gkd.py b/tests/loss/test_grpo_gkd.py index fe6d49341..3d0f55120 100644 --- a/tests/loss/test_grpo_gkd.py +++ b/tests/loss/test_grpo_gkd.py @@ -101,10 +101,6 @@ def test_grpo_weights_sequences_equally(self): assert result['loss'].item() == pytest.approx(-2.0) - def test_grpo_does_not_accept_normalization(self): - with pytest.raises(TypeError, match='normalization'): - GRPOLoss(normalization='token_mean') - def test_grpo_entropy_coef(self): loss_fn = GRPOLoss(epsilon=0.2, entropy_coef=0.01) inputs, outputs, old_logps, _, advantages = _make_rl_batch() diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py index 8a07529f4..8d21c1dbb 100644 --- a/tests/model/test_micro_batch.py +++ b/tests/model/test_micro_batch.py @@ -5,7 +5,7 @@ from twinkle.loss import CrossEntropyLoss, GRPOLoss from twinkle.loss.base import Loss -from twinkle.model.micro_batch import MicroBatchConfig, collect_micro_batch_outputs, plan_micro_batches +from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches from twinkle.model.transformers.transformers import TransformersModel from twinkle.processor import InputProcessor from twinkle.utils.nccl_safe import safe_loss @@ -179,8 +179,6 @@ def backward(self, *, sync_gradients, **_kwargs): assert model.forward_batches == [[0, 1], [2, 3]] assert model.backward_calls == [(False, 1.0), (True, 1.0)] assert model.optimizer_group['adapter'].train_status.num_tokens == 1.0 - assert outputs['micro_batch_count'] == 2 - assert outputs['micro_batch_samples_mean'] == 2.0 def test_fixed_micro_batch_plan_can_match_a_larger_dp_micro_batch_count(): @@ -229,7 +227,6 @@ def all_gather(states, local_state, *, group): OptimizerConfig(), ) - def test_dp_micro_batch_planning_rejects_common_count_on_all_ranks(monkeypatch): from twinkle.model.transformers import transformers as module @@ -262,35 +259,3 @@ def all_gather(states, local_state, *, group): MicroBatchConfig(micro_batch_size=1), OptimizerConfig(), ) - - -def test_dp_collection_reduces_micro_batch_statistics(): - class Mesh: - @staticmethod - def get_collect_ranks(): - return [0, 1] - - result = collect_micro_batch_outputs( - [ - { - 'micro_batch_count': 3, - 'micro_batch_samples_mean': 2.0, - 'micro_batch_tokens_mean': 100.0, - 'micro_batch_tokens_max': 150, - }, - { - 'micro_batch_count': 3, - 'micro_batch_samples_mean': 3.0, - 'micro_batch_tokens_mean': 120.0, - 'micro_batch_tokens_max': 180, - }, - ], - Mesh(), - ) - - assert result == { - 'micro_batch_count': 3, - 'micro_batch_samples_mean': 2.5, - 'micro_batch_tokens_mean': 110.0, - 'micro_batch_tokens_max': 180, - } diff --git a/tests/model/test_multi_lora_target_parameters.py b/tests/model/test_multi_lora_target_parameters.py index 7edc4c6d8..d26b32147 100644 --- a/tests/model/test_multi_lora_target_parameters.py +++ b/tests/model/test_multi_lora_target_parameters.py @@ -85,24 +85,6 @@ def _make_target_cfg(r=2): ) -def test_standard_lora_has_no_target_parameter_trainable_parameters(): - from twinkle.model.multi_lora import LoraTenant, MultiLora - - config = LoraConfig(r=2, lora_alpha=4, target_modules=['linear']) - multi_lora = MultiLora(max_loras=1, max_r=4) - multi_lora.loras = [ - LoraTenant( - index=0, - adapter_name='lora_0', - config=config, - tenant_adapter_name='adapter_a', - tenant_config=config, - ) - ] - - assert multi_lora.get_target_parameter_trainable_parameters('adapter_a') == {} - - def test_target_parameter_multi_lora_updates_only_active_adapter(): from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index dc49b82a8..7a13b73c5 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -5,14 +5,12 @@ from starlette.requests import Request import twinkle_client.types as types -from twinkle.server.model.twinkle_handlers import ( - _model_result_rows, - _register_twinkle_routes, -) +from twinkle.server.model.twinkle_handlers import _register_twinkle_routes +from twinkle.server.model.utils import model_result_rows def test_model_result_rows_keeps_one_output_row_per_sample() -> None: - assert _model_result_rows( + assert model_result_rows( {'logps': [[-1.0], [-2.0]], 'loss': 0.25}, batch_size=2, ) == [ @@ -66,28 +64,6 @@ async def schedule_task_and_wait(self, task, **kwargs): return await task() -@pytest.mark.asyncio -async def test_forward_backward_inline_route_keeps_original_request_shape() -> None: - management = _SchedulingManagement() - app = FastAPI() - _register_twinkle_routes(app, lambda: management) - route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') - request = Request({'type': 'http', 'headers': []}) - request.state.session_id = 'session' - body = types.ForwardRequest( - adapter_name='adapter', - inputs=[{'input_ids': [1, 2, 3]}], - micro_batch_size=1, - ) - - await route.endpoint(request, body, management) - - inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' - assert [row['input_ids'] for row in inputs] == [[1, 2, 3]] - assert forwarded_kwargs == {'micro_batch_size': 1} - - @pytest.mark.asyncio async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() -> None: management = _SchedulingManagement() diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index 9a9e628f0..30b7016b2 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -8,7 +8,7 @@ from twinkle.data_format import SampledSequence, SampleResponse from twinkle.server.sampler.twinkle_handlers import ( _register_twinkle_sampler_routes, - _sample_models_to_rows, + _build_rollout_rows_and_tags, ) @@ -28,7 +28,7 @@ def _response(tokens: list[int]) -> types.SampleResponseModel: def test_async_sampler_flattens_generations_to_tagged_tq_rows() -> None: - rows, tags = _sample_models_to_rows( + rows, tags = _build_rollout_rows_and_tags( [_response([10, 11]), _response([20, 21])], group_ids=['group-a', 'group-b'], policy_version=7, @@ -51,7 +51,7 @@ def test_async_sampler_flattens_generations_to_tagged_tq_rows() -> None: def test_async_sampler_rejects_group_id_count_mismatch() -> None: with pytest.raises(ValueError, match='group_ids contains 1 values for 2'): - _sample_models_to_rows( + _build_rollout_rows_and_tags( [_response([10]), _response([20])], group_ids=['only-one'], policy_version=0, diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py index 28b6f06a9..92d55f110 100644 --- a/tests/twinkle_agentic/test_async_rl_native_tq.py +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -189,12 +189,7 @@ def __init__(self): def forward_backward(self, **kwargs): self.calls.append(kwargs) - return lambda: { - 'micro_batch_count': 4, - 'micro_batch_samples_mean': 1.0, - 'micro_batch_tokens_mean': 1.0, - 'micro_batch_tokens_max': 1, - } + return lambda: {} def clip_grad_and_step(self, **_kwargs): self.optimizer_steps += 1 @@ -229,7 +224,6 @@ def calculate_metric(self, **_kwargs): assert model.calls[0]['loss_scale'] == 1.0 assert model.optimizer_steps == 1 assert metrics['micro_batch_size_per_rank'] == 1 - assert 'micro_batch_count' not in metrics def test_dynamic_micro_batch_planner_honors_per_rank_sample_and_token_limits(): @@ -869,7 +863,7 @@ def __init__(self, loop): self.submitted = asyncio.Event() self.loop = loop - def sample(self, _groups, _sampling_params, _allow_partial_rollout): + def submit_prompt_groups(self, _groups, _sampling_params, _allow_partial_rollout): self.loop.call_soon_threadsafe(self.submitted.set) loaded_batches = [] diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index 7e5749261..f01c0a545 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -18,7 +18,7 @@ def _bare_sampler() -> VLLMSamplerTQ: def test_generation_dispatch_allows_one_prompt_with_multiple_dp_workers() -> None: assert VLLMSamplerTQ.submit_generation._dispatch is _dispatch_generation - assert VLLMSamplerTQ.sample._dispatch == 'slice_dp' + assert VLLMSamplerTQ.submit_prompt_groups._dispatch == 'slice_dp' shards = [ _dispatch_generation( 3, @@ -129,7 +129,7 @@ def test_native_prompt_group_sampling_requires_context_manager() -> None: sampler.context_manager = None with pytest.raises(RuntimeError, match='context_manager is required'): - sampler.sample([], SamplingParams(max_tokens=4)) + sampler.submit_prompt_groups([], SamplingParams(max_tokens=4)) def test_server_waiter_admits_later_submission_before_first_finishes() -> None: diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index 31af25757..1ff69f5f4 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -309,32 +309,6 @@ def _build_from_scripts(scripts_spec: List[Dict[str, Any]]): return trajectories, sampler, template -@pytest.mark.asyncio -async def test_async_rollout_passes_explicit_policy_to_sampler() -> None: - trajectories, sampler, template = _build_from_scripts([{ - 'num_tools': 0, - 'terminal': 'stop', - 'logprobs': True, - }]) - calls = [] - - async def asample(inputs, **kwargs): - calls.append(kwargs) - return sampler.sample(inputs, sampling_params=kwargs.get('sampling_params')) - - sampler.asample = asample - rollout = ClientMultiTurnRollout(sampler, template, max_turns=2) - outputs = await rollout.arun( - trajectories, - adapter_name='math-lora', - adapter_uri='twinkle://run/weights/policy-3', - ) - - assert len(outputs) == 1 - assert calls[0]['adapter_name'] == 'math-lora' - assert calls[0]['adapter_uri'] == 'twinkle://run/weights/policy-3' - - # ============================================================================= # Hypothesis strategies # ============================================================================= From 435e5cfd023c71696755987791dd97811fcd90f5 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 21 Aug 2026 18:09:17 +0800 Subject: [PATCH 14/20] wip --- .../rl/compare_single_lora_gsm8k_async.yaml | 128 +++++++++++++++++ .../rl/compare_single_lora_gsm8k_sync.yaml | 130 ++++++++++++++++++ 2 files changed, 258 insertions(+) create mode 100644 cookbook/rl/compare_single_lora_gsm8k_async.yaml create mode 100644 cookbook/rl/compare_single_lora_gsm8k_sync.yaml diff --git a/cookbook/rl/compare_single_lora_gsm8k_async.yaml b/cookbook/rl/compare_single_lora_gsm8k_async.yaml new file mode 100644 index 000000000..c8144a6d0 --- /dev/null +++ b/cookbook/rl/compare_single_lora_gsm8k_async.yaml @@ -0,0 +1,128 @@ +# Asynchronous comparison run. It is intentionally identical to the strict +# baseline except for max_staleness, run_id, and output paths. +runtime: + run_id: compare_single_lora_gsm8k_async + model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} + mode: ray + model_gpus: 1 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 1 + seed: 1 + max_staleness: 2 + max_steps: 125 + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/compare_single_lora_gsm8k_async + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/compare_single_lora_gsm8k_async_metrics.jsonl + summary_path: outputs/async_rl/compare_single_lora_gsm8k_async_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: true + +model: + strategy: native_fsdp + attn_implementation: flash_attention_2 + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 1 + padding_free: false + max_length: 2048 + +sampler: + max_model_len: 2048 + max_num_batched_tokens: 4096 + gpu_memory_utilization: 0.7 + max_num_seqs: 64 + enforce_eager: false + +rollout_output: + enabled: true + output_dir: ${runtime.output_dir}/rollouts + include_token_ids: false + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +evaluation: + enabled: true + interval: 10 + batch_size: 16 + sampling_params: + max_tokens: 1024 + temperature: 0.6 + top_p: 1.0 + +lora: + target_modules: all-linear + r: 16 + alpha: 16 + dropout: 0.0 + learning_rate: 1.7e-5 + +loss: + cls: GRPOLoss + epsilon: 0.2 + +lora_contexts: + - tenant_id: tenant_single + training_run_id: gsm8k_compare + adapter_name: gsm8k_compare_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + dataset: + dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID,ms://modelscope/gsm8k} + subset_name: main + split: train + data_num: 2000 + max_length: 1024 + processor: GSM8KProcessor + system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' + eval_dataset: + name: gsm8k/test + dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID,ms://modelscope/gsm8k} + subset_name: main + split: test + data_num: null + max_length: 1024 + processor: GSM8KProcessor + system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 1024 + temperature: 1.0 + top_p: 1.0 + train: + mini_batch_size: 64 + micro_batch_size: 64 + dynamic_batching: true + max_tokens_per_micro_batch: 4096 + packing_algorithm: ffd diff --git a/cookbook/rl/compare_single_lora_gsm8k_sync.yaml b/cookbook/rl/compare_single_lora_gsm8k_sync.yaml new file mode 100644 index 000000000..d77d86caf --- /dev/null +++ b/cookbook/rl/compare_single_lora_gsm8k_sync.yaml @@ -0,0 +1,130 @@ +# Strict on-policy baseline for comparison with compare_single_lora_gsm8k_async.yaml. +# max_staleness=0 permits only one live rollout partition. The workers may still +# pipeline groups inside that partition, but no rollout partition can cross a +# policy-version boundary. +runtime: + run_id: compare_single_lora_gsm8k_sync + model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} + mode: ray + model_gpus: 1 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 1 + seed: 1 + max_staleness: 0 + max_steps: 125 + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/compare_single_lora_gsm8k_sync + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/compare_single_lora_gsm8k_sync_metrics.jsonl + summary_path: outputs/async_rl/compare_single_lora_gsm8k_sync_summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Template + enable_thinking: true + +model: + strategy: native_fsdp + attn_implementation: flash_attention_2 + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 1 + padding_free: false + max_length: 2048 + +sampler: + max_model_len: 2048 + max_num_batched_tokens: 4096 + gpu_memory_utilization: 0.7 + max_num_seqs: 64 + enforce_eager: false + +rollout_output: + enabled: true + output_dir: ${runtime.output_dir}/rollouts + include_token_ids: false + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +evaluation: + enabled: true + interval: 10 + batch_size: 16 + sampling_params: + max_tokens: 1024 + temperature: 0.6 + top_p: 1.0 + +lora: + target_modules: all-linear + r: 16 + alpha: 16 + dropout: 0.0 + learning_rate: 1.7e-5 + +loss: + cls: GRPOLoss + epsilon: 0.2 + +lora_contexts: + - tenant_id: tenant_single + training_run_id: gsm8k_compare + adapter_name: gsm8k_compare_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + dataset: + dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID,ms://modelscope/gsm8k} + subset_name: main + split: train + data_num: 2000 + max_length: 1024 + processor: GSM8KProcessor + system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' + eval_dataset: + name: gsm8k/test + dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID,ms://modelscope/gsm8k} + subset_name: main + split: test + data_num: null + max_length: 1024 + processor: GSM8KProcessor + system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 1024 + temperature: 1.0 + top_p: 1.0 + train: + mini_batch_size: 64 + micro_batch_size: 64 + dynamic_batching: true + max_tokens_per_micro_batch: 4096 + packing_algorithm: ffd From 86a2e09cf869af5683dff1fb3a9a2c59b2bacda2 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 21 Aug 2026 19:35:03 +0800 Subject: [PATCH 15/20] wip --- src/twinkle_agentic/async_rl/utils.py | 2 -- .../test_async_rl_native_tq.py | 34 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/src/twinkle_agentic/async_rl/utils.py b/src/twinkle_agentic/async_rl/utils.py index 1a6f98bdc..2c8384aaa 100644 --- a/src/twinkle_agentic/async_rl/utils.py +++ b/src/twinkle_agentic/async_rl/utils.py @@ -37,8 +37,6 @@ def sample_responses_to_rollout_rows( for sequence in response.sequences: row = dict(source) row.update(sequence.new_input_feature or {}) - row.setdefault('group_id', source['group_id']) - row.setdefault('generation_idx', source['generation_idx']) row['logprobs'] = _extract_sampled_token_logps(sequence.logprobs) row['stop_reason'] = sequence.stop_reason row['completion_length'] = len(sequence.tokens) diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py index 92d55f110..e2e65db25 100644 --- a/tests/twinkle_agentic/test_async_rl_native_tq.py +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -34,6 +34,7 @@ resolve_context_loss_config, resolve_model_attention_implementation, resolve_sequence_parallel_size, + sample_responses_to_rollout_rows, sampler_data_parallel_size, validate_context_batch_config, ) @@ -300,6 +301,39 @@ def _sample_response(tokens, stop_reason, input_ids): ) +def test_evaluation_rows_do_not_require_rollout_group_metadata(): + prompt = {'input_ids': [1, 2], 'labels': [-100, -100]} + + rows = sample_responses_to_rollout_rows( + [prompt], + [_sample_response([3], 'stop', [1, 2, 3])], + policy_version=10, + ) + + assert len(rows) == 1 + assert 'group_id' not in rows[0] + assert 'generation_idx' not in rows[0] + assert rows[0]['rollout_policy_version'] == 10 + + +def test_training_rows_preserve_rollout_group_metadata(): + source = { + 'input_ids': [1, 2], + 'labels': [-100, -100], + 'group_id': 'partition/group_0', + 'generation_idx': 2, + } + + rows = sample_responses_to_rollout_rows( + [source], + [_sample_response([3], 'stop', [1, 2, 3])], + policy_version=4, + ) + + assert rows[0]['group_id'] == 'partition/group_0' + assert rows[0]['generation_idx'] == 2 + + def test_sampler_data_parallel_size_is_derived_from_gpu_and_tp_sizes(): assert sampler_data_parallel_size(8, 2) == 4 assert sampler_data_parallel_size(1, 1) == 1 From 984b34b8fa199d4e6a0647e57e251938dbaa658a Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 21 Aug 2026 23:53:45 +0800 Subject: [PATCH 16/20] wip --- .../rl/compare_single_lora_gsm8k_async.yaml | 6 +- .../rl/compare_single_lora_gsm8k_sync.yaml | 10 +- cookbook/rl/sync_barrier_multi_lora_grpo.py | 689 ++++++++++++++++++ 3 files changed, 696 insertions(+), 9 deletions(-) create mode 100644 cookbook/rl/sync_barrier_multi_lora_grpo.py diff --git a/cookbook/rl/compare_single_lora_gsm8k_async.yaml b/cookbook/rl/compare_single_lora_gsm8k_async.yaml index c8144a6d0..b0a2015e0 100644 --- a/cookbook/rl/compare_single_lora_gsm8k_async.yaml +++ b/cookbook/rl/compare_single_lora_gsm8k_async.yaml @@ -1,5 +1,5 @@ -# Asynchronous comparison run. It is intentionally identical to the strict -# baseline except for max_staleness, run_id, and output paths. +# Asynchronous comparison run. It uses the Worker/TQ pipeline with up to three +# live partitions and otherwise matches the synchronous barrier baseline. runtime: run_id: compare_single_lora_gsm8k_async model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} @@ -70,7 +70,7 @@ scheduler: train: {policy: sticky, max_consecutive_units: null} evaluation: - enabled: true + enabled: false interval: 10 batch_size: 16 sampling_params: diff --git a/cookbook/rl/compare_single_lora_gsm8k_sync.yaml b/cookbook/rl/compare_single_lora_gsm8k_sync.yaml index d77d86caf..60dcdcee8 100644 --- a/cookbook/rl/compare_single_lora_gsm8k_sync.yaml +++ b/cookbook/rl/compare_single_lora_gsm8k_sync.yaml @@ -1,7 +1,6 @@ -# Strict on-policy baseline for comparison with compare_single_lora_gsm8k_async.yaml. -# max_staleness=0 permits only one live rollout partition. The workers may still -# pipeline groups inside that partition, but no rollout partition can cross a -# policy-version boundary. +# Fully synchronous barrier baseline for comparison with +# compare_single_lora_gsm8k_async.yaml. Run this file with +# sync_barrier_multi_lora_grpo.py, not async_multi_lora_grpo.py. runtime: run_id: compare_single_lora_gsm8k_sync model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} @@ -11,7 +10,6 @@ runtime: sampler_tp: 1 sampler_max_loras: 1 seed: 1 - max_staleness: 0 max_steps: 125 allow_partial_rollout: false rollout_max_retries: 2 @@ -72,7 +70,7 @@ scheduler: train: {policy: sticky, max_consecutive_units: null} evaluation: - enabled: true + enabled: false interval: 10 batch_size: 16 sampling_params: diff --git a/cookbook/rl/sync_barrier_multi_lora_grpo.py b/cookbook/rl/sync_barrier_multi_lora_grpo.py new file mode 100644 index 000000000..55f2a5ce4 --- /dev/null +++ b/cookbook/rl/sync_barrier_multi_lora_grpo.py @@ -0,0 +1,689 @@ +"""Synchronous barrier baseline for native async multi-LoRA GRPO. + +The model, sampler, datasets, rewards, batch semantics, and checkpoint cadence +match ``async_multi_lora_grpo.py``. The only intentional difference is the +execution schedule: every round finishes rollout for all active contexts +before any context starts training. +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import time +from dataclasses import dataclass +from typing import Any, Iterator, Sequence + +from omegaconf import OmegaConf + +from twinkle.metric import MetricRecord, create_metrics_reporter +from twinkle_agentic.async_rl.metrics import advantage_signal_metrics, rollout_metrics +from twinkle_agentic.async_rl.pipeline import (_prompt_batches, _reward_for_context, _train_batch) +from twinkle_agentic.async_rl.tq_utils import REQUIRED_MODEL_INPUT_FIELDS, columns_to_tq_fields +from twinkle_agentic.async_rl.types import LoraContext, PartitionAdmission +from twinkle_agentic.async_rl.utils import ( + TrainBatchConfig, + build_native_fsdp_model_kwargs, + configure_lora_lr_scheduler, + resolve_context_learning_rate, + resolve_context_lora_target_modules, + resolve_context_loss_config, + resolve_model_attention_implementation, + resolve_sequence_parallel_size, + sample_responses_to_rollout_rows, + sampler_data_parallel_size, + validate_context_batch_config, +) +from twinkle_agentic.async_rl.vllm_sampler_tq import _compute_reward_metrics + + +@dataclass +class SyncContextState: + context: LoraContext + prompt_batches: Iterator[Sequence[dict[str, Any]]] + rollout_batch_size: int + num_generations: int + sampling_params: Any + mini_batch_size: int + reward_fn: Any + adapter_path: str + adapter_history: list[str] + partition_step: int = 0 + optimizer_steps: int = 0 + policy_version: int = 0 + exhausted: bool = False + + +@dataclass +class SyncPartition: + admission: PartitionAdmission + state: SyncContextState + rows: list[dict[str, Any]] + rewards: list[float] + advantages: list[float] | None = None + + +class SyncBarrierMultiLoraGRPO: + + def __init__(self, raw_config: dict[str, Any]): + import twinkle + from peft import LoraConfig + from twinkle import DeviceGroup, DeviceMesh + from twinkle.data_format import SamplingParams + from twinkle.model import MultiLoraTransformersModel + from twinkle.processor import InputProcessor + from twinkle.sampler import vLLMSampler + + raw_config = OmegaConf.to_container(OmegaConf.create(raw_config), resolve=True) + if not isinstance(raw_config, dict): + raise TypeError('sync RL config must resolve to a mapping') + + runtime = raw_config['runtime'] + model_config = raw_config['model'] + lora_data = raw_config['lora'] + loss_data = raw_config.get('loss') + template_data = raw_config.get('template', {}) + template_cls = template_data.get('cls', 'Qwen3_5Template') + enable_thinking = bool(template_data.get('enable_thinking', False)) + model_gpus = int(runtime['model_gpus']) + sampler_gpus = int(runtime['sampler_gpus']) + sampler_tp = int(runtime['sampler_tp']) + sampler_dp = sampler_data_parallel_size(sampler_gpus, sampler_tp) + sequence_parallel_size = resolve_sequence_parallel_size( + model_gpus, + int(model_config['sequence_parallel_size']), + ) + padding_free = bool(model_config['padding_free']) + attn_implementation = resolve_model_attention_implementation( + model_config, + padding_free=padding_free, + sequence_parallel_size=sequence_parallel_size, + ) + model_max_length = int(model_config['max_length']) + sampler_config = raw_config['sampler'] + total_gpus = model_gpus + sampler_gpus + + twinkle.initialize( + mode='ray', + nproc_per_node=total_gpus, + groups=[ + DeviceGroup('model', list(range(model_gpus)), device_type='GPU'), + DeviceGroup( + 'sampler', + list(range(model_gpus, total_gpus)), + device_type='GPU', + gpus_per_worker=sampler_tp, + ), + ], + lazy_collect=False, + ) + model_mesh = DeviceMesh.from_sizes( + world_size=model_gpus, + dp_size=model_gpus, + ulysses_size=sequence_parallel_size, + ) + model_data_parallel_size = model_mesh.data_world_size + self.model_data_parallel_size = model_data_parallel_size + sampler_mesh = DeviceMesh.from_sizes( + world_size=sampler_gpus, + dp_size=sampler_dp, + tp_size=sampler_tp, + ) + model_kwargs = build_native_fsdp_model_kwargs(model_config) + if attn_implementation is not None: + model_kwargs['attn_implementation'] = attn_implementation + self.model = MultiLoraTransformersModel( + model_id=runtime['model_id'], + device_mesh=model_mesh, + remote_group='model', + max_length=model_max_length, + **model_kwargs, + ) + self.train_batch_configs: dict[str, TrainBatchConfig] = {} + self.states: list[SyncContextState] = [] + self.evaluation_configs: dict[str, dict[str, Any]] = {} + self._evaluation_batches: dict[str, list[Sequence[dict[str, Any]]]] = {} + global_evaluation = dict(raw_config.get('evaluation') or {}) + for item in raw_config['lora_contexts']: + context = LoraContext( + item['tenant_id'], + item['training_run_id'], + runtime['model_id'], + item['adapter_name'], + ) + rollout = item['rollout'] + train = item['train'] + rollout_batch_size = int(rollout['batch_size']) + num_generations = int(rollout['num_generations']) + train_batch_config = TrainBatchConfig( + mini_batch_size=int(train['mini_batch_size']), + micro_batch_size=int(train['micro_batch_size']), + dynamic_batching=bool(train.get('dynamic_batching', False)), + max_tokens_per_micro_batch=( + int(train['max_tokens_per_micro_batch']) + if train.get('max_tokens_per_micro_batch') is not None else None + ), + packing_algorithm=str(train.get('packing_algorithm', 'ffd')), + ) + validate_context_batch_config( + context.key, + rollout_groups=rollout_batch_size, + num_generations=num_generations, + train=train_batch_config, + sampler_dp=sampler_dp, + model_dp=model_data_parallel_size, + ) + adapter_lora_config = LoraConfig( + target_modules=resolve_context_lora_target_modules(item, lora_data), + r=lora_data['r'], + lora_alpha=lora_data['alpha'], + lora_dropout=lora_data['dropout'], + ) + self.model.add_adapter_to_model( + context.adapter_name, + adapter_lora_config, + gradient_accumulation_steps=1, + ) + self.model.set_optimizer( + 'AdamW', + lr=resolve_context_learning_rate(train, lora_data), + adapter_name=context.adapter_name, + ) + configure_lora_lr_scheduler(self.model, context.adapter_name, lora_data) + loss_cls, loss_kwargs = resolve_context_loss_config(item, loss_data) + self.model.set_loss( + loss_cls, + adapter_name=context.adapter_name, + **loss_kwargs, + ) + self.model.set_processor( + InputProcessor, + adapter_name=context.adapter_name, + padding_free=padding_free, + ) + self.model.set_template( + template_cls, + model_id=runtime['model_id'], + adapter_name=context.adapter_name, + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + initial_path = self.model.save( + f'sync-{context.adapter_name}-initial', + output_dir=runtime['output_dir'], + adapter_name=context.adapter_name, + ) + state = SyncContextState( + context=context, + prompt_batches=iter( + _prompt_batches( + item['dataset'], + model_id=runtime['model_id'], + batch_size=rollout_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + )), + rollout_batch_size=rollout_batch_size, + num_generations=num_generations, + sampling_params=SamplingParams( + max_tokens=rollout['max_tokens'], + temperature=rollout['temperature'], + top_p=rollout['top_p'], + repetition_penalty=float(rollout.get('repetition_penalty', 1.0)), + logprobs=1, + num_samples=1, + ), + mini_batch_size=train_batch_config.mini_batch_size, + reward_fn=_reward_for_context( + item.get('reward'), + context_key=context.key, + ), + adapter_path=initial_path, + adapter_history=[initial_path], + ) + self.states.append(state) + self.train_batch_configs[context.key] = train_batch_config + if bool(global_evaluation.get('enabled', False)): + eval_dataset = item.get('eval_dataset') + if eval_dataset is None: + raise ValueError(f'eval_dataset is required for periodic evaluation of {context.key}') + eval_batch_size = int(global_evaluation.get('batch_size', 16)) + eval_interval = int(global_evaluation.get('interval', 1)) + if eval_batch_size <= 0 or eval_interval <= 0: + raise ValueError('evaluation.batch_size and evaluation.interval must be positive') + eval_sampling = dict(global_evaluation.get('sampling_params') or {}) + self.evaluation_configs[context.key] = { + 'interval': eval_interval, + 'dataset_name': eval_dataset.get('name', eval_dataset['dataset_id']), + 'prompt_batches': _prompt_batches( + eval_dataset, + model_id=runtime['model_id'], + batch_size=eval_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + full_batches_only=False, + ), + 'sampling_params': SamplingParams( + max_tokens=int(eval_sampling.get('max_tokens', rollout['max_tokens'])), + temperature=float(eval_sampling.get('temperature', 0.0)), + top_p=float(eval_sampling.get('top_p', 1.0)), + repetition_penalty=float(eval_sampling.get('repetition_penalty', 1.0)), + logprobs=0, + num_samples=1, + ), + 'reward_fn': _reward_for_context( + eval_dataset.get('reward'), + context_key=f'{context.key} evaluation', + ), + } + + sampler_engine_args = { + 'tensor_parallel_size': sampler_tp, + 'enable_lora': True, + 'max_loras': int(runtime['sampler_max_loras']), + 'max_lora_rank': lora_data['r'], + 'max_model_len': int(sampler_config['max_model_len']), + 'gpu_memory_utilization': float(sampler_config['gpu_memory_utilization']), + 'max_num_seqs': int(sampler_config['max_num_seqs']), + 'enforce_eager': bool(sampler_config['enforce_eager']), + } + if sampler_config.get('max_num_batched_tokens') is not None: + sampler_engine_args['max_num_batched_tokens'] = int(sampler_config['max_num_batched_tokens']) + self.sampler = vLLMSampler( + model_id=runtime['model_id'], + remote_group='sampler', + device_mesh=sampler_mesh, + engine_args=sampler_engine_args, + ) + self.sampler.set_template( + template_cls, + model_id=runtime['model_id'], + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + self.output_dir = runtime['output_dir'] + self.max_steps = runtime.get('max_steps') + self.max_steps = None if self.max_steps is None else int(self.max_steps) + self.keep_adapter_versions = max(0, int(runtime.get('keep_adapter_versions', 0))) + self.metrics = create_metrics_reporter( + raw_config.get('metrics'), + run_id=str(runtime.get('run_id', 'sync_barrier_multi_lora_grpo')), + ) + self.completed_partitions = 0 + self._creation_order = 0 + + def _record_metric( + self, + stage: str, + *, + admission: PartitionAdmission | None = None, + context: LoraContext | None = None, + values: dict[str, Any] | None = None, + status: str = 'completed', + attributes: dict[str, Any] | None = None, + optimizer_step: int | None = None, + policy_version: int | None = None, + ) -> None: + if self.metrics is None: + return + self.metrics.record(MetricRecord( + stage=stage, + values=dict(values or {}), + context_key=( + admission.context.key if admission is not None + else context.key if context is not None else None + ), + partition_id=admission.partition_id if admission is not None else None, + partition_index=admission.step if admission is not None else None, + optimizer_step=optimizer_step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) + + def run(self) -> dict[str, Any]: + started = time.perf_counter() + try: + while self.max_steps is None or self.completed_partitions < self.max_steps: + partitions = self._rollout_round() + if not partitions: + break + self._advantage_round(partitions) + self._train_round(partitions) + except Exception as exc: + self._record_metric( + 'run', + status='failed', + values={'wall_time_s': time.perf_counter() - started}, + attributes={'error': f'{type(exc).__name__}: {exc}'}, + ) + if self.metrics is not None: + self.metrics.close() + raise + result = { + 'trained_partitions': self.completed_partitions, + 'wall_time_s': time.perf_counter() - started, + 'per_context': { + state.context.key: { + 'optimizer_steps': state.optimizer_steps, + 'policy_version': state.policy_version, + 'adapter_path': state.adapter_path, + } + for state in self.states + }, + } + self._record_metric( + 'run', + values={ + 'trained_partitions': result['trained_partitions'], + 'wall_time_s': result['wall_time_s'], + }, + ) + if self.metrics is not None: + self.metrics.flush() + result['metrics_health'] = self.metrics.health() + self.metrics.close() + return result + + def _rollout_round(self) -> list[SyncPartition]: + partitions = [] + for state in self.states: + if state.exhausted: + continue + if self.max_steps is not None and self.completed_partitions + len(partitions) >= self.max_steps: + break + prompts = next(state.prompt_batches, None) + if prompts is None or len(prompts) != state.rollout_batch_size: + state.exhausted = True + continue + admission = PartitionAdmission( + context=state.context, + partition_id=state.context.partition_id(state.partition_step), + step=state.partition_step, + target_groups=state.rollout_batch_size, + num_generations=state.num_generations, + created_order=self._creation_order, + ) + self._creation_order += 1 + self._record_metric( + 'rollout', + admission=admission, + status='submitted', + policy_version=state.policy_version, + values={ + 'prompt_count': admission.target_groups, + 'sample_count': admission.sample_count, + 'num_generations': admission.num_generations, + }, + attributes={'scope': 'partition'}, + ) + rollout_started = time.perf_counter() + sources = [{ + **dict(prompt), + 'group_id': f'{admission.partition_id}/group_{group_index}', + 'generation_idx': generation_index, + } for group_index, prompt in enumerate(prompts) + for generation_index in range(state.num_generations)] + responses = self.sampler.sample( + [dict(prompt) for prompt in prompts for _ in range(state.num_generations)], + state.sampling_params, + adapter_name=state.context.adapter_name, + adapter_path=state.adapter_path, + ) + rows = sample_responses_to_rollout_rows( + sources, + responses, + policy_version=state.policy_version, + ) + if len(rows) != admission.sample_count: + raise ValueError( + f'{admission.partition_id} expected {admission.sample_count} samples, got {len(rows)}') + for row in rows: + row.update({ + 'rollout_adapter_path': state.adapter_path, + 'rollout_policy_versions': [state.policy_version], + 'initial_policy_version': state.policy_version, + 'final_policy_version': state.policy_version, + 'policy_version_span': 0, + }) + rewards = [float(value) for value in state.reward_fn(rows, context=state.context)] + if len(rewards) != len(rows): + raise ValueError(f'{admission.partition_id} reward count does not match sample count') + rollout_latency_s = time.perf_counter() - rollout_started + self._record_rollout_groups(state, admission, rows, rewards) + self._record_metric( + 'rollout', + admission=admission, + policy_version=state.policy_version, + values=rollout_metrics( + completion_lengths=[int(row['completion_length']) for row in rows], + stop_reasons=[row.get('stop_reason') for row in rows], + rollout_latency_s=rollout_latency_s, + ), + attributes={'scope': 'partition'}, + ) + partitions.append(SyncPartition(admission, state, rows, rewards)) + state.partition_step += 1 + return partitions + + def _record_rollout_groups( + self, + state: SyncContextState, + admission: PartitionAdmission, + rows: list[dict[str, Any]], + rewards: list[float], + ) -> None: + for group_index in range(admission.target_groups): + start = group_index * admission.num_generations + end = start + admission.num_generations + group_rows = rows[start:end] + group_rewards = rewards[start:end] + metrics = { + **_compute_reward_metrics( + {state.context.key: state.reward_fn}, + state.context, + group_rows, + group_rewards, + ), + **rollout_metrics( + rewards={'reward': group_rewards}, + completion_lengths=[int(row['completion_length']) for row in group_rows], + stop_reasons=[row.get('stop_reason') for row in group_rows], + ), + } + self._record_metric( + 'rollout', + admission=admission, + policy_version=state.policy_version, + values=metrics, + attributes={ + 'scope': 'group', + 'group_id': f'{admission.partition_id}/group_{group_index}', + }, + ) + + def _advantage_round(self, partitions: list[SyncPartition]) -> None: + from twinkle.advantage import GRPOAdvantage + + advantage_fn = GRPOAdvantage() + for partition in partitions: + admission = partition.admission + partition.advantages = advantage_fn( + partition.rewards, + num_generations=admission.num_generations, + scale='group', + ).tolist() + samples_per_batch = admission.num_generations + for start in range(0, len(partition.rows), samples_per_batch): + end = min(start + samples_per_batch, len(partition.rows)) + self._record_metric( + 'advantage', + admission=admission, + policy_version=partition.state.policy_version, + values={ + 'sample_count': end - start, + **advantage_signal_metrics( + partition.rewards[start:end], + partition.advantages[start:end], + num_generations=admission.num_generations, + ), + }, + ) + + def _train_round(self, partitions: list[SyncPartition]) -> None: + for partition in partitions: + admission = partition.admission + state = partition.state + assert partition.advantages is not None + samples_per_batch = state.mini_batch_size + for start in range(0, len(partition.rows), samples_per_batch): + end = start + samples_per_batch + batch = self._training_batch( + partition.rows[start:end], + partition.rewards[start:end], + partition.advantages[start:end], + ) + train_started = time.perf_counter() + metrics = _train_batch( + self.model, + self.train_batch_configs, + batch, + admission, + model_data_parallel_size=self.model_data_parallel_size, + ) + state.optimizer_steps += 1 + metrics.update({ + 'sample_count': end - start, + 'train_latency_s': time.perf_counter() - train_started, + 'policy_version_gap_mean': 0.0, + 'policy_version_gap_p95': 0.0, + 'policy_version_gap_max': 0, + 'rollout_policy_span_mean': 0.0, + 'rollout_policy_span_max': 0, + }) + self._record_metric( + 'train', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values=metrics, + ) + finalize_started = time.perf_counter() + next_policy_version = state.policy_version + 1 + save_started = time.perf_counter() + state.adapter_path = self.model.save( + f'sync-{state.context.adapter_name}-v{next_policy_version}', + output_dir=self.output_dir, + adapter_name=state.context.adapter_name, + ) + adapter_save_latency_s = time.perf_counter() - save_started + publish_started = time.perf_counter() + state.policy_version = next_policy_version + policy_publish_latency_s = time.perf_counter() - publish_started + state.adapter_history.append(state.adapter_path) + self._record_metric( + 'policy', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + }, + attributes={'operation': 'publish', 'adapter_path': state.adapter_path}, + ) + self._evaluate_policy(state, admission) + prune_started = time.perf_counter() + self._prune_adapter_history(state) + adapter_prune_latency_s = time.perf_counter() - prune_started + self.completed_partitions += 1 + self._record_metric( + 'partition', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + 'adapter_prune_latency_s': adapter_prune_latency_s, + 'partition_finalize_latency_s': time.perf_counter() - finalize_started, + }, + ) + + def _evaluate_policy(self, state: SyncContextState, admission: PartitionAdmission) -> None: + config = self.evaluation_configs.get(state.context.key) + if config is None or state.policy_version % int(config['interval']): + return + if state.context.key not in self._evaluation_batches: + self._evaluation_batches[state.context.key] = list(config['prompt_batches']) + + started = time.perf_counter() + rewards: list[float] = [] + completion_lengths: list[int] = [] + prompt_count = 0 + for batch in self._evaluation_batches[state.context.key]: + prompts = list(batch) + responses = self.sampler.sample( + prompts, + config['sampling_params'], + adapter_name=state.context.adapter_name, + adapter_path=state.adapter_path, + ) + rows = sample_responses_to_rollout_rows( + prompts, + responses, + policy_version=state.policy_version, + ) + rewards.extend(float(value) for value in config['reward_fn'](rows, context=state.context)) + completion_lengths.extend(int(row['completion_length']) for row in rows) + prompt_count += len(prompts) + if not rewards: + raise ValueError(f'evaluation dataset is empty for {state.context.key}') + self._record_metric( + 'evaluation', + admission=admission, + optimizer_step=state.optimizer_steps, + policy_version=state.policy_version, + values={ + 'accuracy': sum(rewards) / len(rewards), + 'sample_count': len(rewards), + 'prompt_count': prompt_count, + 'completion_length': sum(completion_lengths) / len(completion_lengths), + 'eval_latency_s': time.perf_counter() - started, + }, + attributes={'eval_dataset': config['dataset_name']}, + ) + + @staticmethod + def _training_batch(rows: list[dict[str, Any]], rewards: list[float], advantages: list[float]): + fields = { + name: [row[name] for row in rows] + for name in (*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs') + } + fields.update({'rewards': rewards, 'advantages': advantages}) + return columns_to_tq_fields(fields, len(rows)) + + def _prune_adapter_history(self, state: SyncContextState) -> None: + retained_count = max(1, self.keep_adapter_versions) + stale = state.adapter_history[:-retained_count] + state.adapter_history = state.adapter_history[-retained_count:] + for path in stale: + if os.path.isdir(path): + shutil.rmtree(path) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument('--config', default='cookbook/rl/sync_barrier_multi_lora_grpo.yaml') + args = parser.parse_args() + config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) + print(SyncBarrierMultiLoraGRPO(config).run()) + + +if __name__ == '__main__': + main() + +# MODEL_ID=/path/to/model \ +# DATASET_ID=/path/to/gsm8k \ +# python cookbook/rl/async_multi_lora_grpo.py From 33c28c4e8d7a86d67aa246d9f5642d9e89e667ba Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Mon, 24 Aug 2026 14:42:35 +0800 Subject: [PATCH 17/20] refactor --- cookbook/client/async_rl/README.md | 198 +++-- .../async_rl/client_orchestrated_dpo.py | 222 ------ .../async_rl/client_orchestrated_grpo.py | 18 +- cookbook/client/async_rl/run_server.sh | 32 + .../server_config.yaml} | 19 +- cookbook/rl/async_multi_lora_dapo_grpo.yaml | 175 ----- .../async_multi_lora_dapo_hparam_sweep.yaml | 203 ----- cookbook/rl/async_rl/README.md | 129 +++ .../{ => async_rl}/async_multi_lora_grpo.py | 2 +- .../{ => async_rl}/async_multi_lora_grpo.yaml | 5 +- .../rl/async_rl/run_async_multi_lora_grpo.sh | 30 + cookbook/rl/async_single_lora_dapo_grpo.yaml | 133 ---- .../rl/async_single_lora_gsm8k_areal.yaml | 190 ----- cookbook/rl/async_single_lora_gsm8k_verl.yaml | 126 --- .../rl/compare_single_lora_gsm8k_async.yaml | 128 --- .../rl/compare_single_lora_gsm8k_sync.yaml | 128 --- cookbook/rl/sync_barrier_multi_lora_grpo.py | 689 ---------------- tests/model/test_multi_lora.py | 6 + tests/twinkle_agentic/test_async_rl_config.py | 226 ++++++ .../test_async_rl_data_plane.py | 130 +++ .../test_async_rl_native_tq.py | 737 +----------------- .../test_vllm_sampler_tq_generation.py | 315 +++++++- .../test_client_orchestrated_dpo.py | 108 --- .../test_client_orchestrated_grpo.py | 12 +- 24 files changed, 1050 insertions(+), 2911 deletions(-) delete mode 100644 cookbook/client/async_rl/client_orchestrated_dpo.py create mode 100755 cookbook/client/async_rl/run_server.sh rename cookbook/client/{server/transformer/server_config_local.yaml => async_rl/server_config.yaml} (87%) delete mode 100644 cookbook/rl/async_multi_lora_dapo_grpo.yaml delete mode 100644 cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml create mode 100644 cookbook/rl/async_rl/README.md rename cookbook/rl/{ => async_rl}/async_multi_lora_grpo.py (84%) rename cookbook/rl/{ => async_rl}/async_multi_lora_grpo.yaml (96%) create mode 100755 cookbook/rl/async_rl/run_async_multi_lora_grpo.sh delete mode 100644 cookbook/rl/async_single_lora_dapo_grpo.yaml delete mode 100644 cookbook/rl/async_single_lora_gsm8k_areal.yaml delete mode 100644 cookbook/rl/async_single_lora_gsm8k_verl.yaml delete mode 100644 cookbook/rl/compare_single_lora_gsm8k_async.yaml delete mode 100644 cookbook/rl/compare_single_lora_gsm8k_sync.yaml delete mode 100644 cookbook/rl/sync_barrier_multi_lora_grpo.py create mode 100644 tests/twinkle_agentic/test_async_rl_config.py create mode 100644 tests/twinkle_agentic/test_async_rl_data_plane.py delete mode 100644 tests/twinkle_client/test_client_orchestrated_dpo.py diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md index 2e1fcb2f0..8441c3ddb 100644 --- a/cookbook/client/async_rl/README.md +++ b/cookbook/client/async_rl/README.md @@ -1,72 +1,156 @@ -# Client-orchestrated asynchronous RL +# Client-Orchestrated Async GRPO -This directory demonstrates direct orchestration of the server's Model, -Sampler, and TransferQueue DataPlane components. +The client owns the Dataset, Reward, Advantage, rollout partitions, staleness, +training schedule, and policy publication. The server exposes shared +Multi-LoRA Model, vLLM Sampler, and TransferQueue DataPlane components without +owning the GRPO loop. -The client owns its Dataset, multi-turn rollout, Reward, Advantage, policy -versioning, staleness, and algorithm. There is no central async-RL management, -runtime, tenant submission API, or server-side RL worker involved. +This example deploys a dedicated Qwen3.5-4B GRPO server. It does not combine +different RL algorithms in one deployment. For the YAML-managed multi-tenant +runtime, see [`cookbook/rl`](../../rl/README.md). -- client_orchestrated_grpo.py maps each DataLoader batch to a private client-side - rollout partition. Its Rollout, Advantage, and Trainer workers run as - independent asyncio tasks. Prompt groups stream through TQ independently, so - ready groups can train while the remaining groups are still sampling. A - policy is published once, after the whole partition has trained. -- client_orchestrated_dpo.py uses Dataset, Reference, and Trainer workers. It is - a runnable offline DPO loop and shows that Worker is a role lifecycle rather - than a fixed RL stage graph. +## Resources -Start the component server with -cookbook/client/server/transformer/server_config.yaml. +The server configuration uses two GPUs: + +| Component | GPUs | Purpose | +|---|---:|---| +| Multi-LoRA Model | 1 | GRPO forward, backward, optimizer step, and checkpoint save | +| Async vLLM Sampler | 1 | Shared continuous-batched rollout generation | + +The DataPlane keeps token tensors and sampled log-probabilities server-side. +The client reads decoded completions for Reward and Advantage, appends those +values to the same `DataRef`, and sends only references to the Model. + +## Quick start + +Install the server, client, and async-RL dependencies: ```bash -pip install -e '.[async-rl,client]' -twinkle-server launch -c cookbook/client/server/transformer/server_config.yaml +pip install -e '.[async-rl,client,server]' ``` -Then start one or more independent client orchestrators: +Terminal 1 — start the dedicated component server: ```bash -python cookbook/client/async_rl/client_orchestrated_grpo.py +export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +export CUDA_VISIBLE_DEVICES=0,1 +bash cookbook/client/async_rl/run_server.sh ``` -The training loop composes only the low-level component methods: - -- `sampler.sample_to_data_plane(...)` / `sampler.asample_to_data_plane(...)` -- `model.forward_only_from_data_plane(...)` -- `model.forward_backward_from_data_plane(...)` -- `model.clip_grad_and_step(...)` -- `model.save(...)` -- `data_plane.put/get/append/release(...)` and - `aput/aget/aget_batch/aappend/arelease(...)` - -There is no additional RL runtime or orchestration protocol. The GRPO example -uses `asample_to_data_plane()` so the Sampler's output `DataRef` remains in TQ. Each -generation is one row tagged with its group, generation index, rollout policy, -and status. The Advantage worker reads only decoded completions and appends -reward and advantage to the same keys. Token tensors and sampled log-probabilities -remain server-side. The Trainer passes one or more `DataRef` values to -`forward_backward_from_data_plane()` and releases them in a local `finally` block. `asample()` -remains the materialized-response convenience API. - -`_RolloutPartition` is a private client record, not a server resource or SDK -API. The local FIFO limits live DataLoader batches before rollout, ready prompt -groups immediately use the Model primitives above, and the client calls -`save()` once after a whole batch has trained. `WorkerPipeline` only -starts, joins, and fail-fast cancels concrete roles; queues and algorithm state -remain ordinary client Python code. - -Different client processes may run GRPO and DPO against the same component -server. Model adapters are session-scoped. DataRefs are opaque capabilities -whose UUID identifies an independent physical TQ partition; DataPlane storage -does not know about tokens or sessions. The client remains the single writer -and algorithm owner for its adapter. Async Sampler requests share vLLM -continuous batching; this example does not promise strict round-robin fairness -between sampler tenants. - -The original YAML-managed runtime is still separate and can be started with: +The server script starts a local Ray cluster when needed, validates +`server_config.yaml`, and launches Ray Serve on port 8000. + +Terminal 2 — start one GRPO client: ```bash -python cookbook/rl/async_multi_lora_grpo.py \ - --config cookbook/rl/async_multi_lora_grpo.yaml +export TWINKLE_SERVER_URL=http://127.0.0.1:8000 +export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN +export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B +export TWINKLE_DATASET_ID=/absolute/path/to/gsm8k +python cookbook/client/async_rl/client_orchestrated_grpo.py +``` + +`TWINKLE_TEMPLATE_MODEL_ID` is the tokenizer/template model path used in the +client process. The template class is fixed to `Qwen3_5Template`; a model ID +such as `Qwen/Qwen3.5-4B` is not a template class name. + +## Execution model + +Each DataLoader batch becomes a private client-side `_RolloutPartition` bound +to one immutable adapter checkpoint: + +```text +DataLoader + -> Rollout Worker + -> async vLLM Sampler + -> DataPlane DataRef + -> Advantage Worker + -> Trainer Worker + -> save and publish policy +``` + +The workers are independent asyncio tasks: + +- Prompt groups are submitted separately through + `sampler.asample_to_data_plane()`. +- A complete group is rewarded and made trainable without waiting for the rest + of its partition. +- The Trainer consumes `TRAIN_MINI_BATCH_SIZE / NUM_GENERATIONS` ready groups + at a time. +- Partitions train and publish in FIFO order. +- A policy version advances only after every group in its partition has + trained and the new checkpoint has been saved. + +`MAX_STALENESS` controls the number of live partitions. Setting it to zero +still permits rollout/training overlap inside one partition; values above zero +also permit cross-partition overlap. + +## Training metrics + +After every `clip_grad_and_step()`, the Trainer calls: + +```python +model.calculate_metric(is_training=True) +``` + +and prints the returned metrics to client stdout: + +```text +optimizer_step=1 grad_norm=0.42 learning_rate=2e-05 loss=0.031 ``` + +Partition publication is logged separately: + +```text +partition=0 policy=1 optimizer_step=8 staleness=0 +``` + +## Client settings + +| Environment variable | Default | Purpose | +|---|---|---| +| `TWINKLE_SERVER_URL` | `http://localhost:8000` | Server base URL | +| `TWINKLE_SERVER_TOKEN` | `EMPTY_TOKEN` | Request authentication token | +| `TWINKLE_TEMPLATE_MODEL_ID` | `ms://Qwen/Qwen3.5-4B` | Client tokenizer/template source | +| `TWINKLE_DATASET_ID` | `ms://modelscope/gsm8k` | GSM8K dataset source | +| `TWINKLE_ADAPTER_NAME` | `client-grpo` | LoRA adapter name | +| `TWINKLE_MAX_PARTITIONS` | `100` | Maximum DataLoader batches admitted | +| `TWINKLE_MAX_STALENESS` | `2` | Maximum extra live partitions | +| `TWINKLE_ROLLOUT_CONCURRENCY` | `8` | Concurrent prompt-group submissions | +| `TWINKLE_NUM_GENERATIONS` | `4` | Generations per prompt group | +| `TWINKLE_BATCH_SIZE` | `8` | Prompt groups per partition | +| `TWINKLE_TRAIN_MINI_BATCH_SIZE` | `8` | Samples per optimizer step | +| `TWINKLE_MICRO_BATCH_SIZE` | `4` | Model micro-batch size | +| `TWINKLE_MAX_TOKENS_PER_MICRO_BATCH` | `4096` | Dynamic batching token limit | + +The following must hold: + +```text +TWINKLE_TRAIN_MINI_BATCH_SIZE % TWINKLE_NUM_GENERATIONS == 0 +TWINKLE_BATCH_SIZE % ( + TWINKLE_TRAIN_MINI_BATCH_SIZE / TWINKLE_NUM_GENERATIONS +) == 0 +``` + +## Files + +| File | Role | +|---|---| +| `client_orchestrated_grpo.py` | Client-side async GRPO orchestration | +| `run_server.sh` | Start Ray and the dedicated component server | +| `server_config.yaml` | Qwen3.5-4B Model, Sampler, DataPlane, Gateway, and Processor deployment | + +## Troubleshooting + +- **Client receives HTTP 404 for the model** — use the provided server config; + its public route is fixed to `Qwen/Qwen3.5-4B`, matching the client. +- **The process tries to access ModelScope while offline** — set both + `TWINKLE_LOCAL_MODEL_PATH` on the server and + `TWINKLE_TEMPLATE_MODEL_ID`/`TWINKLE_DATASET_ID` on the client to local + paths. +- **No loss is printed** — confirm the Model request completed and look for an + `optimizer_step=...` line. Metrics are calculated after every optimizer + step, not after each rollout group. +- **Only one GPU is active** — verify Ray sees two GPUs and that Model and + Sampler placement groups were assigned to different devices. diff --git a/cookbook/client/async_rl/client_orchestrated_dpo.py b/cookbook/client/async_rl/client_orchestrated_dpo.py deleted file mode 100644 index eb9a4bf1a..000000000 --- a/cookbook/client/async_rl/client_orchestrated_dpo.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Runnable offline DPO composed from the low-level async component clients.""" -from __future__ import annotations - -import asyncio -import inspect -import os -from typing import Any - -from peft import LoraConfig - -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.preprocessor import EmojiDPOProcessor -from twinkle_client import DataPlaneClient, init_twinkle_client -from twinkle_client.async_rl import Worker, WorkerPipeline -from twinkle_client.common.json_utils import json_safe -from twinkle_client.model import MultiLoraTransformersModel -from twinkle_client.types import DataRef - -BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') -MODEL_ID = f'ms://{BASE_MODEL}' -TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) -TEMPLATE_CLS = os.environ.get( - 'TWINKLE_TEMPLATE_CLS', - 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', -) -DATASET_ID = os.environ.get('TWINKLE_DPO_DATASET_ID', 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji') -ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-dpo') -MAX_STEPS = int(os.environ.get('TWINKLE_MAX_STEPS', '100')) -BATCH_SIZE = int(os.environ.get('TWINKLE_BATCH_SIZE', '4')) -MAX_LENGTH = int(os.environ.get('TWINKLE_MAX_LENGTH', '2048')) - - -def create_dataset() -> Dataset: - """Load and encode preference pairs in the client process.""" - dataset = Dataset(DatasetMeta(DATASET_ID, data_slice=range(MAX_STEPS * BATCH_SIZE))) - dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=MAX_LENGTH) - dataset.map(EmojiDPOProcessor, init_args={'system': 'You are a helpful assistant.'}) - dataset.encode() - return dataset - - -def prepare_dpo_batch(batch: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Flatten pairs as ``chosen_0, rejected_0, ...`` for DP-safe slicing.""" - rows: list[dict[str, Any]] = [] - for pair in batch: - common = {key: value for key, value in pair.items() if key not in ('positive', 'negative')} - rows.append({**common, **pair['positive']}) - rows.append({**common, **pair['negative']}) - return json_safe(rows) - - -async def _put_rows(data_plane, rows, *, kind, tags): - try: - return await data_plane.aput(rows, kind=kind, tags=tags) - except TypeError as error: - if 'tags' not in str(error): - raise - return await data_plane.aput(rows, kind=kind) - - -async def _submit(method, *args, **kwargs): - if inspect.iscoroutinefunction(method): - return await method(*args, **kwargs) - task = await asyncio.to_thread(method, *args, **kwargs) - if inspect.isawaitable(task): - return await task - return task - - -def _response_payload(value: Any) -> dict[str, Any]: - return value if isinstance(value, dict) else value.model_dump() - - -class _DatasetWorker(Worker): - - def __init__(self, dataloader, data_plane, output): - super().__init__('dataset') - self.dataloader = dataloader - self.data_plane = data_plane - self.output = output - - async def run(self) -> None: - completed = 0 - for batch in self.dataloader: - if completed >= MAX_STEPS: - break - rows = prepare_dpo_batch(batch) - tags = [] - for index in range(0, len(rows), 2): - source_pair_id = rows[index].get('pair_id', f'pair-{completed}-{index // 2}') - tags.extend(( - { - 'record_type': 'preference', - 'pair_id': str(source_pair_id), - 'pair_role': 'chosen', - 'pair_status': 'DATA_READY', - }, - { - 'record_type': 'preference', - 'pair_id': str(source_pair_id), - 'pair_role': 'rejected', - 'pair_status': 'DATA_READY', - }, - )) - ref = await _put_rows( - self.data_plane, rows, kind='dpo-preference', tags=tags) - await self.output.put(ref) - completed += 1 - await self.output.put(None) - - -class _ReferenceWorker(Worker): - - def __init__(self, model, source, output): - super().__init__('reference') - self.model = model - self.source = source - self.output = output - - async def run(self) -> None: - while True: - item = await self.source.get() - if item is None: - await self.output.put(None) - return - ref = await _reference_forward(self.model, item) - await self.output.put(ref) - - -class _TrainerWorker(Worker): - - def __init__(self, model, data_plane, source): - super().__init__('trainer') - self.model = model - self.data_plane = data_plane - self.source = source - self.completed_steps = 0 - self.saved = None - - async def run(self) -> None: - while True: - item = await self.source.get() - if item is None: - self.saved = _response_payload(await _submit( - self.model.save, - f'dpo-policy-{self.completed_steps}', - save_optimizer=True, - )) - return - ref = item - try: - await _submit( - self.model.forward_backward_from_data_plane, - ref, - kwarg_fields={'ref_outputs.logps': 'ref_logps'}, - ) - await _submit(self.model.clip_grad_and_step, max_grad_norm=1.0) - finally: - await self.data_plane.arelease(ref) - self.completed_steps += 1 - - -async def _reference_forward( - model: MultiLoraTransformersModel, - batch_ref: DataRef, -) -> DataRef: - """Run the frozen base model and append reference logps to the same rows.""" - return await _submit( - model.forward_only_from_data_plane, - batch_ref, - disable_lora=True, - output_ref=batch_ref, - output_fields={'logps': 'ref_logps'}, - ) - - -async def run_dpo( - dataloader: DataLoader, - model: MultiLoraTransformersModel, - data_plane: DataPlaneClient, -) -> dict[str, Any]: - """Run client-owned DPO roles over the shared Model and DataPlane services.""" - preference_ready = asyncio.Queue(maxsize=2) - reference_ready = asyncio.Queue(maxsize=2) - trainer = _TrainerWorker(model, data_plane, reference_ready) - await WorkerPipeline(( - _DatasetWorker(dataloader, data_plane, preference_ready), - _ReferenceWorker(model, preference_ready, reference_ready), - trainer, - )).run() - return trainer.saved - - -async def train() -> None: - client = init_twinkle_client( - base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), - api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), - ) - model = MultiLoraTransformersModel(MODEL_ID) - data_plane = DataPlaneClient() - - model.add_adapter_to_model( - ADAPTER_NAME, - LoraConfig(target_modules='all-linear', r=8, lora_alpha=32, lora_dropout=0.05), - ) - model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) - model.set_processor('InputProcessor', padding_side='right') - model.set_loss('DPOLoss', beta=0.1, loss_type='sigmoid', reference_free=False) - model.add_metric('DPOMetric', beta=0.1) - model.set_optimizer('AdamW', lr=1e-5) - - try: - dataloader = DataLoader(dataset=create_dataset(), batch_size=BATCH_SIZE, num_workers=0) - saved = await run_dpo(dataloader, model, data_plane) - print(f"saved DPO adapter to {saved['twinkle_path']}") - finally: - client.close() - - -if __name__ == '__main__': - asyncio.run(train()) diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py index 2a1ed0458..6f250cfab 100644 --- a/cookbook/client/async_rl/client_orchestrated_grpo.py +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -21,13 +21,11 @@ from twinkle_client.model import MultiLoraTransformersModel from twinkle_client.sampler import vLLMSampler -BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +BASE_MODEL = 'Qwen/Qwen3.5-4B' MODEL_ID = f'ms://{BASE_MODEL}' TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) -TEMPLATE_CLS = os.environ.get( - 'TWINKLE_TEMPLATE_CLS', - 'Qwen3_5Template' if ('Qwen3.5' in BASE_MODEL or 'Qwen3.6' in BASE_MODEL) else 'Template', -) +TEMPLATE_CLS = 'Qwen3_5Template' +DATASET_ID = os.environ.get('TWINKLE_DATASET_ID', 'ms://modelscope/gsm8k') ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-grpo') MAX_PARTITIONS = int(os.environ.get('TWINKLE_MAX_PARTITIONS', '100')) MAX_STALENESS = int(os.environ.get('TWINKLE_MAX_STALENESS', '2')) @@ -121,7 +119,7 @@ async def publish(self, partition: _RolloutPartition, policy: _Policy) -> None: def create_dataset() -> Dataset: - dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) + dataset = Dataset(DatasetMeta(DATASET_ID, subset_name='main', split='train')) dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=2048, enable_thinking=False) dataset.map(GSM8KProcessor(system='Put the final answer within \\boxed{}.')) dataset.encode(add_generation_prompt=True) @@ -319,6 +317,14 @@ async def _train(self, groups: list[_ReadyGroup]) -> None: ) await _submit(self.model.clip_grad_and_step, max_grad_norm=1.0) self.optimizer_step += 1 + metric_response = await _submit(self.model.calculate_metric, is_training=True) + metrics = dict( + metric_response['result'] + if isinstance(metric_response, dict) + else metric_response.result + ) + values = ' '.join(f'{name}={value}' for name, value in sorted(metrics.items())) + print(f'optimizer_step={self.optimizer_step} {values}'.rstrip()) finally: await asyncio.gather(*(self.data_plane.arelease(ref) for ref in refs)) diff --git a/cookbook/client/async_rl/run_server.sh b/cookbook/client/async_rl/run_server.sh new file mode 100755 index 000000000..d3a235d7c --- /dev/null +++ b/cookbook/client/async_rl/run_server.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${repo_root}" + +: "${TWINKLE_LOCAL_MODEL_PATH:?Set TWINKLE_LOCAL_MODEL_PATH to the local Qwen3.5-4B directory}" + +ray_port="${RAY_PORT:-6379}" +ray_address="${RAY_ADDRESS:-127.0.0.1:${ray_port}}" + +if ! command -v ray >/dev/null 2>&1; then + echo "ray is not installed; install the async-RL dependencies first" >&2 + exit 1 +fi +if ! command -v twinkle-server >/dev/null 2>&1; then + echo "twinkle-server is not installed; run: pip install -e '.[async-rl,client]'" >&2 + exit 1 +fi + +if ! ray status --address="${ray_address}" >/dev/null 2>&1; then + ray start \ + --head \ + --port="${ray_port}" \ + --num-gpus="${RAY_NUM_GPUS:-2}" \ + --include-dashboard=false \ + --disable-usage-stats +fi + +config=cookbook/client/async_rl/server_config.yaml +twinkle-server check-config -c "${config}" +exec twinkle-server launch -c "${config}" diff --git a/cookbook/client/server/transformer/server_config_local.yaml b/cookbook/client/async_rl/server_config.yaml similarity index 87% rename from cookbook/client/server/transformer/server_config_local.yaml rename to cookbook/client/async_rl/server_config.yaml index 833445bf3..103ea523f 100644 --- a/cookbook/client/server/transformer/server_config_local.yaml +++ b/cookbook/client/async_rl/server_config.yaml @@ -1,11 +1,10 @@ -# Twinkle Server Configuration - local Qwen3-4B for client-orchestrated async RL +# Twinkle Server Configuration - local Qwen3.5-4B for client-orchestrated async RL # Set the absolute Hugging Face-compatible model directory before loading this file: -# export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3-4B -# export TWINKLE_MODEL_ID=Qwen/Qwen3-4B # in every client process -# export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3-4B # if clients share that path +# export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +# export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B # in every client process # -# The public HTTP model name remains Qwen/Qwen3-4B. Both the training Model +# The public HTTP model name remains Qwen/Qwen3.5-4B. Both the training Model # and the vLLM Sampler load the same local directory, so neither component # downloads the base model independently. @@ -33,7 +32,7 @@ applications: server_config: per_token_model_limit: 3 supported_models: - - Qwen/Qwen3-4B + - Qwen/Qwen3.5-4B deployments: - name: TinkerCompatServer max_ongoing_requests: 50 @@ -63,8 +62,8 @@ applications: num_cpus: 1 # One GPU hosts the training base model and multiple LoRA adapters. - - name: models-Qwen3-4B - route_prefix: /api/v1/model/Qwen/Qwen3-4B + - name: models-Qwen3.5-4B + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B import_path: model args: backend: transformers @@ -99,8 +98,8 @@ applications: TWINKLE_FAIL_FAST: "0" # A second GPU hosts vLLM and loads the same local base model. - - name: sampler-Qwen3-4B - route_prefix: /api/v1/sampler/Qwen/Qwen3-4B + - name: sampler-Qwen3.5-4B + route_prefix: /api/v1/sampler/Qwen/Qwen3.5-4B import_path: sampler args: model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} diff --git a/cookbook/rl/async_multi_lora_dapo_grpo.yaml b/cookbook/rl/async_multi_lora_dapo_grpo.yaml deleted file mode 100644 index f49bf1ccb..000000000 --- a/cookbook/rl/async_multi_lora_dapo_grpo.yaml +++ /dev/null @@ -1,175 +0,0 @@ -runtime: - run_id: async_multi_lora_dapo_grpo - model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3-4B} - mode: ray - model_gpus: 2 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 4 - max_staleness: 2 - max_steps: null - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/async_multi_lora_dapo_grpo - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/dapo_metrics.jsonl - summary_path: outputs/async_rl/dapo_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: false - -model: - strategy: native_fsdp - fsdp_config: - reshard_after_forward: true - # Two trainer GPUs form one Ulysses SP group. The effective model DP size is 1. - sequence_parallel_size: 2 - padding_free: true - attn_implementation: flash_attention_2 - max_length: 12288 - -sampler: - max_model_len: 12288 - gpu_memory_utilization: 0.8 - max_num_seqs: 32 - max_num_batched_tokens: 16384 - # Keep the long-sequence sampler on eager execution until CUDA-graph stability is verified. - enforce_eager: true - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - train: {policy: sticky, max_consecutive_units: null} - -evaluation: - enabled: true - interval: 5 - batch_size: 16 - sampling_params: - max_tokens: 8192 - temperature: 0.0 - top_p: 1.0 - -lora: - target_modules: all-linear - r: 16 - alpha: 32 - dropout: 0.0 - # Match verl's stable GRPO-on-DAPO learning rate. No scheduler means constant LR. - learning_rate: 1.0e-6 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - # Set both environment variables to distinct local parquet splits for a disjoint-tenant experiment. - - tenant_id: tenant_a - training_run_id: dapo_async - adapter_name: tenant_a_dapo_math_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:TENANT_A_DAPO_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 2000 - # Keeps prompt + max rollout tokens within model.max_length. - max_length: 4096 - processor: DAPOMathProcessor - eval_dataset: - name: aime2024 - dataset_id: ${oc.env:AIME2024_DATASET_ID,ms://Maxwell-Jia/AIME_2024} - subset_name: default - split: train - data_num: null - max_length: 4096 - processor: AIME2024Processor - reward: - class_path: twinkle.reward.BoxedMathAccuracyReward - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - # Four prompt groups x eight generations = 32 samples per optimizer step. - mini_batch_size: 32 - micro_batch_size: 1 - dynamic_batching: false - - - tenant_id: tenant_b - training_run_id: dapo_async - adapter_name: tenant_b_dapo_math_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:TENANT_B_DAPO_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 2000 - max_length: 4096 - processor: DAPOMathProcessor - eval_dataset: - name: aime2024 - dataset_id: ${oc.env:AIME2024_DATASET_ID,ms://Maxwell-Jia/AIME_2024} - subset_name: default - split: train - data_num: null - max_length: 4096 - processor: AIME2024Processor - reward: - class_path: twinkle.reward.BoxedMathAccuracyReward - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - mini_batch_size: 32 - micro_batch_size: 1 - dynamic_batching: false - - -# export MODEL_ID=/nas/disk1/Qwen3.5-4B - -# # 可选。不设置时，两个租户默认都读取远程 DAPO-Math-17k。 -# export TENANT_A_DAPO_DATASET_ID=/path/to/tenant_a_dapo.parquet -# export TENANT_B_DAPO_DATASET_ID=/path/to/tenant_b_dapo.parquet - -# python cookbook/rl/async_multi_lora_grpo.py \ -# --config cookbook/rl/async_multi_lora_dapo_grpo.yaml diff --git a/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml b/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml deleted file mode 100644 index 166f96332..000000000 --- a/cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml +++ /dev/null @@ -1,203 +0,0 @@ -runtime: - run_id: async_multi_lora_dapo_hparam_sweep - model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3-4B} - mode: ray - # Three-GPU layout: two trainer GPUs in one SP group plus one sampler GPU. - model_gpus: 2 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 4 - max_staleness: 2 - max_steps: null - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/async_multi_lora_dapo_hparam_sweep - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/dapo_hparam_sweep_metrics.jsonl - summary_path: outputs/async_rl/dapo_hparam_sweep_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: false - -model: - strategy: native_fsdp - fsdp_config: - reshard_after_forward: true - sequence_parallel_size: 2 - padding_free: true - attn_implementation: flash_attention_2 - max_length: 12288 - -sampler: - max_model_len: 12288 - gpu_memory_utilization: 0.8 - max_num_seqs: 32 - max_num_batched_tokens: 16384 - enforce_eager: true - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - # Interleave hyperparameter candidates so one adapter cannot monopolize training. - train: {policy: round_robin, max_consecutive_units: 1} - -lora: - target_modules: all-linear - r: 16 - alpha: 32 - dropout: 0.05 - # Used only when a context omits train.learning_rate. - learning_rate: 1.0e-6 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - # A/B/C isolate learning rate while keeping mini_batch_size fixed at 32 samples. - - tenant_id: tenant_lr1e6_g4 - training_run_id: dapo_hparam_sweep - adapter_name: dapo_lr1e6_g4_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 500 - max_length: 4096 - processor: DAPOMathProcessor - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - learning_rate: 1.0e-6 - mini_batch_size: 32 - micro_batch_size: 1 - dynamic_batching: false - - - tenant_id: tenant_lr5e6_g4 - training_run_id: dapo_hparam_sweep - adapter_name: dapo_lr5e6_g4_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 500 - max_length: 4096 - processor: DAPOMathProcessor - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - learning_rate: 5.0e-6 - mini_batch_size: 32 - micro_batch_size: 1 - dynamic_batching: false - - - tenant_id: tenant_lr1e5_g4 - training_run_id: dapo_hparam_sweep - adapter_name: dapo_lr1e5_g4_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 500 - max_length: 4096 - processor: DAPOMathProcessor - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - learning_rate: 1.0e-5 - mini_batch_size: 32 - micro_batch_size: 1 - dynamic_batching: false - - # B/D isolate train batch size while keeping learning_rate fixed at 5e-6. - - tenant_id: tenant_lr5e6_g16 - training_run_id: dapo_hparam_sweep - adapter_name: dapo_lr5e6_g16_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:DAPO_HPARAM_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 500 - max_length: 4096 - processor: DAPOMathProcessor - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - learning_rate: 5.0e-6 - mini_batch_size: 128 - micro_batch_size: 1 - dynamic_batching: false - - -# CUDA_VISIBLE_DEVICES=0,1,2 \ -# MODEL_ID=/nas/disk1/Qwen3-4B \ -# DAPO_HPARAM_DATASET_ID=/path/to/dapo_math_500.parquet \ -# TWINKLE_SEED=42 \ -# python3 cookbook/rl/async_multi_lora_grpo.py \ -# --config cookbook/rl/async_multi_lora_dapo_hparam_sweep.yaml diff --git a/cookbook/rl/async_rl/README.md b/cookbook/rl/async_rl/README.md new file mode 100644 index 000000000..4c8188269 --- /dev/null +++ b/cookbook/rl/async_rl/README.md @@ -0,0 +1,129 @@ +# Async Multi-LoRA GRPO + +One YAML configuration launches two LoRA tenants over a shared training model, +vLLM sampler, and TransferQueue data plane. Rollout, advantage calculation, and +training run as independent workers, while each tenant keeps its own dataset, +reward, optimizer, scheduler, partitions, and policy versions. + +This directory contains the YAML-managed CLI workflow. For client-orchestrated +GRPO over HTTP, see +[`cookbook/client/async_rl`](../../client/async_rl/README.md). + +## Resources + +The default configuration uses three GPUs: + +| Component | GPUs | Purpose | +|---|---:|---| +| Training model | 2 | Native FSDP training for all LoRA tenants | +| vLLM sampler | 1 | Shared rollout generation | + +The example has two GSM8K tenants. Each tenant consumes 128 prompts in eight +partitions: + +```text +16 prompts × 4 generations = 64 samples per partition +64 samples ÷ mini_batch_size 4 = 16 optimizer steps per partition +8 partitions × 16 optimizer steps = 128 optimizer steps per tenant +``` + +## Quick start + +Install the async-RL dependencies: + +```bash +pip install -e '.[async-rl]' +``` + +Set a model and datasets that both training processes and Ray workers can +access: + +```bash +export MODEL_ID=/absolute/path/to/Qwen3.5-4B +export TENANT_A_DATASET_ID=/absolute/path/to/gsm8k +export TENANT_B_DATASET_ID=/absolute/path/to/gsm8k +export CUDA_VISIBLE_DEVICES=0,1,2 +``` + +`MODEL_ID` must be a Hugging Face-compatible directory readable by both +Transformers and vLLM. Dataset values may be local ModelScope-compatible paths +or `ms://...` identifiers. Use local paths when running offline. + +Start a local Ray cluster when needed and launch training: + +```bash +bash cookbook/rl/async_rl/run_async_multi_lora_grpo.sh +``` + +If Ray is already running, the script reuses it. To launch the Python entry +point directly: + +```bash +python cookbook/rl/async_rl/async_multi_lora_grpo.py \ + --config cookbook/rl/async_rl/async_multi_lora_grpo.yaml +``` + +## Scheduling and staleness + +The workers exchange complete prompt groups through TransferQueue: + +```text +RolloutWorker -> TransferQueue -> AdvantageWorker -> TrainerWorker +``` + +`max_staleness` limits live partitions, not mini-batches within one partition: + +- `max_staleness=0` allows one live partition. Training may still overlap with + rollout inside that partition once enough complete prompt groups form a + mini-batch. +- `max_staleness=1` allows two live partitions, so training an older partition + may also overlap with rollout for the next partition. + +The default `mini_batch_size=4` equals one complete four-generation prompt +group, allowing the Trainer to consume a group without waiting for all 16 +groups in the partition. A policy version is published only after every +mini-batch in the partition has trained. + +Batch settings must satisfy: + +```text +partition_samples = rollout.batch_size × rollout.num_generations +partition_samples % train.mini_batch_size == 0 +train.mini_batch_size % rollout.num_generations == 0 +train.mini_batch_size % model_dp == 0 +``` + +## Outputs + +| Output | Default path | +|---|---| +| LoRA checkpoints | `output/async_multi_lora_grpo/` | +| Rollout JSONL | `output/async_multi_lora_grpo/rollouts/` | +| Metrics | `outputs/async_rl/metrics.jsonl` | +| Metrics summary | `outputs/async_rl/summary.json` | + +Set `rollout_output.enabled: false` when benchmarking throughput to avoid +writing one JSONL file per completed prompt group. + +## Files + +| File | Role | +|---|---| +| `async_multi_lora_grpo.py` | Load the YAML and run `AsyncMultiLoraGRPOPipeline` | +| `async_multi_lora_grpo.yaml` | Model, sampler, scheduler, LoRA, dataset, reward, and tenant configuration | +| `run_async_multi_lora_grpo.sh` | Validate required environment variables, start local Ray if needed, and launch training | + +## Troubleshooting + +- **A placement group stays pending** — the default run needs three visible + GPUs. Check `ray status` and make sure Ray was started with at least three + GPUs. +- **The process tries to download a model or dataset** — replace every + `MODEL_ID` and `TENANT_*_DATASET_ID` value with an absolute local path visible + to all Ray workers. +- **Async is slower with `max_staleness=0`** — this setting disables + cross-partition overlap but retains Worker, Ray RPC, and TransferQueue + overhead. Partition-internal overlap also depends on prompt completion + distribution. +- **Rollout files dominate runtime** — disable `rollout_output.enabled` for + performance measurements. diff --git a/cookbook/rl/async_multi_lora_grpo.py b/cookbook/rl/async_rl/async_multi_lora_grpo.py similarity index 84% rename from cookbook/rl/async_multi_lora_grpo.py rename to cookbook/rl/async_rl/async_multi_lora_grpo.py index 92b7f6598..f6c19b8c4 100644 --- a/cookbook/rl/async_multi_lora_grpo.py +++ b/cookbook/rl/async_rl/async_multi_lora_grpo.py @@ -11,7 +11,7 @@ def main() -> None: parser = argparse.ArgumentParser() - parser.add_argument('--config', default='cookbook/rl/async_multi_lora_grpo.yaml') + parser.add_argument('--config', default='cookbook/rl/async_rl/async_multi_lora_grpo.yaml') args = parser.parse_args() config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) print(AsyncMultiLoraGRPOPipeline.from_config(config).run()) diff --git a/cookbook/rl/async_multi_lora_grpo.yaml b/cookbook/rl/async_rl/async_multi_lora_grpo.yaml similarity index 96% rename from cookbook/rl/async_multi_lora_grpo.yaml rename to cookbook/rl/async_rl/async_multi_lora_grpo.yaml index cca7775c1..36a5c303b 100644 --- a/cookbook/rl/async_multi_lora_grpo.yaml +++ b/cookbook/rl/async_rl/async_multi_lora_grpo.yaml @@ -6,6 +6,7 @@ runtime: sampler_gpus: 1 sampler_tp: 1 sampler_max_loras: 4 + seed: 1 max_staleness: 1 max_steps: null allow_partial_rollout: false @@ -73,8 +74,8 @@ lora: learning_rate: 5.0e-5 lr_scheduler: cls: CosineAnnealingLR - # One optimizer step per prompt group with mini_batch_size=4. - T_max: 2000 + # 128 prompts / 16 prompts per partition * 16 groups per partition. + T_max: 128 eta_min: 0.0 loss: diff --git a/cookbook/rl/async_rl/run_async_multi_lora_grpo.sh b/cookbook/rl/async_rl/run_async_multi_lora_grpo.sh new file mode 100755 index 000000000..7149a2d84 --- /dev/null +++ b/cookbook/rl/async_rl/run_async_multi_lora_grpo.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${repo_root}" + +: "${MODEL_ID:?Set MODEL_ID to a Hugging Face-compatible model directory or model ID}" +: "${TENANT_A_DATASET_ID:?Set TENANT_A_DATASET_ID to a GSM8K dataset path or ID}" +: "${TENANT_B_DATASET_ID:?Set TENANT_B_DATASET_ID to a GSM8K dataset path or ID}" + +ray_port="${RAY_PORT:-6379}" +ray_address="${RAY_ADDRESS:-127.0.0.1:${ray_port}}" +ray_num_gpus="${RAY_NUM_GPUS:-3}" + +if ! command -v ray >/dev/null 2>&1; then + echo "ray is not installed; install the async-RL dependencies first" >&2 + exit 1 +fi + +if ! ray status --address="${ray_address}" >/dev/null 2>&1; then + ray start \ + --head \ + --port="${ray_port}" \ + --num-gpus="${ray_num_gpus}" \ + --include-dashboard=false \ + --disable-usage-stats +fi + +exec python cookbook/rl/async_rl/async_multi_lora_grpo.py \ + --config cookbook/rl/async_rl/async_multi_lora_grpo.yaml diff --git a/cookbook/rl/async_single_lora_dapo_grpo.yaml b/cookbook/rl/async_single_lora_dapo_grpo.yaml deleted file mode 100644 index 701b750c6..000000000 --- a/cookbook/rl/async_single_lora_dapo_grpo.yaml +++ /dev/null @@ -1,133 +0,0 @@ -runtime: - run_id: async_single_lora_dapo_grpo - model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3-4B} - mode: ray - # Same three-GPU layout as the four-context sweep: two trainer GPUs plus one sampler GPU. - model_gpus: 2 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 4 - max_staleness: 0 - max_steps: null - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/async_single_lora_dapo_grpo - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/single_lora_dapo_metrics.jsonl - summary_path: outputs/async_rl/single_lora_dapo_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: false - -model: - strategy: native_fsdp - fsdp_config: - reshard_after_forward: true - sequence_parallel_size: 2 - padding_free: true - attn_implementation: flash_attention_2 - max_length: 12288 - -sampler: - max_model_len: 12288 - gpu_memory_utilization: 0.8 - max_num_seqs: 32 - max_num_batched_tokens: 16384 - enforce_eager: true - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - # Keep scheduler settings identical to the four-context experiment. - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - train: {policy: round_robin, max_consecutive_units: 1} - -evaluation: - enabled: true - interval: 5 - batch_size: 16 - sampling_params: - max_tokens: 8192 - temperature: 0.0 - top_p: 1.0 - -lora: - target_modules: all-linear - r: 16 - alpha: 32 - dropout: 0.05 - learning_rate: 1.0e-6 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - - tenant_id: tenant_single_dapo - training_run_id: dapo_single_tenant - adapter_name: dapo_single_lr1e6_g4_lora - reward: - class_path: twinkle.reward.DAPOMathReward - kwargs: - max_response_length: ${...rollout.max_tokens} - # Intentionally identical to tenant_lr1e6_g4 in the four-context sweep. - overlong_buffer_length: 4096 - overlong_penalty_factor: 1.0 - score_tail_chars: 300 - dataset: - dataset_id: ${oc.env:DAPO_SINGLE_DATASET_ID,ms://BytedTsinghua-SIA/DAPO-Math-17k} - subset_name: default - split: train - data_num: 500 - max_length: 4096 - processor: DAPOMathProcessor - eval_dataset: - name: aime2024 - dataset_id: ${oc.env:AIME2024_DATASET_ID,ms://Maxwell-Jia/AIME_2024} - subset_name: default - split: train - data_num: null - max_length: 4096 - processor: AIME2024Processor - reward: - class_path: twinkle.reward.BoxedMathAccuracyReward - rollout: - batch_size: 16 - num_generations: 8 - max_tokens: 8192 - temperature: 1.0 - top_p: 0.95 - train: - learning_rate: 1.0e-6 - mini_batch_size: 32 - micro_batch_size: 1 - dynamic_batching: false - - -# CUDA_VISIBLE_DEVICES=0,1,2 \ -# MODEL_ID=/nas/disk1/Qwen3-4B \ -# DAPO_SINGLE_DATASET_ID=/path/to/the_same_dapo_math_500.parquet \ -# TWINKLE_SEED=42 \ -# python3 cookbook/rl/async_multi_lora_grpo.py \ -# --config cookbook/rl/async_single_lora_dapo_grpo.yaml diff --git a/cookbook/rl/async_single_lora_gsm8k_areal.yaml b/cookbook/rl/async_single_lora_gsm8k_areal.yaml deleted file mode 100644 index d6becbbd5..000000000 --- a/cookbook/rl/async_single_lora_gsm8k_areal.yaml +++ /dev/null @@ -1,190 +0,0 @@ -runtime: - run_id: async_single_lora_gsm8k_areal - model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} - mode: ray - model_gpus: 1 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 1 - seed: 1 - max_staleness: 0 - max_steps: null - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/async_single_lora_gsm8k_areal - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/single_lora_gsm8k_areal_metrics.jsonl - summary_path: outputs/async_rl/single_lora_gsm8k_areal_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: true - -model: - strategy: native_fsdp - attn_implementation: flash_attention_2 - fsdp_config: - reshard_after_forward: true - sequence_parallel_size: 1 - padding_free: false - max_length: 2048 - -sampler: - max_model_len: 2048 - gpu_memory_utilization: 0.8 - max_num_seqs: 64 - enforce_eager: false - -rollout_output: - enabled: true - output_dir: ${runtime.output_dir}/rollouts - include_token_ids: false - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - train: {policy: sticky, max_consecutive_units: null} - -evaluation: - enabled: false - interval: 10 - batch_size: 16 - sampling_params: - max_tokens: 2048 - temperature: 0.6 - top_p: 1.0 - -lora: - target_modules: all-linear - r: 16 - alpha: 16 - dropout: 0.0 - learning_rate: 1.7e-5 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - - tenant_id: tenant_single - training_run_id: gsm8k_areal_pk - adapter_name: gsm8k_areal_pk_lora - reward: - class_path: twinkle.reward.MathVerifyAccuracyReward - dataset: - dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID} - subset_name: main - split: train - data_num: 2000 - max_length: 1024 - processor: AReaLGSM8KProcessor - eval_dataset: - name: gsm8k/test - dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID} - subset_name: main - split: test - data_num: null - max_length: 1024 - processor: AReaLGSM8KProcessor - reward: - class_path: twinkle.reward.MathVerifyAccuracyReward - rollout: - batch_size: 16 - num_generations: 4 - max_tokens: 1024 - temperature: 1.0 - top_p: 1.0 - train: - mini_batch_size: 64 - micro_batch_size: 64 - dynamic_batching: true - max_tokens_per_micro_batch: 4096 - packing_algorithm: ffd - - -# export VIRTUAL_ENV=/opt/.venv -# export PATH=/opt/.venv/bin:$PATH -# export CUDA_VISIBLE_DEVICES=2,3 -# export AREAL_MODEL_PATH=/nas/disk1/Qwen3-4B -# export AREAL_ADMIN_API_KEY="$( -# python3 -c 'import secrets; print(secrets.token_urlsafe(32))' -# )" -# export AREAL_TRIAL="pk-$(date +%Y%m%d-%H%M%S)" - -# python3 examples/math/gsm8k_rl.py \ -# --config examples/math/gsm8k_grpo_lora.yaml \ -# scheduler.type=local \ -# +rollout.agent.admin_api_key="$AREAL_ADMIN_API_KEY" \ -# seed=1 \ -# cluster.n_nodes=1 \ -# cluster.n_gpus_per_node=2 \ -# actor.path="$AREAL_MODEL_PATH" \ -# +actor.attn_impl=flash_attention_2 \ -# rollout.backend=vllm:d1 \ -# actor.backend=fsdp:d1 \ -# train_dataset.path=/model/ljl/project/data/gsm8k \ -# valid_dataset.path=/model/ljl/project/data/gsm8k \ -# train_dataset.batch_size=16 \ -# train_dataset.shuffle=false \ -# total_train_epochs=1 \ -# +total_train_steps=125 \ -# rollout.consumer_batch_size=16 \ -# rollout.max_concurrent_rollouts=16 \ -# rollout.max_head_offpolicyness=0 \ -# gconfig.n_samples=4 \ -# gconfig.max_new_tokens=1024 \ -# gconfig.max_tokens=2048 \ -# gconfig.temperature=1.0 \ -# +gconfig.top_p=1.0 \ -# +actor.mb_spec.n_mbs=64 \ -# ++actor.mb_spec.max_tokens_per_mb=null \ -# actor.use_lora=true \ -# actor.lora_rank=16 \ -# actor.lora_alpha=16 \ -# actor.optimizer.lr=1.7e-4 \ -# actor.optimizer.weight_decay=0.01 \ -# actor.optimizer.lr_scheduler_type=constant \ -# actor.eps_clip=0.2 \ -# actor.ppo_n_minibatches=1 \ -# actor.reward_scaling=1.0 \ -# actor.reward_bias=0.0 \ -# actor.adv_norm=null \ -# actor.kl_ctl=0.0 \ -# actor.recompute_logprob=false \ -# actor.use_decoupled_loss=false \ -# actor.rejection_sampling=null \ -# ++vllm.max_model_len=2048 \ -# ++vllm.gpu_memory_utilization=0.8 \ -# ++vllm.max_num_seqs=64 \ -# ++vllm.max_loras=1 \ -# ++vllm.enforce_eager=false \ -# evaluator.freq_epochs=null \ -# evaluator.freq_steps=null \ -# evaluator.freq_secs=null \ -# trial_name="$AREAL_TRIAL" \ -# +stats_logger.tensorboard.path="/tmp/areal/tensorboard/gsm8k-grpo/$AREAL_TRIAL" \ -# cluster.fileroot=/nas/disk1/areal-experiments \ -# saver.freq_epochs=null \ -# saver.freq_steps=125 \ -# saver.freq_secs=null diff --git a/cookbook/rl/async_single_lora_gsm8k_verl.yaml b/cookbook/rl/async_single_lora_gsm8k_verl.yaml deleted file mode 100644 index 7dae0f1ee..000000000 --- a/cookbook/rl/async_single_lora_gsm8k_verl.yaml +++ /dev/null @@ -1,126 +0,0 @@ -runtime: - run_id: async_single_lora_gsm8k_accuracy - model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} - mode: ray - model_gpus: 1 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 1 - seed: 1 - max_staleness: 0 - max_steps: 125 - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/async_single_lora_gsm8k_accuracy - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/single_lora_gsm8k_accuracy_metrics.jsonl - summary_path: outputs/async_rl/single_lora_gsm8k_accuracy_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: true - -model: - strategy: native_fsdp - attn_implementation: flash_attention_2 - fsdp_config: - reshard_after_forward: true - sequence_parallel_size: 1 - padding_free: false - max_length: 2048 - -sampler: - max_model_len: 2048 - max_num_batched_tokens: 4096 - gpu_memory_utilization: 0.7 - max_num_seqs: 64 - enforce_eager: false - -rollout_output: - enabled: true - output_dir: ${runtime.output_dir}/rollouts - include_token_ids: false - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - train: {policy: sticky, max_consecutive_units: null} - -evaluation: - enabled: true - interval: 10 - batch_size: 16 - sampling_params: - max_tokens: 1024 - temperature: 0.6 - top_p: 1.0 - -lora: - target_modules: all-linear - r: 16 - alpha: 16 - dropout: 0.0 - learning_rate: 1.7e-5 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - - tenant_id: tenant_single - training_run_id: gsm8k_accuracy - adapter_name: gsm8k_accuracy_lora - reward: - class_path: twinkle.reward.GSM8KAccuracyReward - dataset: - dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID} - subset_name: main - split: train - data_num: 2000 - max_length: 1024 - processor: GSM8KProcessor - system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' - eval_dataset: - name: gsm8k/test - dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID} - subset_name: main - split: test - data_num: null - max_length: 1024 - processor: GSM8KProcessor - system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' - reward: - class_path: twinkle.reward.GSM8KAccuracyReward - rollout: - batch_size: 16 - num_generations: 4 - max_tokens: 1024 - temperature: 1.0 - top_p: 1.0 - train: - mini_batch_size: 64 - micro_batch_size: 64 - dynamic_batching: true - max_tokens_per_micro_batch: 4096 - packing_algorithm: ffd diff --git a/cookbook/rl/compare_single_lora_gsm8k_async.yaml b/cookbook/rl/compare_single_lora_gsm8k_async.yaml deleted file mode 100644 index b0a2015e0..000000000 --- a/cookbook/rl/compare_single_lora_gsm8k_async.yaml +++ /dev/null @@ -1,128 +0,0 @@ -# Asynchronous comparison run. It uses the Worker/TQ pipeline with up to three -# live partitions and otherwise matches the synchronous barrier baseline. -runtime: - run_id: compare_single_lora_gsm8k_async - model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} - mode: ray - model_gpus: 1 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 1 - seed: 1 - max_staleness: 2 - max_steps: 125 - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/compare_single_lora_gsm8k_async - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/compare_single_lora_gsm8k_async_metrics.jsonl - summary_path: outputs/async_rl/compare_single_lora_gsm8k_async_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: true - -model: - strategy: native_fsdp - attn_implementation: flash_attention_2 - fsdp_config: - reshard_after_forward: true - sequence_parallel_size: 1 - padding_free: false - max_length: 2048 - -sampler: - max_model_len: 2048 - max_num_batched_tokens: 4096 - gpu_memory_utilization: 0.7 - max_num_seqs: 64 - enforce_eager: false - -rollout_output: - enabled: true - output_dir: ${runtime.output_dir}/rollouts - include_token_ids: false - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - train: {policy: sticky, max_consecutive_units: null} - -evaluation: - enabled: false - interval: 10 - batch_size: 16 - sampling_params: - max_tokens: 1024 - temperature: 0.6 - top_p: 1.0 - -lora: - target_modules: all-linear - r: 16 - alpha: 16 - dropout: 0.0 - learning_rate: 1.7e-5 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - - tenant_id: tenant_single - training_run_id: gsm8k_compare - adapter_name: gsm8k_compare_lora - reward: - class_path: twinkle.reward.GSM8KAccuracyReward - dataset: - dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID,ms://modelscope/gsm8k} - subset_name: main - split: train - data_num: 2000 - max_length: 1024 - processor: GSM8KProcessor - system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' - eval_dataset: - name: gsm8k/test - dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID,ms://modelscope/gsm8k} - subset_name: main - split: test - data_num: null - max_length: 1024 - processor: GSM8KProcessor - system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' - reward: - class_path: twinkle.reward.GSM8KAccuracyReward - rollout: - batch_size: 16 - num_generations: 4 - max_tokens: 1024 - temperature: 1.0 - top_p: 1.0 - train: - mini_batch_size: 64 - micro_batch_size: 64 - dynamic_batching: true - max_tokens_per_micro_batch: 4096 - packing_algorithm: ffd diff --git a/cookbook/rl/compare_single_lora_gsm8k_sync.yaml b/cookbook/rl/compare_single_lora_gsm8k_sync.yaml deleted file mode 100644 index 60dcdcee8..000000000 --- a/cookbook/rl/compare_single_lora_gsm8k_sync.yaml +++ /dev/null @@ -1,128 +0,0 @@ -# Fully synchronous barrier baseline for comparison with -# compare_single_lora_gsm8k_async.yaml. Run this file with -# sync_barrier_multi_lora_grpo.py, not async_multi_lora_grpo.py. -runtime: - run_id: compare_single_lora_gsm8k_sync - model_id: ${oc.env:MODEL_ID,/nas/disk1/Qwen3-4B} - mode: ray - model_gpus: 1 - sampler_gpus: 1 - sampler_tp: 1 - sampler_max_loras: 1 - seed: 1 - max_steps: 125 - allow_partial_rollout: false - rollout_max_retries: 2 - rollout_retry_delay_s: 0.5 - keep_adapter_versions: 2 - output_dir: output/compare_single_lora_gsm8k_sync - -metrics: - enabled: true - drain_interval_s: 1.0 - queue_capacity: 10000 - close_timeout_s: 10 - jsonl: - enabled: true - path: outputs/async_rl/compare_single_lora_gsm8k_sync_metrics.jsonl - summary_path: outputs/async_rl/compare_single_lora_gsm8k_sync_summary.json - batch_size: 64 - flush_interval_s: 2.0 - swanlab: - enabled: false - mode: local - project: twinkle-rl - name: ${runtime.run_id} - log_dir: outputs/swanlab - -template: - cls: Template - enable_thinking: true - -model: - strategy: native_fsdp - attn_implementation: flash_attention_2 - fsdp_config: - reshard_after_forward: true - sequence_parallel_size: 1 - padding_free: false - max_length: 2048 - -sampler: - max_model_len: 2048 - max_num_batched_tokens: 4096 - gpu_memory_utilization: 0.7 - max_num_seqs: 64 - enforce_eager: false - -rollout_output: - enabled: true - output_dir: ${runtime.output_dir}/rollouts - include_token_ids: false - -tq: - polling_mode: true - storage_units: 2 - -scheduler: - rollout: {policy: round_robin, max_consecutive_units: 1} - advantage: {policy: oldest_partition, max_consecutive_units: 1} - train: {policy: sticky, max_consecutive_units: null} - -evaluation: - enabled: false - interval: 10 - batch_size: 16 - sampling_params: - max_tokens: 1024 - temperature: 0.6 - top_p: 1.0 - -lora: - target_modules: all-linear - r: 16 - alpha: 16 - dropout: 0.0 - learning_rate: 1.7e-5 - -loss: - cls: GRPOLoss - epsilon: 0.2 - -lora_contexts: - - tenant_id: tenant_single - training_run_id: gsm8k_compare - adapter_name: gsm8k_compare_lora - reward: - class_path: twinkle.reward.GSM8KAccuracyReward - dataset: - dataset_id: ${oc.env:GSM8K_TRAIN_DATASET_ID,ms://modelscope/gsm8k} - subset_name: main - split: train - data_num: 2000 - max_length: 1024 - processor: GSM8KProcessor - system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' - eval_dataset: - name: gsm8k/test - dataset_id: ${oc.env:GSM8K_TEST_DATASET_ID,ms://modelscope/gsm8k} - subset_name: main - split: test - data_num: null - max_length: 1024 - processor: GSM8KProcessor - system_prompt: 'You are a helpful math assistant. Solve the problem step by step and put your final answer within \boxed{}.' - reward: - class_path: twinkle.reward.GSM8KAccuracyReward - rollout: - batch_size: 16 - num_generations: 4 - max_tokens: 1024 - temperature: 1.0 - top_p: 1.0 - train: - mini_batch_size: 64 - micro_batch_size: 64 - dynamic_batching: true - max_tokens_per_micro_batch: 4096 - packing_algorithm: ffd diff --git a/cookbook/rl/sync_barrier_multi_lora_grpo.py b/cookbook/rl/sync_barrier_multi_lora_grpo.py deleted file mode 100644 index 55f2a5ce4..000000000 --- a/cookbook/rl/sync_barrier_multi_lora_grpo.py +++ /dev/null @@ -1,689 +0,0 @@ -"""Synchronous barrier baseline for native async multi-LoRA GRPO. - -The model, sampler, datasets, rewards, batch semantics, and checkpoint cadence -match ``async_multi_lora_grpo.py``. The only intentional difference is the -execution schedule: every round finishes rollout for all active contexts -before any context starts training. -""" - -from __future__ import annotations - -import argparse -import os -import shutil -import time -from dataclasses import dataclass -from typing import Any, Iterator, Sequence - -from omegaconf import OmegaConf - -from twinkle.metric import MetricRecord, create_metrics_reporter -from twinkle_agentic.async_rl.metrics import advantage_signal_metrics, rollout_metrics -from twinkle_agentic.async_rl.pipeline import (_prompt_batches, _reward_for_context, _train_batch) -from twinkle_agentic.async_rl.tq_utils import REQUIRED_MODEL_INPUT_FIELDS, columns_to_tq_fields -from twinkle_agentic.async_rl.types import LoraContext, PartitionAdmission -from twinkle_agentic.async_rl.utils import ( - TrainBatchConfig, - build_native_fsdp_model_kwargs, - configure_lora_lr_scheduler, - resolve_context_learning_rate, - resolve_context_lora_target_modules, - resolve_context_loss_config, - resolve_model_attention_implementation, - resolve_sequence_parallel_size, - sample_responses_to_rollout_rows, - sampler_data_parallel_size, - validate_context_batch_config, -) -from twinkle_agentic.async_rl.vllm_sampler_tq import _compute_reward_metrics - - -@dataclass -class SyncContextState: - context: LoraContext - prompt_batches: Iterator[Sequence[dict[str, Any]]] - rollout_batch_size: int - num_generations: int - sampling_params: Any - mini_batch_size: int - reward_fn: Any - adapter_path: str - adapter_history: list[str] - partition_step: int = 0 - optimizer_steps: int = 0 - policy_version: int = 0 - exhausted: bool = False - - -@dataclass -class SyncPartition: - admission: PartitionAdmission - state: SyncContextState - rows: list[dict[str, Any]] - rewards: list[float] - advantages: list[float] | None = None - - -class SyncBarrierMultiLoraGRPO: - - def __init__(self, raw_config: dict[str, Any]): - import twinkle - from peft import LoraConfig - from twinkle import DeviceGroup, DeviceMesh - from twinkle.data_format import SamplingParams - from twinkle.model import MultiLoraTransformersModel - from twinkle.processor import InputProcessor - from twinkle.sampler import vLLMSampler - - raw_config = OmegaConf.to_container(OmegaConf.create(raw_config), resolve=True) - if not isinstance(raw_config, dict): - raise TypeError('sync RL config must resolve to a mapping') - - runtime = raw_config['runtime'] - model_config = raw_config['model'] - lora_data = raw_config['lora'] - loss_data = raw_config.get('loss') - template_data = raw_config.get('template', {}) - template_cls = template_data.get('cls', 'Qwen3_5Template') - enable_thinking = bool(template_data.get('enable_thinking', False)) - model_gpus = int(runtime['model_gpus']) - sampler_gpus = int(runtime['sampler_gpus']) - sampler_tp = int(runtime['sampler_tp']) - sampler_dp = sampler_data_parallel_size(sampler_gpus, sampler_tp) - sequence_parallel_size = resolve_sequence_parallel_size( - model_gpus, - int(model_config['sequence_parallel_size']), - ) - padding_free = bool(model_config['padding_free']) - attn_implementation = resolve_model_attention_implementation( - model_config, - padding_free=padding_free, - sequence_parallel_size=sequence_parallel_size, - ) - model_max_length = int(model_config['max_length']) - sampler_config = raw_config['sampler'] - total_gpus = model_gpus + sampler_gpus - - twinkle.initialize( - mode='ray', - nproc_per_node=total_gpus, - groups=[ - DeviceGroup('model', list(range(model_gpus)), device_type='GPU'), - DeviceGroup( - 'sampler', - list(range(model_gpus, total_gpus)), - device_type='GPU', - gpus_per_worker=sampler_tp, - ), - ], - lazy_collect=False, - ) - model_mesh = DeviceMesh.from_sizes( - world_size=model_gpus, - dp_size=model_gpus, - ulysses_size=sequence_parallel_size, - ) - model_data_parallel_size = model_mesh.data_world_size - self.model_data_parallel_size = model_data_parallel_size - sampler_mesh = DeviceMesh.from_sizes( - world_size=sampler_gpus, - dp_size=sampler_dp, - tp_size=sampler_tp, - ) - model_kwargs = build_native_fsdp_model_kwargs(model_config) - if attn_implementation is not None: - model_kwargs['attn_implementation'] = attn_implementation - self.model = MultiLoraTransformersModel( - model_id=runtime['model_id'], - device_mesh=model_mesh, - remote_group='model', - max_length=model_max_length, - **model_kwargs, - ) - self.train_batch_configs: dict[str, TrainBatchConfig] = {} - self.states: list[SyncContextState] = [] - self.evaluation_configs: dict[str, dict[str, Any]] = {} - self._evaluation_batches: dict[str, list[Sequence[dict[str, Any]]]] = {} - global_evaluation = dict(raw_config.get('evaluation') or {}) - for item in raw_config['lora_contexts']: - context = LoraContext( - item['tenant_id'], - item['training_run_id'], - runtime['model_id'], - item['adapter_name'], - ) - rollout = item['rollout'] - train = item['train'] - rollout_batch_size = int(rollout['batch_size']) - num_generations = int(rollout['num_generations']) - train_batch_config = TrainBatchConfig( - mini_batch_size=int(train['mini_batch_size']), - micro_batch_size=int(train['micro_batch_size']), - dynamic_batching=bool(train.get('dynamic_batching', False)), - max_tokens_per_micro_batch=( - int(train['max_tokens_per_micro_batch']) - if train.get('max_tokens_per_micro_batch') is not None else None - ), - packing_algorithm=str(train.get('packing_algorithm', 'ffd')), - ) - validate_context_batch_config( - context.key, - rollout_groups=rollout_batch_size, - num_generations=num_generations, - train=train_batch_config, - sampler_dp=sampler_dp, - model_dp=model_data_parallel_size, - ) - adapter_lora_config = LoraConfig( - target_modules=resolve_context_lora_target_modules(item, lora_data), - r=lora_data['r'], - lora_alpha=lora_data['alpha'], - lora_dropout=lora_data['dropout'], - ) - self.model.add_adapter_to_model( - context.adapter_name, - adapter_lora_config, - gradient_accumulation_steps=1, - ) - self.model.set_optimizer( - 'AdamW', - lr=resolve_context_learning_rate(train, lora_data), - adapter_name=context.adapter_name, - ) - configure_lora_lr_scheduler(self.model, context.adapter_name, lora_data) - loss_cls, loss_kwargs = resolve_context_loss_config(item, loss_data) - self.model.set_loss( - loss_cls, - adapter_name=context.adapter_name, - **loss_kwargs, - ) - self.model.set_processor( - InputProcessor, - adapter_name=context.adapter_name, - padding_free=padding_free, - ) - self.model.set_template( - template_cls, - model_id=runtime['model_id'], - adapter_name=context.adapter_name, - enable_thinking=enable_thinking, - max_length=model_max_length, - ) - initial_path = self.model.save( - f'sync-{context.adapter_name}-initial', - output_dir=runtime['output_dir'], - adapter_name=context.adapter_name, - ) - state = SyncContextState( - context=context, - prompt_batches=iter( - _prompt_batches( - item['dataset'], - model_id=runtime['model_id'], - batch_size=rollout_batch_size, - template_cls=template_cls, - enable_thinking=enable_thinking, - )), - rollout_batch_size=rollout_batch_size, - num_generations=num_generations, - sampling_params=SamplingParams( - max_tokens=rollout['max_tokens'], - temperature=rollout['temperature'], - top_p=rollout['top_p'], - repetition_penalty=float(rollout.get('repetition_penalty', 1.0)), - logprobs=1, - num_samples=1, - ), - mini_batch_size=train_batch_config.mini_batch_size, - reward_fn=_reward_for_context( - item.get('reward'), - context_key=context.key, - ), - adapter_path=initial_path, - adapter_history=[initial_path], - ) - self.states.append(state) - self.train_batch_configs[context.key] = train_batch_config - if bool(global_evaluation.get('enabled', False)): - eval_dataset = item.get('eval_dataset') - if eval_dataset is None: - raise ValueError(f'eval_dataset is required for periodic evaluation of {context.key}') - eval_batch_size = int(global_evaluation.get('batch_size', 16)) - eval_interval = int(global_evaluation.get('interval', 1)) - if eval_batch_size <= 0 or eval_interval <= 0: - raise ValueError('evaluation.batch_size and evaluation.interval must be positive') - eval_sampling = dict(global_evaluation.get('sampling_params') or {}) - self.evaluation_configs[context.key] = { - 'interval': eval_interval, - 'dataset_name': eval_dataset.get('name', eval_dataset['dataset_id']), - 'prompt_batches': _prompt_batches( - eval_dataset, - model_id=runtime['model_id'], - batch_size=eval_batch_size, - template_cls=template_cls, - enable_thinking=enable_thinking, - full_batches_only=False, - ), - 'sampling_params': SamplingParams( - max_tokens=int(eval_sampling.get('max_tokens', rollout['max_tokens'])), - temperature=float(eval_sampling.get('temperature', 0.0)), - top_p=float(eval_sampling.get('top_p', 1.0)), - repetition_penalty=float(eval_sampling.get('repetition_penalty', 1.0)), - logprobs=0, - num_samples=1, - ), - 'reward_fn': _reward_for_context( - eval_dataset.get('reward'), - context_key=f'{context.key} evaluation', - ), - } - - sampler_engine_args = { - 'tensor_parallel_size': sampler_tp, - 'enable_lora': True, - 'max_loras': int(runtime['sampler_max_loras']), - 'max_lora_rank': lora_data['r'], - 'max_model_len': int(sampler_config['max_model_len']), - 'gpu_memory_utilization': float(sampler_config['gpu_memory_utilization']), - 'max_num_seqs': int(sampler_config['max_num_seqs']), - 'enforce_eager': bool(sampler_config['enforce_eager']), - } - if sampler_config.get('max_num_batched_tokens') is not None: - sampler_engine_args['max_num_batched_tokens'] = int(sampler_config['max_num_batched_tokens']) - self.sampler = vLLMSampler( - model_id=runtime['model_id'], - remote_group='sampler', - device_mesh=sampler_mesh, - engine_args=sampler_engine_args, - ) - self.sampler.set_template( - template_cls, - model_id=runtime['model_id'], - enable_thinking=enable_thinking, - max_length=model_max_length, - ) - self.output_dir = runtime['output_dir'] - self.max_steps = runtime.get('max_steps') - self.max_steps = None if self.max_steps is None else int(self.max_steps) - self.keep_adapter_versions = max(0, int(runtime.get('keep_adapter_versions', 0))) - self.metrics = create_metrics_reporter( - raw_config.get('metrics'), - run_id=str(runtime.get('run_id', 'sync_barrier_multi_lora_grpo')), - ) - self.completed_partitions = 0 - self._creation_order = 0 - - def _record_metric( - self, - stage: str, - *, - admission: PartitionAdmission | None = None, - context: LoraContext | None = None, - values: dict[str, Any] | None = None, - status: str = 'completed', - attributes: dict[str, Any] | None = None, - optimizer_step: int | None = None, - policy_version: int | None = None, - ) -> None: - if self.metrics is None: - return - self.metrics.record(MetricRecord( - stage=stage, - values=dict(values or {}), - context_key=( - admission.context.key if admission is not None - else context.key if context is not None else None - ), - partition_id=admission.partition_id if admission is not None else None, - partition_index=admission.step if admission is not None else None, - optimizer_step=optimizer_step, - policy_version=policy_version, - status=status, - attributes=dict(attributes or {}), - )) - - def run(self) -> dict[str, Any]: - started = time.perf_counter() - try: - while self.max_steps is None or self.completed_partitions < self.max_steps: - partitions = self._rollout_round() - if not partitions: - break - self._advantage_round(partitions) - self._train_round(partitions) - except Exception as exc: - self._record_metric( - 'run', - status='failed', - values={'wall_time_s': time.perf_counter() - started}, - attributes={'error': f'{type(exc).__name__}: {exc}'}, - ) - if self.metrics is not None: - self.metrics.close() - raise - result = { - 'trained_partitions': self.completed_partitions, - 'wall_time_s': time.perf_counter() - started, - 'per_context': { - state.context.key: { - 'optimizer_steps': state.optimizer_steps, - 'policy_version': state.policy_version, - 'adapter_path': state.adapter_path, - } - for state in self.states - }, - } - self._record_metric( - 'run', - values={ - 'trained_partitions': result['trained_partitions'], - 'wall_time_s': result['wall_time_s'], - }, - ) - if self.metrics is not None: - self.metrics.flush() - result['metrics_health'] = self.metrics.health() - self.metrics.close() - return result - - def _rollout_round(self) -> list[SyncPartition]: - partitions = [] - for state in self.states: - if state.exhausted: - continue - if self.max_steps is not None and self.completed_partitions + len(partitions) >= self.max_steps: - break - prompts = next(state.prompt_batches, None) - if prompts is None or len(prompts) != state.rollout_batch_size: - state.exhausted = True - continue - admission = PartitionAdmission( - context=state.context, - partition_id=state.context.partition_id(state.partition_step), - step=state.partition_step, - target_groups=state.rollout_batch_size, - num_generations=state.num_generations, - created_order=self._creation_order, - ) - self._creation_order += 1 - self._record_metric( - 'rollout', - admission=admission, - status='submitted', - policy_version=state.policy_version, - values={ - 'prompt_count': admission.target_groups, - 'sample_count': admission.sample_count, - 'num_generations': admission.num_generations, - }, - attributes={'scope': 'partition'}, - ) - rollout_started = time.perf_counter() - sources = [{ - **dict(prompt), - 'group_id': f'{admission.partition_id}/group_{group_index}', - 'generation_idx': generation_index, - } for group_index, prompt in enumerate(prompts) - for generation_index in range(state.num_generations)] - responses = self.sampler.sample( - [dict(prompt) for prompt in prompts for _ in range(state.num_generations)], - state.sampling_params, - adapter_name=state.context.adapter_name, - adapter_path=state.adapter_path, - ) - rows = sample_responses_to_rollout_rows( - sources, - responses, - policy_version=state.policy_version, - ) - if len(rows) != admission.sample_count: - raise ValueError( - f'{admission.partition_id} expected {admission.sample_count} samples, got {len(rows)}') - for row in rows: - row.update({ - 'rollout_adapter_path': state.adapter_path, - 'rollout_policy_versions': [state.policy_version], - 'initial_policy_version': state.policy_version, - 'final_policy_version': state.policy_version, - 'policy_version_span': 0, - }) - rewards = [float(value) for value in state.reward_fn(rows, context=state.context)] - if len(rewards) != len(rows): - raise ValueError(f'{admission.partition_id} reward count does not match sample count') - rollout_latency_s = time.perf_counter() - rollout_started - self._record_rollout_groups(state, admission, rows, rewards) - self._record_metric( - 'rollout', - admission=admission, - policy_version=state.policy_version, - values=rollout_metrics( - completion_lengths=[int(row['completion_length']) for row in rows], - stop_reasons=[row.get('stop_reason') for row in rows], - rollout_latency_s=rollout_latency_s, - ), - attributes={'scope': 'partition'}, - ) - partitions.append(SyncPartition(admission, state, rows, rewards)) - state.partition_step += 1 - return partitions - - def _record_rollout_groups( - self, - state: SyncContextState, - admission: PartitionAdmission, - rows: list[dict[str, Any]], - rewards: list[float], - ) -> None: - for group_index in range(admission.target_groups): - start = group_index * admission.num_generations - end = start + admission.num_generations - group_rows = rows[start:end] - group_rewards = rewards[start:end] - metrics = { - **_compute_reward_metrics( - {state.context.key: state.reward_fn}, - state.context, - group_rows, - group_rewards, - ), - **rollout_metrics( - rewards={'reward': group_rewards}, - completion_lengths=[int(row['completion_length']) for row in group_rows], - stop_reasons=[row.get('stop_reason') for row in group_rows], - ), - } - self._record_metric( - 'rollout', - admission=admission, - policy_version=state.policy_version, - values=metrics, - attributes={ - 'scope': 'group', - 'group_id': f'{admission.partition_id}/group_{group_index}', - }, - ) - - def _advantage_round(self, partitions: list[SyncPartition]) -> None: - from twinkle.advantage import GRPOAdvantage - - advantage_fn = GRPOAdvantage() - for partition in partitions: - admission = partition.admission - partition.advantages = advantage_fn( - partition.rewards, - num_generations=admission.num_generations, - scale='group', - ).tolist() - samples_per_batch = admission.num_generations - for start in range(0, len(partition.rows), samples_per_batch): - end = min(start + samples_per_batch, len(partition.rows)) - self._record_metric( - 'advantage', - admission=admission, - policy_version=partition.state.policy_version, - values={ - 'sample_count': end - start, - **advantage_signal_metrics( - partition.rewards[start:end], - partition.advantages[start:end], - num_generations=admission.num_generations, - ), - }, - ) - - def _train_round(self, partitions: list[SyncPartition]) -> None: - for partition in partitions: - admission = partition.admission - state = partition.state - assert partition.advantages is not None - samples_per_batch = state.mini_batch_size - for start in range(0, len(partition.rows), samples_per_batch): - end = start + samples_per_batch - batch = self._training_batch( - partition.rows[start:end], - partition.rewards[start:end], - partition.advantages[start:end], - ) - train_started = time.perf_counter() - metrics = _train_batch( - self.model, - self.train_batch_configs, - batch, - admission, - model_data_parallel_size=self.model_data_parallel_size, - ) - state.optimizer_steps += 1 - metrics.update({ - 'sample_count': end - start, - 'train_latency_s': time.perf_counter() - train_started, - 'policy_version_gap_mean': 0.0, - 'policy_version_gap_p95': 0.0, - 'policy_version_gap_max': 0, - 'rollout_policy_span_mean': 0.0, - 'rollout_policy_span_max': 0, - }) - self._record_metric( - 'train', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values=metrics, - ) - finalize_started = time.perf_counter() - next_policy_version = state.policy_version + 1 - save_started = time.perf_counter() - state.adapter_path = self.model.save( - f'sync-{state.context.adapter_name}-v{next_policy_version}', - output_dir=self.output_dir, - adapter_name=state.context.adapter_name, - ) - adapter_save_latency_s = time.perf_counter() - save_started - publish_started = time.perf_counter() - state.policy_version = next_policy_version - policy_publish_latency_s = time.perf_counter() - publish_started - state.adapter_history.append(state.adapter_path) - self._record_metric( - 'policy', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values={ - 'adapter_save_latency_s': adapter_save_latency_s, - 'policy_publish_latency_s': policy_publish_latency_s, - }, - attributes={'operation': 'publish', 'adapter_path': state.adapter_path}, - ) - self._evaluate_policy(state, admission) - prune_started = time.perf_counter() - self._prune_adapter_history(state) - adapter_prune_latency_s = time.perf_counter() - prune_started - self.completed_partitions += 1 - self._record_metric( - 'partition', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values={ - 'adapter_save_latency_s': adapter_save_latency_s, - 'policy_publish_latency_s': policy_publish_latency_s, - 'adapter_prune_latency_s': adapter_prune_latency_s, - 'partition_finalize_latency_s': time.perf_counter() - finalize_started, - }, - ) - - def _evaluate_policy(self, state: SyncContextState, admission: PartitionAdmission) -> None: - config = self.evaluation_configs.get(state.context.key) - if config is None or state.policy_version % int(config['interval']): - return - if state.context.key not in self._evaluation_batches: - self._evaluation_batches[state.context.key] = list(config['prompt_batches']) - - started = time.perf_counter() - rewards: list[float] = [] - completion_lengths: list[int] = [] - prompt_count = 0 - for batch in self._evaluation_batches[state.context.key]: - prompts = list(batch) - responses = self.sampler.sample( - prompts, - config['sampling_params'], - adapter_name=state.context.adapter_name, - adapter_path=state.adapter_path, - ) - rows = sample_responses_to_rollout_rows( - prompts, - responses, - policy_version=state.policy_version, - ) - rewards.extend(float(value) for value in config['reward_fn'](rows, context=state.context)) - completion_lengths.extend(int(row['completion_length']) for row in rows) - prompt_count += len(prompts) - if not rewards: - raise ValueError(f'evaluation dataset is empty for {state.context.key}') - self._record_metric( - 'evaluation', - admission=admission, - optimizer_step=state.optimizer_steps, - policy_version=state.policy_version, - values={ - 'accuracy': sum(rewards) / len(rewards), - 'sample_count': len(rewards), - 'prompt_count': prompt_count, - 'completion_length': sum(completion_lengths) / len(completion_lengths), - 'eval_latency_s': time.perf_counter() - started, - }, - attributes={'eval_dataset': config['dataset_name']}, - ) - - @staticmethod - def _training_batch(rows: list[dict[str, Any]], rewards: list[float], advantages: list[float]): - fields = { - name: [row[name] for row in rows] - for name in (*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs') - } - fields.update({'rewards': rewards, 'advantages': advantages}) - return columns_to_tq_fields(fields, len(rows)) - - def _prune_adapter_history(self, state: SyncContextState) -> None: - retained_count = max(1, self.keep_adapter_versions) - stale = state.adapter_history[:-retained_count] - state.adapter_history = state.adapter_history[-retained_count:] - for path in stale: - if os.path.isdir(path): - shutil.rmtree(path) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument('--config', default='cookbook/rl/sync_barrier_multi_lora_grpo.yaml') - args = parser.parse_args() - config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) - print(SyncBarrierMultiLoraGRPO(config).run()) - - -if __name__ == '__main__': - main() - -# MODEL_ID=/path/to/model \ -# DATASET_ID=/path/to/gsm8k \ -# python cookbook/rl/async_multi_lora_grpo.py diff --git a/tests/model/test_multi_lora.py b/tests/model/test_multi_lora.py index f2795a5ad..8ae68c5ac 100644 --- a/tests/model/test_multi_lora.py +++ b/tests/model/test_multi_lora.py @@ -3,6 +3,12 @@ from twinkle.model.multi_lora import MultiLora +def test_multi_lora_transformers_save_disables_lazy_collect(): + from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + assert MultiLoraTransformersModel.save._lazy_collect is False + + def test_check_length_checks_each_sample_independently(): multi_lora = MultiLora(max_length=4) diff --git a/tests/twinkle_agentic/test_async_rl_config.py b/tests/twinkle_agentic/test_async_rl_config.py new file mode 100644 index 000000000..64d6f5760 --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_config.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import pytest + +from twinkle_agentic.async_rl.pipeline import _reward_for_context +from twinkle_agentic.async_rl.utils import ( + TrainBatchConfig, + build_native_fsdp_model_kwargs, + configure_lora_lr_scheduler, + resolve_context_learning_rate, + resolve_context_lora_target_modules, + resolve_context_loss_config, + resolve_model_attention_implementation, + resolve_sequence_parallel_size, + sampler_data_parallel_size, + validate_context_batch_config, +) + + +def test_sequence_parallel_size_must_divide_model_gpus(): + assert resolve_sequence_parallel_size(2, 1) == 1 + assert resolve_sequence_parallel_size(2, 2) == 2 + + with pytest.raises(ValueError, match='must be divisible'): + resolve_sequence_parallel_size(2, 3) + + +def test_padding_free_sequence_parallel_requires_flash_attention(): + assert resolve_model_attention_implementation( + {'attn_implementation': 'flash_attention_2'}, + padding_free=True, + sequence_parallel_size=2, + ) == 'flash_attention_2' + + with pytest.raises(ValueError, match='model.attn_implementation'): + resolve_model_attention_implementation({}, padding_free=True, sequence_parallel_size=2) + + +@pytest.mark.parametrize( + ('sampler_gpus', 'sampler_tp', 'expected_dp'), + [(8, 2, 4), (1, 1, 1)], +) +def test_sampler_data_parallel_size(sampler_gpus, sampler_tp, expected_dp): + assert sampler_data_parallel_size(sampler_gpus, sampler_tp) == expected_dp + + +def test_sampler_parallelism_rejects_incomplete_tp_group(): + with pytest.raises(ValueError, match='must be divisible'): + sampler_data_parallel_size(3, 2) + + +def test_lora_lr_scheduler_uses_shared_adapter_config(): + calls = [] + + class Model: + def set_lr_scheduler(self, scheduler_cls, **kwargs): + calls.append((scheduler_cls, kwargs)) + + configure_lora_lr_scheduler( + Model(), + 'tenant_lora', + { + 'lr_scheduler': { + 'cls': 'CosineAnnealingLR', + 'T_max': 2000, + 'eta_min': 0.0, + }, + }, + ) + + assert calls == [('CosineAnnealingLR', { + 'adapter_name': 'tenant_lora', + 'T_max': 2000, + 'eta_min': 0.0, + })] + + +def test_context_learning_rate_overrides_global_default(): + assert resolve_context_learning_rate({'learning_rate': 5e-6}, {'learning_rate': 1e-6}) == pytest.approx(5e-6) + assert resolve_context_learning_rate({}, {'learning_rate': 1e-6}) == pytest.approx(1e-6) + + +@pytest.mark.parametrize('value', [0, -1e-6, float('inf')]) +def test_context_learning_rate_rejects_invalid_values(value): + with pytest.raises(ValueError, match='positive finite'): + resolve_context_learning_rate({'learning_rate': value}, {'learning_rate': 1e-6}) + + +def test_context_lora_target_modules_override_global_default(): + defaults = {'target_modules': 'all-linear'} + + assert resolve_context_lora_target_modules({}, defaults) == 'all-linear' + assert resolve_context_lora_target_modules( + {'lora': {'target_modules': ['q_proj', 'v_proj']}}, + defaults, + ) == ['q_proj', 'v_proj'] + + +@pytest.mark.parametrize('value', ['', [], [None], {'q_proj': True}]) +def test_context_lora_target_modules_reject_invalid_values(value): + with pytest.raises(ValueError, match='target_modules'): + resolve_context_lora_target_modules( + {'lora': {'target_modules': value}}, + {'target_modules': 'all-linear'}, + ) + + +def test_context_loss_config_overrides_global_defaults(): + loss_cls, loss_kwargs = resolve_context_loss_config( + {'loss': {'cls': 'GSPOLoss', 'epsilon_high': 0.3}}, + {'cls': 'GRPOLoss', 'epsilon': 0.2}, + ) + + assert loss_cls == 'GSPOLoss' + assert loss_kwargs == {'epsilon': 0.2, 'epsilon_high': 0.3} + + +def test_context_loss_config_uses_grpo_defaults(): + assert resolve_context_loss_config({}) == ('GRPOLoss', {'epsilon': 0.2}) + + +def test_context_loss_config_rejects_empty_class_name(): + with pytest.raises(ValueError, match='loss.cls'): + resolve_context_loss_config({'loss': {'cls': ''}}) + + +def test_rl_model_kwargs_enforce_native_fsdp(): + assert build_native_fsdp_model_kwargs({}) == { + 'strategy': 'native_fsdp', + 'fsdp_config': {}, + } + assert build_native_fsdp_model_kwargs({ + 'strategy': 'native_fsdp', + 'fsdp_config': {'reshard_after_forward': False}, + }) == { + 'strategy': 'native_fsdp', + 'fsdp_config': {'reshard_after_forward': False}, + } + with pytest.raises(ValueError, match='must be native_fsdp'): + build_native_fsdp_model_kwargs({'strategy': 'accelerate'}) + + +def test_reward_factory_loads_class_and_resolved_kwargs(): + reward = _reward_for_context( + { + 'class_path': 'twinkle.reward.DAPOMathReward', + 'kwargs': { + 'max_response_length': 8192, + 'overlong_buffer_length': 4096, + 'overlong_penalty_factor': 1.0, + 'score_tail_chars': 300, + }, + }, + context_key='tenant/run/adapter', + ) + + assert reward.max_response_length == 8192 + assert reward.overlong_buffer_length == 4096 + + +def test_reward_factory_rejects_non_reward_class(): + with pytest.raises(TypeError, match='Reward subclass'): + _reward_for_context( + {'class_path': 'collections.Counter'}, + context_key='tenant/run/adapter', + ) + + +def test_context_batch_config_accepts_group_aligned_dp_batches(): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=8, micro_batch_size=2), + sampler_dp=2, + model_dp=2, + ) + + +def test_context_batch_config_allows_training_group_to_span_model_dp_ranks(): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=16, micro_batch_size=2), + sampler_dp=1, + model_dp=8, + ) + + +def test_context_batch_config_rejects_partition_tail_and_undersized_rank_batch(): + with pytest.raises(ValueError, match='complete prompt groups'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=6, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=6, micro_batch_size=1), + sampler_dp=2, + model_dp=2, + ) + + with pytest.raises(ValueError, match='per-rank train batch'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=2, + train=TrainBatchConfig(mini_batch_size=2, micro_batch_size=2), + sampler_dp=2, + model_dp=2, + ) + + +def test_context_batch_config_requires_token_limit_for_dynamic_batching(): + with pytest.raises(ValueError, match='max_tokens_per_micro_batch'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig( + mini_batch_size=8, + micro_batch_size=2, + dynamic_batching=True, + ), + sampler_dp=1, + model_dp=1, + ) diff --git a/tests/twinkle_agentic/test_async_rl_data_plane.py b/tests/twinkle_agentic/test_async_rl_data_plane.py new file mode 100644 index 000000000..05208849b --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_data_plane.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle_agentic.async_rl import LoraContext, TQDataPlane +from twinkle_agentic.async_rl.data_plane import build_rollout_group_sample_write +from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup + + +def _context() -> LoraContext: + return LoraContext('tenant', 'run_adapter', 'model', 'adapter') + + +def test_rollout_sample_tags_use_new_context_descriptor_only(): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, batch_meta=None) + fields, tags = build_rollout_group_sample_write( + group, + [ + { + 'generation_idx': 0, + 'labels': [-100, 1], + 'logprobs': [-.1], + 'rollout_policy_version': 3, + 'rollout_adapter_path': 'adapter-v3', + }, + { + 'generation_idx': 1, + 'labels': [-100, 2], + 'logprobs': [-.2], + 'rollout_policy_version': 4, + 'rollout_adapter_path': 'adapter-v4', + }, + ], + rewards=[1., 0.], + expected_num_generations=2, + ) + assert [row['rewards'] for row in fields] == [1., 0.] + assert [tag['generation_idx'] for tag in tags] == [0, 1] + assert all(tag['context_key'] == context.key for tag in tags) + assert [tag['rollout_policy_version'] for tag in tags] == [3, 4] + + +def test_data_plane_completes_rollout_with_full_training_trajectory(): + class Metadata: + def __init__(self): + self.size = 2 + self.custom_meta = [{}, {}] + + def update_custom_meta(self, updates): + for tag, update in zip(self.custom_meta, updates): + tag.update(update) + + class Client: + def __init__(self): + self.written = None + self.calls = [] + + async def async_put(self, data, metadata=None, partition_id=None): + self.calls.append('fields') + self.written = data + return metadata + + async def async_set_custom_meta(self, _metadata): + self.calls.append('tags') + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) + metadata = Metadata() + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) + client = Client() + rows = [{ + 'input_ids': [1, 2, token], + 'labels': [-100, -100, token], + 'attention_mask': [1, 1, 1], + 'position_ids': [0, 1, 2], + 'logprobs': [-.1], + 'generation_idx': generation_idx, + 'rollout_policy_version': 3, + 'rollout_policy_versions': [3], + 'initial_policy_version': 3, + 'final_policy_version': 3, + 'policy_version_span': 0, + 'rollout_adapter_path': 'adapter-v3', + 'completion_length': 1, + } for generation_idx, token in enumerate((7, 8))] + + asyncio.run( + TQDataPlane(client).complete_rollout_group( + group, + rollout_rows=rows, + rewards=[1., 0.], + submission_id='submission', + )) + + assert set(client.written.keys()) == { + 'input_ids', 'labels', 'attention_mask', 'position_ids', 'logprobs', 'rewards' + } + assert client.calls == ['tags', 'fields'] + assert [tag['rollout_status'] for tag in metadata.custom_meta] == ['ROLLOUT_DONE', 'ROLLOUT_DONE'] + assert [tag['submission_id'] for tag in metadata.custom_meta] == ['submission', 'submission'] + + +def test_data_plane_rejects_rollout_without_complete_model_inputs(): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 1, 0) + metadata = type('Metadata', (), {'size': 1})() + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) + row = { + 'input_ids': [1, 2], + 'labels': [-100, 2], + 'logprobs': [-.1], + 'generation_idx': 0, + 'rollout_policy_version': 0, + } + + with pytest.raises(ValueError) as error: + asyncio.run( + TQDataPlane(object()).complete_rollout_group( + group, + rollout_rows=[row], + rewards=[1.], + submission_id='submission', + )) + + assert 'attention_mask' in str(error.value) + assert 'position_ids' in str(error.value) diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py index e2e65db25..8f6479b7d 100644 --- a/tests/twinkle_agentic/test_async_rl_native_tq.py +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -2,46 +2,25 @@ import asyncio import inspect -import json -import time import pytest -from twinkle import DeviceMesh -from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams -from twinkle.infra import _dispatch_args +from twinkle.data_format import SampledSequence, SampleResponse from twinkle.metric import MetricRecord from twinkle_agentic.async_rl import (AsyncMultiLoraGRPOPipeline, ContextSchedulePolicy, ContextScheduler, ContextStatus, LoraContext, LoraContextManager, ScheduleCandidate, SchedulerConfig, TQDataPlane, TrainerWorker) -from twinkle_agentic.async_rl.data_plane import build_rollout_group_sample_write from twinkle_agentic.async_rl.metrics import training_policy_metrics from twinkle_agentic.async_rl.native_tq import ContextGRPOGroupNSampler -from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches from twinkle_agentic.async_rl.pipeline import ( _require_adapter_path, - _reward_for_context, _train_batch, create_cpu_actor, ) -from twinkle_agentic.async_rl.types import (PartitionAdmission, PreparedPartition, PromptGroup, RolloutPolicy) +from twinkle_agentic.async_rl.types import PartitionAdmission, PreparedPartition from twinkle_agentic.async_rl.utils import ( TrainBatchConfig, - build_native_fsdp_model_kwargs, - configure_lora_lr_scheduler, - resolve_context_learning_rate, - resolve_context_lora_target_modules, - resolve_context_loss_config, - resolve_model_attention_implementation, - resolve_sequence_parallel_size, sample_responses_to_rollout_rows, - sampler_data_parallel_size, - validate_context_batch_config, -) -from twinkle_agentic.async_rl.vllm_sampler_tq import ( - VLLMSamplerTQ, - _GeneratedSample, - _PromptGroupRolloutStats, ) from twinkle_agentic.async_rl.workers import RolloutWorker @@ -93,53 +72,6 @@ def fake_remote(**options): assert captured['actor_kwargs'] == {'enabled': True} -def test_unload_lora_paths_does_not_require_pruned_checkpoint(tmp_path): - removed_paths = [] - - class Engine: - - async def unload_lora_paths(self, paths): - removed_paths.extend(paths) - - class Completed: - - @staticmethod - def result(): - return None - - sampler = object.__new__(VLLMSamplerTQ) - sampler.engine = Engine() - - def submit(coro): - asyncio.run(coro) - return Completed() - - sampler._submit_in_loop = submit - pruned_path = tmp_path / 'already-pruned' - sampler.unload_lora_paths([str(pruned_path)]) - - assert removed_paths == [str(pruned_path.resolve())] - - -def test_sequence_parallel_size_must_divide_model_gpus(): - assert resolve_sequence_parallel_size(2, 1) == 1 - assert resolve_sequence_parallel_size(2, 2) == 2 - - with pytest.raises(ValueError, match='must be divisible'): - resolve_sequence_parallel_size(2, 3) - - -def test_padding_free_sequence_parallel_requires_flash_attention(): - assert resolve_model_attention_implementation( - {'attn_implementation': 'flash_attention_2'}, - padding_free=True, - sequence_parallel_size=2, - ) == 'flash_attention_2' - - with pytest.raises(ValueError, match='model.attn_implementation'): - resolve_model_attention_implementation({}, padding_free=True, sequence_parallel_size=2) - - def test_train_batch_preserves_position_ids_from_tq(): class Batch(dict): batch_size = (1, ) @@ -227,59 +159,6 @@ def calculate_metric(self, **_kwargs): assert metrics['micro_batch_size_per_rank'] == 1 -def test_dynamic_micro_batch_planner_honors_per_rank_sample_and_token_limits(): - lengths = [10, 9, 8, 7, 4, 3, 2, 1] - inputs = [{'input_ids': list(range(length))} for length in lengths] - config = MicroBatchConfig( - micro_batch_size=3, - dynamic_batching=True, - max_tokens_per_micro_batch=18, - ) - - batches = plan_micro_batches(inputs, config, padding_free=False) - - assert sorted(index for batch in batches for index in batch) == list(range(8)) - for batch in batches: - assert len(batch) <= 3 - padded_tokens = max(lengths[index] for index in batch) * len(batch) - assert padded_tokens <= 18 - - -class PolicyProvider: - - def __init__(self, policies): - self.policies = iter(policies) - self.released = [] - - def get_rollout_policy(self, _context): - return next(self.policies) - - def acquire_rollout_policy(self, context): - return self.get_rollout_policy(context) - - def release_rollout_policy(self, policy): - self.released.append(policy) - - -class GenerationHarness: - _merge_partial_responses = VLLMSamplerTQ._merge_partial_responses - - def __init__(self, policies, responses): - self.context_manager = LocalActorHandle(PolicyProvider(policies)) - self.responses = iter(responses) - self.rollout_max_retries = 1 - self.rollout_retry_delay_s = 0 - self.calls = [] - self.template = type('Template', (), {'decode': staticmethod(lambda tokens: str(tokens))})() - - async def _load_lora_for_policy(self, policy): - return policy.version - - async def _sample_single(self, feat, sampling_params, *, lora_request, multi_modal_data, logprobs_only): - self.calls.append((list(feat['input_ids']), sampling_params.max_tokens, lora_request)) - return next(self.responses) - - def _context(name: str = 'adapter') -> LoraContext: return LoraContext('tenant', f'run_{name}', 'model', name) @@ -334,129 +213,6 @@ def test_training_rows_preserve_rollout_group_metadata(): assert rows[0]['generation_idx'] == 2 -def test_sampler_data_parallel_size_is_derived_from_gpu_and_tp_sizes(): - assert sampler_data_parallel_size(8, 2) == 4 - assert sampler_data_parallel_size(1, 1) == 1 - - -def test_sampler_parallelism_rejects_incomplete_tp_group(): - try: - sampler_data_parallel_size(3, 2) - except ValueError as exc: - assert 'must be divisible' in str(exc) - else: - raise AssertionError('expected invalid sampler GPU/TP layout to fail') - - -def test_lora_lr_scheduler_uses_shared_adapter_config(): - calls = [] - - class Model: - def set_lr_scheduler(self, scheduler_cls, **kwargs): - calls.append((scheduler_cls, kwargs)) - - configure_lora_lr_scheduler( - Model(), - 'tenant_lora', - { - 'lr_scheduler': { - 'cls': 'CosineAnnealingLR', - 'T_max': 2000, - 'eta_min': 0.0, - }, - }, - ) - - assert calls == [('CosineAnnealingLR', { - 'adapter_name': 'tenant_lora', - 'T_max': 2000, - 'eta_min': 0.0, - })] - - -def test_context_learning_rate_overrides_global_default(): - assert resolve_context_learning_rate({'learning_rate': 5e-6}, {'learning_rate': 1e-6}) == pytest.approx(5e-6) - assert resolve_context_learning_rate({}, {'learning_rate': 1e-6}) == pytest.approx(1e-6) - - -def test_context_lora_target_modules_override_global_default(): - defaults = {'target_modules': 'all-linear'} - - assert resolve_context_lora_target_modules({}, defaults) == 'all-linear' - assert resolve_context_lora_target_modules( - {'lora': {'target_modules': ['q_proj', 'v_proj']}}, - defaults, - ) == ['q_proj', 'v_proj'] - - -@pytest.mark.parametrize('value', ['', [], [None], {'q_proj': True}]) -def test_context_lora_target_modules_reject_invalid_values(value): - with pytest.raises(ValueError, match='target_modules'): - resolve_context_lora_target_modules( - {'lora': {'target_modules': value}}, - {'target_modules': 'all-linear'}, - ) - - -def test_context_loss_config_overrides_global_defaults(): - loss_cls, loss_kwargs = resolve_context_loss_config( - { - 'loss': { - 'cls': 'GSPOLoss', - 'epsilon_high': 0.3, - } - }, - { - 'cls': 'GRPOLoss', - 'epsilon': 0.2, - }, - ) - - assert loss_cls == 'GSPOLoss' - assert loss_kwargs == { - 'epsilon': 0.2, - 'epsilon_high': 0.3, - } - - -def test_context_loss_config_uses_grpo_defaults(): - assert resolve_context_loss_config({}) == ( - 'GRPOLoss', - { - 'epsilon': 0.2, - }, - ) - - -def test_async_single_lora_gsm8k_verl_config_matches_current_grpo_api(): - from omegaconf import OmegaConf - from twinkle.loss import GRPOLoss - - config_path = 'cookbook/rl/async_single_lora_gsm8k_verl.yaml' - config = OmegaConf.to_container(OmegaConf.load(config_path), resolve=False) - context = config['lora_contexts'][0] - loss_cls, loss_kwargs = resolve_context_loss_config(context, config['loss']) - - assert loss_cls == 'GRPOLoss' - GRPOLoss(**loss_kwargs) - assert context['reward']['class_path'] == 'twinkle.reward.GSM8KAccuracyReward' - assert context['eval_dataset']['reward']['class_path'] == 'twinkle.reward.GSM8KAccuracyReward' - validate_context_batch_config( - f"{context['tenant_id']}/{context['training_run_id']}/{context['adapter_name']}", - rollout_groups=context['rollout']['batch_size'], - num_generations=context['rollout']['num_generations'], - train=TrainBatchConfig( - mini_batch_size=context['train']['mini_batch_size'], - micro_batch_size=context['train']['micro_batch_size'], - dynamic_batching=context['train']['dynamic_batching'], - max_tokens_per_micro_batch=context['train']['max_tokens_per_micro_batch'], - packing_algorithm=context['train']['packing_algorithm'], - ), - sampler_dp=config['runtime']['sampler_gpus'] // config['runtime']['sampler_tp'], - model_dp=config['runtime']['model_gpus'], - ) - - def test_adapter_path_rejects_uncollected_remote_result(): def lazy_result(): return '/tmp/policy' @@ -465,359 +221,6 @@ def lazy_result(): _require_adapter_path(lazy_result, operation='test save') -def test_multi_lora_transformers_save_disables_lazy_collect(): - from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel - - assert MultiLoraTransformersModel.save._lazy_collect is False - - -def test_remote_function_metadata_uses_explicit_lazy_collect_value(): - from twinkle import remote_function - - class Component: - - @remote_function(lazy_collect=False) - def eager(self): - return None - - @remote_function(lazy_collect=True) - def lazy(self): - return None - - assert Component.eager._lazy_collect is False - assert Component.lazy._lazy_collect is True - - -def test_context_loss_config_rejects_empty_class_name(): - with pytest.raises(ValueError, match='loss.cls'): - resolve_context_loss_config({'loss': {'cls': ''}}) - - -@pytest.mark.parametrize('value', [0, -1e-6, float('inf')]) -def test_context_learning_rate_rejects_invalid_values(value): - with pytest.raises(ValueError, match='positive finite'): - resolve_context_learning_rate({'learning_rate': value}, {'learning_rate': 1e-6}) - - -def test_rl_model_kwargs_enforce_native_fsdp(): - assert build_native_fsdp_model_kwargs({}) == { - 'strategy': 'native_fsdp', - 'fsdp_config': {}, - } - assert build_native_fsdp_model_kwargs({ - 'strategy': 'native_fsdp', - 'fsdp_config': {'reshard_after_forward': False}, - }) == { - 'strategy': 'native_fsdp', - 'fsdp_config': {'reshard_after_forward': False}, - } - with pytest.raises(ValueError, match='must be native_fsdp'): - build_native_fsdp_model_kwargs({'strategy': 'accelerate'}) - - -def test_reward_factory_loads_class_and_resolved_kwargs(): - reward = _reward_for_context( - { - 'class_path': 'twinkle.reward.DAPOMathReward', - 'kwargs': { - 'max_response_length': 8192, - 'overlong_buffer_length': 4096, - 'overlong_penalty_factor': 1.0, - 'score_tail_chars': 300, - }, - }, - context_key='tenant/run/adapter', - ) - - assert reward.max_response_length == 8192 - assert reward.overlong_buffer_length == 4096 - - -def test_reward_factory_rejects_non_reward_class(): - with pytest.raises(TypeError, match='Reward subclass'): - _reward_for_context( - {'class_path': 'collections.Counter'}, - context_key='tenant/run/adapter', - ) - - -def test_context_batch_config_accepts_group_aligned_dp_batches(): - validate_context_batch_config( - 'tenant/run/adapter', - rollout_groups=8, - num_generations=4, - train=TrainBatchConfig(mini_batch_size=8, micro_batch_size=2), - sampler_dp=2, - model_dp=2, - ) - - -def test_context_batch_config_allows_training_group_to_span_model_dp_ranks(): - validate_context_batch_config( - 'tenant/run/adapter', - rollout_groups=8, - num_generations=4, - train=TrainBatchConfig(mini_batch_size=16, micro_batch_size=2), - sampler_dp=1, - model_dp=8, - ) - - -def test_context_batch_config_rejects_partition_tail_and_undersized_rank_batch(): - try: - validate_context_batch_config( - 'tenant/run/adapter', - rollout_groups=6, - num_generations=4, - train=TrainBatchConfig(mini_batch_size=6, micro_batch_size=1), - sampler_dp=2, - model_dp=2, - ) - except ValueError as exc: - assert 'complete prompt groups' in str(exc) - else: - raise AssertionError('expected a split prompt group to fail') - - try: - validate_context_batch_config( - 'tenant/run/adapter', - rollout_groups=8, - num_generations=2, - train=TrainBatchConfig(mini_batch_size=2, micro_batch_size=2), - sampler_dp=2, - model_dp=2, - ) - except ValueError as exc: - assert 'per-rank train batch' in str(exc) - else: - raise AssertionError('expected an oversized micro batch to fail') - - -def test_context_batch_config_requires_token_limit_for_dynamic_batching(): - with pytest.raises(ValueError, match='max_tokens_per_micro_batch'): - validate_context_batch_config( - 'tenant/run/adapter', - rollout_groups=8, - num_generations=4, - train=TrainBatchConfig( - mini_batch_size=8, - micro_batch_size=2, - dynamic_batching=True, - ), - sampler_dp=1, - model_dp=1, - ) - - -def test_sampler_dp_dispatch_slices_complete_groups_without_duplication(): - mesh = DeviceMesh.from_sizes(world_size=4, dp_size=2, tp_size=2) - groups = ['group_0', 'group_1', 'group_2', 'group_3'] - dispatched = _dispatch_args( - workers=['dp_0', 'dp_1'], - dispatch='slice_dp', - execute='all', - device_mesh=mesh, - args=(groups, 'sampling_params', False), - kwargs={}, - ) - - assert [worker for worker, _, _ in dispatched] == ['dp_0', 'dp_1'] - assert [args[0] for _, args, _ in dispatched] == [groups[:2], groups[2:]] - assert [group for _, args, _ in dispatched for group in args[0]] == groups - - -@pytest.mark.parametrize(('dp_size', 'expected_scope'), [(1, 'partition'), (2, 'shard')]) -def test_sampler_reports_submission_throughput_at_partition_or_shard_scope(dp_size, expected_scope): - context = _context() - admission = PartitionAdmission(context, context.partition_id(0), 0, 2, 2, 0) - groups = [ - PromptGroup(context, admission, f'{admission.partition_id}/group_{index}', {}, object()) - for index in range(2) - ] - - class RolloutMetricsHarness: - def __init__(self): - self.device_mesh = DeviceMesh.from_sizes(world_size=dp_size, dp_size=dp_size) - self.events = [] - - async def _run_prompt_group(self, *, group, **_kwargs): - index = int(group.group_id.rsplit('_', 1)[1]) - lengths = ((10, 20), (30, 40))[index] - reasons = (('stop', 'length'), ('stop', 'stop'))[index] - return _PromptGroupRolloutStats(lengths, reasons, (index + 1, index + 1)) - - def _record_metrics(self, group, values, **kwargs): - self.events.append((group, values, kwargs)) - - sampler = RolloutMetricsHarness() - asyncio.run( - VLLMSamplerTQ._sample_prompt_groups( - sampler, - 'submission', - groups, - SamplingParams(max_tokens=64), - False, - time.perf_counter() - 1, - )) - - recorded_group, metrics, record_options = sampler.events[-1] - assert recorded_group.context == context - assert recorded_group.partition_id == admission.partition_id - assert record_options['attributes']['scope'] == expected_scope - assert metrics['prompt_group_count'] == 2 - assert metrics['sample_count'] == 4 - assert metrics['output_tokens'] == 100 - assert metrics['completion_length_mean'] == 25 - assert metrics['completion_truncated_count'] == 1 - assert metrics['policy_version_min'] == 1 - assert metrics['policy_version_max'] == 2 - assert metrics['sampler_dp_size'] == dp_size - assert metrics['output_tokens_per_s'] == pytest.approx(100 / metrics['rollout_latency_s']) - - -def test_sampler_writes_one_atomic_rollout_file_per_prompt_group(tmp_path): - context = _context() - admission = PartitionAdmission(context, context.partition_id(3), 3, 1, 2, 0) - group = PromptGroup( - context, - admission, - f'{admission.partition_id}/group_0', - {'user_data': [('ground_truth', '"42"')]}, - object(), - ) - policy = RolloutPolicy(context.key, context.adapter_name, 7, '/tmp/adapter-v7') - generated = [ - _GeneratedSample( - SampleResponse( - sequences=[SampledSequence('stop', [20 + index], decoded=f'completion-{index}')], - prompt_token_ids=[10, 11], - ), - (policy,), - attempts=1, - was_aborted=False, - resumed_partial_output=False, - ) - for index in range(2) - ] - rows = [ - { - 'generation_idx': index, - 'rollout_policy_version': 7, - 'initial_policy_version': 7, - 'final_policy_version': 7, - 'rollout_policy_versions': [7], - 'rollout_adapter_path': '/tmp/adapter-v7', - 'stop_reason': 'stop', - 'logprobs': [-0.1], - } - for index in range(2) - ] - - class Template: - @staticmethod - def decode(token_ids, **_kwargs): - return ' '.join(map(str, token_ids)) - - sampler = object.__new__(VLLMSamplerTQ) - sampler.rollout_output_dir = tmp_path - sampler.rollout_output_include_token_ids = False - sampler.template = Template() - - sampler._write_rollout_group('submission-1', group, generated, rows, [1.0, 0.0]) - sampler._write_rollout_group('submission-2', group, generated, rows, [1.0, 0.0]) - - output_path = ( - tmp_path - / context.tenant_id - / context.training_run_id - / context.adapter_name - / 'policy_7' - / 'train_3-group_0.jsonl' - ) - records = [json.loads(line) for line in output_path.read_text().splitlines()] - assert len(records) == 2 - assert records[0]['submission_id'] == 'submission-2' - assert records[0]['prompt'] == '10 11' - assert records[0]['completion'] == '20' - assert records[0]['ground_truth'] == '42' - assert records[0]['reward'] == 1.0 - assert records[0]['head_version'] == 7 - assert records[0]['tail_version'] == 7 - assert 'prompt_token_ids' not in records[0] - - -def test_aborted_generation_restarts_from_original_prompt_when_partial_is_disabled(): - context = _context() - policies = [ - RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), - RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), - ] - sampler = GenerationHarness( - policies, - [ - _sample_response([7], 'abort', [1, 2, 7]), - _sample_response([8], 'stop', [1, 2, 8]), - ], - ) - generated = asyncio.run( - VLLMSamplerTQ._generate_sample( - sampler, - context, - { - 'input_ids': [1, 2], - 'labels': [-100, -100] - }, - SamplingParams(max_tokens=4, logprobs=1), - multi_modal_data=None, - logprobs_only=False, - allow_partial_rollout=False, - )) - - assert sampler.calls == [([1, 2], 4, 3), ([1, 2], 4, 4)] - assert generated.response.sequences[0].tokens == [8] - assert [policy.version for policy in generated.policies] == [4] - assert generated.retry_count == 1 - assert generated.was_aborted - assert not generated.resumed_partial_output - - -def test_aborted_generation_continues_from_partial_tokens_when_enabled(): - context = _context() - policies = [ - RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), - RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), - ] - sampler = GenerationHarness( - policies, - [ - _sample_response([7], 'abort', [1, 2, 7]), - _sample_response([8], 'stop', [1, 2, 7, 8]), - ], - ) - generated = asyncio.run( - VLLMSamplerTQ._generate_sample( - sampler, - context, - { - 'input_ids': [1, 2], - 'labels': [-100, -100] - }, - SamplingParams(max_tokens=4, logprobs=1), - multi_modal_data=None, - logprobs_only=False, - allow_partial_rollout=True, - )) - - assert sampler.calls == [([1, 2], 4, 3), ([1, 2, 7], 3, 4)] - assert generated.response.sequences[0].tokens == [7, 8] - assert [policy.version for policy in generated.policies] == [3, 4] - assert generated.initial_policy.version == 3 - assert generated.final_policy.version == 4 - assert generated.retry_count == 1 - assert generated.was_aborted - assert generated.resumed_partial_output - - def test_training_policy_metrics_use_final_version_and_partial_span(): metrics = training_policy_metrics(( { @@ -995,6 +398,11 @@ def test_scheduler_supports_round_robin_sticky_and_oldest(): sticky.on_blocked(candidates[1]) assert sticky.choose(candidates).context == a + capped = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.STICKY, 1)) + first = capped.choose(candidates) + capped.on_success(first) + assert capped.choose(candidates).context == b + manager = LoraContextManager(max_staleness=2) manager.register_context(a) manager.register_context(b) @@ -1004,17 +412,6 @@ def test_scheduler_supports_round_robin_sticky_and_oldest(): assert oldest.choose([ScheduleCandidate(b, new), ScheduleCandidate(a, old)]).partition == old -def test_sticky_scheduler_switches_context_at_consecutive_cap(): - a, b = _context('a'), _context('b') - candidates = [ScheduleCandidate(a), ScheduleCandidate(b)] - scheduler = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.STICKY, 1)) - - first = scheduler.choose(candidates) - scheduler.on_success(first) - - assert scheduler.choose(candidates).context == b - - def test_context_group_sampler_uses_request_generation_count(): sampler = ContextGRPOGroupNSampler() @@ -1145,126 +542,6 @@ def test_zero_max_steps_finishes_without_admission(): assert manager.is_run_finished() -def test_rollout_sample_tags_use_new_context_descriptor_only(): - context = _context() - admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) - group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, batch_meta=None) - fields, tags = build_rollout_group_sample_write( - group, - [ - { - 'generation_idx': 0, - 'labels': [-100, 1], - 'logprobs': [-.1], - 'rollout_policy_version': 3, - 'rollout_adapter_path': 'adapter-v3', - }, - { - 'generation_idx': 1, - 'labels': [-100, 2], - 'logprobs': [-.2], - 'rollout_policy_version': 4, - 'rollout_adapter_path': 'adapter-v4', - }, - ], - rewards=[1., 0.], - expected_num_generations=2, - ) - assert [row['rewards'] for row in fields] == [1., 0.] - assert [tag['generation_idx'] for tag in tags] == [0, 1] - assert all(tag['context_key'] == context.key for tag in tags) - assert [tag['rollout_policy_version'] for tag in tags] == [3, 4] - - -def test_data_plane_completes_rollout_with_full_training_trajectory(): - class Metadata: - def __init__(self): - self.size = 2 - self.custom_meta = [{}, {}] - - def update_custom_meta(self, updates): - for tag, update in zip(self.custom_meta, updates): - tag.update(update) - - class Client: - def __init__(self): - self.written = None - self.calls = [] - - async def async_put(self, data, metadata=None, partition_id=None): - self.calls.append('fields') - self.written = data - return metadata - - async def async_set_custom_meta(self, _metadata): - self.calls.append('tags') - return None - - context = _context() - admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) - metadata = Metadata() - group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) - client = Client() - rows = [{ - 'input_ids': [1, 2, token], - 'labels': [-100, -100, token], - 'attention_mask': [1, 1, 1], - 'position_ids': [0, 1, 2], - 'logprobs': [-.1], - 'generation_idx': generation_idx, - 'rollout_policy_version': 3, - 'rollout_policy_versions': [3], - 'initial_policy_version': 3, - 'final_policy_version': 3, - 'policy_version_span': 0, - 'rollout_adapter_path': 'adapter-v3', - 'completion_length': 1, - } for generation_idx, token in enumerate((7, 8))] - - asyncio.run( - TQDataPlane(client).complete_rollout_group( - group, - rollout_rows=rows, - rewards=[1., 0.], - submission_id='submission', - )) - - assert set(client.written.keys()) == { - 'input_ids', 'labels', 'attention_mask', 'position_ids', 'logprobs', 'rewards' - } - assert client.calls == ['tags', 'fields'] - assert [tag['rollout_status'] for tag in metadata.custom_meta] == ['ROLLOUT_DONE', 'ROLLOUT_DONE'] - assert [tag['submission_id'] for tag in metadata.custom_meta] == ['submission', 'submission'] - - -def test_data_plane_rejects_rollout_without_complete_model_inputs(): - context = _context() - admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 1, 0) - metadata = type('Metadata', (), {'size': 1})() - group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) - row = { - 'input_ids': [1, 2], - 'labels': [-100, 2], - 'logprobs': [-.1], - 'generation_idx': 0, - 'rollout_policy_version': 0, - } - - try: - asyncio.run( - TQDataPlane(object()).complete_rollout_group( - group, - rollout_rows=[row], - rewards=[1.], - submission_id='submission', - )) - except ValueError as exc: - assert 'attention_mask' in str(exc) - assert 'position_ids' in str(exc) - else: - raise AssertionError('expected incomplete rollout model fields to fail') - - def test_checkpoint_retention_preserves_current_policy_and_history_window(): context = _context() manager = LoraContextManager() diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index f01c0a545..ce427a810 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -1,13 +1,95 @@ from __future__ import annotations import asyncio +import inspect +import json +import time from concurrent.futures import Future import pytest -from twinkle.data_format import SamplingParams +from twinkle import DeviceMesh +from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams +from twinkle.infra import _dispatch_args from twinkle.server.sampler.twinkle_handlers import _await_generation -from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ, _dispatch_generation +from twinkle_agentic.async_rl import LoraContext +from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup, RolloutPolicy +from twinkle_agentic.async_rl.vllm_sampler_tq import ( + VLLMSamplerTQ, + _GeneratedSample, + _PromptGroupRolloutStats, + _dispatch_generation, +) + + +class LocalActorHandle: + def __init__(self, target): + self.target = target + + def __getattr__(self, name): + method = getattr(self.target, name) + + class RemoteMethod: + async def remote(_, *args, **kwargs): + result = method(*args, **kwargs) + return await result if inspect.isawaitable(result) else result + + return RemoteMethod() + + +class PolicyProvider: + def __init__(self, policies): + self.policies = iter(policies) + self.released = [] + + def get_rollout_policy(self, _context): + return next(self.policies) + + def acquire_rollout_policy(self, context): + return self.get_rollout_policy(context) + + def release_rollout_policy(self, policy): + self.released.append(policy) + + +class GenerationHarness: + _merge_partial_responses = VLLMSamplerTQ._merge_partial_responses + + def __init__(self, policies, responses): + self.context_manager = LocalActorHandle(PolicyProvider(policies)) + self.responses = iter(responses) + self.rollout_max_retries = 1 + self.rollout_retry_delay_s = 0 + self.calls = [] + self.template = type('Template', (), {'decode': staticmethod(lambda tokens: str(tokens))})() + + async def _load_lora_for_policy(self, policy): + return policy.version + + async def _sample_single(self, feat, sampling_params, *, lora_request, multi_modal_data, logprobs_only): + self.calls.append((list(feat['input_ids']), sampling_params.max_tokens, lora_request)) + return next(self.responses) + + +def _context(name: str = 'adapter') -> LoraContext: + return LoraContext('tenant', f'run_{name}', 'model', name) + + +def _sample_response(tokens, stop_reason, input_ids): + return SampleResponse( + prompt_token_ids=[1, 2], + sequences=[ + SampledSequence( + stop_reason=stop_reason, + tokens=tokens, + logprobs=[[(token, -.1)] for token in tokens], + new_input_feature={ + 'input_ids': input_ids, + 'labels': [-100, -100, *tokens], + }, + ) + ], + ) def _bare_sampler() -> VLLMSamplerTQ: @@ -208,3 +290,232 @@ def cancel_generation(self, _submission_id): assert result == ['completed'] assert sampler.status_calls == 2 assert not sampler.cancelled + + +def test_unload_lora_paths_does_not_require_pruned_checkpoint(tmp_path): + removed_paths = [] + + class Engine: + async def unload_lora_paths(self, paths): + removed_paths.extend(paths) + + class Completed: + @staticmethod + def result(): + return None + + sampler = object.__new__(VLLMSamplerTQ) + sampler.engine = Engine() + + def submit(coro): + asyncio.run(coro) + return Completed() + + sampler._submit_in_loop = submit + pruned_path = tmp_path / 'already-pruned' + sampler.unload_lora_paths([str(pruned_path)]) + + assert removed_paths == [str(pruned_path.resolve())] + + +def test_sampler_dp_dispatch_slices_complete_groups_without_duplication(): + mesh = DeviceMesh.from_sizes(world_size=4, dp_size=2, tp_size=2) + groups = ['group_0', 'group_1', 'group_2', 'group_3'] + dispatched = _dispatch_args( + workers=['dp_0', 'dp_1'], + dispatch='slice_dp', + execute='all', + device_mesh=mesh, + args=(groups, 'sampling_params', False), + kwargs={}, + ) + + assert [worker for worker, _, _ in dispatched] == ['dp_0', 'dp_1'] + assert [args[0] for _, args, _ in dispatched] == [groups[:2], groups[2:]] + assert [group for _, args, _ in dispatched for group in args[0]] == groups + + +@pytest.mark.parametrize(('dp_size', 'expected_scope'), [(1, 'partition'), (2, 'shard')]) +def test_sampler_reports_submission_throughput_at_partition_or_shard_scope(dp_size, expected_scope): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 2, 2, 0) + groups = [ + PromptGroup(context, admission, f'{admission.partition_id}/group_{index}', {}, object()) + for index in range(2) + ] + + class RolloutMetricsHarness: + def __init__(self): + self.device_mesh = DeviceMesh.from_sizes(world_size=dp_size, dp_size=dp_size) + self.events = [] + + async def _run_prompt_group(self, *, group, **_kwargs): + index = int(group.group_id.rsplit('_', 1)[1]) + lengths = ((10, 20), (30, 40))[index] + reasons = (('stop', 'length'), ('stop', 'stop'))[index] + return _PromptGroupRolloutStats(lengths, reasons, (index + 1, index + 1)) + + def _record_metrics(self, group, values, **kwargs): + self.events.append((group, values, kwargs)) + + sampler = RolloutMetricsHarness() + asyncio.run( + VLLMSamplerTQ._sample_prompt_groups( + sampler, + 'submission', + groups, + SamplingParams(max_tokens=64), + False, + time.perf_counter() - 1, + )) + + recorded_group, metrics, record_options = sampler.events[-1] + assert recorded_group.context == context + assert recorded_group.partition_id == admission.partition_id + assert record_options['attributes']['scope'] == expected_scope + assert metrics['prompt_group_count'] == 2 + assert metrics['sample_count'] == 4 + assert metrics['output_tokens'] == 100 + assert metrics['completion_length_mean'] == 25 + assert metrics['completion_truncated_count'] == 1 + assert metrics['policy_version_min'] == 1 + assert metrics['policy_version_max'] == 2 + assert metrics['sampler_dp_size'] == dp_size + assert metrics['output_tokens_per_s'] == pytest.approx(100 / metrics['rollout_latency_s']) + + +def test_sampler_writes_one_atomic_rollout_file_per_prompt_group(tmp_path): + context = _context() + admission = PartitionAdmission(context, context.partition_id(3), 3, 1, 2, 0) + group = PromptGroup( + context, + admission, + f'{admission.partition_id}/group_0', + {'user_data': [('ground_truth', '"42"')]}, + object(), + ) + policy = RolloutPolicy(context.key, context.adapter_name, 7, '/tmp/adapter-v7') + generated = [ + _GeneratedSample( + SampleResponse( + sequences=[SampledSequence('stop', [20 + index], decoded=f'completion-{index}')], + prompt_token_ids=[10, 11], + ), + (policy,), + attempts=1, + was_aborted=False, + resumed_partial_output=False, + ) + for index in range(2) + ] + rows = [ + { + 'generation_idx': index, + 'rollout_policy_version': 7, + 'initial_policy_version': 7, + 'final_policy_version': 7, + 'rollout_policy_versions': [7], + 'rollout_adapter_path': '/tmp/adapter-v7', + 'stop_reason': 'stop', + 'logprobs': [-0.1], + } + for index in range(2) + ] + + class Template: + @staticmethod + def decode(token_ids, **_kwargs): + return ' '.join(map(str, token_ids)) + + sampler = object.__new__(VLLMSamplerTQ) + sampler.rollout_output_dir = tmp_path + sampler.rollout_output_include_token_ids = False + sampler.template = Template() + + sampler._write_rollout_group('submission-1', group, generated, rows, [1.0, 0.0]) + sampler._write_rollout_group('submission-2', group, generated, rows, [1.0, 0.0]) + + output_path = ( + tmp_path + / context.tenant_id + / context.training_run_id + / context.adapter_name + / 'policy_7' + / 'train_3-group_0.jsonl' + ) + records = [json.loads(line) for line in output_path.read_text().splitlines()] + assert len(records) == 2 + assert records[0]['submission_id'] == 'submission-2' + assert records[0]['prompt'] == '10 11' + assert records[0]['completion'] == '20' + assert records[0]['ground_truth'] == '42' + assert records[0]['reward'] == 1.0 + assert records[0]['head_version'] == 7 + assert records[0]['tail_version'] == 7 + assert 'prompt_token_ids' not in records[0] + + +def test_aborted_generation_restarts_from_original_prompt_when_partial_is_disabled(): + context = _context() + policies = [ + RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), + RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), + ] + sampler = GenerationHarness( + policies, + [ + _sample_response([7], 'abort', [1, 2, 7]), + _sample_response([8], 'stop', [1, 2, 8]), + ], + ) + generated = asyncio.run( + VLLMSamplerTQ._generate_sample( + sampler, + context, + {'input_ids': [1, 2], 'labels': [-100, -100]}, + SamplingParams(max_tokens=4, logprobs=1), + multi_modal_data=None, + logprobs_only=False, + allow_partial_rollout=False, + )) + + assert sampler.calls == [([1, 2], 4, 3), ([1, 2], 4, 4)] + assert generated.response.sequences[0].tokens == [8] + assert [policy.version for policy in generated.policies] == [4] + assert generated.retry_count == 1 + assert generated.was_aborted + assert not generated.resumed_partial_output + + +def test_aborted_generation_continues_from_partial_tokens_when_enabled(): + context = _context() + policies = [ + RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), + RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), + ] + sampler = GenerationHarness( + policies, + [ + _sample_response([7], 'abort', [1, 2, 7]), + _sample_response([8], 'stop', [1, 2, 7, 8]), + ], + ) + generated = asyncio.run( + VLLMSamplerTQ._generate_sample( + sampler, + context, + {'input_ids': [1, 2], 'labels': [-100, -100]}, + SamplingParams(max_tokens=4, logprobs=1), + multi_modal_data=None, + logprobs_only=False, + allow_partial_rollout=True, + )) + + assert sampler.calls == [([1, 2], 4, 3), ([1, 2, 7], 3, 4)] + assert generated.response.sequences[0].tokens == [7, 8] + assert [policy.version for policy in generated.policies] == [3, 4] + assert generated.initial_policy.version == 3 + assert generated.final_policy.version == 4 + assert generated.retry_count == 1 + assert generated.was_aborted + assert generated.resumed_partial_output diff --git a/tests/twinkle_client/test_client_orchestrated_dpo.py b/tests/twinkle_client/test_client_orchestrated_dpo.py deleted file mode 100644 index fa1bd0823..000000000 --- a/tests/twinkle_client/test_client_orchestrated_dpo.py +++ /dev/null @@ -1,108 +0,0 @@ -import asyncio - -from cookbook.client.async_rl.client_orchestrated_dpo import ( - prepare_dpo_batch, - run_dpo, -) -from twinkle_client.types import DataRef - - -def test_prepare_dpo_batch_interleaves_complete_pairs() -> None: - batch = [ - { - 'pair_id': 'a', - 'positive': {'input_ids': [1, 2], 'labels': [-100, 2]}, - 'negative': {'input_ids': [1, 3], 'labels': [-100, 3]}, - }, - { - 'pair_id': 'b', - 'positive': {'input_ids': [4], 'labels': [4]}, - 'negative': {'input_ids': [5], 'labels': [5]}, - }, - ] - - rows = prepare_dpo_batch(batch) - - assert [row['pair_id'] for row in rows] == ['a', 'a', 'b', 'b'] - assert [row['input_ids'] for row in rows] == [[1, 2], [1, 3], [4], [5]] - - -def test_dpo_roles_overlap_reference_and_training(monkeypatch) -> None: - import cookbook.client.async_rl.client_orchestrated_dpo as module - - monkeypatch.setattr(module, 'MAX_STEPS', 2) - first_train_started = asyncio.Event() - events = [] - - class FakeDataPlane: - - def __init__(self): - self.rows = {} - self.released = [] - - async def aput(self, rows, *, kind, tags=None): - ref = DataRef( - ref_id=f'{kind}-{len(self.rows)}', - size=len(rows), - fields=list(rows[0]), - kind=kind, - ) - self.rows[ref.ref_id] = rows - return ref - - async def aappend(self, ref, updates, *, tags=None): - self.rows[ref.ref_id] = [ - {**row, **update} - for row, update in zip(self.rows[ref.ref_id], updates) - ] - return ref.model_copy(update={'fields': list(self.rows[ref.ref_id][0])}) - - async def arelease(self, ref): - self.released.append(ref.ref_id) - - class FakeModel: - - def __init__(self): - self.references = 0 - self.steps = 0 - self.forward_backward_kwargs = [] - - async def forward_only_from_data_plane(self, ref, **kwargs): - self.references += 1 - name = f'reference-{self.references}' - events.append(f'{name}-start') - if self.references == 2: - await first_train_started.wait() - events.append(f'{name}-done') - assert kwargs['output_ref'] == ref - assert kwargs['output_fields'] == {'logps': 'ref_logps'} - return ref.model_copy(update={'fields': [*ref.fields, 'ref_logps']}) - - async def forward_backward_from_data_plane(self, _ref, **kwargs): - self.forward_backward_kwargs.append(kwargs) - events.append('train-start') - first_train_started.set() - - async def clip_grad_and_step(self, **_kwargs): - self.steps += 1 - - async def save(self, name, **_kwargs): - return {'twinkle_path': name} - - batches = [ - [{'pair_id': 'a', 'positive': {'input_ids': [1]}, 'negative': {'input_ids': [2]}}], - [{'pair_id': 'b', 'positive': {'input_ids': [3]}, 'negative': {'input_ids': [4]}}], - ] - model = FakeModel() - data_plane = FakeDataPlane() - - saved = asyncio.run(run_dpo(batches, model, data_plane)) - - assert events.index('train-start') < events.index('reference-2-done') - assert model.steps == 2 - assert model.forward_backward_kwargs == [ - {'kwarg_fields': {'ref_outputs.logps': 'ref_logps'}}, - {'kwarg_fields': {'ref_outputs.logps': 'ref_logps'}}, - ] - assert saved == {'twinkle_path': 'dpo-policy-2'} - assert len(data_plane.released) == 2 diff --git a/tests/twinkle_client/test_client_orchestrated_grpo.py b/tests/twinkle_client/test_client_orchestrated_grpo.py index 2ba86ef6f..2302e3bfc 100644 --- a/tests/twinkle_client/test_client_orchestrated_grpo.py +++ b/tests/twinkle_client/test_client_orchestrated_grpo.py @@ -22,7 +22,7 @@ def _load_module(): return module -def test_rollout_and_train_overlap_with_fifo_policy_publication(monkeypatch) -> None: +def test_rollout_and_train_overlap_with_fifo_policy_publication(monkeypatch, capsys) -> None: module = _load_module() monkeypatch.setattr(module, 'BATCH_SIZE', 2) monkeypatch.setattr(module, 'NUM_GENERATIONS', 2) @@ -74,6 +74,9 @@ async def forward_backward_from_data_plane(self, _refs, **kwargs): async def clip_grad_and_step(self, **_kwargs): self.steps += 1 + async def calculate_metric(self, **_kwargs): + return {'result': {'loss': 1.0 / self.steps, 'grad_norm': 0.5}} + class FakeDataPlane: def __init__(self): self.released = [] @@ -116,6 +119,10 @@ async def run(): assert snapshots['p1-g0'] == (0, '/checkpoints/policy-0') assert snapshots['p2-g0'][0] in (1, 2) assert snapshots['p2-g0'][1] == f'/checkpoints/policy-{snapshots["p2-g0"][0]}' + output = capsys.readouterr().out + assert 'optimizer_step=1' in output + assert 'loss=1.0' in output + assert 'grad_norm=0.5' in output def test_younger_rollout_failure_stops_admission(monkeypatch) -> None: @@ -158,6 +165,9 @@ async def forward_backward(self, _ref, **_kwargs): async def clip_grad_and_step(self, **_kwargs): return None + async def calculate_metric(self, **_kwargs): + return {'result': {'loss': 1.0}} + class FakeDataPlane: async def aget(self, ref, *, fields=None): assert fields == ['decoded'] From 538289fd0facad0082cba59b4e4fdd3a11e073e8 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Mon, 24 Aug 2026 17:10:10 +0800 Subject: [PATCH 18/20] wip --- cookbook/client/async_rl/client_orchestrated_grpo.py | 2 +- src/twinkle/metric/loss.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py index 6f250cfab..0f3162569 100644 --- a/cookbook/client/async_rl/client_orchestrated_grpo.py +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -403,7 +403,7 @@ async def train() -> None: ) model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) model.set_optimizer('AdamW', lr=2e-5) - model.set_processor('InputProcessor', padding_free=True) + model.set_processor('InputProcessor', padding_free=False) model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) sampler.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) diff --git a/src/twinkle/metric/loss.py b/src/twinkle/metric/loss.py index df0ce15c3..90acb647a 100644 --- a/src/twinkle/metric/loss.py +++ b/src/twinkle/metric/loss.py @@ -71,7 +71,7 @@ def calculate(self): self.reset() results = {} if avg_loss is not None: - results['loss'] = f'{avg_loss:.4f}' + results['loss'] = f'{avg_loss:.5f}' if grad_norm > 0: results['grad_norm'] = f'{grad_norm:.6f}' return results From 031829d416621608b086b55d992a351ee130380e Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 25 Aug 2026 16:04:59 +0800 Subject: [PATCH 19/20] fix ci --- .dev_scripts/ci_container_test.sh | 2 +- .../model/transformers/multi_lora_transformers.py | 7 +------ src/twinkle_agentic/async_rl/pipeline.py | 11 +++++++++-- tests/twinkle_agentic/test_async_rl_native_tq.py | 9 +++++++++ 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index beb5fe530..b0484a71a 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.dev_scripts/ci_container_test.sh @@ -1,5 +1,5 @@ install_twinkle_with_kernels() { - pip install ".[test,client,server]" -i https://mirrors.aliyun.com/pypi/simple/ || pip install ".[test,client,server]" + pip install ".[test,client,server,async-rl]" -i https://mirrors.aliyun.com/pypi/simple/ || pip install ".[test,client,server,async-rl]" } if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 50df7a05f..ea53930de 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -267,12 +267,7 @@ def _get_adapter_state_dict_for_save(self, adapter_name: str) -> dict: adapter_state = self.multi_adapter.get_state_dict(adapter_name) return {key: torch_util.to_local_tensor(value).cpu() for key, value in adapter_state.items()} - # Saving publishes an immutable adapter checkpoint to callers, so the - # checkpoint path must be collected before returning. In particular, - # TrainerWorker calls this model handle from another Ray actor; inheriting - # that actor's default lazy-collect mode would otherwise return a callable - # instead of the path string. - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect='first') def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): self._check_adapter_valid(kwargs.get('adapter_name')) with self.multi_adapter.save_context(kwargs.get('adapter_name')): diff --git a/src/twinkle_agentic/async_rl/pipeline.py b/src/twinkle_agentic/async_rl/pipeline.py index 1bce5fbf4..97a3b9a72 100644 --- a/src/twinkle_agentic/async_rl/pipeline.py +++ b/src/twinkle_agentic/async_rl/pipeline.py @@ -289,7 +289,7 @@ def from_config( eval_dataset.get('reward'), context_key=f'{context.key} evaluation', ) - initial_paths[context.key] = _require_adapter_path( + initial_paths[context.key] = _collect_adapter_path( model.save( f'async-{context.adapter_name}-initial', output_dir=runtime['output_dir'], @@ -672,7 +672,7 @@ def _train_batch_with_config( def _save_adapter(model: Any, output_dir: str, admission: PartitionAdmission) -> str: - return _require_adapter_path( + return _collect_adapter_path( model.save( f'async-{admission.context.adapter_name}-v{admission.step + 1}', output_dir=output_dir, @@ -682,6 +682,13 @@ def _save_adapter(model: Any, output_dir: str, admission: PartitionAdmission) -> ) +def _collect_adapter_path(value: Any, *, operation: str) -> str: + """Collect a lazy model.save result at the async-RL publication boundary.""" + if callable(value) and getattr(value, '_is_lazy_collect', False): + value = value() + return _require_adapter_path(value, operation=operation) + + def _require_adapter_path(value: Any, *, operation: str) -> str: """Fail at the save boundary instead of publishing an invalid policy.""" if not isinstance(value, str) or not value: diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py index 8f6479b7d..ef8ec0174 100644 --- a/tests/twinkle_agentic/test_async_rl_native_tq.py +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -13,6 +13,7 @@ from twinkle_agentic.async_rl.metrics import training_policy_metrics from twinkle_agentic.async_rl.native_tq import ContextGRPOGroupNSampler from twinkle_agentic.async_rl.pipeline import ( + _collect_adapter_path, _require_adapter_path, _train_batch, create_cpu_actor, @@ -221,6 +222,14 @@ def lazy_result(): _require_adapter_path(lazy_result, operation='test save') +def test_adapter_path_collects_lazy_remote_result(): + def lazy_result(): + return '/tmp/policy' + + lazy_result._is_lazy_collect = True + assert _collect_adapter_path(lazy_result, operation='test save') == '/tmp/policy' + + def test_training_policy_metrics_use_final_version_and_partial_span(): metrics = training_policy_metrics(( { From 05d845596b09b75e6e313dc3545095445ba8bd0a Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 25 Aug 2026 16:37:14 +0800 Subject: [PATCH 20/20] fix tests --- src/twinkle/infra/__init__.py | 2 +- src/twinkle/loss/grpo.py | 4 ++- tests/model/test_multi_lora.py | 6 ---- .../test_multi_lora_target_parameters.py | 30 +------------------ tests/server/utils/test_task_queue_mixin.py | 8 ++--- 5 files changed, 8 insertions(+), 42 deletions(-) diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index ae04779bc..a3c90eae3 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -830,7 +830,7 @@ def _notifying_result_func(*rargs, **rkwargs): wrapper._execute = execute wrapper._collect = collect wrapper._dispatch = dispatch - wrapper._lazy_collect = _lazy_collect if lazy_collect is None else lazy_collect + wrapper._lazy_collect = _lazy_collect wrapper._sync = sync return wrapper diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 900a91c09..85d48a82c 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -105,7 +105,8 @@ def _aggregate_loss( """ Aggregate per-token loss to scalar. - Mean over response tokens within each sequence, then over sequences. + Override this method in subclasses for different normalization. + Default: mean over sequences, then mean over batch. Args: per_token_loss: [batch, seq_len] per-token loss values @@ -115,6 +116,7 @@ def _aggregate_loss( Returns: loss: scalar loss value """ + # Per-sequence mean, then batch mean (aligned with Swift/TRL GRPO). # Each sequence contributes equally regardless of length. return ((per_token_loss * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).mean() diff --git a/tests/model/test_multi_lora.py b/tests/model/test_multi_lora.py index 8ae68c5ac..f2795a5ad 100644 --- a/tests/model/test_multi_lora.py +++ b/tests/model/test_multi_lora.py @@ -3,12 +3,6 @@ from twinkle.model.multi_lora import MultiLora -def test_multi_lora_transformers_save_disables_lazy_collect(): - from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel - - assert MultiLoraTransformersModel.save._lazy_collect is False - - def test_check_length_checks_each_sample_independently(): multi_lora = MultiLora(max_length=4) diff --git a/tests/model/test_multi_lora_target_parameters.py b/tests/model/test_multi_lora_target_parameters.py index d26b32147..b28ef6b36 100644 --- a/tests/model/test_multi_lora_target_parameters.py +++ b/tests/model/test_multi_lora_target_parameters.py @@ -45,16 +45,6 @@ def forward(self, x, expert_idx=0): return self.mlp.experts(x, expert_idx=expert_idx) -class FakeLinearModel(nn.Module): - - def __init__(self): - super().__init__() - self.proj = nn.Linear(4, 4) - - def forward(self, x): - return self.proj(x) - - def test_peft_target_parameter_key_shapes_for_3d_experts(): model = FakeModel() cfg = LoraConfig( @@ -223,24 +213,6 @@ def test_multilora_state_dict_round_trips_target_parameters(): assert torch.allclose(actual, expected, atol=1e-6) -def test_multilora_state_dict_without_target_parameters_does_not_require_slot(): - from twinkle.model.multi_lora import LoraTenant, MultiLora - - slot_cfg = LoraConfig(r=4, lora_alpha=8, target_modules=["proj"]) - model = get_peft_model(FakeLinearModel(), slot_cfg, adapter_name="lora_0") - multi_lora = MultiLora(max_loras=1, max_r=4) - multi_lora.module = model - multi_lora.loras = [LoraTenant(index=0, adapter_name="lora_0", config=slot_cfg)] - - tenant_cfg = LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]) - multi_lora.acquire_lora("adapter_a", tenant_cfg) - - assert "adapter_a" not in multi_lora.target_parameter_manager.tenant_to_slot - state = multi_lora.get_state_dict("adapter_a") - assert state - multi_lora.set_state_dict("adapter_a", state) - - def test_multilora_transformers_installs_target_parameters_once(): from twinkle.model.multi_lora import LoraTenant, MultiLora @@ -294,4 +266,4 @@ def test_multilora_transformers_installs_target_parameters_once(): assert test_target_parameter_multi_lora_updates_only_active_adapter() == True assert test_multilora_releases_target_parameter_slot_to_initial_weights() == True assert test_multilora_state_dict_round_trips_target_parameters() == True - assert test_multilora_transformers_installs_target_parameters_once() == True + assert test_multilora_transformers_installs_target_parameters_once() == True \ No newline at end of file diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index 2aa133a7e..f0bdbf963 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -58,8 +58,6 @@ async def test_preflight_rejects_batch_without_per_dp_multiple(): assert result == {'request_id': 'req1', 'model_id': 'model1'} _, kwargs = queue.state.records[-1] assert kwargs['result']['category'] == 'User' - assert 'token' not in kwargs - assert 'session_id' not in kwargs assert 'Batch size 2 must be divisible by 4' in kwargs['result']['error'] @@ -82,7 +80,7 @@ async def test_preflight_accepts_batch_with_per_dp_multiple(): @pytest.mark.asyncio -async def test_background_task_tracks_status_without_owner_or_preflight(): +async def test_background_task_tracks_status(): queue = _DummyQueue() async def work(): @@ -94,8 +92,8 @@ async def work(): ) await asyncio.sleep(0) - assert queue.state.records[0][0][1] == 'running' - assert all('token' not in kwargs and 'session_id' not in kwargs for _, kwargs in queue.state.records) + assert [args[1] for args, _ in queue.state.records] == ['running', 'completed'] + assert queue.state.records[-1][1]['result'] == {'ok': True} @pytest.mark.asyncio