| """Model config for a REAP-pruned checkpoint (apply_selection.py output): the base model's config |
| with pretrained = the checkpoint directory, revision = SHA-1 of its reap_selection.json (a local |
| directory ignores the revision; the hash ties every run to the selection) and run_name = --name. |
| Everything else (serving, sampling, chat template) stays the base config's. |
| |
| cd apex-moe && .venv/bin/python external/reap/experiments/apex/make_config.py \ |
| --base configs/models/qwen3.6-35b-a3b-nvfp4.yaml --checkpoint data/reap-official/ckpt/<dir> \ |
| --name qwen3.6-35b-a3b-nvfp4-reapofficial-keep50 |
| """ |
|
|
| import argparse |
| import hashlib |
| import json |
| import re |
| from pathlib import Path |
|
|
| parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| parser.add_argument("--base", required=True) |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--name", required=True) |
| args = parser.parse_args() |
|
|
| ckpt = Path(args.checkpoint) |
| selection = (ckpt / "reap_selection.json").read_bytes() |
| digest = hashlib.sha1(selection).hexdigest() |
| meta = json.loads(selection) |
| text = Path(args.base).read_text() |
| header = ( |
| f"# REAP-pruned (official REAP, apex port: external/reap) {Path(args.base).name}: " |
| f"{meta['kept_per_layer']} experts per layer kept\n" |
| f"# ({100 * meta['ratio_pruned']:.0f}% pruned by {meta['metric']} saliency on apex-moe's calibration prompts\n" |
| "# with their datasets' answers, REAP style); revision = SHA-1 of the checkpoint's reap_selection.json\n" |
| ) |
| if meta.get("padding_experts"): |
| header += ( |
| f"# plus {meta['padding_experts']} padding experts no token can reach (zero weights, router bias -1e4),\n" |
| "# for vLLM's sqrtsoftplus top-k kernel (64/128/192/256/384 experts)\n" |
| ) |
| text, n1 = re.subn(r"(?m)^( pretrained:\s*)\S+", rf"\g<1>{ckpt}", text) |
| text, n2 = re.subn(r"(?m)^(run_name:\s*)\S+", rf"\g<1>{args.name}", text) |
| assert n1 == 1 and n2 == 1, (n1, n2) |
| if re.search(r"(?m)^ revision:", text): |
| text = re.sub(r"(?m)^( revision:\s*)\S+", rf"\g<1>{digest}", text) |
| else: |
| text = re.sub(r"(?m)^( pretrained:.*\n)", rf"\g<1> revision: {digest}\n", text, count=1) |
| out = Path("configs/models") / f"{args.name}.yaml" |
| out.write_text(header + text) |
| print(f"{out}: pretrained {ckpt}, revision {digest}") |
|
|
