test: refactor tests and fix inference edge cases

- Convert protocol and MoE test classes to plain functions
- Add real server/engine integration and generate_async tests
- Isolate test model per test and use pytest tmp_path
- Reset FastAPI engine state after inference tests
- Fix generate_async StopIteration handling on Python 3.12
- Fix HF adapter MoE dense/shared and Gemma qk_norm mapping
- Correct dev dependency httpx2 to httpx

astrai/inference/engine.py

@@ -156,9 +156,8 @@ class InferenceEngine:
         async def _agen():
             loop = asyncio.get_event_loop()
             while True:
-                try:
-                    token = await loop.run_in_executor(None, next, sync_gen)
-                except StopIteration:
+                token = await loop.run_in_executor(None, next, sync_gen, None)
+                if token is None:
                     break
                 yield token
 

astrai/serialization/hf_adapter.py

@@ -44,6 +44,8 @@ HF_MODEL_TYPES = frozenset(
 
 _EMBED = re.compile(r"^model\.embed_tokens\.weight$")
 _ATTN = re.compile(r"^model\.layers\.(\d+)\.self_attn\.(q|k|v|o)_proj\.(weight|bias)$")
+_Q_NORM = re.compile(r"^model\.layers\.(\d+)\.self_attn\.q_norm\.weight$")
+_K_NORM = re.compile(r"^model\.layers\.(\d+)\.self_attn\.k_norm\.weight$")
 _INPUT_NORM = re.compile(r"^model\.layers\.(\d+)\.input_layernorm\.weight$")
 _POST_NORM = re.compile(r"^model\.layers\.(\d+)\.post_attention_layernorm\.weight$")
 _FINAL_NORM = re.compile(r"^model\.norm\.weight$")
@@ -56,7 +58,7 @@ _MOE_EXPERTS = re.compile(
     r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(weight|bias)$"
 )
 _MOE_SHARED = re.compile(
-    r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(\d+)\."
+    r"^model\.layers\.(\d+)\.mlp\.shared_expert(?:s)?\.(\d+)\."
     r"(gate|up|down)_proj\.(weight|bias)$"
 )
 
@@ -74,6 +76,17 @@ def looks_like_hf_state_dict(state_dict: Mapping[str, Any]) -> bool:
     )
 
 
+def _is_dense_mlp_layer(config: BaseConfig, layer_id: int) -> bool:
+    """Return whether a layer uses dense MLP instead of routed experts."""
+    if getattr(config, "ffn_type", "mlp") != "moe":
+        return True
+    mlp_only = getattr(config, "mlp_only_layers", None) or []
+    if layer_id in mlp_only:
+        return True
+    step = getattr(config, "decoder_sparse_step", 1) or 1
+    return step > 1 and (layer_id + 1) % step != 0
+
+
 def adapt_config(raw: Dict[str, Any]) -> Dict[str, Any]:
     """Translate *raw* for AstrAI if it looks like an HF model config."""
     if raw.get("model_type") in HF_MODEL_TYPES:
@@ -112,6 +125,8 @@ def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]:
         "topk_method",
         "norm_topk_prob",
         "moe_aux_loss_coef",
+        "decoder_sparse_step",
+        "mlp_only_layers",
         "neftune_alpha",
     ):
         if key in raw:
@@ -119,6 +134,13 @@ def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]:
 
     if "qk_norm" in raw and "use_qk_norm" not in cfg:
         cfg["use_qk_norm"] = raw["qk_norm"]
+    if (
+        raw.get("model_type") in ("gemma", "gemma2")
+        and "use_qk_norm" not in cfg
+        and "qk_norm" not in raw
+    ):
+        # Gemma/Gemma2 always apply RMSNorm to Q and K before attention.
+        cfg["use_qk_norm"] = True
 
     n_heads = raw.get("num_attention_heads")
     if cfg.get("num_key_value_heads") is None and n_heads is not None:
@@ -205,6 +227,10 @@ def convert_hf_weights(
                             f"layers.{m.group(1)}.mlp.shared_experts.{m.group(2)}."
                             f"{m.group(3)}.{m.group(4)}"
                         )
