| |
| """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() |
|
|
