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

test: prune low-value and duplicate tests

Browse Source 
- Remove tautological test_trainer assertions that never trained
- Drop grpo isfinite-only smokes and merge frozen-model checks via parametrize
- Merge duplicate tool_parser cases (find/streaming/factory) with parametrize
- Collapse duplicate dataset store/detect_format tests
- Remove misleading scheduler/task tests that asserted the opposite of their names
- Merge signal-handler SIGTERM/SIGINT into one parametrized case
- Drop cross-file grpo strategy duplication kept in online_strategy
This commit is contained in:
ViperEkura
2026-08-01 16:01:20 +08:00
parent 91acaf4b0b
commit a27c8a819d
10 changed files with 137 additions and 342 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/trainer/test_grpo_strategy.py
+4 -55
View File
@@ -71,23 +71,14 @@ def test_grpo_loss_backward(grpo_strategy):
assert has_grad
def test_grpo_ref_model_not_updated(grpo_strategy):
"""Backward should not populate gradients on ref_model."""
@pytest.mark.parametrize("model_name", ["ref_model", "old_model"])
def test_grpo_frozen_models_not_updated(grpo_strategy, model_name):
"""Backward should not populate gradients on ref_model or old_model."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
loss.backward()
for p in strategy.ref_model.parameters():
assert p.grad is None
def test_grpo_old_model_not_updated(grpo_strategy):
"""Backward should not populate gradients on old_model."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
loss.backward()
for p in strategy.old_model.parameters():
for p in getattr(strategy, model_name).parameters():
assert p.grad is None
@@ -133,45 +124,3 @@ def test_grpo_sync_old_model(grpo_strategy):
if k in old_sd_after
)
assert matches
def test_grpo_partial_mask(grpo_strategy):
"""Only the first half of response tokens are valid."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
B, G, R = batch["masks"].shape
half = R // 2
batch["masks"][:, :, half:] = 0.0
loss = strategy.compute_loss(batch)
assert torch.isfinite(loss).item()
def test_grpo_clipping_effect(grpo_strategy):
"""After diverging policy from ref, ratio should be clipped to [1-eps, 1+eps]
on the surrogate. Verify loss is finite and non-zero for distinct rewards."""
strategy, device = grpo_strategy
with torch.no_grad():
for p in strategy.model.parameters():
p.add_(0.3)
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
assert torch.isfinite(loss).item()
assert loss.abs().item() > 1e-4
def test_grpo_no_reduction_param():
"""GRPOStrategy.__init__ must not accept ``reduction`` (removed)."""
import inspect
sig = inspect.signature(GRPOStrategy.__init__)
assert "reduction" not in sig.parameters
def test_grpo_shapes_3d_batch(grpo_strategy):
"""Verify compute_loss handles non-square prompt/response lengths."""
strategy, device = grpo_strategy
batch = _make_batch(
batch_size=3, group_size=4, prompt_len=10, response_len=8, device=device
)
loss = strategy.compute_loss(batch)
assert torch.isfinite(loss).item()
tests/trainer/test_online_strategy.py
+4 -28
View File
@@ -96,12 +96,10 @@ def test_factory_registers_online_aliases():
assert StrategyFactory.get_component_class("online_dpo") is DPOStrategy
def test_grpo_supports_online(device):
assert _make_grpo(device).supports_online() is True
def test_dpo_supports_online(device):
assert _make_dpo(device).supports_online() is True
@pytest.mark.parametrize("make_fn", ["_make_grpo", "_make_dpo"])
def test_online_strategies_support_online(device, make_fn):
maker = {"_make_grpo": _make_grpo, "_make_dpo": _make_dpo}[make_fn]
assert maker(device).supports_online() is True
def test_base_strategy_prepare_from_rollout_raises_by_default(device):
@@ -267,17 +265,6 @@ def test_step_called_when_sync_gradients_true(device):
assert runner.step_calls == 1
def test_loss_is_differentiable_grpo(device):
strat = _make_grpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
loss.backward()
has_grad = any(
p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
)
assert has_grad
def test_loss_is_differentiable_dpo(device):
strat = _make_dpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
@@ -289,17 +276,6 @@ def test_loss_is_differentiable_dpo(device):
assert has_grad
def test_ref_and_old_model_not_updated_by_backward_grpo(device):
strat = _make_grpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
loss.backward()
for p in strat.ref_model.parameters():
assert p.grad is None
for p in strat.old_model.parameters():
assert p.grad is None
def test_ref_model_not_updated_by_backward_dpo(device):
strat = _make_dpo(device)
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
tests/trainer/test_rollout.py
+3
View File
@@ -81,6 +81,9 @@ def test_rollout_result_inherits_raw_rollout_fields():
assert r.prompts.shape == (2, 4)
assert r.responses.shape == (2, 3, 5)
assert r.prompt_mask.shape == (2, 4)
# RolloutResult must carry every RawRollout field.
raw_fields = {f for f in RawRollout.__dataclass_fields__}
assert raw_fields.issubset(set(RolloutResult.__dataclass_fields__))
def test_base_reward_model_is_abstract():
tests/trainer/test_signal_handler.py
+3 -11
View File
@@ -131,18 +131,10 @@ def test_register_signal_handlers():
assert ctx.stop_requested
def test_sigterm_triggers_checkpoint_save(base_test_env):
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGTERM)
assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
meta = load_checkpoint_meta(base_test_env["test_dir"])
assert "consumed_samples" in meta
assert meta["consumed_samples"] >= 0
@pytest.mark.slow
def test_sigint_triggers_checkpoint_save(base_test_env):
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGINT)
@pytest.mark.parametrize("sig", [signal.SIGTERM, signal.SIGINT])
def test_signal_triggers_checkpoint_save(base_test_env, sig):
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], sig)
assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
meta = load_checkpoint_meta(base_test_env["test_dir"])
tests/trainer/test_trainer.py
+22 -48
View File
@@ -1,13 +1,13 @@
import pytest
from astrai.trainer import Trainer
# train_config_factory is injected via fixture
def test_different_batch_sizes(base_test_env, random_dataset, train_config_factory):
"""Test training with different batch sizes"""
batch_sizes = [1, 2, 4, 8]
for batch_per_device in batch_sizes:
def test_training_runs_with_various_batch_sizes(
base_test_env, random_dataset, train_config_factory
):
"""Training should complete for a range of batch sizes without error."""
for batch_per_device in [1, 2, 4]:
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
dataset=random_dataset,
@@ -15,48 +15,22 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto
device=base_test_env["device"],
batch_per_device=batch_per_device,
)
assert train_config.batch_per_device == batch_per_device
def test_gradient_accumulation(base_test_env, random_dataset, train_config_factory):
"""Test training with different gradient accumulation steps"""
grad_accum_steps_list = [1, 2, 4]
for grad_accum_steps in grad_accum_steps_list:
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
dataset=random_dataset,
test_dir=base_test_env["test_dir"],
device=base_test_env["device"],
batch_per_device=2,
grad_accum_steps=grad_accum_steps,
)
trainer = Trainer(train_config)
trainer.train()
assert train_config.grad_accum_steps == grad_accum_steps
def test_memory_efficient_training(base_test_env, random_dataset, train_config_factory):
"""Test training with memory-efficient configurations"""
# Test with smaller batch sizes and gradient checkpointing
small_batch_configs = [
{"batch_per_device": 1, "grad_accum_steps": 8},
{"batch_per_device": 2, "grad_accum_steps": 4},
{"batch_per_device": 4, "grad_accum_steps": 2},
]
for config in small_batch_configs:
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
dataset=random_dataset,
test_dir=base_test_env["test_dir"],
device=base_test_env["device"],
batch_per_device=config["batch_per_device"],
grad_accum_steps=config["grad_accum_steps"],
)
assert train_config.grad_accum_steps == config["grad_accum_steps"]
assert train_config.batch_per_device == config["batch_per_device"]
@pytest.mark.slow
def test_gradient_accumulation_runs(
base_test_env, random_dataset, train_config_factory
):
"""Training with gradient accumulation should complete."""
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
dataset=random_dataset,
test_dir=base_test_env["test_dir"],
device=base_test_env["device"],
batch_per_device=2,
grad_accum_steps=4,
)
trainer = Trainer(train_config)
trainer.train()