+            if new_key is None:
+                m = _DENSE_MLP.match(key)
+                if m and _is_dense_mlp_layer(config, int(m.group(1))):
+                    new_key = f"layers.{m.group(1)}.mlp.{m.group(2)}.{m.group(3)}"
         else:
             m = _DENSE_MLP.match(key)
             if m:
@@ -216,6 +242,10 @@ def convert_hf_weights(
                 new_key = (
                     f"layers.{m.group(1)}.attention.{m.group(2)}_proj.{m.group(3)}"
                 )
+            elif (m := _Q_NORM.match(key)) is not None:
+                new_key = f"layers.{m.group(1)}.attention.q_norm.weight"
+            elif (m := _K_NORM.match(key)) is not None:
+                new_key = f"layers.{m.group(1)}.attention.k_norm.weight"
             elif (m := _INPUT_NORM.match(key)) is not None:
                 new_key = f"layers.{m.group(1)}.input_norm.weight"
             elif (m := _POST_NORM.match(key)) is not None:

pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
 urls = { Homepage = "https://github.com/ViperEkura/AstrAI" }
 
 [project.optional-dependencies]
-dev = ["pytest==9.0.2", "ruff", "httpx2"]
+dev = ["pytest==9.0.2", "ruff", "httpx"]
 flash = ["flash-attn>=2.6"]
 
 [tool.setuptools.packages.find]

tests/conftest.py

@@ -1,7 +1,5 @@
 import json
 import os
-import shutil
-import tempfile
 
 import pytest
 import torch
@@ -44,20 +42,18 @@ def test_tokenizer():
     return create_test_tokenizer()
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def test_model(device):
-    """Session-scoped small AutoRegressiveLM model, created once."""
+    """Function-scoped small AutoRegressiveLM model, isolated per test."""
     config = make_tiny_config()
     model = AutoRegressiveLM(config).to(device=device)
     return {"model": model, "device": device, "config": config}
 
 
 @pytest.fixture
-def temp_dir():
-    """Function-scoped temporary directory, cleaned up after each test."""
-    d = tempfile.mkdtemp()
-    yield d
-    shutil.rmtree(d, ignore_errors=True)
+def temp_dir(tmp_path):
+    """Function-scoped temporary directory, cleaned up by pytest."""
+    return str(tmp_path)
 
 
 @pytest.fixture

tests/inference/conftest.py

@@ -8,6 +8,13 @@ from fastapi.testclient import TestClient
 from astrai.inference import get_app
 
 
+@pytest.fixture(autouse=True)
+def _cleanup_app_engine():
+    """Reset the lazy FastAPI singleton engine after each inference test."""
+    yield
+    get_app().state.engine = None
+
+
 @pytest.fixture
 def client():
     """Provide a test client for the FastAPI app."""

tests/inference/test_engine.py

@@ -1,5 +1,6 @@
 """Unit tests for GenerateResult accumulator and InferenceEngine.generate()."""
 
+import asyncio
 import threading
 from unittest.mock import MagicMock, patch
 
@@ -156,6 +157,35 @@ def test_engine_generate_streaming_yields_tokens():
         assert tokens == ["t1", "t2"]
 
 
+def test_engine_generate_async_yields_tokens_until_stop():
+    mock_model, mock_tokenizer = _make_engine_mocks(decode="tok")
+    callbacks_saved = []
+
+    def capture_cb(prompt, **kw):
+        callbacks_saved.append(kw.get("stream_callback"))
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+        instance.add_task.side_effect = capture_cb
+        instance.remove_task.return_value = []
+
+        eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
+        agen = eng.generate_async("hello")
+
+        async def collect():
+            out = []
+            async for token in agen:
+                out.append(token)
+            return out
+
+        cb = callbacks_saved[0]
+        cb("t1")
+        cb("t2")
+        cb(STOP)
+
+        assert asyncio.run(collect()) == ["t1", "t2"]
+
+
 def test_engine_generate_non_streaming_batch():
     mock_model, mock_tokenizer = _make_engine_mocks(decode="r")
 

tests/inference/test_protocol.py

@@ -3,8 +3,6 @@
 import json
 from unittest.mock import MagicMock
 
-import pytest
-
 from astrai.inference.network.anthropic import AnthropicResponseBuilder
 from astrai.inference.network.openai import OpenAIResponseBuilder
 from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo
@@ -34,48 +32,7 @@ def _sse_payloads(events):
     return payloads
 
 
-class TestStopChecker:
-    def test_check_finds_match(self):
-        sc = StopChecker(["stop", "end"])
-        assert sc.check("hello stop world") == "stop"
-
-    def test_check_returns_none_when_no_match(self):
-        sc = StopChecker(["stop"])
-        assert sc.check("hello world") is None
-
-    def test_check_empty_sequences(self):
-        sc = StopChecker([])
-        assert sc.check("hello") is None
-
-
-class TestGenContext:
-    def test_defaults(self):
-        ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
-        assert ctx.completion_tokens == 0
-
-    def test_fields_mutable(self):
-        ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
-        ctx.completion_tokens = 42
-        assert ctx.completion_tokens == 42
-
-
-class TestStopInfo:
-    def test_defaults(self):
-        s = StopInfo()
-        assert s.matched is None
-        assert s.body == ""
-        assert s.yielded == ""
-
-    def test_with_values(self):
-        s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
-        assert s.matched == "stop"
-        assert s.body == "hello stop"
-        assert s.yielded == "hello "
-
-
-class TestOpenAIResponseBuilder:
-    @pytest.fixture
-    def builder(self):
+def _make_openai_builder():
     builder = OpenAIResponseBuilder()
     req = MagicMock()
     req.messages = [MagicMock(role="user", content="Hello")]
@@ -86,7 +43,61 @@ class TestOpenAIResponseBuilder:
     builder.prepare(req, engine)
     return builder
 
-    def test_prepare_returns_prompt_ctx_stops(self, builder):
+
+def _make_anthropic_builder():
+    builder = AnthropicResponseBuilder()
+    req = MagicMock()
+    req.messages = [MagicMock(role="user", content="Hello")]
+    req.model = "claude"
+    req.system = None
+    engine = MagicMock()
+    engine.tokenizer.apply_chat_template.return_value = "Hello"
+    builder.prepare(req, engine)
+    return builder
+
+
+def test_check_finds_match():
+    sc = StopChecker(["stop", "end"])
+    assert sc.check("hello stop world") == "stop"
+
+
+def test_check_returns_none_when_no_match():
+    sc = StopChecker(["stop"])
+    assert sc.check("hello world") is None
+
+
+def test_check_empty_sequences():
+    sc = StopChecker([])
+    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()
     req.messages = [MagicMock(role="user", content="Hi")]
     req.stop = ["END"]
@@ -99,7 +110,9 @@ class TestOpenAIResponseBuilder:
     assert ctx.prompt_tokens == 0
     assert stops == ["END"]
 
-    def test_prepare_no_stop_returns_empty_list(self, builder):
+
+def test_openai_prepare_no_stop_returns_empty_list():
+    builder = _make_openai_builder()
     req = MagicMock()
     req.messages = []
     req.stop = None
@@ -109,7 +122,9 @@ class TestOpenAIResponseBuilder:
     _, _, stops = builder.prepare(req, engine)
     assert stops == []
 
-    def test_format_stream_start(self, builder):
+
+def test_openai_format_stream_start():
+    builder = _make_openai_builder()
     ctx = _make_ctx()
     events = builder.format_stream_start(ctx)
     payloads = _sse_payloads(events)
@@ -119,13 +134,17 @@ class TestOpenAIResponseBuilder:
     assert p["choices"][0]["delta"]["role"] == "assistant"
     assert p["choices"][0]["finish_reason"] is None
 
-    def test_format_chunk(self, builder):
+
+def test_openai_format_chunk():
+    builder = _make_openai_builder()
     events = builder.format_chunk("hello", body="hello")
     payload = json.loads(events[0].split("data: ", 1)[1])
     assert payload["choices"][0]["delta"]["content"] == "hello"
     assert payload["choices"][0]["finish_reason"] is None
 
-    def test_format_stream_end(self, builder):
+
+def test_openai_format_stream_end():
+    builder = _make_openai_builder()
     ctx = _make_ctx(completion_tokens=5)
     stop = StopInfo(matched="stop")
     events = builder.format_stream_end(ctx, stop)
@@ -136,7 +155,9 @@ class TestOpenAIResponseBuilder:
     assert usage["completion_tokens"] == 5
     assert usage["total_tokens"] == 15
 
-    def test_format_response(self, builder):
+
+def test_openai_format_response():
+    builder = _make_openai_builder()
     ctx = _make_ctx()
     stop = StopInfo()
     resp = builder.format_response(ctx, "hello", stop)
@@ -145,20 +166,8 @@ class TestOpenAIResponseBuilder:
     assert resp["usage"]["prompt_tokens"] == 10
 
 
-class TestAnthropicResponseBuilder:
-    @pytest.fixture
-    def builder(self):
-        builder = AnthropicResponseBuilder()
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hello")]
-        req.model = "claude"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hello"
-        req.system = None
-        builder.prepare(req, engine)
-        return builder
-
-    def test_prepare_messages(self, builder):
+def test_anthropic_prepare_messages():
+    builder = _make_anthropic_builder()
     req = MagicMock()
     req.messages = [MagicMock(role="user", content="Hi")]
     req.model = "claude"
@@ -170,7 +179,9 @@ class TestAnthropicResponseBuilder:
     assert prompt == "Hi"
     assert stops == []
 
-    def test_prepare_with_stop_sequences(self, builder):
+
+def test_anthropic_prepare_with_stop_sequences():
+    builder = _make_anthropic_builder()
     req = MagicMock()
     req.messages = []
     req.model = "x"
@@ -181,7 +192,9 @@ class TestAnthropicResponseBuilder:
     _, _, stops = builder.prepare(req, engine)
     assert stops == ["stop", "end"]
 
-    def test_format_stream_start(self, builder):
+
+def test_anthropic_format_stream_start():
+    builder = _make_anthropic_builder()
     ctx = _make_ctx(prompt_tokens=3)
     events = builder.format_stream_start(ctx)
     payloads = _sse_payloads(events)
@@ -190,23 +203,28 @@ class TestAnthropicResponseBuilder:
     assert payloads[0]["message"]["usage"]["input_tokens"] == 3
     assert payloads[1]["type"] == "content_block_start"
 
-    def test_format_chunk(self, builder):
+
+def test_anthropic_format_chunk():
+    builder = _make_anthropic_builder()
     events = builder.format_chunk("tok", body="tok")
     payload = json.loads(events[0].split("data: ", 1)[1])
     assert payload["type"] == "content_block_delta"
     assert payload["delta"]["text"] == "tok"
 
-    def test_format_stream_end_no_stop(self, builder):
+
+def test_anthropic_format_stream_end_no_stop():
+    builder = _make_anthropic_builder()
     ctx = _make_ctx(completion_tokens=3)
     stop = StopInfo()
     events = builder.format_stream_end(ctx, stop)
     payloads = _sse_payloads(events)
-        # content_block_stop, message_delta, message_stop
     types = [p["type"] for p in payloads]
     assert types == ["content_block_stop", "message_delta", "message_stop"]
     assert payloads[1]["delta"]["stop_reason"] == "end_turn"
 
-    def test_format_stream_end_with_stop_trims_and_emits_remaining(self, builder):
+
+def test_anthropic_format_stream_end_with_stop_trims_and_emits_remaining():
+    builder = _make_anthropic_builder()
     ctx = _make_ctx(completion_tokens=7)
     stop = StopInfo(
         matched="END",
@@ -215,7 +233,6 @@ class TestAnthropicResponseBuilder:
     )
     events = builder.format_stream_end(ctx, stop)
     payloads = _sse_payloads(events)
-        # unyielded delta, content_block_stop, message_delta, message_stop
     types = [p["type"] for p in payloads]
     assert types == [
         "content_block_delta",
@@ -227,7 +244,9 @@ class TestAnthropicResponseBuilder:
     assert payloads[2]["delta"]["stop_reason"] == "stop_sequence"
     assert payloads[2]["delta"]["stop_sequence"] == "END"
 
-    def test_format_stream_end_stop_trimmed_already_yielded(self, builder):
+
+def test_anthropic_format_stream_end_stop_trimmed_already_yielded():
+    builder = _make_anthropic_builder()
     ctx = _make_ctx()
     stop = StopInfo(
         matched="END",
@@ -236,11 +255,12 @@ class TestAnthropicResponseBuilder:
     )
     events = builder.format_stream_end(ctx, stop)
     payloads = _sse_payloads(events)
-        # No unyielded delta (everything already sent)
     types = [p["type"] for p in payloads]
     assert types == ["content_block_stop", "message_delta", "message_stop"]
 
-    def test_format_response_with_stop_trims_content(self, builder):
+
+def test_anthropic_format_response_with_stop_trims_content():
+    builder = _make_anthropic_builder()
     ctx = _make_ctx()
     stop = StopInfo(matched="STOP", body="text STOP extra", yielded="text ")
     resp = builder.format_response(ctx, "text STOP extra", stop)
@@ -248,7 +268,9 @@ class TestAnthropicResponseBuilder:
     assert resp["stop_reason"] == "stop_sequence"
     assert resp["stop_sequence"] == "STOP"
 
-    def test_format_response_no_stop(self, builder):
+
+def test_anthropic_format_response_no_stop():
+    builder = _make_anthropic_builder()
     ctx = _make_ctx()
     stop = StopInfo()
     resp = builder.format_response(ctx, "full text", stop)

tests/inference/test_server.py

@@ -1,8 +1,15 @@
 """Unit tests for the inference HTTP server."""
 
+from pathlib import Path
+
 import pytest
+import torch
 
 from astrai.inference import get_app
+from astrai.inference.network.app import _create_engine
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.serialization import save_model
+from tests.helpers import CHAT_TEMPLATE, build_test_tokenizer, make_tiny_config
 
 
 def test_health_no_model(client):
@@ -212,5 +219,42 @@ def test_chat_completions_stop_sequence_stream(client, loaded_model):
     )
 
 
+def test_chat_completions_real_engine(tmp_path, client):
+    """POST /v1/chat/completions with a real tiny model and tokenizer."""
+    cfg = make_tiny_config(vocab_size=256)
+    model = AutoRegressiveLM(cfg).eval()
+    save_model(cfg.to_dict(), model.state_dict(), str(tmp_path))
+
+    tokenizer = build_test_tokenizer(vocab_size=256, chat_template=CHAT_TEMPLATE)
+    tokenizer.save_pretrained(str(tmp_path))
+
+    engine = _create_engine(
+        Path(tmp_path),
+        device="cpu",
+        dtype=torch.float32,
+        max_batch_size=1,
+        max_seq_len=64,
+    )
+    try:
+        get_app().state.engine = engine
+        response = client.post(
+            "/v1/chat/completions",
+            json={
+                "messages": [{"role": "user", "content": "Hello"}],
+                "max_tokens": 4,
+                "temperature": 0.0,
+                "stream": False,
+            },
+        )
+        assert response.status_code == 200
+        data = response.json()
+        content = data["choices"][0]["message"]["content"]
+        assert isinstance(content, str)
+        assert data["usage"]["completion_tokens"] > 0
+    finally:
+        engine.shutdown()
+        get_app().state.engine = None
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/module/test_model_forward.py

@@ -54,6 +54,20 @@ def _make_batch(config, batch_size=2, seq_len=8, with_extra=False):
     return batch
 
 
+def _make_seq_moe_fixture(device):
+    config = _make_tiny_moe_config()
+    model = _make_model(config).to(device)
+    model.train()
+    return config, model
+
+
+def _make_sft_moe_fixture(device):
+    config = _make_tiny_moe_config()
+    model = _make_model(config).to(device)
+    model.train()
+    return config, model
+
+
 def test_model_forward_contract_uses_dense_training_and_packed_inference():
     from astrai.inference.cache import PagePool, TaskCacheManager
     from astrai.inference.workspace import InferenceWorkspace
@@ -185,35 +199,28 @@ def test_moe_metrics_flow_through_wrapped_model(device):
     assert "router_entropy" in strategy._moe_metrics
 
 
-class TestSEQStrategyMoE:
-    """End‑to‑end tests for SEQStrategy with MoE aux loss."""
-
-    @pytest.fixture(autouse=True)
-    def setup(self, device):
-        self.device = device
-        self.config = _make_tiny_moe_config()
-        self.model = _make_model(self.config).to(device)
-        self.model.train()
-
-    def test_compute_loss_returns_scalar(self):
+def test_seq_compute_loss_returns_scalar(device):
     """compute_loss should return a scalar tensor."""
+    config, model = _make_seq_moe_fixture(device)
     strategy = SEQStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.01,
     )
