tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-04-16 13:58:38 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Garbage collect after step 1 and freeze

Browse Source
This commit is contained in:
Chris McCormick 2026-01-31 21:25:43 -08:00
parent 814475af42
commit 232d1341be
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
scripts/base_train.py
View File

@ -429,8 +429,14 @@ while True:
wandb_run.log(log_data)
# state update
first_step_of_run = (step == 0) or (resuming and step == args.resume_from_step)
step += 1
# After first step of this run, flush torch.compile garbage and freeze long-lived objects
if first_step_of_run:
gc.collect()
gc.freeze() # move survivors to permanent generation, won't be scanned in future GCs
# print a few more stats
print0(f"Peak memory usage: {get_max_memory() / 1024 / 1024:.2f}MiB")
print0(f"Total training time: {total_training_time/60:.2f}m")