File size: 1,546 Bytes
8b7b828 | 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 | #!/usr/bin/env python3
"""Create the full module graph on meta, or explicitly allocate random weights."""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
from configuration_mini_k3 import MiniK3Config
from modeling_mini_k3 import MiniK3ForCausalLM, count_logical_parameters
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=Path, default=Path(__file__).with_name("config.json"))
parser.add_argument(
"--materialize",
action="store_true",
help="Actually allocate and randomly initialize the full model on CPU.",
)
parser.add_argument("--seed", type=int, default=1234, help="Random seed for initialization")
parser.add_argument("--output", type=Path, help="Optional torch.save state-dict path")
args = parser.parse_args()
config = MiniK3Config.from_json(args.config)
model = MiniK3ForCausalLM(config, device="meta")
print(f"meta model defined: {count_logical_parameters(model):,} logical parameters")
if not args.materialize:
print("No storage allocated. Pass --materialize only on a host with enough CPU RAM.")
return
torch.manual_seed(args.seed)
model.materialize_and_initialize("cpu")
print(f"Full model randomly initialized on CPU with seed {args.seed}")
if args.output is not None:
torch.save({"config": config.to_dict(), "model": model.state_dict()}, args.output)
print(f"saved: {args.output}")
if __name__ == "__main__":
main()
|