-        loss = strategy.compute_loss(_make_batch(self.config))
+    loss = strategy.compute_loss(_make_batch(config))
     assert loss.ndim == 0
     assert loss.requires_grad
 
-    def test_compute_loss_output_has_metrics(self):
+
+def test_seq_compute_loss_output_has_metrics(device):
     """compute_loss_output dict with moe_aux_loss_coef > 0 includes MoE metrics."""
+    config, model = _make_seq_moe_fixture(device)
     strategy = SEQStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.01,
     )
-        output = strategy.compute_loss_output(_make_batch(self.config))
+    output = strategy.compute_loss_output(_make_batch(config))
 
     assert "loss" in output
     assert "metrics" in output
@@ -226,14 +233,16 @@ class TestSEQStrategyMoE:
         assert key in metrics, f"Missing metric: {key}"
         assert isinstance(metrics[key], float)
 
-    def test_moe_metrics_populated_after_forward(self):
+
+def test_seq_moe_metrics_populated_after_forward(device):
     """strategy._moe_metrics populated after compute_loss_output."""
+    config, model = _make_seq_moe_fixture(device)
     strategy = SEQStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.01,
     )
-        strategy.compute_loss_output(_make_batch(self.config))
+    strategy.compute_loss_output(_make_batch(config))
 
     moe_metrics = strategy._moe_metrics
     assert moe_metrics, "_moe_metrics should not be empty for MoE model"
