geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

move fs_load and fs_store to nn (#16821)

Browse Source 
similar to safe_load and safe_save
This commit is contained in:
chenyu
2026-07-01 22:01:15 -04:00
committed by GitHub
parent c8aed121cf
commit 2e62dd308d
9 changed files with 95 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/tinyfs/fetch_file.py
+8 -6
View File
@@ -1,4 +1,6 @@
from tinygrad.tensor import Tensor
from tinygrad.helpers import CHUNK_SIZE
from tinygrad.nn.state import fs_load
import argparse, math, hashlib
def _python_hash_1mb(data:bytes|bytearray):
@@ -7,15 +9,15 @@ def _python_hash_1mb(data:bytes|bytearray):
return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
def hash_file(data: bytes|bytearray):
if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
base_chunks = math.ceil(len(data) / Tensor.CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
base_chunks = math.ceil(len(data) / CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
for _ in range(tree_depth + 1):
data_chunks = [data[i:i+Tensor.CHUNK_SIZE] for i in range(0, len(data), Tensor.CHUNK_SIZE)]
data_chunks = [data[i:i+CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
data = b''.join(data_chunk_hashes)
if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
return data[:16]
@@ -27,7 +29,7 @@ if __name__ == "__main__":
parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
args = parser.parse_args()
Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()
fs_load(Tensor(bytes.fromhex(args.hash), device="CPU"), args.len).to(f"disk:{args.dest}").realize()
if args.check:
with open(args.dest, "rb") as f:
extra/tinyfs/fetch_raid.py
+3 -2
View File
@@ -3,6 +3,7 @@ from pathlib import Path
from tinygrad.tensor import Tensor
from tinygrad.helpers import tqdm, getenv
from tinygrad.nn.state import fs_load
raid_root = Path(getenv("RAID_ROOT", "/raid"))
@@ -14,7 +15,7 @@ def fetch_file(item):
path.parent.mkdir(parents=True, exist_ok=True)
try:
pt = Tensor(bytes.fromhex(h), device="CPU").fs_load(size).to(f"disk:{path.as_posix()}").realize()
pt = fs_load(Tensor(bytes.fromhex(h), device="CPU"), size).to(f"disk:{path.as_posix()}").realize()
except Exception as e:
print(f"error fetching {path}, {h}, {size}: {e}")
raise
@@ -22,7 +23,7 @@ def fetch_file(item):
pt.uop.buffer.deallocate()
def fetch_mapping(h, l):
mapping_tensor = Tensor(bytes.fromhex(h)).fs_load(l).realize()
mapping_tensor = fs_load(Tensor(bytes.fromhex(h)), l).realize()
mapping = mapping_tensor.data().tobytes().decode()
mapping = json.loads(mapping)
mapped_files = mapping.items()
extra/tinyfs/upload_raid.py
+3 -2
View File
@@ -3,12 +3,13 @@ import multiprocessing, json
from tinygrad.tensor import Tensor
from tinygrad.helpers import tqdm
from tinygrad.nn.state import fs_store
raid_root = Path("/raid")
def upload_file(path: Path):
pt = Tensor(path).realize()
h = pt.fs_store().realize()
h = fs_store(pt).realize()
pt.uop.realized.deallocate()
return h.data().hex(), path, pt.nbytes()
@@ -26,6 +27,6 @@ if __name__ == "__main__":
mapping = json.dumps(mapping).encode()
mapping_tensor = Tensor(mapping, device="CPU")
h = mapping_tensor.fs_store().realize()
h = fs_store(mapping_tensor).realize()
print(f"final hash: {h.data().hex()}, size: {len(mapping)}")
test/null/test_tinyfs.py
+5 -4
View File
@@ -1,24 +1,25 @@
import unittest
from tinygrad import Tensor
from tinygrad.nn.state import fs_store, fs_load
class TestLoadStore(unittest.TestCase):
def test_load_shape(self):
t = Tensor(bytes(16)).fs_load(1024)
t = fs_load(Tensor(bytes(16)), 1024)
assert t.shape == (1024,), t.shape
t.schedule_linear()
def test_store_shape(self):
t = Tensor.zeros(1024).fs_store()
t = fs_store(Tensor.zeros(1024))
assert t.shape == (16,), t.shape
t.schedule_linear()
def test_load_large_shape(self):
t = Tensor(bytes(16)).fs_load(10_000_000)
t = fs_load(Tensor(bytes(16)), 10_000_000)
assert t.shape == (10_000_000,), t.shape
t.schedule_linear()
def test_store_large_shape(self):
t = Tensor.zeros(10_000_000).fs_store()
t = fs_store(Tensor.zeros(10_000_000))
assert t.shape == (16,), t.shape
t.schedule_linear()
test/unit/test_tinyfs.py
+14 -12
View File
@@ -1,6 +1,8 @@
import json, math, os, socketserver, threading, unittest
import numpy as np
from tinygrad import Tensor, dtypes
from tinygrad.helpers import CHUNK_SIZE
from tinygrad.nn.state import fs_store, fs_load
from extra.tinyfs.fetch_file import hash_file, _python_hash_1mb
_chunks: dict[bytes, bytes] = {}
@@ -14,8 +16,8 @@ class _Handler(socketserver.StreamRequestHandler):
elif cmd.startswith("STORE_IN"):
data = self.rfile.read(int(cmd.split()[1]))
hashes = bytearray()
for i in range(math.ceil(len(data) / Tensor.CHUNK_SIZE)):
chunk = data[i*Tensor.CHUNK_SIZE:(i+1)*Tensor.CHUNK_SIZE].ljust(Tensor.CHUNK_SIZE, b'\0')
for i in range(math.ceil(len(data) / CHUNK_SIZE)):
chunk = data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE].ljust(CHUNK_SIZE, b'\0')
h = _python_hash_1mb(chunk)
_chunks[h] = chunk
hashes.extend(h)
@@ -46,35 +48,35 @@ class TestTinyFS(unittest.TestCase):
cls._server.server_close()
def test_store(self):
h = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
h = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
self.assertEqual(h.shape, (16,))
self.assertEqual(h.dtype, dtypes.uint8)
def test_store_deterministic(self):
a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
a = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
b = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
np.testing.assert_array_equal(a.numpy(), b.numpy())
def test_store_different_data(self):
a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
a = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
b = fs_store(Tensor([5.0, 6.0, 7.0, 8.0])).realize()
self.assertNotEqual(a.tolist(), b.tolist())
def test_roundtrip_uint8(self):
arr = np.arange(256, dtype=np.uint8)
loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
loaded = fs_load(fs_store(Tensor(arr)).realize(), len(arr)).to("CPU")
np.testing.assert_array_equal(loaded.numpy(), arr)
def test_roundtrip_multichunk_uint8(self):
arr = np.random.default_rng(42).integers(0, 256, size=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
arr = np.random.default_rng(42).integers(0, 256, size=CHUNK_SIZE + 1024, dtype=np.uint8)
loaded = fs_load(fs_store(Tensor(arr)).realize(), len(arr)).to("CPU")
np.testing.assert_array_equal(loaded.numpy(), arr)
def test_hash_matches_python_impl(self):
arr = np.arange(256, dtype=np.uint8)
h = Tensor(arr).fs_store().realize()
h = fs_store(Tensor(arr)).realize()
# the hash from fs_store should match the pure-Python hash_file reference
padded = arr.tobytes().ljust(Tensor.CHUNK_SIZE, b'\0')
padded = arr.tobytes().ljust(CHUNK_SIZE, b'\0')
self.assertEqual(h.data().tobytes(), hash_file(padded))
if __name__ == "__main__":
tinygrad/helpers.py
+1
View File
@@ -239,6 +239,7 @@ class _DEV(ContextVar):
DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
CHUNK_SIZE = 2**20 # TinyFS content-addressed store: blob chunk + hash-tree node granularity
WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
TRAINING = ContextVar("TRAINING", 0)
USE_TC, TC_SELECT, TC_OPT = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0)
tinygrad/nn/state.py
+56 -2
View File
@@ -1,9 +1,9 @@
import json, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
import json, math, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
from collections import OrderedDict
from typing import Any, Callable, BinaryIO, Iterable, cast
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes
from tinygrad.helpers import prod, argsort, DEBUG, Timing, GlobalCounters, tqdm, round_up, T, strides_for_shape
from tinygrad.helpers import prod, argsort, DEBUG, Timing, GlobalCounters, tqdm, round_up, T, strides_for_shape, CHUNK_SIZE
class TensorIO(io.RawIOBase, BinaryIO):
def __init__(self, t: Tensor):
@@ -82,6 +82,60 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
t[8:8+len(j)].assign(list(j.encode('utf-8')))
for k,v in safe_load(t).items(): v.assign(tensors[k])
# tinyfs
def fs_store(t:Tensor) -> Tensor:
"""
Store a tensor to storage.
"""
# TODO: this should work locally as well
data = t.contiguous().flatten().bitcast(dtypes.uint8)
# pad to a multiple of 1mb
if (tsize := data.shape[0]) % CHUNK_SIZE != 0: data = data.pad((0, CHUNK_SIZE - tsize % CHUNK_SIZE))
size = data.shape[0]
base_chunks = math.ceil(size / CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
to_device = "CPU" if isinstance(t.device, str) and t.device.startswith("DISK") else t.device
level_chunks = base_chunks
for _ in range(tree_depth + 1):
data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device)
if (tsize := data.shape[0]) % CHUNK_SIZE != 0: data = data.pad((0, CHUNK_SIZE - tsize % CHUNK_SIZE))
level_chunks = math.ceil(data.shape[0] / CHUNK_SIZE)
return data[:16].contiguous()
def fs_load(t:Tensor, size:int) -> Tensor:
"""
Load a tensor from storage.
t should be a tensor of the hash to load
"""
# TODO: this should work locally as well
assert t.dtype == dtypes.uint8, "hash is expected to be uint8"
h = t.contiguous().flatten()
assert h.shape[0] == 16, "expected hash"
base_chunks = math.ceil(size / CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
data, level_chunks = h, 0
for i in reversed(range(tree_depth + 1)):
data = data.to("tinyfs:load")
# if not last level, its still hashes
if i > 0 or tree_depth == 0:
level_chunks = max(1, math.ceil(base_chunks / (CHUNK_SIZE // 16)**(i-1)))
pad_amt = 16 * level_chunks
else: pad_amt = CHUNK_SIZE * level_chunks
if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
data = data[:pad_amt].contiguous()
if i != 0: data = data.to(t.device)
return data[:size]
# state dict
def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:
tinygrad/runtime/ops_tinyfs.py
+4 -5
View File
@@ -1,8 +1,7 @@
import socket, json, asyncio, threading, math
from contextlib import asynccontextmanager
from tinygrad.device import Compiled, Allocator
from tinygrad.helpers import DEBUG, getenv
from tinygrad import Tensor
from tinygrad.helpers import DEBUG, getenv, CHUNK_SIZE
TINYFS_ENDPOINT = getenv("TINYFS_ENDPOINT", "localhost:6767")
TINYFS_TIMEOUT = getenv("TINYFS_TIMEOUT", 60)
@@ -95,7 +94,7 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
dest.copyout_queue = json.loads(locs)
dest.hash_buf = src.tobytes()
elif dest.device.op == "STORE":
expected_hashes = math.ceil(dest.size / Tensor.CHUNK_SIZE)
expected_hashes = math.ceil(dest.size / CHUNK_SIZE)
dest.hash_buf = bytearray(expected_hashes * 16)
self.dev.sfile.readinto(dest.hash_buf)
@@ -109,8 +108,8 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
async def _copyout_async(self, dest:memoryview, src:TinyFSBuffer):
async def _worker(i, loc):
async with self.dev.connection(loc) as (reader, writer):
ptr = i * Tensor.CHUNK_SIZE
size = min(len(dest[ptr:ptr+Tensor.CHUNK_SIZE]), Tensor.CHUNK_SIZE)
ptr = i * CHUNK_SIZE
size = min(len(dest[ptr:ptr+CHUNK_SIZE]), CHUNK_SIZE)
writer.write(f"CHUNK_OUT {size}\r\n".encode())
writer.write(src.hash_buf[i*16:(i+1)*16])
tinygrad/tensor.py
+1 -54
View File
@@ -1,6 +1,6 @@
# inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
from __future__ import annotations
import time, math, functools, sys, inspect, pathlib, hashlib, weakref
import time, functools, sys, inspect, pathlib, hashlib, weakref
from typing import Any, Callable, Sequence, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
if TYPE_CHECKING: import numpy
from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, _from_np_dtype, _to_np_dtype, PyConst
@@ -349,59 +349,6 @@ class Tensor(RandMixin):
if isinstance(y.device, str): return self.to(y.device)
return self if isinstance(self.device, tuple) and (y.device, y.uop.axis) == (self.device, self.uop.axis) else self.shard(y.device, y.uop.axis)
CHUNK_SIZE = 2**20
def fs_load(self, size:int) -> Tensor:
"""
Load a tensor from storage.
self should be a tensor of the hash to load
"""
# TODO: this should work locally as well
assert self.dtype == dtypes.uint8, "hash is expected to be uint8"
h = self.contiguous().flatten()
assert h.shape[0] == 16, "expected hash"
base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
data, level_chunks = h, 0
for i in reversed(range(tree_depth + 1)):
data = data.to("tinyfs:load")
# if not last level, its still hashes
if i > 0 or tree_depth == 0:
level_chunks = max(1, math.ceil(base_chunks / (Tensor.CHUNK_SIZE // 16)**(i-1)))
pad_amt = 16 * level_chunks
else: pad_amt = Tensor.CHUNK_SIZE * level_chunks
if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
data = data[:pad_amt].contiguous()
if i != 0: data = data.to(self.device)
return data[:size]
def fs_store(self) -> Tensor:
"""
Store a tensor to storage.
"""
# TODO: this should work locally as well
data = self.contiguous().flatten().bitcast(dtypes.uint8)
# pad to a multiple of 1mb
if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE))
size = data.shape[0]
base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
to_device = "CPU" if isinstance(self.device, str) and self.device.startswith("DISK") else self.device
level_chunks = base_chunks
for _ in range(tree_depth + 1):
data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device)
if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE))
level_chunks = math.ceil(data.shape[0] / Tensor.CHUNK_SIZE)
return data[:16].contiguous()
# ***** creation entrypoint *****
@staticmethod