vikarti.anatra/blt
1
0
Fork 0
mirror of https://github.com/facebookresearch/blt.git synced 2025-09-01 10:09:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add way to call consolidate

Browse source
This commit is contained in:
Srini Iyer 2025-03-11 22:26:11 +00:00
parent a5ceaaa226
commit 699d0be470
1 changed files with 14 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
bytelatent/checkpoint.py
View file

@ -4,6 +4,7 @@ import json
import logging import logging
import os import os
import re import re
import typer
import fsspec import fsspec
import s3fs import s3fs
@ -323,3 +324,16 @@ class CheckpointManager:
dist.barrier() dist.barrier()
return cls(args) return cls(args)
def main(
command: str,
model_checkpoint_dir: str,
):
if command == "consolidate":
print(f"Consolidating {model_checkpoint_dir}. Output will be in the {CONSOLIDATE_FOLDER} folder.")
consolidate_checkpoints(fsspec.filesystem("file"), model_checkpoint_dir)
else:
raise ValueError("Invalid command")
if __name__ == '__main__':
typer.run(main)