zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-09 18:36:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add checkpoints directory for adapter weights (#431)

Browse Source 
* Add checkpoints directory for adapter weights

The code was modified to create a checkpoints directory if it doesn't exist yet. Adapter weights are now saved to this checkpoints directory during the training iterations.
Corrected indentation of Save adapter weights code because it was part of "if eval"

* Fixing a blank added by mistake
This commit is contained in:
Ivan Fioravanti 2024-02-12 19:50:05 +01:00 committed by GitHub
parent 70465b8cda
commit 4576946151
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
llms/mlx_lm/tuner/trainer.py
View File

@ -127,6 +127,10 @@ def train(
args: TrainingArgs = TrainingArgs(),
loss: callable = default_loss,
):
# Create checkpoints directory if it does not exist
if not os.path.exists('checkpoints'):
os.makedirs('checkpoints')
# Create value and grad function for loss
loss_value_and_grad = nn.value_and_grad(model, loss)
@ -191,12 +195,13 @@ def train(
start = time.perf_counter()
# Save adapter weights if needed
if (it + 1) % args.steps_per_save == 0:
save_adapter(model=model, adapter_file=args.adapter_file)
print(
f"Iter {it + 1}: Saved adapter weights to {os.path.join(args.adapter_file)}."
)
# Save adapter weights if needed
if (it + 1) % args.steps_per_save == 0:
checkpoint_adapter_file = f"checkpoints/{it + 1}_{args.adapter_file}"
save_adapter(model=model, adapter_file=checkpoint_adapter_file)
print(
f"Iter {it + 1}: Saved adapter weights to {os.path.join(checkpoint_adapter_file)}."
)
# save final adapter weights
save_adapter(model=model, adapter_file=args.adapter_file)
print(f"Saved final adapter weights to {os.path.join(args.adapter_file)}.")