| """As smoke_tiny.py, with the tiny Qwen3.5-MoE quantized to block FP8 on load (transformers |
| FP8Experts, as Qwen3.8-Flash-Next-FP8 loads): observer, prune, save, reload, saved names.""" |
|
|
| import json |
| import tempfile |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import torch |
| from safetensors import safe_open |
| from transformers import AutoModelForCausalLM, FineGrainedFP8Config |
| 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=256, hidden_size=256, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, |
| head_dim=64, moe_intermediate_size=128, shared_expert_intermediate_size=128, 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=64, linear_value_head_dim=64, |
| ) |
| with tempfile.TemporaryDirectory() as d: |
| AutoModelForCausalLM.from_config(cfg).save_pretrained(Path(d) / "bf16") |
| model = AutoModelForCausalLM.from_pretrained( |
| Path(d) / "bf16", quantization_config=FineGrainedFP8Config(weight_block_size=[128, 128]), device_map="cuda" |
| ).eval() |
| print(type(model.model.layers[0].mlp.experts).__name__, model.model.layers[0].mlp.experts.gate_up_proj.dtype, |
| tuple(model.model.layers[0].mlp.experts.gate_up_proj_scale_inv.shape)) |
| observer = _setup_observer(model, SimpleNamespace(renormalize_router_weights=True, record_pruning_metrics_only=True)) |
| ids = torch.randint(0, 256, (1, 64)).cuda() |
| with torch.no_grad(), observer.set_attention_mask(torch.ones_like(ids)): |
| ref = model(input_ids=ids).logits |
| 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("reap[0]", [round(float(x), 4) for x in state[0]["reap"]]) |
| observer.close_hooks() |
| 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 |= {k: fh.get_slice(k).get_dtype() for k in fh.keys()} |
| cfg_saved = json.load(open(Path(out) / "config.json")) |
| print("saved num_experts", cfg_saved["num_experts"], "| quantization_config", cfg_saved.get("quantization_config", {}).get("quant_method")) |
| print("saved expert keys", sorted((k, v) for k, v in keys.items() if "layers.0.mlp.experts.0." in k)) |
| re = AutoModelForCausalLM.from_pretrained(out, device_map="cuda").eval() |
| print("reloaded", type(re.model.layers[0].mlp.experts).__name__, tuple(re.model.layers[0].mlp.experts.gate_up_proj.shape)) |
| with torch.no_grad(): |
| print("reloaded forward ok", tuple(re(input_ids=ids).logits.shape)) |
|
|
