File size: 2,516 Bytes
8c38801 | 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 | """CPU smoke test of the apex port on a tiny random Qwen3.5-MoE text model: REAP's observer with the
3D-expert hook, REAP's prune (keep half), save_pretrained, reload, and the saved weight names."""
import json
import tempfile
from pathlib import Path
from types import SimpleNamespace
import torch
from safetensors import safe_open
from transformers import AutoModelForCausalLM
from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeTextConfig
from reap.main import _setup_observer
from reap.prune import prune
torch.manual_seed(0)
cfg = Qwen3_5MoeTextConfig(
vocab_size=128, hidden_size=64, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
head_dim=16, moe_intermediate_size=32, shared_expert_intermediate_size=32, num_experts=8,
num_experts_per_tok=2, layer_types=["linear_attention", "full_attention"], linear_num_key_heads=2,
linear_num_value_heads=2, linear_key_head_dim=16, linear_value_head_dim=16,
)
model = AutoModelForCausalLM.from_config(cfg).eval().cuda()
print(type(model).__name__)
observer = _setup_observer(model, SimpleNamespace(renormalize_router_weights=True, record_pruning_metrics_only=True))
ids = torch.randint(0, 128, (1, 48)).cuda()
mask = torch.ones_like(ids)
mask[0, 20] = 0
with torch.no_grad(), observer.set_attention_mask(mask):
model(input_ids=ids, attention_mask=mask)
state = {k: {m: (v.mean if hasattr(v, "mean") and not torch.is_tensor(v) else v) for m, v in s.items()} for k, s in observer.state.items()}
print("layers observed", list(state), "reap[0]", [round(float(x), 4) for x in state[0]["reap"]], "freq[0]", state[0]["expert_frequency"].tolist())
observer.close_hooks()
with tempfile.TemporaryDirectory() as d:
out = prune(state, model, SimpleNamespace(prune_method="reap", perserve_super_experts=False, perserve_outliers=False), 4, Path(d) / "pruned")
keys = []
for f in Path(out).glob("*.safetensors"):
with safe_open(f, "pt") as fh:
keys += list(fh.keys())
print("saved config num_experts", json.load(open(Path(out) / "config.json"))["num_experts"])
print("saved expert keys", sorted(k for k in keys if "experts" in k)[:6], "...", len([k for k in keys if "experts." in k]))
re = AutoModelForCausalLM.from_pretrained(out).eval().cuda()
print("reloaded experts", re.model.layers[0].mlp.experts.gate_up_proj.shape, re.model.layers[0].mlp.gate.weight.shape)
with torch.no_grad():
print("reloaded forward ok", tuple(re(input_ids=ids).logits.shape))
|