@@ -247,14 +256,16 @@ class TestSEQStrategyMoE:
         assert key in moe_metrics, f"Missing _moe_metrics key: {key}"
         assert isinstance(moe_metrics[key], float)
 
-    def test_zero_coef_zeroes_weighted_aux(self):
+
+def test_seq_zero_coef_zeroes_weighted_aux(device):
     """moe_aux_loss_coef=0 → weighted_aux_loss is zero, task_loss == loss."""
+    config, model = _make_seq_moe_fixture(device)
     strategy = SEQStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.0,
     )
-        output = strategy.compute_loss_output(_make_batch(self.config))
+    output = strategy.compute_loss_output(_make_batch(config))
     metrics = output["metrics"]
 
     # task_loss and loss should be equal (aux weighted by zero)
@@ -269,40 +280,44 @@ class TestSEQStrategyMoE:
     assert strategy._moe_metrics
     assert "router_entropy" in strategy._moe_metrics
 
-    def test_aux_loss_added_to_total_loss(self):
+
+def test_seq_aux_loss_added_to_total_loss(device):
     """Total loss > task_loss when moe_aux_loss_coef > 0."""
+    config, model = _make_seq_moe_fixture(device)
     strategy = SEQStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.01,
     )
-        output = strategy.compute_loss_output(_make_batch(self.config))
+    output = strategy.compute_loss_output(_make_batch(config))
     assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
 
