File size: 5,962 Bytes
16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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 # noqa: E402
from test_model import CONTEXT, SCHEMA, _raw_outputs, _sub # noqa: E402
from deg.infer import fill # noqa: E402
from deg.model import TaskRoute, add_task_moe, load_model, task_moe_params # noqa: E402
DECIDE = ["department", "urgent", "anger"]
EXTRACT = ["email", "order_id", "total"]
@pytest.fixture(scope="module")
def pair():
dense = _gemma_model(2)
dense.meta["gen_layout"] = "chat" # the task MoE writes in the chat turn (G route)
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) # one prefill, G writes in the chat turn
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} # E reads the context
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()) # tied head counted once
assert c["active_without_ple"] + c["ple"] == dense
assert c["total"] == dense + c["ffn_d"] + c["ffn_g"]
|