| """Unstructured Wanda / SparseGPT on SparseMoE expert Linears. |
| |
| Shapes are unchanged. Only expert `up_gate_proj` and `down_proj` are pruned. |
| Layer-sequential: prune block i, then re-forward before scoring block i+1. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import types |
| from typing import Any |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from inference.model.dit.dit_module import SparseMoEMLP |
|
|
| from .sparsegpt_full import SparseGPT |
| from .wanda_unstructured import WrappedGPT, prune_linear_wanda |
|
|
| FORMAT = "davinci_unstructured_v1" |
| CAPTURE_ENV = "MAGI_UNSTRUCTURED_CAPTURE" |
| OUT_ROOT_ENV = "MAGI_UNSTRUCTURED_OUT_ROOT" |
| METHODS_ENV = "MAGI_UNSTRUCTURED_METHODS" |
| SPARSITIES_ENV = "MAGI_UNSTRUCTURED_SPARSITIES" |
| STRIDE_ENV = "MAGI_UNSTRUCTURED_STRIDE" |
| MAX_ENV = "MAGI_UNSTRUCTURED_MAX" |
| SOURCE_ENV = "MAGI_UNSTRUCTURED_SOURCE_CKPT" |
|
|
| __all__ = [ |
| "FORMAT", |
| "moe_expert_linears", |
| "prune_linear_wanda", |
| "sequential_unstructured", |
| "global_mlp_sparsity", |
| "maybe_attach_capture", |
| "maybe_prune_captured", |
| ] |
|
|
|
|
| def moe_expert_linears(mlp: SparseMoEMLP) -> dict[str, nn.Linear]: |
| out: dict[str, nn.Linear] = {} |
| for eid, expert in enumerate(mlp.experts): |
| out[f"expert{eid}.up_gate_proj"] = expert.up_gate_proj |
| out[f"expert{eid}.down_proj"] = expert.down_proj |
| return out |
|
|
|
|
| def global_mlp_sparsity(model) -> float: |
| total = 0 |
| zeros = 0 |
| for layer in model.block.layers: |
| mlp = getattr(layer, "mlp", None) |
| if not isinstance(mlp, SparseMoEMLP): |
| continue |
| for lin in moe_expert_linears(mlp).values(): |
| w = lin.weight.data |
| total += int(w.numel()) |
| zeros += int((w == 0).sum()) |
| return zeros / max(total, 1) |
|
|
|
|
| def _run_layer(layer, sample: dict[str, Any]): |
| with torch.amp.autocast("cuda", dtype=torch.bfloat16): |
| out = layer( |
| sample["h"], |
| sample["rope"], |
| sample["permute_mapping"], |
| sample["inv_permute_mapping"], |
| sample["varlen_handler"], |
| sample["local_attn_handler"], |
| sample["modality_dispatcher"], |
| sample["cp_split_sizes"], |
| ) |
| return out.detach() if torch.is_tensor(out) else out |
|
|
|
|
| def _clone_samples(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| cloned = [] |
| for sample in samples: |
| item = dict(sample) |
| item["h"] = sample["h"].clone() |
| cloned.append(item) |
| return cloned |
|
|
|
|
| @torch.no_grad() |
| def sequential_unstructured( |
| model, |
| samples: list[dict[str, Any]], |
| sparsity: float, |
| method: str, |
| verbose: bool = True, |
| ) -> dict: |
| method = method.strip().lower() |
| if method not in ("wanda", "sparsegpt"): |
| raise ValueError(f"method must be wanda or sparsegpt, got {method}") |
| |
| |
| os.environ["MAGI_MOE_GROUPED"] = "off" |
| work = _clone_samples(samples) |
| layer_meta: dict[str, Any] = {} |
| captured_any = False |
| for layer_idx, layer in enumerate(model.block.layers): |
| mlp = getattr(layer, "mlp", None) |
| is_moe = isinstance(mlp, SparseMoEMLP) |
| linears = moe_expert_linears(mlp) if is_moe else {} |
| wrappers: dict[str, Any] = {} |
| handles = [] |
| if is_moe: |
| if method == "wanda": |
| wrappers = {name: WrappedGPT(lin) for name, lin in linears.items()} |
| else: |
| wrappers = {name: SparseGPT(lin) for name, lin in linears.items()} |
|
|
| def _make_hook(wrapper): |
| def _hook(_, inp, out): |
| wrapper.add_batch(inp[0].data, out.data) |
| return _hook |
|
|
| for name, lin in linears.items(): |
| handles.append(lin.register_forward_hook(_make_hook(wrappers[name]))) |
| for sample in work: |
| _run_layer(layer, sample) |
| for handle in handles: |
| handle.remove() |
| names_meta = {} |
| layer_ns = [int(wrappers[name].nsamples) for name in linears] |
| if sum(layer_ns) == 0: |
| raise RuntimeError( |
| f"unstructured {method} layer {getattr(mlp, 'layer_idx', layer_idx)} " |
| "captured 0 expert activations; MAGI_MOE_GROUPED must be off" |
| ) |
| captured_any = True |
| for name, lin in linears.items(): |
| nsamples = int(wrappers[name].nsamples) |
| if nsamples == 0: |
| names_meta[name] = { |
| "sparsity": 0.0, |
| "shape": list(lin.weight.shape), |
| "nsamples": 0, |
| "skipped": True, |
| } |
| if hasattr(wrappers[name], "free"): |
| wrappers[name].free() |
| continue |
| if method == "wanda": |
| sp = prune_linear_wanda(lin, wrappers[name].scaler_row, sparsity) |
| else: |
| sp = wrappers[name].fasterprune(sparsity, verbose=verbose) |
| wrappers[name].free() |
| names_meta[name] = { |
| "sparsity": float(sp), |
| "shape": list(lin.weight.shape), |
| "nsamples": nsamples, |
| } |
| layer_meta[str(int(getattr(mlp, "layer_idx", layer_idx)))] = names_meta |
| for sample in work: |
| sample["h"] = _run_layer(layer, sample) |
| if verbose and is_moe: |
| print( |
| f"unstructured {method} layer {getattr(mlp, 'layer_idx', layer_idx)} " |
| f"sparsity={sparsity:.3f}", |
| flush=True, |
| ) |
| torch.cuda.empty_cache() |
| if not captured_any: |
| raise RuntimeError(f"unstructured {method} captured 0 expert activations") |
| return { |
| "method": method, |
| "sparsity": float(sparsity), |
| "keep_ratio": float(1.0 - sparsity), |
| "global_mlp_sparsity": global_mlp_sparsity(model), |
| "n_samples": len(samples), |
| "layers": layer_meta, |
| } |
|
|
|
|
| def maybe_attach_capture(model) -> dict | None: |
| flag = os.getenv(CAPTURE_ENV, "").strip() |
| if flag not in ("1", "true", "True", "yes"): |
| return None |
| stride = int(os.getenv(STRIDE_ENV, "8") or "8") |
| max_samples = int(os.getenv(MAX_ENV, "40") or "40") |
| block = model.block |
| orig = block.forward |
| bag: dict[str, Any] = {"samples": [], "calls": 0, "orig": orig} |
|
|
| def wrapped( |
| self, |
| x, |
| rope, |
| permute_mapping, |
| inv_permute_mapping, |
| varlen_handler, |
| local_attn_handler, |
| modality_dispatcher, |
| cp_split_sizes, |
| ): |
| bag["calls"] += 1 |
| if bag["calls"] % stride == 0 and len(bag["samples"]) < max_samples: |
| bag["samples"].append( |
| { |
| "h": x.detach().clone(), |
| "rope": rope, |
| "permute_mapping": permute_mapping, |
| "inv_permute_mapping": inv_permute_mapping, |
| "varlen_handler": varlen_handler, |
| "local_attn_handler": local_attn_handler, |
| "modality_dispatcher": modality_dispatcher, |
| "cp_split_sizes": cp_split_sizes, |
| } |
| ) |
| return orig( |
| x, |
| rope, |
| permute_mapping, |
| inv_permute_mapping, |
| varlen_handler, |
| local_attn_handler, |
| modality_dispatcher, |
| cp_split_sizes, |
| ) |
|
|
| block.forward = types.MethodType(wrapped, block) |
| model._unstr_capture = bag |
| print( |
| f"unstructured capture stride={stride} max={max_samples}", |
| flush=True, |
| ) |
| return bag |
|
|
|
|
| def save_unstructured_checkpoint(path: str, model, meta: dict) -> None: |
| payload = { |
| "format": FORMAT, |
| "model": model.state_dict(), |
| "unstructured_meta": meta, |
| } |
| os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) |
| torch.save(payload, path) |
|
|
|
|
| def maybe_prune_captured(model) -> list[str]: |
| out_root = os.getenv(OUT_ROOT_ENV, "").strip() |
| if not out_root: |
| return [] |
| bag = getattr(model, "_unstr_capture", None) |
| if not bag or not bag.get("samples"): |
| raise RuntimeError( |
| f"{OUT_ROOT_ENV} is set but no TransformerBlock inputs were captured" |
| ) |
| block = model.block |
| orig = bag.get("orig") |
| if orig is not None: |
| block.forward = orig |
| methods = [m.strip() for m in os.getenv(METHODS_ENV, "wanda,sparsegpt").split(",") if m.strip()] |
| sparsities = [ |
| float(x) for x in os.getenv(SPARSITIES_ENV, "0.25,0.50,0.75").split(",") if x.strip() |
| ] |
| source = os.getenv(SOURCE_ENV, "").strip() |
| samples = bag["samples"] |
| print( |
| f"unstructured prune captured={len(samples)} calls={bag.get('calls')} " |
| f"methods={methods} sparsities={sparsities}", |
| flush=True, |
| ) |
| orig_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} |
| written: list[str] = [] |
| os.makedirs(out_root, exist_ok=True) |
| for method in methods: |
| for sparsity in sparsities: |
| keep = 1.0 - float(sparsity) |
| slug = f"{method}_unstr_keep{int(round(keep * 100)):03d}" |
| out_path = os.path.join(out_root, f"{slug}.pt") |
| if os.path.isfile(out_path): |
| print(f"[skip prune] {out_path} exists", flush=True) |
| written.append(out_path) |
| continue |
| print(f"unstructured apply {slug}", flush=True) |
| model.load_state_dict(orig_state, strict=True) |
| meta = sequential_unstructured( |
| model, samples, sparsity=float(sparsity), method=method, verbose=True, |
| ) |
| meta["source"] = source |
| meta["captured_calls"] = int(bag.get("calls") or 0) |
| save_unstructured_checkpoint(out_path, model, meta) |
| meta_json = os.path.splitext(out_path)[0] + "_meta.json" |
| with open(meta_json, "w", encoding="utf-8") as f: |
| json.dump(meta, f, indent=2) |
| f.write("\n") |
| print(f"wrote {out_path} global_mlp_sparsity={meta['global_mlp_sparsity']:.4f}", flush=True) |
| written.append(out_path) |
| torch.cuda.empty_cache() |
| return written |
|
|