-    def test_factory_creates_strategy_with_coef(self):
+
+def test_seq_factory_creates_strategy_with_coef(device):
     """StrategyFactory.create passes moe_aux_loss_coef to strategy."""
+    _, model = _make_seq_moe_fixture(device)
     strategy = StrategyFactory.create(
         "seq",
-            model=self.model,
-            device=self.device,
+        model=model,
+        device=device,
         moe_aux_loss_coef=0.02,
     )
     assert strategy.moe_aux_loss_coef == 0.02
 
-    def test_no_aux_loss_for_mlp_model(self):
-        """Pure MLP model: model outputs no aux_loss → no MoE metrics."""
-        from astrai.config.model_config import AutoRegressiveLMConfig
 
+def test_seq_no_aux_loss_for_mlp_model(device):
+    """Pure MLP model: model outputs no aux_loss → no MoE metrics."""
+    config, _ = _make_seq_moe_fixture(device)
     mlp_config = AutoRegressiveLMConfig(**{**TINY_CONFIG, "ffn_type": "mlp"})
-        mlp_model = AutoRegressiveLM(mlp_config).to(self.device)
+    mlp_model = AutoRegressiveLM(mlp_config).to(device)
     mlp_model.train()
 
     strategy = SEQStrategy(
         mlp_model,
-            self.device,
+        device,
         moe_aux_loss_coef=0.01,
     )
