tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-08 05:12:16 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix race condition in save_checkpoint for non-zero ranks

Browse Source
This commit is contained in:
zzF 2025-11-26 14:05:11 +08:00 committed by GitHub
parent 4a87a0d19f
commit 2c6a007e3c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nanochat/checkpoint_manager.py
View File

@ -21,8 +21,8 @@ def log0(message):
logger.info(message) logger.info(message)
def save_checkpoint(checkpoint_dir, step, model_data, optimizer_data, meta_data, rank=0): def save_checkpoint(checkpoint_dir, step, model_data, optimizer_data, meta_data, rank=0):
if rank == 0:
os.makedirs(checkpoint_dir, exist_ok=True) os.makedirs(checkpoint_dir, exist_ok=True)
if rank == 0:
# Save the model state parameters # Save the model state parameters
model_path = os.path.join(checkpoint_dir, f"model_{step:06d}.pt") model_path = os.path.join(checkpoint_dir, f"model_{step:06d}.pt")
torch.save(model_data, model_path) torch.save(model_data, model_path)