34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
from torch.distributed.checkpoint.metadata import Metadata
|
|
from torch.distributed.checkpoint.filesystem import FileSystemReader
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("checkpoint_dir")
|
|
parser.add_argument("--limit", type=int, default=240)
|
|
args = parser.parse_args()
|
|
|
|
reader = FileSystemReader(args.checkpoint_dir)
|
|
metadata: Metadata = reader.read_metadata()
|
|
count = 0
|
|
skip_prefixes = ("optimizer.", "opt_param_scheduler", "rng_state", "iteration")
|
|
for key in sorted(metadata.state_dict_metadata):
|
|
if key.startswith(skip_prefixes):
|
|
continue
|
|
meta = metadata.state_dict_metadata[key]
|
|
size = getattr(meta, "size", None)
|
|
print(f"{key}\t{tuple(size) if size is not None else '?'}")
|
|
count += 1
|
|
if args.limit and count >= args.limit:
|
|
break
|
|
print(f"printed={count}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|