-        output = strategy.compute_loss_output(_make_batch(self.config))
+    output = strategy.compute_loss_output(_make_batch(config))
     metrics = output["metrics"]
 
     assert "moe_aux_loss" not in metrics
@@ -311,24 +326,15 @@ class TestSEQStrategyMoE:
     assert strategy._moe_metrics == {}
 
 
-class TestSFTStrategyMoE:
-    """End‑to‑end tests for SFTStrategy with MoE aux loss."""
-
-    @pytest.fixture(autouse=True)
-    def setup(self, device):
-        self.device = device
-        self.config = _make_tiny_moe_config()
-        self.model = _make_model(self.config).to(device)
-        self.model.train()
-
-    def test_compute_loss_output_with_aux_loss(self):
+def test_sft_compute_loss_output_with_aux_loss(device):
     """SFTStrategy produces MoE metrics when coef > 0."""
+    config, model = _make_sft_moe_fixture(device)
     strategy = SFTStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.01,
     )
-        output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
+    output = strategy.compute_loss_output(_make_batch(config, with_extra=True))
 
     metrics = output["metrics"]
     assert "moe_aux_loss" in metrics
@@ -339,14 +345,16 @@ class TestSFTStrategyMoE:
     assert "router_entropy" in moe_metrics
     assert "dead_expert_fraction" in moe_metrics
 
