From 2e62dd308d3d5c6d0df83aa43c5380c27ff37ce2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Jul 2026 22:01:15 -0400 Subject: [PATCH] move fs_load and fs_store to nn (#16821) similar to safe_load and safe_save --- extra/tinyfs/fetch_file.py | 14 ++++---- extra/tinyfs/fetch_raid.py | 5 +-- extra/tinyfs/upload_raid.py | 5 +-- test/null/test_tinyfs.py | 9 +++--- test/unit/test_tinyfs.py | 26 ++++++++------- tinygrad/helpers.py | 1 + tinygrad/nn/state.py | 58 ++++++++++++++++++++++++++++++++-- tinygrad/runtime/ops_tinyfs.py | 9 +++--- tinygrad/tensor.py | 55 +------------------------------- 9 files changed, 95 insertions(+), 87 deletions(-) diff --git a/extra/tinyfs/fetch_file.py b/extra/tinyfs/fetch_file.py index fd527c06f6..7354e24308 100644 --- a/extra/tinyfs/fetch_file.py +++ b/extra/tinyfs/fetch_file.py @@ -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: diff --git a/extra/tinyfs/fetch_raid.py b/extra/tinyfs/fetch_raid.py index c628c7efb6..4cf00fd193 100644 --- a/extra/tinyfs/fetch_raid.py +++ b/extra/tinyfs/fetch_raid.py @@ -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() diff --git a/extra/tinyfs/upload_raid.py b/extra/tinyfs/upload_raid.py index 5557ffdd26..5dfed65792 100644 --- a/extra/tinyfs/upload_raid.py +++ b/extra/tinyfs/upload_raid.py @@ -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)}") diff --git a/test/null/test_tinyfs.py b/test/null/test_tinyfs.py index aeb215a9fb..080b7709e0 100644 --- a/test/null/test_tinyfs.py +++ b/test/null/test_tinyfs.py @@ -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() diff --git a/test/unit/test_tinyfs.py b/test/unit/test_tinyfs.py index 8866a1a441..aa7c0f619f 100644 --- a/test/unit/test_tinyfs.py +++ b/test/unit/test_tinyfs.py @@ -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__": diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 8a55f51687..fe30002c50 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -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) diff --git a/tinygrad/nn/state.py b/tinygrad/nn/state.py index 7b340c4ec7..cd47bc988e 100644 --- a/tinygrad/nn/state.py +++ b/tinygrad/nn/state.py @@ -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]: diff --git a/tinygrad/runtime/ops_tinyfs.py b/tinygrad/runtime/ops_tinyfs.py index c8e41113da..8e0cba3ecf 100644 --- a/tinygrad/runtime/ops_tinyfs.py +++ b/tinygrad/runtime/ops_tinyfs.py @@ -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]) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 6ae49c5a77..20539f105e 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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