refactor: eliminate test duplication via shared helpers

- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers
- Replace 15 copies of device one-liner with session-scoped fixture
- Collapse 5 near-identical Dataset subclasses into RandomTokenDataset
- Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions
- Make test_callbacks and test_early_stopping use existing train_config_factory
- Replace 6 duplicate meta.json read blocks with load_shard_meta
- Fix mkdtemp leaks in test_lora.py with TemporaryDirectory

tests/trainer/test_grpo_strategy.py

@@ -1,35 +1,9 @@
 import pytest
 import torch
 
-from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.model.transformer import AutoRegressiveLM
 from astrai.trainer.strategy import GRPOStrategy
-
-
-class _FakeExecutor:
-    """Minimal executor stub providing ``unwrap_model`` for ref model creation."""
-
-    def unwrap_model(self, model):
-        return model.state_dict()
-
-
-def _make_config(vocab_size=200, max_position_embeddings=64):
-    return AutoRegressiveLMConfig(
-        vocab_size=vocab_size,
-        hidden_size=16,
-        num_attention_heads=2,
-        num_key_value_heads=1,
-        intermediate_size=32,
-        max_position_embeddings=max_position_embeddings,
-        num_hidden_layers=2,
-        rms_norm_eps=1e-5,
-    )
-
-
-def _make_model(device):
-    config = _make_config()
-    model = AutoRegressiveLM(config).to(device=device)
-    return model, config
+from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config
 
 
 def _make_batch(
@@ -44,9 +18,7 @@ def _make_batch(
     responses = torch.randint(
         0, 200, (batch_size, group_size, response_len), device=device
     )
-    # All response tokens valid.
     masks = torch.ones(batch_size, group_size, response_len, device=device)
-    # Distinct rewards per group member so std > 0.
     rewards = torch.randn(batch_size, group_size, device=device)
     return {
         "prompts": prompts,
@@ -56,23 +28,12 @@ def _make_batch(
     }
 
 
-def _make_frozen_copy(model, device):
-    """Create a frozen copy of ``model`` with independent weights loaded."""
-    config = _make_config()
-    copy = AutoRegressiveLM(config).to(device=device)
-    copy.load_state_dict(model.state_dict())
-    copy.requires_grad_(False)
-    copy.eval()
-    return copy
-
-
 @pytest.fixture
-def grpo_strategy():
+def grpo_strategy(device):
     """Build a GRPOStrategy with a small real model and fake executor."""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model, config = _make_model(device)
-    old_model = _make_frozen_copy(model, device)
-    ref_model = _make_frozen_copy(model, device)
+    model, config = make_model(device)
+    old_model = make_frozen(model, device)
+    ref_model = make_frozen(model, device)
 
     strategy = GRPOStrategy(
         model=model,
@@ -83,7 +44,7 @@ def grpo_strategy():
         kl_coef=0.01,
         group_size=4,
         model_fn=lambda c=config: AutoRegressiveLM(c).to(device=device),
-        executor=_FakeExecutor(),
+        executor=FakeExecutor(),
     )
     return strategy, device
 
@@ -103,7 +64,6 @@ def test_grpo_loss_backward(grpo_strategy):
     batch = _make_batch(device=device)
     loss = strategy.compute_loss(batch)
     loss.backward()
-    # At least some parameter should receive a gradient.
     has_grad = any(
         p.grad is not None and p.grad.abs().sum().item() > 0
         for p in strategy.model.parameters()
@@ -136,32 +96,27 @@ def test_grpo_prompt_tokens_masked(grpo_strategy):
     the policy loss should be zero (no valid tokens contribute)."""
     strategy, device = grpo_strategy
     batch = _make_batch(device=device)
-    # Zero out all response masks → no response token contributes.
     batch["masks"] = torch.zeros_like(batch["masks"])
     loss = strategy.compute_loss(batch)
-    # With no valid tokens, policy_loss term is 0 and KL term is 0.
     assert loss.item() == pytest.approx(0.0, abs=1e-6)
 
 
 def test_grpo_identical_rewards_zero_advantage(grpo_strategy):
-    """When all group rewards are identical, advantage is 0 → policy_loss is 0.
+    """When all group rewards are identical, advantage is 0 -> policy_loss is 0.
     Only the KL term remains (which is 0 when policy == ref at init)."""
     strategy, device = grpo_strategy
     batch = _make_batch(device=device)
     batch["rewards"] = torch.ones(batch["rewards"].shape, device=device)
     loss = strategy.compute_loss(batch)
-    # At init policy == old == ref, so ratio == 1, KL == 0; advantage == 0.
     assert loss.item() == pytest.approx(0.0, abs=1e-5)
 
 
 def test_grpo_sync_old_model(grpo_strategy):
     """sync_old_model copies current policy weights into old_model."""
     strategy, device = grpo_strategy
-    # Perturb policy model so it differs from old.
     with torch.no_grad():
         for p in strategy.model.parameters():
             p.add_(0.05)
-    # old_model should still hold original weights (differ from policy).
     policy_sd = strategy.model.state_dict()
     old_sd = strategy.old_model.state_dict()
     differs_before = any(
@@ -195,14 +150,12 @@ 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
-    # Diverge policy from ref.
     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()
-    # With distinct rewards and diverged policy, loss should be non-trivial.
     assert loss.abs().item() > 1e-4