| """D/E/G task MoE on the Gemma 4 backbone (:func:`deg.model.add_task_moe`): FFN experts |
| per capability, deterministic routing by fork kind, shared attention / embeddings / PLE |
| (tiny random Gemma 4, CPU). Decide forks run D, extract / count / anchor forks and the |
| context run E (the layer's own FFN), x-generate (KNOWN VALUES, header, decoded tokens) |
| runs G; each expert's parameters get gradients only from its own losses.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import random |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("transformers") |
|
|
| from test_gemma_backbone import _gemma_model, _tok, tiny_gemma |
| from test_model import CONTEXT, SCHEMA, _raw_outputs, _sub |
|
|
| from deg.infer import fill |
| from deg.model import TaskRoute, add_task_moe, load_model, task_moe_params |
|
|
| DECIDE = ["department", "urgent", "anger"] |
| EXTRACT = ["email", "order_id", "total"] |
|
|
|
|
| @pytest.fixture(scope="module") |
| def pair(): |
| dense = _gemma_model(2) |
| dense.meta["gen_layout"] = "chat" |
| moe = _gemma_model(2) |
| add_task_moe(moe) |
| return dense, moe |
|
|
|
|
| def _perturb(module, seed): |
| g = torch.Generator().manual_seed(seed) |
| with torch.no_grad(): |
| for p in module.parameters(): |
| p.add_(torch.randn(p.shape, generator=g) * 0.05) |
|
|
|
|
| def test_fresh_experts_are_the_dense_model(pair): |
| dense, moe = pair |
| want = _raw_outputs(dense, _sub(DECIDE + EXTRACT)) |
| got = _raw_outputs(moe, _sub(DECIDE + EXTRACT)) |
| for k in want: |
| assert torch.equal(want[k], got[k]), k |
| a, b = fill(dense, CONTEXT, SCHEMA), fill(moe, CONTEXT, SCHEMA) |
| assert a["value"] == b["value"] |
| assert moe.generator() == (moe, True, True) |
| assert moe.gen_route(3) == TaskRoute("g") |
|
|
|
|
| def test_routing_isolates_the_experts(): |
| """Perturbing D changes only decide outputs; perturbing G only the generated text; |
| perturbing E (the reader) changes everything that reads the context.""" |
| m = _gemma_model(2) |
| add_task_moe(m) |
| base = _raw_outputs(m, _sub(DECIDE + EXTRACT)) |
| reply = fill(m, CONTEXT, SCHEMA)["value"]["reply"] |
| for layer in m.backbone.text.layers: |
| _perturb(layer.mlp_d, 1) |
| after_d = _raw_outputs(m, _sub(DECIDE + EXTRACT)) |
| assert all(torch.equal(after_d[f"/{n}"], base[f"/{n}"]) for n in EXTRACT) |
| assert all(not torch.equal(after_d[f"/{n}"], base[f"/{n}"]) for n in DECIDE) |
| m.graphs.clear() |
| for layer in m.backbone.text.layers: |
| _perturb(layer.mlp_g, 2) |
| after_g = _raw_outputs(m, _sub(DECIDE + EXTRACT)) |
| for k in after_g: |
| assert torch.equal(after_g[k], after_d[k]), k |
| assert fill(m, CONTEXT, SCHEMA)["value"]["reply"] != reply |
| for layer in m.backbone.text.layers: |
| _perturb(layer.mlp, 3) |
| after_e = _raw_outputs(m, _sub(DECIDE + EXTRACT)) |
| assert all(not torch.equal(after_e[k], after_g[k]) for k in after_e) |
|
|
|
|
| def test_each_expert_trains_only_on_its_own_losses(): |
| from deg.train import deg_example, forward_losses |
|
|
| m = _gemma_model(2) |
| m.meta["gen_layout"] = "chat" |
| add_task_moe(m) |
| a = CONTEXT.index("AB-1234") |
| rec = {"uid": "t1", "schema": _sub(["department", "order_id", "reply"]), |
| "context": CONTEXT, |
| "value": {"department": "billing", "order_id": "AB-1234", "reply": "Refund issued."}, |
| "spans": {"/order_id": [a, a + 7]}, "span_alternatives": {}, "span_probs": {}, |
| "decide_probs": {}} |
| ex = deg_example(m, rec, random.Random(0), shuffle_p=0.0) |
| layers = m.backbone.text.layers |
|
|
| def grads(keep): |
| ex2 = type(ex)(**{**ex.__dict__, "items": [it for it in ex.items if it.kind in keep]}) |
| for p in m.backbone.parameters(): |
| p.grad = None |
| m.backbone.train() |
| try: |
| losses, _ = forward_losses(m, [ex2], 4096, 16) |
| finally: |
| m.backbone.eval() |
| sum(v for k, v in losses.items() if k in ("decide", "extract", "lm")).backward() |
| return {x: any(p.grad is not None and bool(p.grad.abs().sum() > 0) |
| for lay in layers for p in getattr(lay, f"mlp{x}").parameters()) |
| for x in ("", "_d", "_g")} |
|
|
| assert grads({"choice"}) == {"": True, "_d": True, "_g": False} |
| assert grads({"extract"}) == {"": True, "_d": False, "_g": False} |
| assert grads({"generate"}) == {"": True, "_d": False, "_g": True} |
| for p in m.backbone.parameters(): |
| p.grad = None |
|
|
|
|
| def test_task_moe_checkpoint_roundtrip_and_counts(tmp_path): |
| m = _gemma_model(6) |
| add_task_moe(m) |
| with torch.no_grad(): |
| for p in m.backbone.lm.state_dict().values(): |
| if p.is_floating_point(): |
| p.copy_(p.bfloat16().float()) |
| for layer in m.backbone.text.layers: |
| _perturb(layer.mlp_d, 4) |
| for p in layer.mlp_d.parameters(): |
| p.copy_(p.bfloat16().float()) |
| m.save(tmp_path / "ck") |
| meta = json.loads((tmp_path / "ck" / "deg.json").read_text()) |
| assert meta["moe"]["experts"] == ["d", "g"] |
| other = tiny_gemma(3) |
| with torch.no_grad(): |
| other.model.embed_tokens_per_layer.weight.copy_( |
| m.backbone.text.embed_tokens_per_layer.weight) |
| back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=other, |
| tokenizer=_tok()) |
| assert back.moe is not None and back.moe["experts"] == ["d", "g"] |
| want = _raw_outputs(m, _sub(DECIDE + EXTRACT)) |
| got = _raw_outputs(back, _sub(DECIDE + EXTRACT)) |
| for k in want: |
| assert torch.equal(want[k], got[k]), k |
| c = task_moe_params(back) |
| assert c["ffn_d"] == c["ffn_g"] == c["ffn_e"] > 0 |
| dense = sum(p.numel() for p in tiny_gemma(0).parameters()) |
| assert c["active_without_ple"] + c["ple"] == dense |
| assert c["total"] == dense + c["ffn_d"] + c["ffn_g"] |
|
|
