Model loading and saving for tinygrad

329efb2381bdda412ff812645a43acff8771301d · 2024-12-11 03:21:17 -0800 · Nel Nibcord

Files touched

M exo/inference/tinygrad/inference.py

Diff

commit 329efb2381bdda412ff812645a43acff8771301d
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Wed Dec 11 03:21:17 2024 -0800

    Model loading and saving for tinygrad
---
 exo/inference/tinygrad/inference.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index 8d67c6a2..86abd76b 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -4,7 +4,7 @@ import os
 from exo.inference.tinygrad.models.llama import Transformer, TransformerShard, convert_from_huggingface, fix_bf16, sample_logits
 from exo.inference.shard import Shard
 from exo.inference.tokenizers import resolve_tokenizer
-from tinygrad.nn.state import load_state_dict
+from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict
 from tinygrad import Tensor, nn, Context, TinyJit
 from exo.inference.inference_engine import InferenceEngine
 import numpy as np
@@ -96,9 +96,13 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
   
   async def load_checkpoint(self, shard: Shard, path: str):
     await self.ensure_shard(shard)
+    state_dict = safe_load(path)
+    await asyncio.get_running_loop().run_in_executor(self.executor, load_state_dict, self.model, state_dict)
   
   async def save_checkpoint(self, shard: Shard, path: str):
     await self.ensure_shard(shard)
+    state_dict = await asyncio.get_running_loop().run_in_executor(self.executor, get_state_dict, self.model)
+    safe_save(state_dict, path) 
   
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
     await self.ensure_shard(shard)

← b1397b49 Proper sharding in tinygrad · back to Exo · t 8269b4b1 →