Made models save properly

0d3abfca95590cbaf5f1cb1bd57b84c3be9ab6d9 · 2024-11-21 21:14:53 -0800 · Nel Nibcord

Files touched

M exo/main.py
M exo/orchestration/standard_node.py

Diff

commit 0d3abfca95590cbaf5f1cb1bd57b84c3be9ab6d9
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Thu Nov 21 21:14:53 2024 -0800

    Made models save properly
---
 exo/main.py                        |  3 +--
 exo/orchestration/standard_node.py | 14 +++++++++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/exo/main.py b/exo/main.py
index 30f463e2..c8bb72e5 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -267,8 +267,7 @@ async def train_model_cli(node: Node, inference_engine: InferenceEngine, model_n
     loss, tokens = await run_iter(node, shard, True, train, batch_size)
     print(f"epoch {epoch + 1}/{iters}\t| loss: {loss}, tokens: {tokens}")
     if save_interval > 0 and epoch > 0 and (epoch % save_interval) == 0 and checkpoint_dir is not None:
-      node.coordinate_save(checkpoint_dir, shard, epoch)
-      print("Hold up let's save a checkpoint")
+      await node.coordinate_save(shard, epoch, checkpoint_dir)
       await hold_outstanding(node)
   await hold_outstanding(node)
 
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index b0b62874..5def42ea 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -227,16 +227,20 @@ class StandardNode(Node):
   ):
     shard = self.get_current_shard(base_shard)
     model = shard.model_id
+    sid = shard.__hash__()
+    path = f"{destination}/{model}/{sid}-{iteration}.safetensors"
     self.outstanding_requests[f"{sid}::{iteration}"] = "Checking"
     if model not in self.checkpoints:
-      self.checkpoints[model_id] = {}
-    sid = shard.__hash__()
+      self.checkpoints[model] = {}
     if sid not in self.checkpoints[model]:
       self.checkpoints[model][sid] = []
-    if len(self.checkpoints[model][sid]) and self.checkpoints[model][sid][-1] < iteration:
+    if len(self.checkpoints[model][sid]) < 1 or self.checkpoints[model][sid][-1] < iteration:
+      print(f"Saving checkpoint to {path}")
       self.outstanding_requests[f"{sid}::{iteration}"] = "Saving"
-      await self.inference_engine.save_checkpoint(f"{destination}/{model}/{hash}-{iteration}")
-      self.checkpoints[model][sid] = sorted(self.checkpoints.model.sid + [iteration])
+      import os
+      os.makedirs("/".join(path.split("/")[:-1]), exist_ok=True)
+      await self.inference_engine.save_checkpoint(shard, path)
+      self.checkpoints[model][sid] = sorted(self.checkpoints[model][sid] + [iteration])
     self.outstanding_requests.pop(f"{sid}::{iteration}")
 
   async def process_example(

← 9283f6d7 Correct loss propagation so we can see the actual loss inste · back to Exo · Fixing tinygrad model 37a75d6b →