ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: deduplicate and restructure test suite

Browse Source 
- 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.
This commit is contained in:
ViperEkura
2026-08-20 01:53:28 +08:00
parent 53a7149577
commit 84753d3e08
13 changed files with 312 additions and 454 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/module/test_lora.py
+30 -40
View File
@@ -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():
tests/module/test_model_forward.py
+23 -28
View File
@@ -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)