From 84753d3e083590a630a24fd5dacead88aabd5b65 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 20 Aug 2026 01:52:49 +0800 Subject: [PATCH] refactor: deduplicate and restructure test suite - extract preprocessing config factories into tests/data/factories.py - keep conftest.py fixtures-only; stop importing builders from it - promote temp_dir fixture to root conftest for cross-directory reuse - unify duplicate BPE tokenizer builders into build_test_tokenizer - merge grpo/dpo online e2e tests into one parametrized integration test - extract engine mock factory and shared model batch builders - drop local tempfile usage in favor of shared fixtures No behavior change: 519 tests pass. --- tests/conftest.py | 41 +++-- tests/data/conftest.py | 213 +++++-------------------- tests/data/factories.py | 86 ++++++++++ tests/data/test_dataset.py | 4 +- tests/data/test_preprocess_builder.py | 26 +-- tests/data/test_preprocess_config.py | 14 +- tests/data/test_preprocess_pipeline.py | 16 +- tests/helpers.py | 40 +++++ tests/inference/test_engine.py | 43 +++-- tests/module/test_lora.py | 70 ++++---- tests/module/test_model_forward.py | 51 +++--- tests/trainer/test_grpo_e2e.py | 123 -------------- tests/trainer/test_online_e2e.py | 39 ++--- 13 files changed, 312 insertions(+), 454 deletions(-) create mode 100644 tests/data/factories.py delete mode 100644 tests/trainer/test_grpo_e2e.py diff --git a/tests/conftest.py b/tests/conftest.py index 47f302d..d957683 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,12 +5,15 @@ import tempfile import pytest import torch -from tokenizers import Tokenizer, models, pre_tokenizers, trainers from astrai.extension import KERNEL_NAMES, is_available from astrai.model.transformer import AutoRegressiveLM -from astrai.tokenize import AutoTokenizer -from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config +from tests.helpers import ( + TINY_CONFIG, + RandomTokenDataset, + build_test_tokenizer, + make_tiny_config, +) CUDA_AVAIL = torch.cuda.is_available() KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES) @@ -30,18 +33,9 @@ def device(): return "cuda" if torch.cuda.is_available() else "cpu" -def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer: +def create_test_tokenizer(vocab_size: int = 1000): """Create a simple tokenizer for testing purposes.""" - tokenizer = Tokenizer(models.BPE()) - tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel() - trainer = trainers.BpeTrainer( - vocab_size=vocab_size, min_frequency=1, special_tokens=["", ""] - ) - tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer) - auto_tokenizer = AutoTokenizer() - auto_tokenizer._tokenizer = tokenizer - auto_tokenizer._special_token_map = {"unk_token": "", "pad_token": ""} - return auto_tokenizer + return build_test_tokenizer(vocab_size) @pytest.fixture(scope="session") @@ -59,24 +53,29 @@ def test_model(device): @pytest.fixture -def base_test_env(test_model, test_tokenizer): +def temp_dir(): + """Function-scoped temporary directory, cleaned up after each test.""" + d = tempfile.mkdtemp() + yield d + shutil.rmtree(d, ignore_errors=True) + + +@pytest.fixture +def base_test_env(test_model, test_tokenizer, temp_dir): """Function-scoped test environment with isolated temp directory.""" - test_dir = tempfile.mkdtemp() - config_path = os.path.join(test_dir, "config.json") + config_path = os.path.join(temp_dir, "config.json") with open(config_path, "w") as f: json.dump(TINY_CONFIG, f) - yield { + return { "device": test_model["device"], - "test_dir": str(test_dir), + "test_dir": temp_dir, "config_path": config_path, "transformer_config": test_model["config"], "model": test_model["model"], "tokenizer": test_tokenizer, } - shutil.rmtree(test_dir) - @pytest.fixture def random_dataset(): diff --git a/tests/data/conftest.py b/tests/data/conftest.py index 848ac9c..a8ab651 100644 --- a/tests/data/conftest.py +++ b/tests/data/conftest.py @@ -1,21 +1,15 @@ import json import os -import tempfile import pytest -from tokenizers import Tokenizer, models, pre_tokenizers, trainers -from astrai.config.preprocess_config import ( - InputConfig, - PipelineConfig, - ProcessingConfig, -) from astrai.preprocessing.builder import ( MultiOutputMaskBuilder, SectionedMaskBuilder, SingleOutputMaskBuilder, ) -from astrai.tokenize import AutoTokenizer +from tests.data.factories import make_grpo_config +from tests.helpers import build_test_tokenizer _SPECIAL_TOKENS_CONFIG = { "bos_token": "<|begin_of_sentence|>", @@ -41,55 +35,42 @@ _CHAT_TEMPLATE = ( "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}" ) -_CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}] -_INSTRUCTION_SECTIONS = [ - {"field": "prompt", "action": "mask", "add_special_tokens": True}, - {"field": "response", "action": "train"}, +_CHAT_TOKENIZER_DATA = [ + "hello world", + "Hi there!", + "You are helpful.", + "What is 2+2?", + "Tell me a story about dragons and knights.", + "Sure, here is a tale.", + "Translate to French: Hello", + "Bonjour", + "Artificial Intelligence is a field of computer science.", + "system", + "user", + "assistant", + "<|im_start|>", + "<|im_end|>", + *[chr(i) for i in range(32, 127)], ] -_TEXT_SECTIONS = [{"field": "text", "action": "train"}] - -_GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}] +_CHAT_TOKENIZER_MAP = { + "bos_token": "<|begin_of_sentence|>", + "eos_token": "<|end_of_sentence|>", + "pad_token": "<|_pad_|>", + "unk_token": "<|_unk_|>", +} def _build_chat_tokenizer(): - tok = Tokenizer(models.BPE()) - tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) - tr = trainers.BpeTrainer( + return build_test_tokenizer( vocab_size=512, - min_frequency=1, special_tokens=_SPECIAL_TOKENS, + special_token_map=_CHAT_TOKENIZER_MAP, + add_prefix_space=False, + train_data=_CHAT_TOKENIZER_DATA, + chat_template=_CHAT_TEMPLATE, ) - train_data = [ - "hello world", - "Hi there!", - "You are helpful.", - "What is 2+2?", - "Tell me a story about dragons and knights.", - "Sure, here is a tale.", - "Translate to French: Hello", - "Bonjour", - "Artificial Intelligence is a field of computer science.", - "system", - "user", - "assistant", - "<|im_start|>", - "<|im_end|>", - *[chr(i) for i in range(32, 127)], - ] - tok.train_from_iterator(train_data, tr) - - auto_tok = AutoTokenizer() - auto_tok._tokenizer = tok - auto_tok._special_token_map = { - "bos_token": "<|begin_of_sentence|>", - "eos_token": "<|end_of_sentence|>", - "pad_token": "<|_pad_|>", - "unk_token": "<|_unk_|>", - } - auto_tok.set_chat_template(_CHAT_TEMPLATE) - return auto_tok @pytest.fixture(scope="session") @@ -97,116 +78,11 @@ def chat_tokenizer(): return _build_chat_tokenizer() -@pytest.fixture -def temp_dir(): - d = tempfile.mkdtemp() - yield d - import shutil - - shutil.rmtree(d, ignore_errors=True) - - -def make_chat_config(): - return PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), - mask={"system": "mask", "user": "mask", "assistant": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) - - -def make_instruction_config(): - return PipelineConfig( - input=InputConfig(sections=_INSTRUCTION_SECTIONS), - mask={"prompt": "mask", "response": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) - - -def make_text_config(): - return PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), - preprocessing=ProcessingConfig( - max_seq_len=2048, min_chars=1, max_chars=2_000_000 - ), - ) - - -def make_dpo_chat_config(): - return PipelineConfig( - input=InputConfig( - sources={ - "chosen": { - "sections": [ - {"field": "chosen", "action": "$role", "template": True} - ] - }, - "rejected": { - "sections": [ - {"field": "rejected", "action": "$role", "template": True} - ] - }, - } - ), - mask={"user": "mask", "assistant": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) - - -def make_grpo_config(): - return PipelineConfig( - input=InputConfig( - sources={ - "prompts": { - "sections": [ - {"field": "prompt", "action": "mask", "template": True} - ] - }, - "responses": { - "sections": _GRPO_RESPONSE_SECTIONS, - "list_field": True, - "mask_key": "masks", - }, - "rewards": { - "sections": [{"field": "rewards", "action": "value"}], - }, - } - ), - mask={"user": "mask", "assistant": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) - - -def make_grpo_no_template_config(): - return PipelineConfig( - input=InputConfig( - sources={ - "prompts": { - "sections": [ - { - "field": "prompt", - "action": "mask", - "add_special_tokens": True, - } - ] - }, - "responses": { - "sections": _GRPO_RESPONSE_SECTIONS, - "list_field": True, - "mask_key": "masks", - }, - "rewards": { - "sections": [{"field": "rewards", "action": "value"}], - }, - } - ), - mask={"user": "mask", "assistant": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) +def _write_tokenizer_dir(dir_path, tokenizer, tokenizer_config): + """Persist a tokenizer plus ``tokenizer_config.json`` into *dir_path*.""" + tokenizer._tokenizer.save(os.path.join(dir_path, "tokenizer.json")) + with open(os.path.join(dir_path, "tokenizer_config.json"), "w") as f: + json.dump(tokenizer_config, f) @pytest.fixture @@ -228,11 +104,11 @@ def multi_builder(): def tokenizer_dir(temp_dir, test_tokenizer): d = os.path.join(temp_dir, "tok") os.makedirs(d, exist_ok=True) - test_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json")) - with open(os.path.join(d, "tokenizer_config.json"), "w") as f: - json.dump( - {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, f - ) + _write_tokenizer_dir( + d, + test_tokenizer, + {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, + ) return d @@ -240,10 +116,9 @@ def tokenizer_dir(temp_dir, test_tokenizer): def chat_tokenizer_dir(temp_dir, chat_tokenizer): d = os.path.join(temp_dir, "tok") os.makedirs(d, exist_ok=True) - chat_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json")) - with open(os.path.join(d, "tokenizer_config.json"), "w") as f: - json.dump( - {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE}, - f, - ) + _write_tokenizer_dir( + d, + chat_tokenizer, + {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE}, + ) return d diff --git a/tests/data/factories.py b/tests/data/factories.py new file mode 100644 index 0000000..4a5d12b --- /dev/null +++ b/tests/data/factories.py @@ -0,0 +1,86 @@ +"""Test data builders for preprocessing and dataset scenarios.""" + +from astrai.config.preprocess_config import ( + InputConfig, + PipelineConfig, + ProcessingConfig, +) + +CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}] +INSTRUCTION_SECTIONS = [ + {"field": "prompt", "action": "mask", "add_special_tokens": True}, + {"field": "response", "action": "train"}, +] +TEXT_SECTIONS = [{"field": "text", "action": "train"}] +GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}] + + +def make_pipeline_config(sections, *, mask=None, preprocessing=None, sources=None): + """Build a pipeline config with the common test defaults.""" + return PipelineConfig( + input=InputConfig(sections=sections, sources=sources), + mask={} if mask is None else mask, + mask_default="mask", + preprocessing=preprocessing or ProcessingConfig(max_seq_len=2048), + ) + + +def make_chat_config(): + return make_pipeline_config( + CHAT_SECTIONS, + mask={"system": "mask", "user": "mask", "assistant": "train"}, + ) + + +def make_instruction_config(): + return make_pipeline_config( + INSTRUCTION_SECTIONS, + mask={"prompt": "mask", "response": "train"}, + ) + + +def make_text_config(): + return make_pipeline_config( + TEXT_SECTIONS, + preprocessing=ProcessingConfig( + max_seq_len=2048, min_chars=1, max_chars=2_000_000 + ), + ) + + +def make_dpo_chat_config(): + sources = { + name: {"sections": [{"field": name, "action": "$role", "template": True}]} + for name in ("chosen", "rejected") + } + return make_pipeline_config( + None, + mask={"user": "mask", "assistant": "train"}, + sources=sources, + ) + + +def make_grpo_config(*, template=True): + prompt_section = {"field": "prompt", "action": "mask"} + if template: + prompt_section["template"] = True + else: + prompt_section["add_special_tokens"] = True + sources = { + "prompts": {"sections": [prompt_section]}, + "responses": { + "sections": GRPO_RESPONSE_SECTIONS, + "list_field": True, + "mask_key": "masks", + }, + "rewards": {"sections": [{"field": "rewards", "action": "value"}]}, + } + return make_pipeline_config( + None, + mask={"user": "mask", "assistant": "train"}, + sources=sources, + ) + + +def make_grpo_no_template_config(): + return make_grpo_config(template=False) diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index b2b74e7..e926427 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -24,7 +24,7 @@ from astrai.serialization import ( load_bin, save_bin, ) -from tests.data.conftest import make_grpo_no_template_config +from tests.data.factories import make_grpo_config def _rand_seq(length, vocab=1000): @@ -797,7 +797,7 @@ def test_grpo_builder_preserves_response_boundaries(base_test_env): _save_test_tokenizer(base_test_env["test_dir"], tokenizer) builder = SectionedMaskBuilder() - config = make_grpo_no_template_config() + config = make_grpo_config(template=False) config.preprocessing.max_seq_len = 128 item = { diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index 3d0e83c..cd7db16 100644 --- a/tests/data/test_preprocess_builder.py +++ b/tests/data/test_preprocess_builder.py @@ -12,10 +12,10 @@ from astrai.preprocessing.builder import ( SectionedMaskBuilder, SingleOutputMaskBuilder, ) -from tests.data.conftest import ( - _CHAT_SECTIONS, - _INSTRUCTION_SECTIONS, - _TEXT_SECTIONS, +from tests.data.factories import ( + CHAT_SECTIONS, + INSTRUCTION_SECTIONS, + TEXT_SECTIONS, make_chat_config, make_dpo_chat_config, make_grpo_config, @@ -101,7 +101,7 @@ def test_chat_uniform_masking( mask_rules, mask_default, expect_nonzero, chat_tokenizer, builder ): config = PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), + input=InputConfig(sections=CHAT_SECTIONS), mask=mask_rules, mask_default=mask_default, preprocessing=ProcessingConfig(max_seq_len=2048), @@ -128,7 +128,7 @@ def test_chat_empty_messages(chat_tokenizer, builder): def test_chat_domain_extraction(chat_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), + input=InputConfig(sections=CHAT_SECTIONS), mask={"assistant": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), @@ -147,7 +147,7 @@ def test_chat_domain_extraction(chat_tokenizer, builder): def test_chat_truncation(chat_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), + input=InputConfig(sections=CHAT_SECTIONS), mask={"assistant": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=10), @@ -237,7 +237,7 @@ def test_text_empty(test_tokenizer, builder): def test_text_too_short(test_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), + input=InputConfig(sections=TEXT_SECTIONS), preprocessing=ProcessingConfig(min_chars=100), ) assert builder.build({"text": "short"}, config, test_tokenizer) is None @@ -245,7 +245,7 @@ def test_text_too_short(test_tokenizer, builder): def test_text_truncation(test_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), + input=InputConfig(sections=TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1), ) item = {"text": "This is a very long text that should be truncated"} @@ -255,7 +255,7 @@ def test_text_truncation(test_tokenizer, builder): def test_sectioned_chat(chat_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), + input=InputConfig(sections=CHAT_SECTIONS), mask={"system": "mask", "user": "mask", "assistant": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), @@ -275,7 +275,7 @@ def test_sectioned_chat(chat_tokenizer, builder): def test_sectioned_instruction(test_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_INSTRUCTION_SECTIONS), + input=InputConfig(sections=INSTRUCTION_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0), ) item = {"prompt": "Q: Why?", "response": "A: Because."} @@ -288,7 +288,7 @@ def test_sectioned_instruction(test_tokenizer, builder): def test_sectioned_text(test_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), + input=InputConfig(sections=TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1), ) item = {"text": "Hello world, this is a test."} @@ -299,7 +299,7 @@ def test_sectioned_text(test_tokenizer, builder): def test_sectioned_text_too_short(test_tokenizer, builder): config = PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), + input=InputConfig(sections=TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100), ) assert builder.build({"text": "short"}, config, test_tokenizer) is None diff --git a/tests/data/test_preprocess_config.py b/tests/data/test_preprocess_config.py index e1a470a..1500595 100644 --- a/tests/data/test_preprocess_config.py +++ b/tests/data/test_preprocess_config.py @@ -4,9 +4,9 @@ from astrai.config.preprocess_config import ( InputConfig, PipelineConfig, ) -from tests.data.conftest import ( - _INSTRUCTION_SECTIONS, - _TEXT_SECTIONS, +from tests.data.factories import ( + INSTRUCTION_SECTIONS, + TEXT_SECTIONS, make_dpo_chat_config, ) @@ -43,26 +43,26 @@ def test_from_dict_flat(): def test_to_dict_roundtrip(): config = PipelineConfig( - input=InputConfig(sections=_INSTRUCTION_SECTIONS), + input=InputConfig(sections=INSTRUCTION_SECTIONS), mask={"prompt": "mask", "response": "train"}, mask_default="mask", ) d = config.to_dict() config2 = PipelineConfig.from_dict(d) - assert config2.input.sections == _INSTRUCTION_SECTIONS + assert config2.input.sections == INSTRUCTION_SECTIONS assert config2.mask == {"prompt": "mask", "response": "train"} def test_to_file_from_file(temp_dir): config = PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), + input=InputConfig(sections=TEXT_SECTIONS), mask={"text": "train"}, mask_default="mask", ) path = os.path.join(temp_dir, "config.json") config.to_file(path) loaded = PipelineConfig.from_file(path) - assert loaded.input.sections == _TEXT_SECTIONS + assert loaded.input.sections == TEXT_SECTIONS assert loaded.mask == {"text": "train"} diff --git a/tests/data/test_preprocess_pipeline.py b/tests/data/test_preprocess_pipeline.py index 3857a79..125ae8e 100644 --- a/tests/data/test_preprocess_pipeline.py +++ b/tests/data/test_preprocess_pipeline.py @@ -9,10 +9,10 @@ from astrai.config.preprocess_config import ( ) from astrai.preprocessing.packing import PackingStrategyFactory from astrai.preprocessing.pipeline import Pipeline, filter_by_length -from tests.data.conftest import ( - _CHAT_SECTIONS, - _INSTRUCTION_SECTIONS, - _TEXT_SECTIONS, +from tests.data.factories import ( + CHAT_SECTIONS, + INSTRUCTION_SECTIONS, + TEXT_SECTIONS, make_dpo_chat_config, make_grpo_no_template_config, ) @@ -54,7 +54,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir): ) config = PipelineConfig( - input=InputConfig(sections=_CHAT_SECTIONS), + input=InputConfig(sections=CHAT_SECTIONS), mask={"system": "mask", "user": "mask", "assistant": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), @@ -97,7 +97,7 @@ def test_full_text_pipeline(temp_dir, tokenizer_dir): ) config = PipelineConfig( - input=InputConfig(sections=_TEXT_SECTIONS), + input=InputConfig(sections=TEXT_SECTIONS), preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=10), output=OutputConfig(storage_format="bin"), ) @@ -138,7 +138,7 @@ def test_full_instruction_pipeline(temp_dir, tokenizer_dir): ) config = PipelineConfig( - input=InputConfig(sections=_INSTRUCTION_SECTIONS), + input=InputConfig(sections=INSTRUCTION_SECTIONS), mask={"prompt": "mask", "response": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), @@ -164,7 +164,7 @@ def test_dtype_override(temp_dir, tokenizer_dir): f.write(json.dumps({"prompt": "Q", "response": "A"}) + "\n") config = PipelineConfig( - input=InputConfig(sections=_INSTRUCTION_SECTIONS), + input=InputConfig(sections=INSTRUCTION_SECTIONS), mask={"prompt": "mask", "response": "train"}, mask_default="mask", preprocessing=ProcessingConfig(max_seq_len=2048), diff --git a/tests/helpers.py b/tests/helpers.py index 474977e..a0bdf2a 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -4,10 +4,12 @@ import json import os import torch +from tokenizers import Tokenizer, models, pre_tokenizers, trainers from torch.utils.data import Dataset from astrai.config.model_config import AutoRegressiveLMConfig from astrai.model.transformer import AutoRegressiveLM +from astrai.tokenize import AutoTokenizer TINY_CONFIG = dict( vocab_size=1000, @@ -57,6 +59,44 @@ def make_model(device, **cfg_overrides): return model, cfg +def build_test_tokenizer( + vocab_size: int = 1000, + *, + special_tokens=("", ""), + special_token_map=None, + add_prefix_space: bool = True, + train_data=None, + chat_template: str | None = None, +) -> AutoTokenizer: + """Build a lightweight BPE ``AutoTokenizer`` for tests. + + ``special_token_map`` defaults to ``{"unk_token", "pad_token"}`` + pointing at the first two special tokens. + """ + tokenizer = Tokenizer(models.BPE()) + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel( + add_prefix_space=add_prefix_space + ) + trainer = trainers.BpeTrainer( + vocab_size=vocab_size, + min_frequency=1, + special_tokens=list(special_tokens), + ) + tokenizer.train_from_iterator( + train_data if train_data is not None else [chr(i) for i in range(256)], + trainer, + ) + auto_tokenizer = AutoTokenizer() + auto_tokenizer._tokenizer = tokenizer + auto_tokenizer._special_token_map = special_token_map or { + "unk_token": special_tokens[0], + "pad_token": special_tokens[1], + } + if chat_template is not None: + auto_tokenizer.set_chat_template(chat_template) + return auto_tokenizer + + def make_frozen(model, device): """Create a frozen, eval-mode copy of *model* with identical weights.""" cfg = make_rollout_config() diff --git a/tests/inference/test_engine.py b/tests/inference/test_engine.py index 18a2026..dab8c1d 100644 --- a/tests/inference/test_engine.py +++ b/tests/inference/test_engine.py @@ -8,6 +8,17 @@ from astrai.inference import STOP from astrai.inference.engine import GenerateResult, InferenceEngine +def _make_engine_mocks(decode=None): + """Build the standard mock model/tokenizer pair used by engine tests.""" + mock_model = MagicMock() + mock_tokenizer = MagicMock() + mock_tokenizer.encode.return_value = [1, 2, 3] + mock_tokenizer.stop_ids = [0] + if decode is not None: + mock_tokenizer.decode.return_value = decode + return mock_model, mock_tokenizer + + def test_result_append_single(): r = GenerateResult(count=1) r.append("hello", 0) @@ -102,11 +113,7 @@ def test_result_get_results(): def test_engine_generate_non_streaming_single(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() - mock_tokenizer.encode.return_value = [1, 2, 3] - mock_tokenizer.decode.return_value = "response" - mock_tokenizer.stop_ids = [0] + mock_model, mock_tokenizer = _make_engine_mocks(decode="response") with patch("astrai.inference.engine.InferenceScheduler") as MockSched: instance = MockSched.return_value @@ -125,11 +132,7 @@ def test_engine_generate_non_streaming_single(): def test_engine_generate_streaming_yields_tokens(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() - mock_tokenizer.encode.return_value = [1, 2, 3] - mock_tokenizer.decode.return_value = "tok" - mock_tokenizer.stop_ids = [0] + mock_model, mock_tokenizer = _make_engine_mocks(decode="tok") callbacks_saved = [] @@ -154,11 +157,7 @@ def test_engine_generate_streaming_yields_tokens(): def test_engine_generate_non_streaming_batch(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() - mock_tokenizer.encode.return_value = [1, 2, 3] - mock_tokenizer.decode.return_value = "r" - mock_tokenizer.stop_ids = [0] + mock_model, mock_tokenizer = _make_engine_mocks(decode="r") with patch("astrai.inference.engine.InferenceScheduler") as MockSched: instance = MockSched.return_value @@ -177,10 +176,7 @@ def test_engine_generate_non_streaming_batch(): def test_engine_generate_zero_max_tokens_returns_empty(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() - mock_tokenizer.encode.return_value = [1, 2, 3] - mock_tokenizer.stop_ids = [0] + mock_model, mock_tokenizer = _make_engine_mocks() with patch("astrai.inference.engine.InferenceScheduler") as MockSched: instance = MockSched.return_value @@ -192,8 +188,7 @@ def test_engine_generate_zero_max_tokens_returns_empty(): def test_engine_generate_zero_max_tokens_stream_is_empty(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() + mock_model, mock_tokenizer = _make_engine_mocks() with patch("astrai.inference.engine.InferenceScheduler") as MockSched: instance = MockSched.return_value @@ -203,8 +198,7 @@ def test_engine_generate_zero_max_tokens_stream_is_empty(): def test_engine_passes_backend_to_scheduler(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() + mock_model, mock_tokenizer = _make_engine_mocks() with patch("astrai.inference.engine.InferenceScheduler") as MockSched: InferenceEngine( @@ -218,8 +212,7 @@ def test_engine_passes_backend_to_scheduler(): def test_generate_captures_calling_backend_context(): - mock_model = MagicMock() - mock_tokenizer = MagicMock() + mock_model, mock_tokenizer = _make_engine_mocks() captured = [] with patch("astrai.inference.engine.InferenceScheduler") as MockSched: diff --git a/tests/module/test_lora.py b/tests/module/test_lora.py index 39d7b27..ceb5f4b 100644 --- a/tests/module/test_lora.py +++ b/tests/module/test_lora.py @@ -1,9 +1,6 @@ -import tempfile - import pytest import torch -from astrai.config.model_config import AutoRegressiveLMConfig from astrai.model import AutoRegressiveLM from astrai.model.components.linear import Linear from astrai.model.components.lora import ( @@ -16,22 +13,20 @@ from astrai.model.components.lora import ( merge_lora, save_lora, ) +from tests.helpers import make_tiny_config -MODEL_KWARGS = dict( +LORA_MODEL_KWARGS = dict( vocab_size=1000, hidden_size=64, num_attention_heads=4, num_key_value_heads=2, intermediate_size=128, - num_hidden_layers=2, max_position_embeddings=32, - rms_norm_eps=1e-5, ) def _make_model(**kwargs): - kw = {**MODEL_KWARGS, **kwargs} - config = AutoRegressiveLMConfig(**kw) + config = make_tiny_config(**{**LORA_MODEL_KWARGS, **kwargs}) model = AutoRegressiveLM(config) model.eval() return model @@ -227,7 +222,7 @@ def test_state_dict_after_inject_consistent_with_original(): assert len(lora_keys) > 0 -def test_save_load_roundtrip(): +def test_save_load_roundtrip(temp_dir): model = _make_model() cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"}) @@ -240,20 +235,19 @@ def test_save_load_roundtrip(): with torch.no_grad(): out_src = model(x)["logits"].clone() - with tempfile.TemporaryDirectory() as tmpdir: - save_lora(model, tmpdir, cfg) + save_lora(model, temp_dir, cfg) - model2 = _make_model() - model2.load_state_dict(model.state_dict(), strict=False) - load_lora(model2, tmpdir) + model2 = _make_model() + model2.load_state_dict(model.state_dict(), strict=False) + load_lora(model2, temp_dir) - with torch.no_grad(): - out_dst = model2(x)["logits"] + with torch.no_grad(): + out_dst = model2(x)["logits"] - torch.testing.assert_close(out_src, out_dst) + torch.testing.assert_close(out_src, out_dst) -def test_save_after_merge_raises(): +def test_save_after_merge_raises(temp_dir): model = _make_model() cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"}) @@ -262,16 +256,14 @@ def test_save_after_merge_raises(): if isinstance(m, LoRALinear): m.lora_B.fill_(0.5) - with tempfile.TemporaryDirectory() as tmpdir: - save_lora(model, tmpdir, cfg) - merge_lora(model) + save_lora(model, temp_dir, cfg) + merge_lora(model) - with tempfile.TemporaryDirectory() as tmpdir2: - with pytest.raises(RuntimeError, match="No LoRA parameters"): - save_lora(model, tmpdir2, cfg) + with pytest.raises(RuntimeError, match="No LoRA parameters"): + save_lora(model, temp_dir, cfg) -def test_load_lora_on_already_injected(): +def test_load_lora_on_already_injected(temp_dir): model = _make_model() inject_lora(model, r=4, alpha=8, target_modules={"q_proj"}) @@ -280,18 +272,17 @@ def test_load_lora_on_already_injected(): if isinstance(m, LoRALinear): m.lora_B.fill_(0.5) - with tempfile.TemporaryDirectory() as tmpdir: - save_lora(model, tmpdir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",))) + save_lora(model, temp_dir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",))) - model2 = _make_model() - model2.load_state_dict(model.state_dict(), strict=False) - inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"}) + model2 = _make_model() + model2.load_state_dict(model.state_dict(), strict=False) + inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"}) - load_lora(model2, tmpdir) - assert _get_lora_count(model2) > 0 + load_lora(model2, temp_dir) + assert _get_lora_count(model2) > 0 -def test_load_lora_mismatched_r_raises(): +def test_load_lora_mismatched_r_raises(temp_dir): model = _make_model() cfg = inject_lora(model, r=8, alpha=16, target_modules={"q_proj"}) @@ -300,15 +291,14 @@ def test_load_lora_mismatched_r_raises(): if isinstance(m, LoRALinear): m.lora_B.fill_(0.5) - with tempfile.TemporaryDirectory() as tmpdir: - save_lora(model, tmpdir, cfg) + save_lora(model, temp_dir, cfg) - model2 = _make_model() - model2.load_state_dict(model.state_dict(), strict=False) - inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"}) + model2 = _make_model() + model2.load_state_dict(model.state_dict(), strict=False) + inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"}) - with pytest.raises(RuntimeError, match="size mismatch"): - load_lora(model2, tmpdir) + with pytest.raises(RuntimeError, match="size mismatch"): + load_lora(model2, temp_dir) def test_merge_preserves_output(): diff --git a/tests/module/test_model_forward.py b/tests/module/test_model_forward.py index a88726d..2fe7aec 100644 --- a/tests/module/test_model_forward.py +++ b/tests/module/test_model_forward.py @@ -39,6 +39,21 @@ def _make_model(config=None) -> AutoRegressiveLM: return AutoRegressiveLM(config) +def _make_batch(config, batch_size=2, seq_len=8, with_extra=False): + """Build a random token batch, optionally with position ids and loss mask.""" + vocab = config.vocab_size + batch = { + "input_ids": torch.randint(0, vocab, (batch_size, seq_len)), + "target_ids": torch.randint(0, vocab, (batch_size, seq_len)), + } + if with_extra: + batch["position_ids"] = ( + torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + ) + batch["loss_mask"] = torch.ones(batch_size, seq_len, dtype=torch.bool) + return batch + + def test_model_forward_contract_uses_dense_training_and_packed_inference(): from astrai.inference.cache import PagePool, TaskCacheManager from astrai.inference.workspace import InferenceWorkspace @@ -180,13 +195,6 @@ class TestSEQStrategyMoE: self.model = _make_model(self.config).to(device) self.model.train() - def _make_batch(self, batch_size=2, seq_len=8): - vocab = self.config.vocab_size - input_ids = torch.randint(0, vocab, (batch_size, seq_len)) - # target = input shifted right - target_ids = torch.randint(0, vocab, (batch_size, seq_len)) - return {"input_ids": input_ids, "target_ids": target_ids} - def test_compute_loss_returns_scalar(self): """compute_loss should return a scalar tensor.""" strategy = SEQStrategy( @@ -194,7 +202,7 @@ class TestSEQStrategyMoE: self.device, moe_aux_loss_coef=0.01, ) - loss = strategy.compute_loss(self._make_batch()) + loss = strategy.compute_loss(_make_batch(self.config)) assert loss.ndim == 0 assert loss.requires_grad @@ -205,7 +213,7 @@ class TestSEQStrategyMoE: self.device, moe_aux_loss_coef=0.01, ) - output = strategy.compute_loss_output(self._make_batch()) + output = strategy.compute_loss_output(_make_batch(self.config)) assert "loss" in output assert "metrics" in output @@ -225,7 +233,7 @@ class TestSEQStrategyMoE: self.device, moe_aux_loss_coef=0.01, ) - strategy.compute_loss_output(self._make_batch()) + strategy.compute_loss_output(_make_batch(self.config)) moe_metrics = strategy._moe_metrics assert moe_metrics, "_moe_metrics should not be empty for MoE model" @@ -246,7 +254,7 @@ class TestSEQStrategyMoE: self.device, moe_aux_loss_coef=0.0, ) - output = strategy.compute_loss_output(self._make_batch()) + output = strategy.compute_loss_output(_make_batch(self.config)) metrics = output["metrics"] # task_loss and loss should be equal (aux weighted by zero) @@ -268,7 +276,7 @@ class TestSEQStrategyMoE: self.device, moe_aux_loss_coef=0.01, ) - output = strategy.compute_loss_output(self._make_batch()) + output = strategy.compute_loss_output(_make_batch(self.config)) assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12 def test_factory_creates_strategy_with_coef(self): @@ -294,7 +302,7 @@ class TestSEQStrategyMoE: self.device, moe_aux_loss_coef=0.01, ) - output = strategy.compute_loss_output(self._make_batch()) + output = strategy.compute_loss_output(_make_batch(self.config)) metrics = output["metrics"] assert "moe_aux_loss" not in metrics @@ -313,19 +321,6 @@ class TestSFTStrategyMoE: self.model = _make_model(self.config).to(device) self.model.train() - def _make_batch(self, batch_size=2, seq_len=8): - vocab = self.config.vocab_size - input_ids = torch.randint(0, vocab, (batch_size, seq_len)) - target_ids = torch.randint(0, vocab, (batch_size, seq_len)) - position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) - loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool) - return { - "input_ids": input_ids, - "target_ids": target_ids, - "position_ids": position_ids, - "loss_mask": loss_mask, - } - def test_compute_loss_output_with_aux_loss(self): """SFTStrategy produces MoE metrics when coef > 0.""" strategy = SFTStrategy( @@ -333,7 +328,7 @@ class TestSFTStrategyMoE: self.device, moe_aux_loss_coef=0.01, ) - output = strategy.compute_loss_output(self._make_batch()) + output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True)) metrics = output["metrics"] assert "moe_aux_loss" in metrics @@ -351,7 +346,7 @@ class TestSFTStrategyMoE: self.device, moe_aux_loss_coef=0.0, ) - output = strategy.compute_loss_output(self._make_batch()) + output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True)) metrics = output["metrics"] assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6) diff --git a/tests/trainer/test_grpo_e2e.py b/tests/trainer/test_grpo_e2e.py deleted file mode 100644 index d44b006..0000000 --- a/tests/trainer/test_grpo_e2e.py +++ /dev/null @@ -1,123 +0,0 @@ -"""End-to-end integration test for online GRPO rollout.""" - -import os -from functools import partial - -import pytest -import torch -from torch.utils.data import Dataset - -from astrai.config import TrainConfig -from astrai.model.transformer import AutoRegressiveLM -from astrai.trainer.rollout import BaseRewardModel -from astrai.trainer.schedule import SchedulerFactory -from astrai.trainer.trainer import Trainer -from tests.helpers import CHAT_TEMPLATE - - -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 len(self._SAMPLES) - - def __getitem__(self, idx): - return dict(self._SAMPLES[idx]) - - -class LengthRewardModel(BaseRewardModel): - """Rewards each response by its (non-pad) token count. - - Gives the group-normalized advantage a non-degenerate signal. - """ - - def score(self, prompts, responses): - B = len(prompts) - G = len(responses[0]) if B else 0 - rewards = torch.zeros(B, G) - for i in range(B): - for g in range(G): - rewards[i, g] = float(len(responses[i][g])) - 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) - - -def _optimizer_fn(m): - return torch.optim.AdamW(m.parameters(), lr=1e-4) - - -def _scheduler_fn(optim): - return SchedulerFactory.create( - "cosine", optim, warmup_steps=1, lr_decay_steps=4, min_rate=0.05 - ) - - -@pytest.mark.integration -def test_online_grpo_end_to_end(base_test_env): - """Run one epoch of online GRPO with KV-cache-backed rollout.""" - test_dir = base_test_env["test_dir"] - device = base_test_env["device"] - tokenizer = base_test_env["tokenizer"] - model_config = base_test_env["transformer_config"] - - 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 = InstructionDataset() - - train_config = TrainConfig( - strategy="online_grpo", - model_fn=model_fn, - dataset=dataset, - optimizer_fn=optimizer_fn, - scheduler_fn=scheduler_fn, - ckpt_dir=os.path.join(test_dir, "ckpt"), - n_epoch=1, - batch_per_device=2, - ckpt_interval=100, - grad_accum_steps=1, - random_seed=42, - device_type=device, - nprocs=1, - parallel_mode="none", - strategy_kwargs={"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2}, - rollout_interval=1, - rollout_temperature=1.0, - rollout_top_k=0, - rollout_top_p=1.0, - rollout_max_tokens=4, - reward_model_fn=LengthRewardModel, - collate_fn=instruction_collate_fn, - ) - - trainer = Trainer(train_config) - trainer.train(param_path=test_dir) - - assert os.path.isdir(os.path.join(test_dir, "ckpt")) diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index 9439ac1..9f1eac8 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -1,4 +1,4 @@ -"""End-to-end integration test for online DPO rollout.""" +"""End-to-end integration tests for online GRPO/DPO rollout.""" import os from functools import partial @@ -40,7 +40,7 @@ class InstructionDataset(Dataset): class LengthRewardModel(BaseRewardModel): """Rewards each response by its (non-pad) token count. - Enough for DPO to distinguish chosen/rejected from the rollout group. + Gives the group-normalized advantage a non-degenerate signal. """ def score(self, prompts, responses): @@ -75,31 +75,34 @@ def _scheduler_fn(optim): ) +_ONLINE_STRATEGIES = [ + pytest.param( + "online_grpo", + {"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2}, + id="grpo", + ), + pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, id="dpo"), +] + + @pytest.mark.integration -def test_online_dpo_end_to_end(base_test_env): - """Run one epoch of online DPO with KV-cache-backed rollout.""" +@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES) +def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs): + """Run one epoch of online RL rollout with KV-cache-backed generation.""" test_dir = base_test_env["test_dir"] device = base_test_env["device"] tokenizer = base_test_env["tokenizer"] model_config = base_test_env["transformer_config"] - # 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 = InstructionDataset() - train_config = TrainConfig( - strategy="online_dpo", - model_fn=model_fn, - dataset=dataset, - optimizer_fn=optimizer_fn, - scheduler_fn=scheduler_fn, + strategy=strategy, + model_fn=partial(_model_fn, model_config), + dataset=InstructionDataset(), + optimizer_fn=_optimizer_fn, + scheduler_fn=_scheduler_fn, ckpt_dir=os.path.join(test_dir, "ckpt"), n_epoch=1, batch_per_device=2, @@ -109,7 +112,7 @@ def test_online_dpo_end_to_end(base_test_env): device_type=device, nprocs=1, parallel_mode="none", - strategy_kwargs={"beta": 0.1, "group_size": 2}, + strategy_kwargs=strategy_kwargs, rollout_interval=1, rollout_temperature=1.0, rollout_top_k=0,