Spaces:
Running
Running
| import json | |
| import os | |
| import pytest | |
| import torch | |
| from mouh.jev.types import JevMessageType as T | |
| from mouh.model import ( | |
| ByteTokenizer, MouhConfig, MouhForCausalLM, estimate_params, preset_config, type_ids_for, PRESETS, | |
| ) | |
| from mouh.model.layers import GatedSSM, _segsum, ssd_chunked, ssm_step | |
| from mouh.model.moe import DeterministicMoE | |
| torch.manual_seed(0) | |
| def tiny(): | |
| torch.manual_seed(0) | |
| return MouhForCausalLM(preset_config("mouh-tiny")).eval() | |
| def test_param_estimate_matches_real_count(tiny): | |
| real = sum(p.numel() for p in tiny.parameters()) | |
| assert real == estimate_params(tiny.config).total | |
| def test_param_estimate_matches_meta_instantiation(name): | |
| cfg = preset_config(name) | |
| if name != "mouh-tiny": | |
| with torch.device("meta"): | |
| m = MouhForCausalLM(cfg) | |
| else: | |
| m = MouhForCausalLM(cfg) | |
| assert sum(p.numel() for p in m.parameters()) == estimate_params(cfg).total | |
| def test_layer_pattern_is_five_ssm_per_attention(): | |
| cfg = preset_config("mouh-1.5b") | |
| pattern = ["A" if cfg.is_attention_layer(i) else "S" for i in range(cfg.num_hidden_layers)] | |
| assert "".join(pattern) == ("SSSSSA" * 4) | |
| with pytest.raises(ValueError): | |
| MouhConfig(num_hidden_layers=10) | |
| def test_expert_count_and_topk(): | |
| for name in PRESETS: | |
| cfg = preset_config(name) | |
| assert cfg.num_experts == 8 | |
| assert preset_config("mouh-1.5b").num_experts_per_tok == 4 | |
| def test_spec_budget_arithmetic_documented(): | |
| # top-4/8 : actifs >= 50 % des paramètres experts => total 1,5 Md => actifs >= ~750 M | |
| r = estimate_params(preset_config("mouh-1.5b")) | |
| assert 1.45e9 < r.total < 1.55e9 and r.active > 0.75e9 | |
| r = estimate_params(preset_config("mouh-a400m")) | |
| assert abs(r.active - 400e6) / 400e6 < 0.01 | |
| r = estimate_params(preset_config("mouh-1.5b-top1")) | |
| assert 1.45e9 < r.total < 1.55e9 and abs(r.active - 400e6) / 400e6 < 0.02 | |
| def test_ssd_chunked_matches_sequential_recurrence(): | |
| torch.manual_seed(1) | |
| b, T, H, P, N = 2, 37, 3, 4, 5 | |
| x = torch.randn(b, T, H, P) | |
| a = -torch.rand(b, T, H) | |
| B, C = torch.randn(b, T, N), torch.randn(b, T, N) | |
| y, fin = ssd_chunked(x, a, B, C, chunk=8) | |
| state = torch.zeros(b, H, P, N) | |
| ys = [] | |
| for t in range(T): | |
| yt, state = ssm_step(state, x[:, t], a[:, t], B[:, t], C[:, t]) | |
| ys.append(yt) | |
| ref = torch.stack(ys, 1) | |
| assert torch.allclose(y, ref, atol=1e-4, rtol=1e-4) | |
| assert torch.allclose(fin, state, atol=1e-4, rtol=1e-4) | |
| # état initial non nul | |
| init = torch.randn(b, H, P, N) | |
| y2, fin2 = ssd_chunked(x, a, B, C, chunk=8, init_state=init) | |
| state = init.clone(); ys = [] | |
| for t in range(T): | |
| yt, state = ssm_step(state, x[:, t], a[:, t], B[:, t], C[:, t]); ys.append(yt) | |
| assert torch.allclose(y2, torch.stack(ys, 1), atol=1e-4, rtol=1e-4) | |
| assert torch.allclose(fin2, state, atol=1e-4, rtol=1e-4) | |
| def test_ssm_block_streaming_equals_full(tiny): | |
| ssm = [b.mixer for b in tiny.blocks if isinstance(b.mixer, GatedSSM)][0].eval() | |
| x = torch.randn(2, 30, tiny.config.hidden_size) | |
| full, _ = ssm(x) | |
| cache, outs = None, [] | |
| for t in range(30): | |
| o, cache = ssm(x[:, t:t + 1], cache) | |
| outs.append(o) | |
| assert torch.allclose(full, torch.cat(outs, 1), atol=1e-4, rtol=1e-4) | |
| # préremplissage puis flux | |
| o1, cache = ssm(x[:, :17]) | |
| o2, _ = ssm(x[:, 17:], cache) | |
| assert torch.allclose(full, torch.cat([o1, o2], 1), atol=1e-4, rtol=1e-4) | |
| def test_cached_decoding_matches_full_forward(tiny): | |
| ids = torch.randint(0, 256, (2, 25)) | |
| with torch.no_grad(): | |
| full = tiny(ids).logits | |
| o = tiny(ids[:, :10], use_cache=True) | |
| pieces = [o.logits] | |
| cache = o.cache | |
| for t in range(10, 25): | |
| o = tiny(ids[:, t:t + 1], cache=cache) | |
| cache = o.cache | |
| pieces.append(o.logits) | |
| stepped = torch.cat(pieces, 1) | |
| assert torch.allclose(full, stepped, atol=2e-3, rtol=2e-3), (full - stepped).abs().max() | |
| def test_causality(tiny): | |
| a = torch.randint(0, 256, (1, 20)) | |
| b = a.clone(); b[0, 15:] = torch.randint(0, 256, (5,)) | |
| with torch.no_grad(): | |
| la, lb = tiny(a).logits, tiny(b).logits | |
| assert torch.allclose(la[:, :15], lb[:, :15], atol=1e-4) | |
| def test_forward_is_deterministic(tiny): | |
| ids = torch.randint(0, 256, (2, 16)) | |
| with torch.no_grad(): | |
| l1, l2 = tiny(ids, message_type_ids=type_ids_for(T.CODE_PROPOSAL, 2, ids.device)).logits, \ | |
| tiny(ids, message_type_ids=type_ids_for(T.CODE_PROPOSAL, 2, ids.device)).logits | |
| assert torch.equal(l1, l2) | |
| def test_moe_router_topk_deterministic_and_type_prior(): | |
| cfg = preset_config("mouh-tiny") | |
| moe = DeterministicMoE(cfg).eval() | |
| x = torch.randn(1, 12, cfg.hidden_size) | |
| moe(x, torch.tensor([0])); t1 = moe.last_trace | |
| moe(x, torch.tensor([0])); t2 = moe.last_trace | |
| assert torch.equal(t1.experts, t2.experts) and torch.equal(t1.weights, t2.weights) | |
| assert t1.experts.shape == (12, 4) | |
| assert torch.allclose(t1.weights.sum(-1), torch.ones(12), atol=1e-5) | |
| # un prior de type très fort impose les experts de la table JEV (CORRECTION_REQUEST -> n°5,1,7,3) | |
| strong = preset_config("mouh-tiny", router_type_prior_strength=1000.0) | |
| moe2 = DeterministicMoE(strong).eval() | |
| moe2(x, type_ids_for(T.CORRECTION_REQUEST, 1, x.device)) | |
| chosen = {tuple(r) for r in moe2.last_trace.experts.tolist()} | |
| assert chosen == {(4, 0, 6, 2)} | |
| def test_moe_ties_break_to_lowest_index(): | |
| cfg = preset_config("mouh-tiny") | |
| moe = DeterministicMoE(cfg).eval() | |
| with torch.no_grad(): | |
| moe.router.weight.zero_() | |
| moe(torch.randn(1, 3, cfg.hidden_size)) | |
| assert moe.last_trace.experts.tolist() == [[0, 1, 2, 3]] * 3 | |
| def test_moe_capacity_limits_load_in_training(): | |
| cfg = preset_config("mouh-tiny", capacity_factor=1.0) | |
| moe = DeterministicMoE(cfg).train() | |
| with torch.no_grad(): | |
| moe.router.weight.zero_() # tous les tokens visent les mêmes experts -> débordement garanti | |
| N = 32 | |
| moe(torch.randn(1, N, cfg.hidden_size)) | |
| cap = -(-int(1.0 * N * 4) // 8) | |
| assert moe.last_trace.load.max().item() <= cap | |
| assert (~moe.last_trace.kept).any() | |
| def test_moe_balancing_bias_moves_against_overload(): | |
| cfg = preset_config("mouh-tiny") | |
| moe = DeterministicMoE(cfg).train() | |
| with torch.no_grad(): | |
| moe.router.weight.zero_() | |
| moe(torch.randn(1, 64, cfg.hidden_size)) | |
| moe.step_routing_bias() | |
| assert moe.routing_bias[0] < 0 and moe.routing_bias[7] >= 0 | |
| def test_moe_expert_offload_roundtrip(): | |
| cfg = preset_config("mouh-tiny") | |
| moe = DeterministicMoE(cfg).eval() | |
| x = torch.randn(1, 8, cfg.hidden_size) | |
| ref = moe(x) | |
| moe.offload_experts([0, 1], "cpu"); moe.load_experts([0, 1], "cpu") | |
| assert torch.equal(moe(x), ref) | |
| def test_jev_decisions_are_typed(tiny): | |
| ids = torch.randint(0, 256, (1, 12)) | |
| msg_choice, tool_choice = tiny.decide(ids, T.TASK_SUBMISSION) | |
| assert msg_choice.chosen in {t.value for t in T} and abs(sum(msg_choice.score.values) - 1) < 1e-9 | |
| assert tool_choice.chosen in ("none", "PythonExecutionSkill", "PytestRunnerSkill", "FileSystemSkill", "Z3SolverSkill") | |
| def test_generate_greedy_deterministic_and_cache_consistent(tiny): | |
| ids = torch.tensor([ByteTokenizer().encode("def f():")]) | |
| a = tiny.generate_greedy(ids, max_new_tokens=12) | |
| b = tiny.generate_greedy(ids, max_new_tokens=12) | |
| assert torch.equal(a, b) and a.shape[1] >= ids.shape[1] + 1 | |
| # le décodage avec cache équivaut à un décodage glouton sans cache | |
| seq = ids | |
| with torch.no_grad(): | |
| for _ in range(12): | |
| nxt = tiny(seq).logits[:, -1].argmax(-1, keepdim=True) | |
| seq = torch.cat([seq, nxt], 1) | |
| if nxt.item() == tiny.config.eos_token_id: | |
| break | |
| assert torch.equal(a[:, : seq.shape[1]], seq) | |
| def test_gradients_flow_to_all_parameter_groups(): | |
| m = MouhForCausalLM(preset_config("mouh-tiny")).train() | |
| ids = torch.randint(0, 256, (2, 64)) | |
| out = m(ids, labels=ids) | |
| out.loss.backward() | |
| missing = [n for n, p in m.named_parameters() if p.grad is None] | |
| assert not missing, missing[:5] | |
| assert all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None) | |
| def test_tiny_model_overfits_a_batch(): | |
| torch.manual_seed(0) | |
| m = MouhForCausalLM(preset_config("mouh-tiny", router_aux_loss_coef=0.0)).train() | |
| opt = torch.optim.AdamW(m.parameters(), lr=3e-3) | |
| ids = torch.tensor([ByteTokenizer().encode("def add(a, b):\n return a + b\n") * 1]) | |
| first = None | |
| for step in range(60): | |
| out = m(ids, labels=ids) | |
| out.loss.backward(); opt.step(); opt.zero_grad() | |
| first = first or out.loss.item() | |
| assert out.loss.item() < first * 0.35 | |
| def test_safetensors_roundtrip_and_config_json(tmp_path, tiny): | |
| tiny.save_pretrained(tmp_path) | |
| assert (tmp_path / "model.safetensors").exists() | |
| cfg = json.loads((tmp_path / "config.json").read_text()) | |
| assert cfg["model_type"] == "mouh" and cfg["architectures"] == ["MouhForCausalLM"] | |
| assert cfg["num_experts"] == 8 and cfg["num_experts_per_tok"] == 4 and cfg["attn_every"] == 6 | |
| loaded = MouhForCausalLM.from_pretrained(tmp_path).eval() | |
| ids = torch.randint(0, 256, (1, 20)) | |
| with torch.no_grad(): | |
| assert torch.allclose(tiny(ids).logits, loaded(ids).logits, atol=1e-5) | |
| from transformers import AutoModelForCausalLM | |
| auto = AutoModelForCausalLM.from_pretrained(tmp_path) | |
| assert isinstance(auto, MouhForCausalLM) | |
| def test_root_config_json_matches_default_preset(): | |
| root = os.path.join(os.path.dirname(__file__), "..", "config.json") | |
| saved = json.load(open(root)) | |
| cfg = preset_config() | |
| for k in ("hidden_size", "num_hidden_layers", "num_experts", "num_experts_per_tok", | |
| "expert_intermediate_size", "attn_every", "vocab_size", "model_type"): | |
| assert saved[k] == (getattr(cfg, k) if k != "model_type" else "mouh"), k | |
| def test_routing_trace_is_exported_as_validated_bus_messages(tiny): | |
| from mouh.jev import JevBus | |
| from mouh.jev.types import new_id | |
| from mouh.model import routing_messages | |
| ids = torch.randint(0, 256, (1, 24)) | |
| with torch.no_grad(): | |
| tiny(ids, message_type_ids=type_ids_for(T.CODE_PROPOSAL, 1, ids.device)) | |
| bus = JevBus() | |
| msgs = routing_messages(tiny, new_id(), T.CODE_PROPOSAL) | |
| assert len(msgs) == tiny.config.num_hidden_layers | |
| for m in msgs: | |
| out = bus.publish(m) # validation stricte du schéma par le bus | |
| assert out[0].payload.experts and len(out[0].payload.experts) == 4 | |