| """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)) |
|
|
