HyperPrune-gpt-oss-20b-2to4

openai/gpt-oss-20b pruned to 2:4 semi-structured sparsity with HyperPrune (Sun & Sakuma, Learning Semi-Structured Sparsity for LLMs via Shared and Context-Aware Hypernetwork, ICLR 2026, OpenReview).

This is a reproduction run produced at Elastix as part of the BLADE sparsity-method comparison. It is plain sparse bf16/fp16 safetensors and loads with stock transformers:

from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained("elastix-ai/HyperPrune-gpt-oss-20b-2to4")
t = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

What differs from the paper's own recipe

	paper / repo default	this checkpoint
calibration corpus	allenai/c4	DKYoon/SlimPajama-6B, validation (BLADE's corpus)
pruned modules	see below	see below

Everything else — hypernet architecture, both training stages, all learning rates, step counts, temperature, prior, row selection — is HyperPrune's own shipped setting.

Configuration

{
  "model": {
    "name_or_path": "openai/gpt-oss-20b",
    "dtype": "bfloat16",
    "moe_routing": "routed"
  },
  "data": {
    "dataset_name": "elastiml:elastix-ai/elastiml-calib-gpt-oss-20b",
    "num_samples": 128,
    "seq_len": 2048,
    "seed": 42
  },
  "hypernet": {
    "type": "mlp",
    "hidden_dim": 256,
    "emb_dim": 64,
    "use_layer_emb": false,
    "use_comp_emb": false,
    "use_hessian_diag": true
  },
  "training": {
    "sup_steps": 12000,
    "sup_lr": 0.001,
    "ft_lr": 0.0003,
    "ft_nsamples": 4,
    "rows_per_step": 400,
    "cascade_inner_steps": 300,
    "ft_mode": "cascade",
    "tau": 0.5,
    "prior_source": "sparsegpt",
    "wanda_residual_alpha": 2.0,
    "compensated_propagation": true,
    "use_weight_compensation": true,
    "train_on_compensated": true,
    "fixed_rows_count": 200,
    "fixed_rows_pos": "first",
    "dense_layers_list": []
  },
  "output": {
    "save_dir": "/home/ubuntu/hyperprune_work/outputs/hp-gptoss_20b-2to4",
    "wanda_dir": "/home/ubuntu/hyperprune_work/outputs/hp-gptoss_20b-2to4_ref",
    "preserve_wanda_dir": false
  }
}

Measured

metric	value
overall decoder sparsity (check_sparsity)	0.5035 over the 19.1B pruned expert weights (whole-model 0.4600: attention and router stay dense per BLADE's two_four_experts spec); strict 2:4 verified, 0 violations in 4.78B groups
WikiText-2 PPL (HyperPrune eval_ppl.py, seqlen 2048)	nan (fp16 overflow in the layer-wise eval; see the BLADE-protocol column, measured in bf16)
WikiText-2 word PPL (lm-eval-harness, BLADE's protocol)	247.45
training wall-clock	59.9 min
peak GPU during cascade FT	11.24 GB
GPU	1 x NVIDIA RTX PRO 6000 Blackwell (97 GB), CUDA 13.0, torch 2.13.0+cu130

Two things to know before comparing this number to the paper

1. Every decoder layer is pruned here. HyperPrune's own shipped configs set dense_layers_list: [0, 1], leaving 2 layers fully dense and yielding ~46.9 % sparsity rather than 50 %. This checkpoint prunes every layer, matching BLADE's two_four_experts spec, so it is a true 2:4 model in the modules BLADE prunes.

2. Only a few percent of this mask was chosen by the hypernet. The shipped recipe sets fixed_rows_count: 200, so the hypernet decides the mask for the first 200 output rows of each projection and every remaining row keeps the SparseGPT prior's mask verbatim. This is HyperPrune's own default, kept here deliberately because the brief was to change nothing but the calibration corpus.

The two perplexity rows are different quantities and are not comparable to each other. The first is token-level PPL over concatenated WikiText-2 at seqlen 2048 (the Wanda/SparseGPT convention). The second is lm-evaluation-harness word_perplexity at max_length=2048, which is BLADE's protocol — pinned empirically by reproducing BLADE's dense LLaMA-2-7B value of 9.19 (measured 9.1915).

Calibration

ElastiML self-distilled (elastix-ai/elastiml-calib-gpt-oss-20b), BLADE's corpus for this model, consumed from its pre-tokenized input_ids verbatim — necessary because these rows carry retokenizes_exactly: False, so re-encoding the text would give a different token stream.

This checkpoint required working around an upstream transformers defect

save_pretrained in transformers 5.15.1 silently drops gpt-oss-20b's fused expert weights. Established by a two-arm probe on a 2-layer truncation of the released checkpoint: both arms lost them, including the arm with no expert expansion at all — i.e. load the released MXFP4 checkpoint (which transformers dequantizes when the kernels package is absent) and save it straight back. The first attempt here produced a 5.2 GB file for a 21 B model: expert biases present, all 19.1 B expert weights missing, no exception raised.

The tensors are ordinary registered Parameters — in state_dict() and named_parameters(), not buffers, not tied, not in _keys_to_ignore_on_save, no shared storage, no _checkpoint_conversion_mapping — so the loss happens inside save_pretrained's writer. This checkpoint was written by a manual sharded safetensors save and then verified tensor-by-tensor against the live model (411/411).

Scope and behaviour

Experts only (attention and mlp.router dense), matching BLADE's two_four_experts. Calibration used moe_routing: routed; measured coverage at 128 samples is a minimum of 7,280 tokens per expert with none starved.

Reconstruction improved 73-78 % over the SparseGPT prior through layer 14, then decayed with depth (-58 %, -29 %, -33 %) and three layers (6, 22, 23) ended worse than the prior and were kept — HyperPrune has no keep-best guard. Greedy generation is nonetheless coherent.

Provenance

Produced from HyperPrune commit 6d093d7 with a small set of documented patches (bias-dtype autocast, calibration loader, disk-peak reduction, and — for 4:8 checkpoints — the N:M generalization, which the reference implementation does not ship). See the reproduction report for the full diff.