| |
| """Export every parameter name, shape, dtype and gradient flag from a meta model.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from configuration_mini_k3 import MiniK3Config |
| from modeling_mini_k3 import MiniK3ForCausalLM |
|
|
|
|
| def main() -> None: |
| here = Path(__file__).resolve().parent |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", type=Path, default=here / "config.json") |
| parser.add_argument("--output", type=Path, default=here / "parameter_schema.jsonl") |
| args = parser.parse_args() |
|
|
| config = MiniK3Config.from_json(args.config) |
| model = MiniK3ForCausalLM(config, device="meta") |
| logical = 0 |
| tensors = 0 |
| with args.output.open("w") as stream: |
| for name, parameter in model.named_parameters(): |
| record = { |
| "name": name, |
| "shape": list(parameter.shape), |
| "dtype": str(parameter.dtype).removeprefix("torch."), |
| "requires_grad": parameter.requires_grad, |
| "numel": parameter.numel(), |
| } |
| stream.write(json.dumps(record, separators=(",", ":")) + "\n") |
| logical += parameter.numel() |
| tensors += 1 |
| assert logical == config.expected_logical_parameters |
| print(f"wrote {tensors:,} tensors / {logical:,} elements to {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
