| """deg-MoE (:func:`deg.model.add_moe`, :mod:`deg.moetrain`) on a tiny random Qwen3.5 (CPU): |
| |
| * decide/extract outputs are bitwise the plain model's (W attached, random, never read); |
| * routing: R-only calls never read W (W poisoned with NaN changes nothing) and W calls |
| never read R's FFN (R poisoned after the trunk changes nothing in generation); |
| * generate rows are batch-invariant; save/load round-trips; one train step moves only W. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("transformers") |
|
|
| from test_gentrain import RECORD |
| from test_hybrid_generate import SCHEMA |
| from test_model import CONTEXT, SNAPSHOT, _with_note, tiny_lm |
|
|
| from deg.gentrain import gen_items |
| from deg.infer import fill |
| from deg.model import add_moe, load_model, load_tokenizer, moe_params |
| from deg.moetrain import OffloadAdamW, batch_losses, micro_batches |
| from deg.request import FillOptions |
|
|
|
|
| def _plain(): |
| if SNAPSHOT is None: |
| pytest.skip("Qwen3.5-0.8B tokenizer not cached") |
| m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(0), |
| tokenizer=load_tokenizer(SNAPSHOT)) |
| torch.manual_seed(1) |
| with torch.no_grad(): |
| for p in m.heads.parameters(): |
| p.normal_(0.0, 0.5) |
| return m |
|
|
|
|
| def _moe(): |
| m = _plain() |
| params = add_moe(m, tiny_lm(7), writer_rank=4) |
| g = torch.Generator().manual_seed(5) |
| with torch.no_grad(): |
| for p in params: |
| if p.dim() == 2 and p.shape[1] == 4: |
| p.copy_(torch.randn(p.shape, generator=g) * 0.2) |
| return m, params |
|
|
|
|
| def _de(resp): |
| return {k: (resp["value"][k], resp["meta"][f"/{k}"]) for k in ("department", "order_id")} |
|
|
|
|
| def test_decide_extract_bitwise_and_r_never_reads_w(): |
| plain = fill(_plain(), CONTEXT, SCHEMA) |
| m, params = _moe() |
| got = fill(m, CONTEXT, SCHEMA) |
| assert _de(got) == _de(plain) |
| with torch.no_grad(): |
| for layer in m.backbone.text.layers: |
| for p in layer.mlp_w.parameters(): |
| p.fill_(float("nan")) |
| poisoned = fill(m, CONTEXT, {**SCHEMA, "properties": { |
| k: v for k, v in SCHEMA["properties"].items() if k != "reply"}, |
| "required": ["department", "order_id"]}) |
| assert _de(poisoned) == _de(plain) |
| c = moe_params(m) |
| assert c["w_ffn"] == c["r_ffn"] and c["active_writer"] < c["total"] |
|
|
|
|
| def test_w_never_reads_r_ffn_and_rows_are_batch_invariant(): |
| m, _ = _moe() |
| ref = fill(m, CONTEXT, SCHEMA) |
| assert isinstance(ref["value"]["reply"], str) |
| |
| from deg import model as dm |
|
|
| real = dm._expert |
|
|
| def guarded(layer, sel, name="mlp"): |
| e = real(layer, sel, name) |
| if isinstance(sel, dm.WriteRoute): |
| assert e is layer.mlp_w |
| else: |
| assert e is layer.mlp |
| return e |
|
|
| dm._expert = guarded |
| try: |
| assert fill(m, CONTEXT, SCHEMA) == ref |
| finally: |
| dm._expert = real |
| for opts in (FillOptions(), FillOptions(temperature=0.9, seed=3)): |
| alone = fill(m, CONTEXT, SCHEMA, opts) |
| both = fill(m, CONTEXT, _with_note(SCHEMA), opts) |
| assert both["value"]["reply"] == alone["value"]["reply"] |
| assert both["meta"]["/reply"] == alone["meta"]["/reply"] |
|
|
|
|
| def test_save_load_roundtrip(tmp_path): |
| m, _ = _moe() |
| with torch.no_grad(): |
| for k, v in m.backbone.lm.state_dict().items(): |
| if ".lora_" not in k: |
| v.copy_(v.to(torch.bfloat16).float()) |
| ref = fill(m, CONTEXT, SCHEMA) |
| m.save(tmp_path / "ck") |
| back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lm(3), |
| tokenizer=load_tokenizer(SNAPSHOT)) |
| assert back.moe == m.moe and back.generate_with == "moe" |
| assert fill(back, CONTEXT, SCHEMA) == ref |
| with pytest.raises(ValueError): |
| load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lm(3), |
| tokenizer=load_tokenizer(SNAPSHOT), generate="base") |
|
|
|
|
| def test_one_train_step_moves_only_w(): |
| m, params = _moe() |
| for p in params: |
| p.requires_grad_(True) |
| m.backbone.grad_ckpt = True |
| frozen = {k: v.clone() for k, v in m.backbone.lm.state_dict().items() |
| if ".mlp_w." not in k and ".lora_" not in k} |
| w0 = [p.detach().clone() for p in params] |
| items = gen_items(m, [RECORD], {"judge_drop": 0, "budget_drop": 0}) |
| (group,) = micro_batches(items, 4096, None) |
| lora = [p.dim() == 2 and p.shape[1] == 4 or p.shape[0] == 4 for p in params] |
| opt = OffloadAdamW(params, lora, 1e-3, 1e-3, torch.device("cpu")) |
| nll, _, _, n = batch_losses(m, group) |
| assert n == items[0].target_len and torch.isfinite(nll) |
| (nll / n).backward() |
| opt.step() |
| for k, v in m.backbone.lm.state_dict().items(): |
| if k in frozen: |
| assert torch.equal(v, frozen[k]), k |
| assert any(not torch.equal(p, q) for p, q in zip(params, w0)) |
|
|
