diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index 585a530..92cf89f 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -4,7 +4,6 @@ import uuid from typing import Any, Dict, List, Optional, Tuple import torch -from torch import Tensor from astrai.inference.core.cache import ContiguousCache, KVCache from astrai.inference.core.executor import Executor diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index cfcf8e1..3d15a3f 100644 --- a/astrai/preprocessing/pipeline.py +++ b/astrai/preprocessing/pipeline.py @@ -23,7 +23,6 @@ import tqdm from astrai.config.preprocess_config import PipelineConfig from astrai.preprocessing.core import ( build_preprocessing_components, - iter_raw_records, primary_ids, ) from astrai.preprocessing.packing import PackingStrategyFactory diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index f11bdaf..713e22d 100644 --- a/astrai/trainer/rollout.py +++ b/astrai/trainer/rollout.py @@ -18,7 +18,6 @@ from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple import torch -import torch.nn as nn from torch import Tensor from astrai.inference.core.scheduler import InferenceScheduler @@ -33,25 +32,26 @@ class RawRollout: Fields are designed to cover all common RL algorithms: GRPO, PPO, Online DPO, Rejection Sampling, etc. + + Fields: + prompts: Tokenized prompts, shape ``[B, P_len]``. + responses: Generated response token IDs, shape ``[B, G, R_max]``. + response_mask: Boolean mask for real (non-pad) response tokens, + shape ``[B, G, R_max]``. + logprobs_old: Per-token log-probs under the behaviour policy, + shape ``[B, G, R_max]``. + prompt_texts: Decoded prompt strings (for reward models that + need text). + response_texts: Decoded response strings, shape ``[B, G]`` + (for reward models). """ prompts: Tensor - """Tokenized prompts, shape ``[B, P_len]``.""" - responses: Tensor - """Generated response token IDs, shape ``[B, G, R_max]``.""" - response_mask: Tensor - """Boolean mask for real (non-pad) response tokens, shape ``[B, G, R_max]``.""" - logprobs_old: Tensor - """Per-token log-probs under the behaviour policy, shape ``[B, G, R_max]``.""" - prompt_texts: List[str] = field(default_factory=list) - """Decoded prompt strings (for reward models that need text).""" - response_texts: List[List[str]] = field(default_factory=list) - """Decoded response strings, shape ``[B, G]`` (for reward models).""" @dataclass(kw_only=True) @@ -60,10 +60,12 @@ class RolloutResult(RawRollout): Produced by :class:`RolloutRunner` once the :class:`BaseRewardModel` has scored the decoded responses. + + Fields: + rewards: Reward per response, shape ``[B, G]``. """ rewards: Tensor - """Reward per response, shape ``[B, G]``.""" class BaseRewardModel(ABC): @@ -126,26 +128,31 @@ class RolloutGenerator: self.rep_window = rep_window @torch.no_grad() - def generate(self, batch: Dict[str, Tensor]) -> RawRollout: - """Expand prompts by ``group_size`` and generate one response each.""" - prompt_ids = batch["input_ids"] if "input_ids" in batch else batch["prompts"] - prompt_mask = ( - batch["attention_mask"] if "attention_mask" in batch else (prompt_ids != 0) - ) - B, _ = prompt_ids.shape - G = self.group_size + def generate(self, batch: Dict) -> RawRollout: + """Expand prompts by ``group_size`` and generate one response each. - prompt_texts: List[str] = [] - flat_prompt_ids: List[List[int]] = [] - for i in range(B): - ids = prompt_ids[i, prompt_mask[i]].tolist() - text = self.tokenizer.decode(ids, skip_special_tokens=True) - for _ in range(G): - flat_prompt_ids.append(list(ids)) - prompt_texts.append(text) + Accepted batch formats (per sample, repeated B times): + + - **messages**: ``{"messages": [{"role": "user", "content": "..."}, ...]}`` + - **instruction + input + output**: ``{"instruction": "...", + "input": "...", "output": "..."}`` — mapped to ``system`` / + ``user`` / ``assistant`` messages; ``input`` and ``output`` + are optional and skipped when empty. + + Both are rendered through the tokenizer's chat template with + ``add_generation_prompt=True`` so rollout prompts match the + format the policy was SFT-trained on. + """ + prompt_texts, flat_prompt_ids = self._prepare_prompts(batch) + B = len(prompt_texts) + G = self.group_size + # Re-expand flat list to G copies per prompt for run_batch. + expanded_prompt_ids: List[List[int]] = [] + for ids in flat_prompt_ids: + expanded_prompt_ids.extend([list(ids)] * G) results = self.scheduler.run_batch( - flat_prompt_ids, + expanded_prompt_ids, max_tokens=self.max_tokens, temperature=self.temperature, top_k=self.top_k, @@ -161,7 +168,14 @@ class RolloutGenerator: max_len = max(max_len, len(token_ids)) max_len = max(max_len, 1) - device = prompt_ids.device + device = self.scheduler.device + P_len = max(len(ids) for ids in flat_prompt_ids) + prompts_tensor = torch.zeros(B, P_len, dtype=torch.long, device=device) + for i, ids in enumerate(flat_prompt_ids): + prompts_tensor[i, : len(ids)] = torch.tensor( + ids, dtype=torch.long, device=device + ) + responses = torch.full((B, G, max_len), _PAD, dtype=torch.long, device=device) response_mask = torch.zeros((B, G, max_len), dtype=torch.bool, device=device) logprobs_old = torch.zeros((B, G, max_len), dtype=torch.float, device=device) @@ -186,7 +200,7 @@ class RolloutGenerator: ) return RawRollout( - prompts=prompt_ids, + prompts=prompts_tensor, responses=responses, response_mask=response_mask, logprobs_old=logprobs_old, @@ -194,6 +208,72 @@ class RolloutGenerator: response_texts=response_texts, ) + def _prepare_prompts(self, batch: Dict) -> Tuple[List[str], List[List[int]]]: + """Render batch prompts to ``(texts, token_id_lists)``. + + Returns two parallel lists of length B (number of prompts in + the batch). Dispatches by batch keys: + + - ``"messages"``: treated as a pre-built message list per sample. + - ``"instruction"`` (optionally ``"input"`` and ``"output"``): mapped + to ``system`` / ``user`` / ``assistant`` messages respectively. + + Both paths go through the tokenizer's chat template with + ``add_generation_prompt=True``. + """ + if "messages" in batch: + messages_list = batch["messages"] + elif "instruction" in batch: + instructions = batch["instruction"] + B = len(instructions) + inputs = batch.get("input") or [""] * B + outputs = batch.get("output") or [""] * B + messages_list = [ + self._instruction_to_messages(i, u, o) + for i, u, o in zip(instructions, inputs, outputs) + ] + else: + raise ValueError( + "Rollout batch must contain either 'messages' or " + "'instruction' (optionally 'input'/'output'); got keys: " + f"{list(batch.keys())}" + ) + + prompt_texts: List[str] = [] + flat_prompt_ids: List[List[int]] = [] + for messages in messages_list: + text = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + ids = self.tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + ) + prompt_texts.append(text) + flat_prompt_ids.append(list(ids)) + return prompt_texts, flat_prompt_ids + + @staticmethod + def _instruction_to_messages( + instruction: str, inp: str = "", output: str = "" + ) -> List[Dict[str, str]]: + """Map instruction/input/output to chat messages. + + Role mapping follows the convention used throughout the + preprocessing pipeline: ``instruction`` → system, ``input`` → + user, ``output`` → assistant. Empty fields are skipped so a + bare instruction produces a ``[system]`` list and the chat + template's ``add_generation_prompt`` adds the assistant header + for sampling. + """ + messages: List[Dict[str, str]] = [] + if instruction: + messages.append({"role": "system", "content": instruction}) + if inp: + messages.append({"role": "user", "content": inp}) + if output: + messages.append({"role": "assistant", "content": output}) + return messages + class RolloutRunner: """Produces :class:`RolloutResult` from a prompt batch. diff --git a/scripts/eval/evaluate_rouge.py b/scripts/eval/evaluate_rouge.py index 4b5714d..ec2852a 100644 --- a/scripts/eval/evaluate_rouge.py +++ b/scripts/eval/evaluate_rouge.py @@ -15,7 +15,7 @@ Usage:: import argparse import json from collections import Counter -from typing import Dict, List, Tuple +from typing import Dict, List def _tokenize(text: str) -> List[str]: diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index 827781e..516e6a7 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -13,23 +13,40 @@ from astrai.trainer.rollout import BaseRewardModel from astrai.trainer.schedule import SchedulerFactory from astrai.trainer.trainer import Trainer +_CHAT_TEMPLATE = ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}" + "SYSTEM: {{ message['content'] }}\n" + "{% elif message['role'] == 'user' %}" + "USER: {{ message['content'] }}\n" + "{% elif message['role'] == 'assistant' %}" + "ASSISTANT: {{ message['content'] }}\n" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}ASSISTANT: {% endif %}" +) -class PromptDataset(Dataset): - """Toy prompt-only dataset for online RL rollout.""" - def __init__(self, n=4, seq_len=8, vocab_size=1000): - self.n = n - self.seq_len = seq_len - self.vocab_size = vocab_size +class InstructionDataset(Dataset): + """Toy instruction/input dataset for online RL rollout. + + Each sample has an ``instruction`` and an optional ``input``; the + RolloutGenerator renders both through the tokenizer's chat template + so the prompt matches the SFT-trained format. + """ + + _SAMPLES = [ + {"instruction": "Hello", "input": ""}, + {"instruction": "Tell me a story", "input": "about dragons"}, + {"instruction": "Summarize", "input": "the article"}, + {"instruction": "Translate", "input": "to French: hi"}, + ] def __len__(self): - return self.n + return len(self._SAMPLES) def __getitem__(self, idx): - return { - "input_ids": torch.randint(3, self.vocab_size, (self.seq_len,)), - "attention_mask": torch.ones(self.seq_len, dtype=torch.bool), - } + return dict(self._SAMPLES[idx]) class LengthRewardModel(BaseRewardModel): @@ -48,6 +65,14 @@ class LengthRewardModel(BaseRewardModel): return rewards +def instruction_collate_fn(batch): + """Stack a list of instruction/input dicts into a batch dict of lists.""" + return { + "instruction": [b["instruction"] for b in batch], + "input": [b.get("input", "") for b in batch], + } + + def _model_fn(model_config): return AutoRegressiveLM(model_config).to(dtype=torch.float32) @@ -70,15 +95,16 @@ def test_online_dpo_end_to_end(base_test_env): tokenizer = base_test_env["tokenizer"] model_config = base_test_env["transformer_config"] - # base_test_env already wrote config.json into test_dir; we only need - # to drop the tokenizer files so AutoTokenizer.from_pretrained works. + # Equip tokenizer with a chat template so RolloutGenerator can + # render instruction/input via apply_chat_template. + tokenizer.set_chat_template(_CHAT_TEMPLATE) tokenizer.save_pretrained(test_dir) model_fn = partial(_model_fn, model_config) optimizer_fn = _optimizer_fn scheduler_fn = _scheduler_fn - dataset = PromptDataset(n=4, seq_len=8, vocab_size=model_config.vocab_size) + dataset = InstructionDataset() train_config = TrainConfig( strategy="online_dpo", @@ -103,7 +129,7 @@ def test_online_dpo_end_to_end(base_test_env): rollout_top_p=1.0, rollout_max_tokens=4, reward_model_fn=LengthRewardModel, - collate_fn=None, + collate_fn=instruction_collate_fn, ) trainer = Trainer(train_config) diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index 72c660c..05677f0 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -1,10 +1,4 @@ -"""Unit tests for the online rollout module. - -Covers :class:`RolloutResult` / :class:`RawRollout`, :class:`BaseRewardModel`, -:class:`RolloutGenerator` (KV-cache-backed via :class:`InferenceScheduler.run_batch`) -and :class:`RolloutRunner` including its internal cache and rollout-interval -trigger logic. -""" +"""Unit tests for the online rollout module.""" import pytest import torch @@ -20,26 +14,49 @@ from astrai.trainer.rollout import ( RolloutRunner, ) +_CHAT_TEMPLATE = ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}SYSTEM: {{ message['content'] }}\n{% endif %}" + "{% if message['role'] == 'user' %}USER: {{ message['content'] }}\n{% endif %}" + "{% if message['role'] == 'assistant' %}ASSISTANT: {{ message['content'] }}\n{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}ASSISTANT: {% endif %}" +) + class FakeTokenizer: - """Minimal char-level tokenizer stub for rollout tests. - - Vocab: 0 = pad, 1..255 = byte values. ``stop_ids = [2]`` (a fake - EOS) so tests can verify early-stopping behaviour. - """ + """Minimal stub tokenizer with a chat template for rollout tests.""" stop_ids = [2] - def encode(self, texts, out_ids=True, **_): + def __init__(self): + from astrai.tokenize.chat_template import ChatTemplate + + self._chat_template = ChatTemplate.from_string(_CHAT_TEMPLATE) + + def encode(self, texts, **_): if isinstance(texts, str): texts = [texts] return [[b for b in t.encode("utf-8")] for t in texts] def decode(self, ids, skip_special_tokens=True): - out = bytes(b for b in ids if b > 2 or not skip_special_tokens).decode( - "utf-8", errors="ignore" + if isinstance(ids, list): + return bytes(b for b in ids if b > 2).decode("utf-8", errors="ignore") + return str(ids) + + def apply_chat_template( + self, messages, tokenize=True, add_generation_prompt=True, **_ + ): + rendered = self._chat_template.render( + messages=messages, add_generation_prompt=add_generation_prompt ) - return out + if tokenize: + return ( + self.encode(rendered)[0] + if isinstance(rendered, str) + else [self.encode(t)[0] for t in rendered] + ) + return rendered class ConstantRewardModel(BaseRewardModel): @@ -84,10 +101,11 @@ def _make_scheduler(model, tokenizer, max_batch_size=8, max_len=128): ) -def _make_prompt_batch(batch_size=2, prompt_len=6, device="cpu"): - ids = torch.randint(3, 200, (batch_size, prompt_len), device=device) - mask = torch.ones(batch_size, prompt_len, dtype=torch.bool, device=device) - return {"input_ids": ids, "attention_mask": mask} +def _make_instruction_batch(n=2): + """Build a batch of instruction+input prompts as lists of strings.""" + instructions = [f"Tell me about topic {i}" for i in range(n)] + inputs = [f"context {i}" for i in range(n)] + return {"instruction": instructions, "input": inputs} def test_raw_rollout_fields(): @@ -114,8 +132,6 @@ def test_rollout_result_inherits_raw_rollout_fields(): assert r.rewards.shape == (2, 3) assert r.prompts.shape == (2, 4) assert r.responses.shape == (2, 3, 5) - assert r.prompt_texts == [] - assert r.response_texts == [] def test_base_reward_model_is_abstract(): @@ -158,9 +174,8 @@ def _make_generator(device, **kw): def test_rollout_generator_shapes(device): gen, _ = _make_generator(device, group_size=3, max_tokens=5) - batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device) + batch = _make_instruction_batch(n=2) r = gen.generate(batch) - assert r.prompts.shape == (2, 4) assert r.responses.shape == (2, 3, 5) assert r.response_mask.shape == (2, 3, 5) assert r.logprobs_old.shape == (2, 3, 5) @@ -172,14 +187,12 @@ def test_rollout_generator_shapes(device): def test_rollout_generator_mask_matches_responses(device): """Positions beyond a response's length are pad (mask False).""" gen, _ = _make_generator(device, group_size=2, max_tokens=6) - batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device) + batch = _make_instruction_batch(n=2) r = gen.generate(batch) for i in range(2): for g in range(2): real = r.response_mask[i, g].sum().item() - # Pad positions should be 0 assert r.responses[i, g, real:].sum() == 0 - # logprobs after the real tokens are 0 (padding) if real < r.logprobs_old.size(-1): assert torch.all(r.logprobs_old[i, g, real:] == 0) @@ -187,7 +200,7 @@ def test_rollout_generator_mask_matches_responses(device): def test_rollout_generator_logprobs_are_nonpositive(device): """Behaviour-policy logprobs of sampled tokens should be ≤ 0.""" gen, _ = _make_generator(device, group_size=2, max_tokens=4) - batch = _make_prompt_batch(batch_size=1, prompt_len=3, device=device) + batch = _make_instruction_batch(n=1) r = gen.generate(batch) for i in range(1): for g in range(2): @@ -196,6 +209,45 @@ def test_rollout_generator_logprobs_are_nonpositive(device): assert torch.all(lp <= 1e-5) +def test_rollout_generator_instruction_role_mapping(device): + """instruction → system, input → user, output → assistant.""" + gen, _ = _make_generator(device, group_size=1, max_tokens=4) + batch = { + "instruction": ["Be helpful"], + "input": ["What is 2+2?"], + "output": ["Four"], + } + r = gen.generate(batch) + text = r.prompt_texts[0] + assert "SYSTEM: Be helpful" in text + assert "USER: What is 2+2?" in text + assert "ASSISTANT: Four" in text + + +def test_rollout_generator_messages_format(device): + """Rollout also accepts pre-built messages.""" + gen, _ = _make_generator(device, group_size=2, max_tokens=4) + batch = { + "messages": [ + [{"role": "user", "content": "Hello"}], + [{"role": "user", "content": "Goodbye"}], + ] + } + r = gen.generate(batch) + assert r.responses.shape[0] == 2 + assert len(r.prompt_texts) == 2 + assert "Hello" in r.prompt_texts[0] or "USER" in r.prompt_texts[0] + + +def test_rollout_generator_bad_batch_raises(device): + """Batch without messages or instruction raises a clear error.""" + gen, _ = _make_generator(device) + with pytest.raises( + ValueError, match="must contain either 'messages' or 'instruction'" + ): + gen.generate({"input_ids": torch.zeros(2, 4, dtype=torch.long)}) + + def _make_runner(device, **kw): generator, model = _make_generator( device, @@ -217,10 +269,9 @@ def _make_runner(device, **kw): def test_rollout_runner_shapes(device): runner, _ = _make_runner(device, group_size=3, max_tokens=5) - batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device) + batch = _make_instruction_batch(n=2) r, is_fresh = runner(batch) assert is_fresh - assert r.prompts.shape == (2, 4) assert r.responses.shape == (2, 3, 5) assert r.response_mask.shape == (2, 3, 5) assert r.rewards.shape == (2, 3) @@ -232,7 +283,7 @@ def test_rollout_runner_shapes(device): def test_rollout_runner_cache_returns_stale_flag(device): runner, _ = _make_runner(device, rollout_interval=10) - batch = _make_prompt_batch(device=device) + batch = _make_instruction_batch() r1, fresh1 = runner(batch) r2, fresh2 = runner(batch) assert r1 is r2 @@ -242,7 +293,7 @@ def test_rollout_runner_cache_returns_stale_flag(device): def test_rollout_runner_step_triggers_new_rollout(device): runner, _ = _make_runner(device, rollout_interval=2) - batch = _make_prompt_batch(device=device) + batch = _make_instruction_batch() r1, fresh1 = runner(batch) assert fresh1 is True runner.step() @@ -259,7 +310,7 @@ def test_rollout_runner_step_triggers_new_rollout(device): def test_rollout_runner_clear_cache_forces_rerun(device): runner, _ = _make_runner(device, rollout_interval=100) - batch = _make_prompt_batch(device=device) + batch = _make_instruction_batch() r1, _ = runner(batch) runner.clear_cache() r2, fresh2 = runner(batch) @@ -269,7 +320,7 @@ def test_rollout_runner_clear_cache_forces_rerun(device): def test_rollout_runner_step_resets_counter(device): runner, _ = _make_runner(device, rollout_interval=1) - batch = _make_prompt_batch(device=device) + batch = _make_instruction_batch() r1, _ = runner(batch) runner.step() r2, fresh2 = runner(batch)