File size: 3,771 Bytes
8e0794f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """Load the ShellElastic supernet of this repository on the CPU, run one sub-network and extract it.
Usage, from the downloaded repository:
python example.py # the Low sub-network, float32
python example.py --tier max --dtype bfloat16
python example.py --tier low --save low_model # also write the extracted Low model to ./low_model
Needs only torch and safetensors (see requirements.txt). Nothing is written to disk unless --save is given.
"""
import argparse
import os
import sys
sys.dont_write_bytecode = True # leave the downloaded folder unchanged
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import torch # noqa: E402
from modeling_shellelastic import ShellElasticForCausalLM # noqa: E402
# GPT-NeoX-20B token ids of "The capital of France is". Any copy of that tokenizer gives them, for example
# tokenizers.Tokenizer.from_pretrained("EleutherAI/gpt-neox-20b").encode("The capital of France is").ids
PROMPT_IDS = [510, 5347, 273, 6181, 310]
def main() -> None:
parser = argparse.ArgumentParser(description="Run and extract one sub-network of this supernet on the CPU.")
parser.add_argument("--tier", default="low", choices=["low", "medium", "high", "extra", "max"],
help="sub-network (default: low)")
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"],
help="compute precision (the released weights are float32)")
parser.add_argument("--new-tokens", type=int, default=8, help="greedy tokens to generate")
parser.add_argument("--save", metavar="DIR", help="write the extracted sub-network (config.json and "
"model.safetensors, loadable with modeling_shellelastic.py) to DIR")
args = parser.parse_args()
model = ShellElasticForCausalLM.from_pretrained(HERE, device="cpu", dtype=getattr(torch, args.dtype))
cfg = model.config
print(f"supernet: {model.num_parameters():,} stored parameters, {args.dtype}")
print(f"{args.tier}: width {cfg.width(args.tier)}, {cfg.width(args.tier) // cfg.head_dim} heads, "
f"{model.num_parameters(args.tier):,} parameters")
ids = torch.tensor([PROMPT_IDS])
with torch.no_grad():
logits, loss = model(ids[:, :-1], tier=args.tier, labels=ids[:, 1:])
all_tiers = model(ids, tiers=cfg.tiers) # every sub-network from one full-width trunk pass
out = model.generate(ids, max_new_tokens=args.new_tokens, tier=args.tier)
assert logits.shape == (1, len(PROMPT_IDS) - 1, cfg.vocab_size)
assert torch.isfinite(logits.float()).all() and torch.isfinite(loss)
print(f"prompt ids: {PROMPT_IDS}")
print(f"mean next-token loss on the prompt: {loss.item():.4f}")
print("next-token id after the prompt, per sub-network: "
+ ", ".join(f"{t} {all_tiers[t][0, -1].argmax().item()}" for t in cfg.tiers))
print(f"greedy continuation ids: {out[0, len(PROMPT_IDS):].tolist()}")
# The sub-network as a standalone model: its embedding columns, the core and the shells it contains, and its own
# terminal norm and readout. By Proposition 1 of the paper its logits equal the supernet's for this tier.
small = model.extract(args.tier)
with torch.no_grad():
diff = (small(ids, tier=args.tier).float() - model(ids, tier=args.tier).float()).abs().max().item()
print(f"extracted {args.tier}: {small.num_parameters():,} parameters, max |logit difference| to the "
f"supernet {diff:g}")
assert diff == 0.0
if args.save:
small.save_pretrained(args.save)
print(f"wrote {os.path.join(args.save, 'config.json')} and {os.path.join(args.save, 'model.safetensors')}")
if __name__ == "__main__":
main()
|