diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index e926427..cba63df 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -26,11 +26,22 @@ from astrai.serialization import ( ) from tests.data.factories import make_grpo_config +SIMPLE_CHAT_TEMPLATE = ( + "{% for message in messages %}" + "{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" +) + def _rand_seq(length, vocab=1000): return torch.randint(0, vocab, (length,), dtype=torch.int64) +def _dump_jsonl(path, records): + with open(path, "w", encoding="utf-8") as f: + for record in records: + f.write(json.dumps(record, ensure_ascii=False) + "\n") + + def _save_test_tokenizer(test_dir, tokenizer): tokenizer_path = os.path.join(test_dir, "tokenizer") os.makedirs(tokenizer_path, exist_ok=True) @@ -38,19 +49,19 @@ def _save_test_tokenizer(test_dir, tokenizer): return tokenizer_path -def _write_jsonl_dataset(test_dir, tokenizer_path, records, config_overrides=None): - data_dir = os.path.join(test_dir, "jsonl_data") +def _write_text_dataset( + test_dir, dirname, tokenizer_path, records, config_overrides=None +): + """Write a JSONL dataset directory with a text-section default config.""" + data_dir = os.path.join(test_dir, dirname) os.makedirs(data_dir, exist_ok=True) - - with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: - for record in records: - f.write(json.dumps(record, ensure_ascii=False) + "\n") + _dump_jsonl(os.path.join(data_dir, "data.jsonl"), records) config = { "tokenizer_path": tokenizer_path, "version": 1, "input": {"sections": [{"field": "text", "action": "train"}]}, - "preprocessing": {"max_seq_len": 128}, + "preprocessing": {"max_seq_len": 128, "min_chars": 0}, "output": {"position_ids_mode": "continuous"}, } if config_overrides: @@ -71,6 +82,27 @@ def _fake_fetch_record(self, idx, keys): return {k: self._data[k][idx] for k in keys} +def _grpo_fake_store(prompts, responses, masks, rewards): + """Fake GRPO record store matching real Store semantics.""" + return type( + "FakeStore", + (), + { + "keys": ["prompts", "responses", "masks", "rewards"], + "num_records": len(prompts), + "token_count": 0, + "_data": { + "prompts": prompts, + "responses": responses, + "masks": masks, + "rewards": rewards, + }, + "fetch_record": _fake_fetch_record, + "__len__": lambda self: self.num_records, + }, + )() + + def _make_seq_dataset( test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs ): @@ -85,32 +117,6 @@ def _make_seq_dataset( ) -def test_dataset_loader_random_paths(base_test_env): - """Test dataset loader with multiple random paths""" - test_dir = base_test_env["test_dir"] - - loaded_dataset = None - num_files = np.random.randint(2, 5) - for i in range(num_files): - seq_length = np.random.randint(200, 400) - dummy_data = {"sequence": [_rand_seq(seq_length) for _ in range(10)]} - sub_dir = os.path.join(test_dir, f"sub_{i}") - os.makedirs(sub_dir, exist_ok=True) - loaded_dataset = _make_seq_dataset( - sub_dir, f"data_{i}", seq_length, data=dummy_data - ) - assert loaded_dataset is not None - assert len(loaded_dataset) > 0 - - # Test that we can get items without errors - for i in range(len(loaded_dataset)): - item = loaded_dataset[i] - assert "input_ids" in item - assert "target_ids" in item - assert item["input_ids"].shape == item["target_ids"].shape - assert item["input_ids"].shape[0] == 64 - - def test_dpo_strategy_with_random_data(base_test_env): """Test DPO strategy with randomized preference data""" test_dir = base_test_env["test_dir"] @@ -345,25 +351,12 @@ def test_normalize_mixed_empty_key(): def test_grpo_dataset_dtype(base_test_env): """GRPO dataset returns correct dtypes for per-record structured data.""" G = 4 - store = type( - "FakeStore", - (), - { - "keys": ["prompts", "responses", "masks", "rewards"], - "num_records": 1, - "token_count": 0, - "_data": { - "prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)], - "responses": [ - [torch.randint(0, 100, (5,), dtype=torch.int32) for _ in range(G)] - ], - "masks": [[torch.ones(5, dtype=torch.int32) for _ in range(G)]], - "rewards": [torch.rand(G, dtype=torch.float32)], - }, - "fetch_record": _fake_fetch_record, - "__len__": lambda self: self.num_records, - }, - )() + store = _grpo_fake_store( + prompts=[torch.randint(0, 100, (10,), dtype=torch.int32)], + responses=[[torch.randint(0, 100, (5,), dtype=torch.int32) for _ in range(G)]], + masks=[[torch.ones(5, dtype=torch.int32) for _ in range(G)]], + rewards=[torch.rand(G, dtype=torch.float32)], + ) dataset = GRPODataset(store=store) item = dataset[0] @@ -378,23 +371,12 @@ def test_grpo_dataset_load(base_test_env): G = 3 prompt_len = 8 resp_lens = [5, 7, 4] - store = type( - "FakeStore", - (), - { - "keys": ["prompts", "responses", "masks", "rewards"], - "num_records": 1, - "token_count": 0, - "_data": { - "prompts": [torch.randint(0, 100, (prompt_len,))], - "responses": [[torch.randint(0, 100, (rl,)) for rl in resp_lens]], - "masks": [[torch.ones(rl, dtype=torch.int64) for rl in resp_lens]], - "rewards": [torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)], - }, - "fetch_record": _fake_fetch_record, - "__len__": lambda self: self.num_records, - }, - )() + store = _grpo_fake_store( + prompts=[torch.randint(0, 100, (prompt_len,))], + responses=[[torch.randint(0, 100, (rl,)) for rl in resp_lens]], + masks=[[torch.ones(rl, dtype=torch.int64) for rl in resp_lens]], + rewards=[torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)], + ) dataset = GRPODataset(store=store) assert len(dataset) == 1 @@ -465,62 +447,28 @@ def test_dataset_load_explicit_storage_type(base_test_env): assert dataset.token_count == 200 -def _write_json_dataset(test_dir, tokenizer_path, records, config_overrides=None): - """Write JSONL dataset — one JSON object per line.""" - data_dir = os.path.join(test_dir, "json_data") - os.makedirs(data_dir, exist_ok=True) - - with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: - for rec in records: - f.write(json.dumps(rec, ensure_ascii=False) + "\n") - - config = { - "tokenizer_path": tokenizer_path, - "version": 1, - "input": {"sections": [{"field": "text", "action": "train"}]}, - "preprocessing": {"max_seq_len": 128, "min_chars": 0}, - "output": {"position_ids_mode": "continuous"}, - } - if config_overrides: - config.update(config_overrides) - - with open( - os.path.join(data_dir, "dataset_config.json"), "w", encoding="utf-8" - ) as f: - json.dump(config, f, ensure_ascii=False, indent=2) - - return data_dir - - -@pytest.mark.parametrize( - "use_jsonl", - [True, False], -) +@pytest.mark.parametrize("use_jsonl", [True, False]) def test_detect_format_data_dir(base_test_env, use_jsonl): """detect_format returns 'jsonl' for dirs of .jsonl or .json files.""" test_dir = base_test_env["test_dir"] tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) - if use_jsonl: - data_dir = _write_jsonl_dataset( - test_dir, - tokenizer_path, - [{"text": "hello world"}, {"text": "foo bar baz"}], - ) - else: - data_dir = _write_json_dataset( - test_dir, - tokenizer_path, - [{"text": "hello world"}, {"text": "foo bar baz qux"}], - ) + data_dir = _write_text_dataset( + test_dir, + "jsonl_data" if use_jsonl else "json_data", + tokenizer_path, + [{"text": "hello world"}, {"text": "foo bar baz"}], + ) assert detect_format(data_dir) == "jsonl" -def test_json_store_seq(base_test_env): - """JsonlStore loads .json array correctly.""" +@pytest.mark.parametrize("dirname", ["json_data", "jsonl_data"]) +def test_json_store_seq(base_test_env, dirname): + """JsonlStore loads a text JSONL dataset and feeds seq training.""" test_dir = base_test_env["test_dir"] tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) - data_dir = _write_json_dataset( + data_dir = _write_text_dataset( test_dir, + dirname, tokenizer_path, [{"text": "hello world"}, {"text": "foo bar baz qux"}], ) @@ -535,15 +483,14 @@ def test_json_store_seq(base_test_env): item = dataset[0] assert "input_ids" in item assert "target_ids" in item + assert item["input_ids"].dtype == torch.long def test_json_store_no_tokenizer_path(base_test_env): """JsonlStore uses dataset dir as tokenizer_path when omitted.""" test_dir = base_test_env["test_dir"] tokenizer = base_test_env["tokenizer"] - tokenizer.set_chat_template( - "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" - ) + tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE) data_dir = os.path.join(test_dir, "self_contained") os.makedirs(data_dir, exist_ok=True) @@ -560,9 +507,7 @@ def test_json_store_no_tokenizer_path(base_test_env): ] } ] - with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: - for rec in records: - f.write(json.dumps(rec, ensure_ascii=False) + "\n") + _dump_jsonl(os.path.join(data_dir, "data.jsonl"), records) # dataset_config.json WITHOUT tokenizer_path config = { @@ -587,38 +532,14 @@ def test_json_store_no_tokenizer_path(base_test_env): assert "loss_mask" in store.keys -def test_jsonl_store_seq(base_test_env): - test_dir = base_test_env["test_dir"] - tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) - data_dir = _write_jsonl_dataset( - test_dir, - tokenizer_path, - [{"text": "hello world"}, {"text": "foo bar baz qux"}], - config_overrides={"preprocessing": {"max_seq_len": 128, "min_chars": 0}}, - ) - - store = StoreFactory.create("jsonl") - store.load(data_dir, transform=_build_jsonl_transform(data_dir)) - assert len(store) > 0 - assert "sequence" in store.keys - - dataset = DatasetFactory.load("seq", data_dir, window_size=8) - assert len(dataset) > 0 - item = dataset[0] - assert "input_ids" in item - assert "target_ids" in item - assert item["input_ids"].dtype == torch.long - - def test_jsonl_store_sft(base_test_env): test_dir = base_test_env["test_dir"] tokenizer = base_test_env["tokenizer"] - tokenizer.set_chat_template( - "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" - ) + tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE) tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) - data_dir = _write_jsonl_dataset( + data_dir = _write_text_dataset( test_dir, + "sft_jsonl", tokenizer_path, [ { @@ -661,9 +582,7 @@ def test_sft_jsonl_default_messages_config(base_test_env): """ test_dir = base_test_env["test_dir"] tokenizer = base_test_env["tokenizer"] - tokenizer.set_chat_template( - "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" - ) + tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE) tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) data_dir = os.path.join(test_dir, "jsonl_data") @@ -683,9 +602,7 @@ def test_sft_jsonl_default_messages_config(base_test_env): ] }, ] - with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: - for record in records: - f.write(json.dumps(record, ensure_ascii=False) + "\n") + _dump_jsonl(os.path.join(data_dir, "data.jsonl"), records) dataset = DatasetFactory.load( "sft", data_dir, window_size=8, tokenizer_path=tokenizer_path @@ -706,13 +623,12 @@ def test_sft_jsonl_explicit_config_takes_priority(base_test_env): """When dataset_config.json exists, it overrides the default messages config.""" test_dir = base_test_env["test_dir"] tokenizer = base_test_env["tokenizer"] - tokenizer.set_chat_template( - "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}" - ) + tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE) tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) - data_dir = _write_jsonl_dataset( + data_dir = _write_text_dataset( test_dir, + "sft_explicit", tokenizer_path, [ { @@ -748,10 +664,7 @@ def _write_grpo_jsonl(test_dir, tokenizer_path, records): """Write a GRPO JSONL dataset directory with config.""" data_dir = os.path.join(test_dir, "grpo_jsonl") os.makedirs(data_dir, exist_ok=True) - - with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: - for rec in records: - f.write(json.dumps(rec, ensure_ascii=False) + "\n") + _dump_jsonl(os.path.join(data_dir, "data.jsonl"), records) config = { "tokenizer_path": tokenizer_path, @@ -927,28 +840,15 @@ def test_grpo_multiple_records(base_test_env): [torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)] for _ in range(n_records) ] - store = type( - "FakeStore", - (), - { - "keys": ["prompts", "responses", "masks", "rewards"], - "num_records": n_records, - "token_count": 0, - "_data": { - "prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)], - "responses": dummy_responses, - "masks": [ - [torch.ones(r.shape[0], dtype=torch.int64) for r in resps] - for resps in dummy_responses - ], - "rewards": [ - torch.rand(G, dtype=torch.float32) for _ in range(n_records) - ], - }, - "fetch_record": _fake_fetch_record, - "__len__": lambda self: self.num_records, - }, - )() + store = _grpo_fake_store( + prompts=[torch.randint(0, 100, (10,)) for _ in range(n_records)], + responses=dummy_responses, + masks=[ + [torch.ones(r.shape[0], dtype=torch.int64) for r in resps] + for resps in dummy_responses + ], + rewards=[torch.rand(G, dtype=torch.float32) for _ in range(n_records)], + ) dataset = GRPODataset(store=store) assert len(dataset) == n_records @@ -965,9 +865,7 @@ def test_grpo_multiple_records(base_test_env): def _write_dpo_jsonl(test_dir, records): """Write a raw DPO JSONL file (no dataset_config.json).""" path = os.path.join(test_dir, "dpo.jsonl") - with open(path, "w", encoding="utf-8") as f: - for rec in records: - f.write(json.dumps(rec, ensure_ascii=False) + "\n") + _dump_jsonl(path, records) return path @@ -1075,12 +973,12 @@ def test_jsonl_store_eager_len_returns_token_count(base_test_env): """JsonlStore in eager mode: num_records reflects per-record count.""" test_dir = base_test_env["test_dir"] tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) - data_dir = _write_jsonl_dataset( + data_dir = _write_text_dataset( test_dir, + "jsonl_data", tokenizer_path, [{"text": "hello world"}, {"text": "foo bar"}], config_overrides={ - "preprocessing": {"max_seq_len": 128, "min_chars": 0}, "output": {"position_ids_mode": "none"}, }, ) diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index cd7db16..251ab66 100644 --- a/tests/data/test_preprocess_builder.py +++ b/tests/data/test_preprocess_builder.py @@ -14,7 +14,6 @@ from astrai.preprocessing.builder import ( ) from tests.data.factories import ( CHAT_SECTIONS, - INSTRUCTION_SECTIONS, TEXT_SECTIONS, make_chat_config, make_dpo_chat_config, @@ -35,9 +34,6 @@ def test_chat_simple(chat_tokenizer, builder): } result = builder.build(item, config, chat_tokenizer) assert result is not None - assert "sequence" in result - assert "loss_mask" in result - assert len(result["sequence"]) == len(result["loss_mask"]) ids = chat_tokenizer.decode(result["sequence"], skip_special_tokens=False) assert "system" in ids.lower() or "<|im_start|>system" in ids @@ -51,21 +47,26 @@ def test_chat_simple(chat_tokenizer, builder): def test_chat_mask_only_assistant(chat_tokenizer, builder): config = make_chat_config() - item = { - "messages": [ - {"role": "user", "content": "What is 2+2?"}, - {"role": "assistant", "content": "4"}, - ] - } - result = builder.build(item, config, chat_tokenizer) + messages = [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + ] + result = builder.build({"messages": messages}, config, chat_tokenizer) mask = result["loss_mask"] - ids = result["sequence"] - assert len(ids) == len(mask) - trained = [i for i, m in enumerate(mask) if m == 1] - masked = [i for i, m in enumerate(mask) if m == 0] - assert len(trained) > 0 - assert len(masked) > 0 + def template_ids(message): + rendered = chat_tokenizer.apply_chat_template( + [message], tokenize=False, add_generation_prompt=False + ) + return chat_tokenizer.encode(rendered, add_special_tokens=False) + + user_len = len(template_ids(messages[0])) + assistant_len = len(template_ids(messages[1])) + bos = 1 if chat_tokenizer.bos_token_id is not None else 0 + + assert len(mask) == bos + user_len + assistant_len + assert all(m == 0 for m in mask[: bos + user_len]) + assert all(m == 1 for m in mask[bos + user_len :]) def test_chat_batch_matches_single(chat_tokenizer, builder): @@ -166,14 +167,6 @@ def test_chat_truncation(chat_tokenizer, builder): assert len(result["loss_mask"]) == len(result["sequence"]) -def test_instruction_basic(test_tokenizer, builder): - config = make_instruction_config() - item = {"prompt": "Translate to French: Hello", "response": "Bonjour"} - result = builder.build(item, config, test_tokenizer) - assert result is not None - assert len(result["sequence"]) == len(result["loss_mask"]) - - def test_instruction_batch_matches_single(test_tokenizer, builder): config = make_instruction_config() items = [ @@ -224,7 +217,6 @@ def test_text_basic(test_tokenizer, builder): item = {"text": "Hello world. This is a test document."} result = builder.build(item, config, test_tokenizer) assert result is not None - assert "sequence" in result assert len(result["sequence"]) > 0 assert "loss_mask" not in result @@ -253,72 +245,17 @@ def test_text_truncation(test_tokenizer, builder): assert len(result["sequence"]) <= 3 -def test_sectioned_chat(chat_tokenizer, builder): - config = PipelineConfig( - input=InputConfig(sections=CHAT_SECTIONS), - mask={"system": "mask", "user": "mask", "assistant": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - ) - item = { - "messages": [ - {"role": "user", "content": "What is 2+2?"}, - {"role": "assistant", "content": "4"}, - ] - } - result = builder.build(item, config, chat_tokenizer) - assert result is not None - assert len(result["sequence"]) == len(result["loss_mask"]) - assert sum(result["loss_mask"]) > 0 - assert 0 in result["loss_mask"] - - -def test_sectioned_instruction(test_tokenizer, builder): - config = PipelineConfig( - input=InputConfig(sections=INSTRUCTION_SECTIONS), - preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0), - ) - item = {"prompt": "Q: Why?", "response": "A: Because."} - result = builder.build(item, config, test_tokenizer) - assert result is not None - mask = result["loss_mask"] - assert mask[0] == 0 - assert mask[-1] == 1 - - -def test_sectioned_text(test_tokenizer, builder): - config = PipelineConfig( - input=InputConfig(sections=TEXT_SECTIONS), - preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1), - ) - item = {"text": "Hello world, this is a test."} - result = builder.build(item, config, test_tokenizer) - assert result is not None - assert "loss_mask" not in result - - -def test_sectioned_text_too_short(test_tokenizer, builder): - config = PipelineConfig( - input=InputConfig(sections=TEXT_SECTIONS), - preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100), - ) - assert builder.build({"text": "short"}, config, test_tokenizer) is None - - -def test_factory_registered(): - names = MaskBuilderFactory.list_registered() - assert "single" in names - assert "multi" in names - assert "sectioned" in names - - -def test_factory_create(): - single = MaskBuilderFactory.create("single") - assert isinstance(single, SingleOutputMaskBuilder) - multi = MaskBuilderFactory.create("multi") - assert isinstance(multi, MultiOutputMaskBuilder) - sectioned = MaskBuilderFactory.create("sectioned") - assert isinstance(sectioned, SectionedMaskBuilder) +@pytest.mark.parametrize( + ("name", "builder_cls"), + [ + ("single", SingleOutputMaskBuilder), + ("multi", MultiOutputMaskBuilder), + ("sectioned", SectionedMaskBuilder), + ], +) +def test_factory_create(name, builder_cls): + assert name in MaskBuilderFactory.list_registered() + assert isinstance(MaskBuilderFactory.create(name), builder_cls) def test_dpo_chat_basic(chat_tokenizer, builder): diff --git a/tests/data/test_preprocess_pipeline.py b/tests/data/test_preprocess_pipeline.py index 125ae8e..8554119 100644 --- a/tests/data/test_preprocess_pipeline.py +++ b/tests/data/test_preprocess_pipeline.py @@ -115,49 +115,6 @@ def test_full_text_pipeline(temp_dir, tokenizer_dir): assert "loss_mask" not in meta -def test_full_instruction_pipeline(temp_dir, tokenizer_dir): - jsonl_path = os.path.join(temp_dir, "instruct.jsonl") - with open(jsonl_path, "w", encoding="utf-8") as f: - f.write( - json.dumps( - { - "prompt": "Tell me a joke", - "response": "Why did the chicken cross the road?", - } - ) - + "\n" - ) - f.write( - json.dumps( - { - "prompt": "What is AI?", - "response": "Artificial Intelligence is a field of computer science.", - } - ) - + "\n" - ) - - config = PipelineConfig( - input=InputConfig(sections=INSTRUCTION_SECTIONS), - mask={"prompt": "mask", "response": "train"}, - mask_default="mask", - preprocessing=ProcessingConfig(max_seq_len=2048), - output=OutputConfig(storage_format="bin"), - ) - - out_dir = os.path.join(temp_dir, "output") - Pipeline( - config=config, - input_paths=[jsonl_path], - output_dir=out_dir, - tokenizer_path=tokenizer_dir, - ).run() - - meta = load_shard_meta(out_dir) - assert "sequence" in meta - assert "loss_mask" in meta - - def test_dtype_override(temp_dir, tokenizer_dir): jsonl_path = os.path.join(temp_dir, "data.jsonl") with open(jsonl_path, "w", encoding="utf-8") as f: diff --git a/tests/extension/test_kernel_mask.py b/tests/extension/test_kernel_mask.py index 1b29d31..fc776d2 100644 --- a/tests/extension/test_kernel_mask.py +++ b/tests/extension/test_kernel_mask.py @@ -1,5 +1,7 @@ """Kernel-level mask dimension support (2D, 3D, 4D).""" +import math + import torch from astrai.extension.ops.attention import attn_prefill @@ -7,9 +9,27 @@ from tests.conftest import skip_no_kernel from tests.extension.conftest import D +def _reference(q, k, v, mask): + """fp32 masked GQA attention reference (True=keep).""" + b, s_q, h, d = q.shape + rep = h // k.shape[2] + qf = q.float().transpose(1, 2) + kf = k.float().repeat_interleave(rep, dim=2).transpose(1, 2) + vf = v.float().repeat_interleave(rep, dim=2).transpose(1, 2) + scores = qf @ kf.transpose(-1, -2) / math.sqrt(d) + if mask.dim() == 2: + mask = mask[:, None, None, :] + elif mask.dim() == 3: + mask = mask[:, None, :, :] + # 4D [batch, 1, q_len, kv_len] broadcasts over heads as-is + scores = scores.masked_fill(~mask, float("-inf")) + return (scores.softmax(dim=-1) @ vf).transpose(1, 2).to(q.dtype) + + @skip_no_kernel def test_kernel_accepts_2d_mask(): - """Kernel should accept 2D mask [batch, kv_len].""" + """2D mask [batch, kv_len] gates the softmax, not just parses.""" + torch.manual_seed(11) batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1 kv_len = 8 q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16) @@ -20,25 +40,31 @@ def test_kernel_accepts_2d_mask(): out = attn_prefill(q, k, v, mask=mask, is_causal=False) assert out.shape == (batch, q_len, n_heads, D) + torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05) @skip_no_kernel def test_kernel_accepts_3d_mask(): - """Kernel should accept 3D mask [batch, q_len, kv_len].""" + """3D mask [batch, q_len, kv_len] applies per-query-row gating.""" + torch.manual_seed(12) batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1 kv_len = 8 q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16) k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16) v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16) mask = torch.ones(batch, q_len, kv_len, dtype=torch.bool, device="cuda") + mask[:, 0, 5:] = False # differs per query row: only the 3D path can apply it + mask[:, 1, 6:] = False out = attn_prefill(q, k, v, mask=mask, is_causal=False) assert out.shape == (batch, q_len, n_heads, D) + torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05) @skip_no_kernel def test_kernel_accepts_4d_mask(): - """Kernel should accept 4D mask [batch, n_heads, q_len, kv_len].""" + """4D mask [batch, 1, q_len, kv_len] broadcasts over heads and gates.""" + torch.manual_seed(13) batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1 kv_len = 8 q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16) @@ -49,6 +75,7 @@ def test_kernel_accepts_4d_mask(): out = attn_prefill(q, k, v, mask=mask, is_causal=False) assert out.shape == (batch, q_len, n_heads, D) + torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05) @skip_no_kernel diff --git a/tests/extension/test_linear_dispatch.py b/tests/extension/test_linear_dispatch.py index 56c0adf..77c3fd6 100644 --- a/tests/extension/test_linear_dispatch.py +++ b/tests/extension/test_linear_dispatch.py @@ -6,7 +6,6 @@ import torch import torch.nn.functional as F from astrai.extension import is_available, linear -from astrai.extension.backend import linear as public_linear from astrai.extension.dispatch import explain, op_backend, resolve # The package attribute ``linear`` is the dispatched function; reach the @@ -36,10 +35,6 @@ def _routes_to_gemv(monkeypatch, x, weight, bias=None) -> bool: return linear(x, weight, bias) is sentinel -def test_linear_backend_is_public(): - assert linear is public_linear - - def test_model_linear_routes_through_backend(monkeypatch): sentinel = torch.randn(2, 4) diff --git a/tests/extension/test_swiglu_dispatch.py b/tests/extension/test_swiglu_dispatch.py index 4611242..3ef8109 100644 --- a/tests/extension/test_swiglu_dispatch.py +++ b/tests/extension/test_swiglu_dispatch.py @@ -38,10 +38,13 @@ def test_cpu_and_training_calls_fall_back_with_gradients(monkeypatch): def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog): monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode") + x = torch.randn(2, 8) + up_weight = torch.randn(4, 8) + gate_weight = torch.randn(4, 8) with caplog.at_level(logging.WARNING): - actual = swiglu(torch.randn(2, 8), torch.randn(4, 8), torch.randn(4, 8)) - assert actual.shape == (2, 4) + actual = swiglu(x, up_weight, gate_weight) assert "using auto" in caplog.text + torch.testing.assert_close(actual, reference_swiglu(x, up_weight, gate_weight)) def test_mlp_routes_through_swiglu_backend(monkeypatch): diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index 6b5bcbd..bafb22d 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -34,13 +34,6 @@ def _make_task_cache(pool: PagePool) -> TaskCacheManager: # ---- page_hash ---- -def test_page_hash_full_page(): - token_ids = list(range(256)) - h = page_hash(token_ids, 0, 64) - assert isinstance(h, int) - assert h >= 0 - - def test_page_hash_different_page_differs(): token_ids = list(range(256)) assert page_hash(token_ids, 0, 64) != page_hash(token_ids, 1, 64) @@ -121,19 +114,13 @@ def test_prefix_cache_ignores_partial_last_page(): def test_prefix_cache_on_evict_clears_mappings(): prefix = RadixCache(64) + assert not prefix.has_page(0) prefix.record(0, list(range(64)), 0) assert prefix.has_page(0) prefix.evict(0) assert not prefix.has_page(0) -def test_prefix_cache_has_page(): - prefix = RadixCache(64) - assert not prefix.has_page(0) - prefix.record(0, list(range(64)), 0) - assert prefix.has_page(0) - - def test_prefix_cache_does_not_reuse_page_without_parent_prefix(): prefix = RadixCache(2) prefix.record(0, [1, 2, 3, 4], 0) diff --git a/tests/inference/test_engine.py b/tests/inference/test_engine.py index 049b728..6d7a68d 100644 --- a/tests/inference/test_engine.py +++ b/tests/inference/test_engine.py @@ -20,12 +20,6 @@ def _make_engine_mocks(decode=None): return mock_model, mock_tokenizer -def test_result_append_single(): - r = GenerateResult(count=1) - r.append("hello", 0) - assert r.results[0] == "hello" - - def test_result_append_multiple_tasks(): r = GenerateResult(count=3) r.append("a", 0) diff --git a/tests/inference/test_protocol.py b/tests/inference/test_protocol.py index 39144e9..2af9a44 100644 --- a/tests/inference/test_protocol.py +++ b/tests/inference/test_protocol.py @@ -71,31 +71,6 @@ def test_check_empty_sequences(): assert sc.check("hello") is None -def test_gen_context_defaults(): - ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10) - assert ctx.completion_tokens == 0 - - -def test_gen_context_fields_mutable(): - ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10) - ctx.completion_tokens = 42 - assert ctx.completion_tokens == 42 - - -def test_stop_info_defaults(): - s = StopInfo() - assert s.matched is None - assert s.body == "" - assert s.yielded == "" - - -def test_stop_info_with_values(): - s = StopInfo(matched="stop", body="hello stop", yielded="hello ") - assert s.matched == "stop" - assert s.body == "hello stop" - assert s.yielded == "hello " - - def test_openai_prepare_returns_prompt_ctx_stops(): builder = _make_openai_builder() req = MagicMock() diff --git a/tests/inference/test_sample.py b/tests/inference/test_sample.py index f950fd3..aee2ac4 100644 --- a/tests/inference/test_sample.py +++ b/tests/inference/test_sample.py @@ -62,8 +62,9 @@ def test_top_p_nucleus_filtering(): logits = torch.tensor([[10.0, 1.0, 1.0, 1.0, 1.0]]) s = TopPStrategy(top_p=0.5) result = s.apply(logits.clone(), filter_value=-1e9) + # The dominant logit alone exceeds the nucleus mass; the rest are filtered. kept = (result > -1e9).sum().item() - assert kept >= 1 + assert kept == 1 def test_top_p_skip_when_one(): diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 82d7496..7f504c0 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -40,18 +40,32 @@ def mock_model_and_tokenizer(): return mock_model, mock_tokenizer +def _make_mock_scheduler(mock_model_and_tokenizer): + """Build a CPU scheduler over mocks, patching scheduler-internal imports.""" + mock_model, mock_tokenizer = mock_model_and_tokenizer + with ( + patch("astrai.inference.scheduler.AutoModel"), + patch("astrai.inference.scheduler.AutoTokenizer"), + ): + return InferenceScheduler( + model=mock_model, + tokenizer=mock_tokenizer, + max_batch_size=4, + device="cpu", + ) + + +def _run_threads(*workers, timeout=10.0): + threads = [threading.Thread(target=worker) for worker in workers] + for t in threads: + t.start() + for t in threads: + t.join(timeout=timeout) + + def test_scheduler_concurrent_add_task(mock_model_and_tokenizer): """Test concurrent add_task operations.""" - mock_model, mock_tokenizer = mock_model_and_tokenizer - - with patch("astrai.inference.scheduler.AutoModel"): - with patch("astrai.inference.scheduler.AutoTokenizer"): - scheduler = InferenceScheduler( - model=mock_model, - tokenizer=mock_tokenizer, - max_batch_size=4, - device="cpu", - ) + scheduler = _make_mock_scheduler(mock_model_and_tokenizer) results = {"task_ids": [], "errors": []} lock = threading.Lock() @@ -65,13 +79,7 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer): except Exception as e: results["errors"].append(str(e)) - threads = [threading.Thread(target=add_task_worker, args=(i,)) for i in range(5)] - - for t in threads: - t.start() - - for t in threads: - t.join() + _run_threads(*(lambda wid=i: add_task_worker(wid) for i in range(5))) scheduler.stop() @@ -205,16 +213,7 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits(): def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer): """Test concurrent add and remove task operations.""" - mock_model, mock_tokenizer = mock_model_and_tokenizer - - with patch("astrai.inference.scheduler.AutoModel"): - with patch("astrai.inference.scheduler.AutoTokenizer"): - scheduler = InferenceScheduler( - model=mock_model, - tokenizer=mock_tokenizer, - max_batch_size=4, - device="cpu", - ) + scheduler = _make_mock_scheduler(mock_model_and_tokenizer) results = {"added": [], "removed": [], "errors": []} add_ready = threading.Event() @@ -238,14 +237,7 @@ def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer): except Exception as e: results["errors"].append(f"Remove: {str(e)}") - add_thread = threading.Thread(target=add_worker) - remove_thread = threading.Thread(target=remove_worker) - - add_thread.start() - remove_thread.start() - - add_thread.join() - remove_thread.join() + _run_threads(add_worker, remove_worker) scheduler.stop() assert len(results["errors"]) == 0, f"Errors: {results['errors']}" @@ -254,16 +246,7 @@ def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer): def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer): """Test concurrent get_stats operations.""" - mock_model, mock_tokenizer = mock_model_and_tokenizer - - with patch("astrai.inference.scheduler.AutoModel"): - with patch("astrai.inference.scheduler.AutoTokenizer"): - scheduler = InferenceScheduler( - model=mock_model, - tokenizer=mock_tokenizer, - max_batch_size=4, - device="cpu", - ) + scheduler = _make_mock_scheduler(mock_model_and_tokenizer) results = {"stats": [], "errors": []} started = threading.Event() @@ -287,17 +270,9 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer): except Exception as e: results["errors"].append(f"Get stats: {str(e)}") - add_thread = threading.Thread(target=add_tasks) - stats_thread = threading.Thread(target=get_stats) - - add_thread.start() - stats_thread.start() - - add_thread.join() - stats_done.wait(timeout=5.0) + _run_threads(add_tasks, get_stats) scheduler.stop() - - stats_thread.join() + stats_done.wait(timeout=5.0) assert len(results["errors"]) == 0, f"Errors: {results['errors']}" assert len(results["stats"]) == 50 @@ -504,16 +479,6 @@ def test_ragged_prefill_matches_sequential_greedy_tokens_and_logprobs(device): scheduler.stop() -def test_run_batch_respects_max_tokens(device): - scheduler, _tok, _model = _make_real_scheduler(device) - try: - prompts = [[10, 20, 30]] - results = scheduler.run_batch(prompts, max_tokens=3, temperature=1.0) - assert len(results[0]) <= 3 - finally: - scheduler.stop() - - def test_run_batch_zero_max_tokens_returns_empty(device): scheduler, _tok, _model = _make_real_scheduler(device) try: @@ -526,12 +491,12 @@ def test_run_batch_stop_id_terminates(device): """A token matching stop_ids terminates generation for that prompt.""" scheduler, _tok, _model = _make_real_scheduler(device) try: + # Make every token a stop id: generation must end after exactly + # one token (the stop token itself) instead of running to max_tokens. + scheduler._task_mgr.tokenizer.stop_ids = list(range(200)) prompts = [[10, 20, 30]] results = scheduler.run_batch(prompts, max_tokens=32, temperature=1.0) - # If stop token 2 was produced, it is the last token - if results[0] and results[0][-1] == 2: - # No tokens after stop should exist (since we terminate) - assert 2 not in results[0][:-1] + assert len(results[0]) == 1 finally: scheduler.stop() diff --git a/tests/inference/test_tool_parser.py b/tests/inference/test_tool_parser.py index 2b63235..82b0258 100644 --- a/tests/inference/test_tool_parser.py +++ b/tests/inference/test_tool_parser.py @@ -252,13 +252,15 @@ def test_feed_with_tools_constructor(): tools = [{"type": "function", "function": {"name": "get_weather"}}] parser = SimpleJsonToolParser(tools=tools, tool_choice="auto") deltas = parser.feed('{"name": "get_weather", "arguments": {"city": "BJ"}}') - assert len(deltas) > 0 + tc_deltas = [d for d in deltas if "tool_calls" in d] + assert tc_deltas[0]["tool_calls"][0]["function"]["name"] == "get_weather" def test_feed_content_after_tool_call_is_not_emitted(): parser = SimpleJsonToolParser() - parser.feed('{"name": "f", "arguments": {}} trailing text') + deltas = parser.feed('{"name": "f", "arguments": {}} trailing text') assert parser.has_tool_calls + assert not any("trailing" in d.get("content", "") for d in deltas) def _simulate_streaming(parser, text): @@ -513,10 +515,6 @@ def test_factory_create_passes_tools(): assert parser.tool_choice == "required" -def test_factory_list_registered(): - assert "simple_json" in ToolParserFactory.list_registered() - - def test_factory_create_with_tools_only(): tools = [ { @@ -544,29 +542,6 @@ def test_feed_token_ids_do_not_affect_parsing(): ) -def test_parser_uses_token_ids_for_detection(): - class TokenIdParser(BaseToolParser): - def __init__(self, tools=None, tool_choice="auto"): - super().__init__(tools, tool_choice) - self._detections = 0 - - def feed(self, body, current_token_ids=None, delta_token_ids=None): - if current_token_ids and 999 in current_token_ids: - self._detections += 1 - return [] - - def parse_complete(self, body): - return None - - @property - def has_tool_calls(self): - return self._detections > 0 - - parser = TokenIdParser() - parser.feed("hello", current_token_ids=[1, 999, 3]) - assert parser.has_tool_calls - - def test_streaming_partial_name_prefix_never_leaks_into_content(): parser = SimpleJsonToolParser() parts = ["Hello ", '{"', '{"n', '{"na', '{"name"'] diff --git a/tests/module/test_forward_configs.py b/tests/module/test_forward_configs.py index 5e876a1..7f55463 100644 --- a/tests/module/test_forward_configs.py +++ b/tests/module/test_forward_configs.py @@ -369,7 +369,9 @@ def test_moe_component_forward_returns_ffn_output(): output = moe(torch.randn(2, 8, 8)) assert output["hidden_states"].shape == (2, 8, 8) - assert output["aux_loss"] is not None + assert output["aux_loss"].ndim == 0 + assert output["aux_loss"].requires_grad + assert torch.isfinite(output["aux_loss"]) @pytest.mark.parametrize("decoder_sparse_step", [0, -1]) diff --git a/tests/module/test_lora.py b/tests/module/test_lora.py index ceb5f4b..afdf46d 100644 --- a/tests/module/test_lora.py +++ b/tests/module/test_lora.py @@ -332,4 +332,6 @@ def test_collect_lora_info(): info = _collect_lora_info(model) assert "q_proj" in info assert "o_proj" in info - assert "q_proj" in info # each layer has one + # every decoder layer contributes one of each attention projection + assert len(info["q_proj"]) == len(info["o_proj"]) == 2 + assert all(name.endswith("attention.q_proj") for name in info["q_proj"]) diff --git a/tests/optim/test_mano_adamw.py b/tests/optim/test_mano_adamw.py index 75bc767..3a54d8c 100644 --- a/tests/optim/test_mano_adamw.py +++ b/tests/optim/test_mano_adamw.py @@ -34,18 +34,32 @@ def test_mano_one_step_projects_to_tangent_space(): def test_mano_alternates_projection_axis(): - param = torch.nn.Parameter(torch.eye(4) * 3.0) - param.grad = torch.ones(4, 4) + """Step 0 projects along dim 0, step 1 along dim 1 (steps % 2).""" + original = torch.tensor([[3.0, 4.0], [0.0, 2.0]]) + param = torch.nn.Parameter(original.clone()) + eps = 1e-8 + optimizer = Mano( + [param], lr=0.1, momentum=0.0, nesterov=False, eps=eps, weight_decay=0.0 + ) - optimizer = Mano([param], lr=0.1, momentum=0.0, nesterov=False) + param.grad = torch.ones(2, 2) optimizer.step() - dim_step0 = 0 + after_step0 = param.detach().clone() - param.grad = torch.ones(4, 4) + param.grad = torch.ones(2, 2) optimizer.step() - dim_step1 = 1 - assert dim_step0 != dim_step1 + grad = torch.ones(2, 2) + + def projected(after, dim): + tangent = grad - (torch.sum(grad * after, dim=dim, keepdim=True) * after) + direction = tangent / (torch.norm(tangent, p=2, dim=dim, keepdim=True) + eps) + adjusted_lr = 0.1 * 0.2 * math.sqrt(direction.shape[dim]) + return after - adjusted_lr * direction + + torch.testing.assert_close(param.detach(), projected(after_step0, dim=1)) + # The dim=1 result is distinct, so the assertion above pins the axis. + assert not torch.allclose(param.detach(), projected(after_step0, dim=0)) def test_mano_rejects_non_2d_parameters(): diff --git a/tests/serialization/test_hf_adapter.py b/tests/serialization/test_hf_adapter.py index cf5d45a..616bc7a 100644 --- a/tests/serialization/test_hf_adapter.py +++ b/tests/serialization/test_hf_adapter.py @@ -113,6 +113,44 @@ def to_hf_keys(state_dict, head_dim=None): return out +MOE_KWARGS = { + "ffn_type": "moe", + "n_routed_experts": 2, + "n_shared_experts": 1, + "n_activated_experts": 1, + "moe_intermediate_size": 16, + "shared_expert_intermediate_size": 16, +} + + +def _hf_keyed_state_dict(model, cfg): + return to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads) + + +def _assert_hf_roundtrip(cfg, convert_cfg=None): + """HF-keyed weights of a fresh model must convert back exactly.""" + model = AutoRegressiveLM(cfg) + sd = model.state_dict() + converted = convert_hf_weights(_hf_keyed_state_dict(model, cfg), convert_cfg or cfg) + assert_state_dicts_equal(converted, sd) + + +def _assert_hf_directory_load(tmp_path, cfg, raw_config): + """Save HF-format weights and check from_pretrained reproduces logits.""" + model = AutoRegressiveLM(cfg).eval() + save_model( + config=raw_config, + state_dict=_hf_keyed_state_dict(model, cfg), + save_directory=str(tmp_path), + ) + loaded = AutoModel.from_pretrained(tmp_path).eval() + input_ids = torch.randint(0, cfg.vocab_size, (1, 8)) + with torch.no_grad(): + torch.testing.assert_close( + loaded(input_ids)["logits"], model(input_ids)["logits"] + ) + + def test_convert_hf_config_llama(): cfg = convert_hf_config(LLAMA_RAW) assert cfg["model_type"] == "autoregressive_lm" @@ -168,31 +206,14 @@ def test_adapt_config_passthrough(): def test_convert_hf_weights_dense_roundtrip(): - cfg = make_tiny_config() - model = AutoRegressiveLM(cfg) - converted = convert_hf_weights( - to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg - ) - assert_state_dicts_equal(converted, model.state_dict()) + _assert_hf_roundtrip(make_tiny_config()) def test_convert_hf_weights_moe_roundtrip(): - cfg = make_tiny_config( - ffn_type="moe", - n_routed_experts=2, - n_shared_experts=1, - n_activated_experts=1, - moe_intermediate_size=16, - shared_expert_intermediate_size=16, - ) - model = AutoRegressiveLM(cfg) + cfg = make_tiny_config(**MOE_KWARGS) hf_raw = convert_hf_config(MOE_RAW) hf_cfg = ConfigFactory.load(hf_raw) - converted = convert_hf_weights( - to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), - hf_cfg, - ) - assert_state_dicts_equal(converted, model.state_dict()) + _assert_hf_roundtrip(cfg, convert_cfg=hf_cfg) def test_convert_hf_weights_keeps_astrai_keys(): @@ -239,66 +260,28 @@ def test_convert_hf_config_gemma_enables_qk_norm(): def test_convert_hf_weights_moe_with_dense_layers_roundtrip(): - cfg = make_tiny_config( - ffn_type="moe", - n_routed_experts=2, - n_shared_experts=1, - n_activated_experts=1, - moe_intermediate_size=16, - shared_expert_intermediate_size=16, - mlp_only_layers=[0], - decoder_sparse_step=1, + _assert_hf_roundtrip( + make_tiny_config(**MOE_KWARGS, mlp_only_layers=[0], decoder_sparse_step=1) ) - model = AutoRegressiveLM(cfg) - converted = convert_hf_weights( - to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg - ) - assert_state_dicts_equal(converted, model.state_dict()) def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip(): - cfg = make_tiny_config( - ffn_type="moe", - n_routed_experts=2, - n_shared_experts=1, - n_activated_experts=1, - moe_intermediate_size=16, - shared_expert_intermediate_size=16, - ) + cfg = make_tiny_config(**MOE_KWARGS) model = AutoRegressiveLM(cfg) - hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads) hf_sd = { - k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items() + k.replace("shared_experts.", "shared_expert.", 1): v + for k, v in _hf_keyed_state_dict(model, cfg).items() } converted = convert_hf_weights(hf_sd, cfg) assert_state_dicts_equal(converted, model.state_dict()) def test_convert_hf_weights_gemma_qk_norm_roundtrip(): - cfg = make_tiny_config(use_qk_norm=True) - model = AutoRegressiveLM(cfg) - converted = convert_hf_weights( - to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg - ) - assert_state_dicts_equal(converted, model.state_dict()) + _assert_hf_roundtrip(make_tiny_config(use_qk_norm=True)) def test_from_pretrained_hf_directory(tmp_path): - cfg = make_tiny_config() - model = AutoRegressiveLM(cfg).eval() - save_model( - config=LLAMA_RAW, - state_dict=to_hf_keys( - model.state_dict(), cfg.hidden_size // cfg.num_attention_heads - ), - save_directory=str(tmp_path), - ) - loaded = AutoModel.from_pretrained(tmp_path).eval() - input_ids = torch.randint(0, cfg.vocab_size, (1, 8)) - with torch.no_grad(): - torch.testing.assert_close( - loaded(input_ids)["logits"], model(input_ids)["logits"] - ) + _assert_hf_directory_load(tmp_path, make_tiny_config(), LLAMA_RAW) def test_from_pretrained_astrai_directory(tmp_path): @@ -332,9 +315,7 @@ def test_from_pretrained_weights_format_astrai_rejects_hf(tmp_path): model = AutoRegressiveLM(cfg) save_model( config=LLAMA_RAW, - state_dict=to_hf_keys( - model.state_dict(), cfg.hidden_size // cfg.num_attention_heads - ), + state_dict=_hf_keyed_state_dict(model, cfg), save_directory=str(tmp_path), ) with pytest.raises(ValueError): @@ -355,7 +336,7 @@ def test_from_pretrained_invalid_weights_format(tmp_path): def test_from_pretrained_hf_directory_sharded(tmp_path): cfg = make_tiny_config() model = AutoRegressiveLM(cfg).eval() - hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads) + hf_sd = _hf_keyed_state_dict(model, cfg) keys = sorted(hf_sd) split = len(keys) // 2 shard_a = {k: hf_sd[k] for k in keys[:split]} @@ -507,25 +488,4 @@ def test_hf_import_qk_norm_matches_norm_before_rope_reference(): def test_from_pretrained_hf_directory_with_moe(tmp_path): - cfg = make_tiny_config( - ffn_type="moe", - n_routed_experts=2, - n_shared_experts=1, - n_activated_experts=1, - moe_intermediate_size=16, - shared_expert_intermediate_size=16, - ) - model = AutoRegressiveLM(cfg).eval() - save_model( - config=MOE_RAW, - state_dict=to_hf_keys( - model.state_dict(), cfg.hidden_size // cfg.num_attention_heads - ), - save_directory=str(tmp_path), - ) - loaded = AutoModel.from_pretrained(tmp_path).eval() - input_ids = torch.randint(0, cfg.vocab_size, (1, 8)) - with torch.no_grad(): - torch.testing.assert_close( - loaded(input_ids)["logits"], model(input_ids)["logits"] - ) + _assert_hf_directory_load(tmp_path, make_tiny_config(**MOE_KWARGS), MOE_RAW) diff --git a/tests/test_serve_cli.py b/tests/test_serve_cli.py deleted file mode 100644 index 46a8443..0000000 --- a/tests/test_serve_cli.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Unit tests for the serving CLI YAML merge logic.""" - -import click -import pytest -import torch -from click.testing import CliRunner - -from scripts.tools.server import ( - _merge_yaml_into_kwargs, - _resolve_server_config, - server_command, -) - - -def _passed() -> dict: - return { - "host": "0.0.0.0", - "port": 8000, - "reload": False, - "param_path": None, - "device": "cuda", - "dtype": "bfloat16", - "max_batch_size": 16, - "max_seq_len": None, - } - - -def test_yaml_overrides_click_defaults_but_not_explicit_cli(tmp_path): - config_path = tmp_path / "serve.yaml" - config_path.write_text( - "server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n", - encoding="utf-8", - ) - merged = _merge_yaml_into_kwargs( - str(config_path), _passed(), explicit_keys={"device"} - ) - - assert merged["device"] == "cuda" - assert merged["dtype"] == "float16" - assert merged["max_batch_size"] == 8 - - -def test_resolve_config_yaml_wins_by_default(tmp_path): - config_path = tmp_path / "serve.yaml" - config_path.write_text( - "server:\n port: 9000\n max_seq_len: 2048\n", - encoding="utf-8", - ) - resolved = _resolve_server_config(str(config_path), _passed()) - - assert resolved["port"] == 9000 - assert resolved["max_seq_len"] == 2048 - assert resolved["device"] == "cuda" - assert resolved["dtype"] == "bfloat16" - - -def test_resolve_config_rejects_bad_dtype(tmp_path): - config_path = tmp_path / "serve.yaml" - config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8") - - with pytest.raises(click.UsageError, match="server.dtype"): - _resolve_server_config(str(config_path), _passed()) - - -def test_server_command_rejects_bad_yaml_dtype(tmp_path): - config_path = tmp_path / "serve.yaml" - config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8") - - result = CliRunner().invoke(server_command, ["--config", str(config_path)]) - - assert result.exit_code == 2 - assert "server.dtype" in result.output - - -def test_server_command_merges_yaml_and_cli(tmp_path, monkeypatch): - """Full CLI path: YAML values apply, explicit CLI flags override, args reach run_server.""" - config_path = tmp_path / "serve.yaml" - config_path.write_text( - "server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n", - encoding="utf-8", - ) - captured = {} - - def fake_run_server(**kwargs): - captured.update(kwargs) - - monkeypatch.setattr("scripts.tools.server.run_server", fake_run_server) - result = CliRunner().invoke( - server_command, - ["--config", str(config_path), "--max_batch_size", "32"], - ) - - assert result.exit_code == 0, result.output - assert captured["device"] == "cpu" - assert captured["dtype"] == torch.float16 - assert captured["max_batch_size"] == 32 - assert captured["port"] == 8000 - - -def test_config_option_rejects_missing_file(tmp_path): - result = CliRunner().invoke( - server_command, ["--config", str(tmp_path / "nope.yaml")] - ) - - assert result.exit_code == 2 diff --git a/tests/test_serve_runtime.py b/tests/test_serve_runtime.py deleted file mode 100644 index fe50c8f..0000000 --- a/tests/test_serve_runtime.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Unit tests for the serving runtime configuration parser.""" - -import pytest - -from scripts.docker.serve_runtime import load_runtime - - -def _write(tmp_path, body: str) -> str: - config_path = tmp_path / "serve.yaml" - config_path.write_text(body, encoding="utf-8") - return str(config_path) - - -def test_runtime_exports_defaults(tmp_path): - config_path = _write( - tmp_path, - "runtime:\n" - " port: 8000\n" - " paths:\n" - " param: ./params\n" - "server:\n" - " device: cuda\n", - ) - runtime = load_runtime(config_path) - assert runtime["SERVE_PORT"] == "8000" - assert runtime["SERVE_CONTAINER_PORT"] == "8000" - assert runtime["SERVE_PARAM_DIR"] == str((tmp_path / "params").resolve()) - assert runtime["SERVE_GPU_ENABLED"] == "true" - assert "CUDA_VISIBLE_DEVICES" not in runtime - assert runtime["SERVE_DEVICE"] == "cuda" - assert runtime["CUDA_TAG"] == "cu128" - assert runtime["SERVE_JOB_NAME"] == "" - - -def test_runtime_gpu_disabled_requires_cpu(tmp_path): - config_path = _write( - tmp_path, - "runtime:\n gpu:\n enabled: false\nserver:\n device: cuda\n", - ) - with pytest.raises(ValueError, match="server.device must be 'cpu'"): - load_runtime(config_path) - - -def test_runtime_gpu_devices_single_and_ports(tmp_path): - config_path = _write( - tmp_path, - "runtime:\n" - " gpu:\n" - " devices: [1]\n" - " port: 8080\n" - "server:\n" - " port: 9000\n" - " device: cuda\n", - ) - runtime = load_runtime(config_path) - assert runtime["SERVE_PORT"] == "8080" - assert runtime["SERVE_CONTAINER_PORT"] == "9000" - assert runtime["CUDA_VISIBLE_DEVICES"] == "1" - - -def test_runtime_gpu_devices_rejects_multi(tmp_path): - config_path = _write( - tmp_path, - "runtime:\n gpu:\n devices: [0, 1]\n", - ) - with pytest.raises(ValueError, match="single-device"): - load_runtime(config_path) - - -def test_runtime_port_out_of_range(tmp_path): - config_path = _write(tmp_path, "runtime:\n port: 70000\n") - with pytest.raises(ValueError, match="between 1 and 65535"): - load_runtime(config_path) - - -def test_runtime_environment_export(tmp_path): - config_path = _write( - tmp_path, - "runtime:\n environment:\n TOKENIZERS_PARALLELISM: 'false'\n", - ) - runtime = load_runtime(config_path) - assert runtime["environment"] == {"TOKENIZERS_PARALLELISM": "false"} diff --git a/tests/test_train_cli.py b/tests/test_train_cli.py deleted file mode 100644 index 2f7574f..0000000 --- a/tests/test_train_cli.py +++ /dev/null @@ -1,62 +0,0 @@ -import re - -from click.testing import CliRunner - -from scripts.tools.train import _merge_yaml_into_kwargs, train_command - - -def test_yaml_overrides_click_defaults_but_not_explicit_cli(tmp_path): - config_path = tmp_path / "train.yaml" - config_path.write_text( - "training:\n" - " optimizer: nora_nadamw\n" - " max_lr: 0.0002\n" - " nora_lr: 0.004\n" - " batch_per_device: 8\n", - encoding="utf-8", - ) - click_values = { - "optimizer": "nora_nadamw", - "max_lr": 3e-4, - "nora_lr": 5e-3, - "batch_per_device": 16, - } - - merged = _merge_yaml_into_kwargs( - str(config_path), click_values, explicit_keys={"batch_per_device"} - ) - - assert merged["max_lr"] == 2e-4 - assert merged["nora_lr"] == 4e-3 - assert merged["batch_per_device"] == 16 - - -def test_train_dry_run_uses_yaml_then_explicit_cli(tmp_path): - data_path = tmp_path / "data" - model_path = tmp_path / "model" - data_path.mkdir() - model_path.mkdir() - config_path = tmp_path / "train.yaml" - config_path.write_text( - "data:\n" - f" data_root_path: {data_path}\n" - "model:\n" - f" param_path: {model_path}\n" - "training:\n" - " train_type: seq\n" - " optimizer: nora_nadamw\n" - " max_lr: 0.0002\n" - " nora_lr: 0.004\n" - " batch_per_device: 8\n", - encoding="utf-8", - ) - - result = CliRunner().invoke( - train_command, - ["--config", str(config_path), "--dry-run", "--batch_per_device", "16"], - ) - - assert result.exit_code == 0, result.output - assert re.search(r"Optimizer\s+: nora_nadamw", result.output) - assert re.search(r"Batch/device\s+: 16", result.output) - assert re.search(r"Max LR\s+: 0.0002", result.output) diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py index 3194755..5fab573 100644 --- a/tests/trainer/test_callbacks.py +++ b/tests/trainer/test_callbacks.py @@ -23,9 +23,9 @@ def test_gradient_checkpointing_enable_disable(test_model): for layer in model.layers: callback._enable(layer) - for layer in model.layers: + for i, layer in enumerate(model.layers): assert hasattr(layer, "_original_forward") - assert layer.forward is not originals[0] + assert layer.forward is not originals[i] for layer in model.layers: callback._disable(layer) @@ -110,6 +110,10 @@ def test_gradient_checkpointing_trainer_integration( ) trainer = Trainer(train_config) + gc_callbacks = [ + c for c in trainer.callbacks if isinstance(c, GradientCheckpointingCallback) + ] + assert gc_callbacks and gc_callbacks[0].modules == (DecoderBlock,) trainer.train() diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index 411bb30..24c30a3 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -58,6 +58,46 @@ def _make_instruction_batch(n=2): return {"instruction": instructions, "input": inputs} +def _blocking_hook(original, started, release): + """Wrap a hook so it signals ``started`` then blocks until ``release``.""" + + def hook(*args, **kwargs): + started.set() + assert release.wait(timeout=5) + return original(*args, **kwargs) + + return hook + + +def _assert_interleaved(first, second, *, started, release, finished): + """Run ``first`` until it blocks, then assert ``second`` cannot finish + while ``first`` holds the lock; release, join both, and surface errors.""" + errors = [] + + def run_safely(fn): + def run(): + try: + fn() + except BaseException as exc: + errors.append(exc) + + return run + + first_thread = threading.Thread(target=run_safely(first)) + second_thread = threading.Thread(target=run_safely(second)) + first_thread.start() + assert started.wait(timeout=5) + second_thread.start() + assert not finished.wait(timeout=0.1) + + release.set() + first_thread.join(timeout=5) + second_thread.join(timeout=5) + assert not first_thread.is_alive() + assert not second_thread.is_alive() + assert errors == [] + + def test_raw_rollout_fields(): r = RawRollout( prompts=torch.zeros(2, 4, dtype=torch.long), @@ -96,13 +136,6 @@ def test_base_reward_model_is_abstract(): BaseRewardModel() -def test_constant_reward_model_shape(): - rm = ConstantRewardModel(0.5) - out = rm.score(["a", "b"], [["x", "y", "z"], ["p", "q", "r"]]) - assert out.shape == (2, 3) - assert torch.all(out == 0.5) - - def _make_generator(device, **kw): model, _ = make_model(device, max_position_embeddings=128) tokenizer = FakeTokenizer(with_chat_template=True) @@ -159,41 +192,18 @@ def test_rollout_generator_serializes_generation_and_policy_update(device): generation_started = threading.Event() allow_generation_to_finish = threading.Event() update_finished = threading.Event() - thread_errors = [] - original = gen._generate_eval + gen._generate_eval = _blocking_hook( + gen._generate_eval, generation_started, allow_generation_to_finish + ) - def blocking_generate(batch, generation_version): - generation_started.set() - assert allow_generation_to_finish.wait(timeout=5) - return original(batch, generation_version) + _assert_interleaved( + lambda: gen.generate(_make_instruction_batch(n=1)), + lambda: gen.apply_weight_update(1, update_finished.set), + started=generation_started, + release=allow_generation_to_finish, + finished=update_finished, + ) - gen._generate_eval = blocking_generate - - def generate(): - try: - gen.generate(_make_instruction_batch(n=1)) - except BaseException as exc: - thread_errors.append(exc) - - def apply_update(): - try: - gen.apply_weight_update(1, update_finished.set) - except BaseException as exc: - thread_errors.append(exc) - - generation_thread = threading.Thread(target=generate) - update_thread = threading.Thread(target=apply_update) - generation_thread.start() - assert generation_started.wait(timeout=5) - update_thread.start() - assert not update_finished.wait(timeout=0.1) - - allow_generation_to_finish.set() - generation_thread.join(timeout=5) - update_thread.join(timeout=5) - assert not generation_thread.is_alive() - assert not update_thread.is_alive() - assert thread_errors == [] assert update_finished.is_set() assert gen.policy_version == 1 @@ -203,43 +213,23 @@ def test_rollout_generator_serializes_direct_scheduler_update(device): generation_started = threading.Event() allow_generation_to_finish = threading.Event() update_finished = threading.Event() - thread_errors = [] - original = gen._generate_eval - - def blocking_generate(batch, generation_version): - generation_started.set() - assert allow_generation_to_finish.wait(timeout=5) - return original(batch, generation_version) - - gen._generate_eval = blocking_generate + gen._generate_eval = _blocking_hook( + gen._generate_eval, generation_started, allow_generation_to_finish + ) rollout = [] - def generate(): - try: - rollout.append(gen.generate(_make_instruction_batch(n=1))) - except BaseException as exc: - thread_errors.append(exc) - def update_scheduler_directly(): - try: - gen.scheduler.update_weights(1) - update_finished.set() - except BaseException as exc: - thread_errors.append(exc) + gen.scheduler.update_weights(1) + update_finished.set() - generation_thread = threading.Thread(target=generate) - update_thread = threading.Thread(target=update_scheduler_directly) - generation_thread.start() - assert generation_started.wait(timeout=5) - update_thread.start() - assert not update_finished.wait(timeout=0.1) + _assert_interleaved( + lambda: rollout.append(gen.generate(_make_instruction_batch(n=1))), + update_scheduler_directly, + started=generation_started, + release=allow_generation_to_finish, + finished=update_finished, + ) - allow_generation_to_finish.set() - generation_thread.join(timeout=5) - update_thread.join(timeout=5) - assert not generation_thread.is_alive() - assert not update_thread.is_alive() - assert thread_errors == [] assert rollout[0].policy_version == 0 assert gen.policy_version == 1 @@ -466,7 +456,6 @@ def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device): allow_final_validation_to_finish = threading.Event() update_finished = threading.Event() rollout_finished = threading.Event() - thread_errors = [] validation_calls = 0 original_validate = runner._validate_policy_version @@ -481,31 +470,17 @@ def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device): runner._validate_policy_version = blocking_validate def produce_rollout(): - try: - runner(_make_instruction_batch(n=1)) - rollout_finished.set() - except BaseException as exc: - thread_errors.append(exc) + runner(_make_instruction_batch(n=1)) + rollout_finished.set() - def apply_update(): - try: - runner.apply_weight_update(1, update_finished.set) - except BaseException as exc: - thread_errors.append(exc) + _assert_interleaved( + produce_rollout, + lambda: runner.apply_weight_update(1, update_finished.set), + started=final_validation_started, + release=allow_final_validation_to_finish, + finished=update_finished, + ) - rollout_thread = threading.Thread(target=produce_rollout) - update_thread = threading.Thread(target=apply_update) - rollout_thread.start() - assert final_validation_started.wait(timeout=5) - update_thread.start() - assert not update_finished.wait(timeout=0.1) - - allow_final_validation_to_finish.set() - rollout_thread.join(timeout=5) - update_thread.join(timeout=5) - assert not rollout_thread.is_alive() - assert not update_thread.is_alive() - assert thread_errors == [] assert rollout_finished.is_set() assert update_finished.is_set() assert runner._cache is not None diff --git a/tests/trainer/test_train_strategy.py b/tests/trainer/test_train_strategy.py index cb83dab..8d43a49 100644 --- a/tests/trainer/test_train_strategy.py +++ b/tests/trainer/test_train_strategy.py @@ -1,120 +1,78 @@ -import numpy as np +import math + +import pytest import torch from astrai.trainer.schedule import CosineScheduler, SchedulerFactory, SGDRScheduler -def test_schedule_factory_random_configs(): - """Test scheduler factory with random configurations""" +def _stepped_lrs(scheduler, optimizer, n_steps): + """Return the lr after construction plus each of *n_steps* steps.""" + lrs = list(scheduler.get_last_lr()) + for _ in range(n_steps): + optimizer.step() + scheduler.step() + lrs.append(scheduler.get_last_lr()[0]) + return lrs - # Create a simple model and optimizer for testing + +def test_cosine_scheduler_warms_up_then_decays_to_floor(): + """lr ramps linearly to base_lr during warmup, cosine-decays after it, + and never drops below min_rate * base_lr.""" + base_lr = 0.001 model = torch.nn.Linear(10, 2) - optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) + optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr) + scheduler = SchedulerFactory.create( + "cosine", optimizer, warmup_steps=2, lr_decay_steps=4, min_rate=0.1 + ) - # Test multiple random configurations - for _ in range(5): # Test 5 random configurations - # Test multiple random configurations - cosine_params = { - "schedule_type": "cosine", - "warmup_steps": np.random.randint(50, 200), - "total_steps": np.random.randint(1000, 5000), - "min_rate": np.random.uniform(0.01, 0.1), - } - sgdr_params = { - "schedule_type": "sgdr", - "warmup_steps": np.random.randint(50, 200), - "cycle_length": np.random.randint(500, 2000), - "t_mult": np.random.randint(1, 3), - "min_rate": np.random.uniform(0.01, 0.1), - } - for params in [cosine_params, sgdr_params]: - schedule_type = params["schedule_type"] - # Convert parameters for scheduler constructor - if schedule_type == "cosine": - warmup_steps = params["warmup_steps"] - total_steps = params["total_steps"] - min_rate = params["min_rate"] - lr_decay_steps = total_steps - warmup_steps - scheduler = SchedulerFactory.create( - schedule_type, - optimizer, - warmup_steps=warmup_steps, - lr_decay_steps=lr_decay_steps, - min_rate=min_rate, - ) - assert isinstance(scheduler, CosineScheduler) - assert scheduler.warmup_steps == warmup_steps - assert scheduler.lr_decay_steps == lr_decay_steps - assert scheduler.min_rate == min_rate - elif schedule_type == "sgdr": - warmup_steps = params["warmup_steps"] - cycle_length = params["cycle_length"] - t_mult = params["t_mult"] - min_rate = params["min_rate"] - scheduler = SchedulerFactory.create( - schedule_type, - optimizer, - warmup_steps=warmup_steps, - cycle_length=cycle_length, - t_mult=t_mult, - min_rate=min_rate, - ) - assert isinstance(scheduler, SGDRScheduler) - assert scheduler.warmup_steps == warmup_steps - assert scheduler.cycle_length == cycle_length - assert scheduler.t_mult == t_mult - assert scheduler.min_rate == min_rate + assert isinstance(scheduler, CosineScheduler) + lrs = _stepped_lrs(scheduler, optimizer, n_steps=7) - # Test scheduler state dict functionality - state_dict = scheduler.state_dict() - assert "warmup_steps" in state_dict - assert "min_rate" in state_dict - - # Test scheduler step functionality - initial_lr = scheduler.get_last_lr() - optimizer.step() - scheduler.step() - new_lr = scheduler.get_last_lr() - - # Learning rate should change after step, or if it's the first step, - # the epoch counter should increment - assert initial_lr != new_lr or scheduler.last_epoch > -1 + assert lrs[0] == pytest.approx(0.1 * base_lr) # warmup starts at the floor + assert lrs[1] == pytest.approx(0.5 * base_lr) # halfway through warmup + assert lrs[2] == pytest.approx(base_lr) # warmup complete + expected_mid = base_lr * 0.5 * (1.0 + math.cos(math.pi * 0.25)) + assert lrs[3] == pytest.approx(expected_mid) # quarter into decay + assert lrs[5] > 0.1 * base_lr # 3/4 into decay: not clamped yet + assert lrs[6] == pytest.approx(0.1 * base_lr) # clamped at min_rate floor + assert lrs[7] == pytest.approx(0.1 * base_lr) # stays at the floor + assert all(lr >= 0.1 * base_lr - 1e-12 for lr in lrs) -def test_schedule_factory_edge_cases(): - """Test scheduler factory with edge cases and boundary conditions""" - +def test_cosine_scheduler_decays_to_zero_with_min_rate_zero(): + """min_rate=0 must reach exactly 0.0 at the end of decay, not NaN.""" + base_lr = 0.001 model = torch.nn.Linear(10, 2) - optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) + optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr) + scheduler = SchedulerFactory.create( + "cosine", optimizer, warmup_steps=1, lr_decay_steps=9, min_rate=0.0 + ) - # Test edge cases for CosineScheduleConfig - edge_cases = [ - # Minimal warmup and steps - {"warmup_steps": 1, "total_steps": 10, "min_rate": 0.01}, - # Large values - {"warmup_steps": 1000, "total_steps": 10000, "min_rate": 0.5}, - # Zero min_rate (edge case) - {"warmup_steps": 100, "total_steps": 1000, "min_rate": 0.0}, - ] + lrs = _stepped_lrs(scheduler, optimizer, n_steps=11) - for params in edge_cases: - warmup_steps = params["warmup_steps"] - total_steps = params["total_steps"] - min_rate = params["min_rate"] - lr_decay_steps = total_steps - warmup_steps - scheduler = SchedulerFactory.create( - "cosine", - optimizer, - warmup_steps=warmup_steps, - lr_decay_steps=lr_decay_steps, - min_rate=min_rate, - ) - assert scheduler is not None + assert lrs[10] == 0.0 + assert lrs[11] == 0.0 + assert all(math.isfinite(lr) for lr in lrs) - # Test multiple steps - for _ in range(10): - optimizer.step() - scheduler.step() + +def test_sgdr_scheduler_restarts_each_cycle(): + """lr anneals within a cycle, then jumps back to base_lr on restart.""" + base_lr = 0.001 + model = torch.nn.Linear(10, 2) + optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr) + scheduler = SchedulerFactory.create( + "sgdr", optimizer, warmup_steps=2, cycle_length=4, t_mult=1, min_rate=0.1 + ) + + assert isinstance(scheduler, SGDRScheduler) + lrs = _stepped_lrs(scheduler, optimizer, n_steps=7) + + assert lrs[2] == pytest.approx(base_lr) # cycle start + expected_mid = base_lr * (0.1 + 0.9 * 0.5) # halfway through the cycle + assert lrs[4] == pytest.approx(expected_mid) + assert lrs[5] < lrs[4] # still annealing at the cycle end + assert lrs[6] == pytest.approx(base_lr) # restart: back to full lr def test_schedule_factory_state_persistence():