-    def test_sft_zero_coef_zeroes_weighted_aux(self):
+
+def test_sft_zero_coef_zeroes_weighted_aux(device):
     """SFTStrategy with zero coef: weighted aux is zero, loss == task_loss."""
+    config, model = _make_sft_moe_fixture(device)
     strategy = SFTStrategy(
-            self.model,
-            self.device,
+        model,
+        device,
         moe_aux_loss_coef=0.0,
     )
-        output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
+    output = strategy.compute_loss_output(_make_batch(config, with_extra=True))
     metrics = output["metrics"]
 
     assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)

tests/serialization/test_hf_adapter.py

@@ -183,6 +183,68 @@ def test_convert_hf_weights_rejects_mla():
         convert_hf_weights(sd, cfg)
 
 
+def test_convert_hf_config_qwen2_moe_preserves_sparse_fields():
+    raw = {
+        **LLAMA_RAW,
+        "model_type": "qwen2_moe",
+        "num_local_experts": 2,
+        "num_experts_per_tok": 1,
+        "n_shared_experts": 1,
+        "decoder_sparse_step": 2,
+        "mlp_only_layers": [0],
+    }
+    cfg = ConfigFactory.load(convert_hf_config(raw))
+    assert cfg.decoder_sparse_step == 2
+    assert cfg.mlp_only_layers == [0]
+
+
+def test_convert_hf_config_gemma_enables_qk_norm():
+    raw = {**LLAMA_RAW, "model_type": "gemma"}
+    cfg = ConfigFactory.load(convert_hf_config(raw))
+    assert cfg.use_qk_norm is True
+
+
+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,
+    )
+    model = AutoRegressiveLM(cfg)
+    converted = convert_hf_weights(to_hf_keys(model.state_dict()), 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,
+    )
+    model = AutoRegressiveLM(cfg)
+    hf_sd = to_hf_keys(model.state_dict())
+    hf_sd = {
+        k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.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)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
 def test_from_pretrained_hf_directory(tmp_path):
     cfg = make_tiny_config()
     model = AutoRegressiveLM(cfg).eval()

tests/trainer/test_callbacks.py

@@ -31,13 +31,11 @@ def test_gradient_checkpointing_empty_modules_noop(test_model):
     model = test_model["model"]
     callback = GradientCheckpointingCallback()
 
-    originals = [layer.forward for layer in model.layers]
-
     for layer in model.layers:
         callback._enable(layer)
 
-    for layer, orig in zip(model.layers, originals):
-        assert layer.forward is orig
+    for layer in model.layers:
+        assert not hasattr(layer, "_original_forward")
 
 
 def test_gradient_checkpointing_forward_unchanged(test_model):