tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 13:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

use BufferedReader (#8032)

Browse Source
This commit is contained in:
leopf
2024-12-04 19:08:54 +08:00
committed by GitHub
parent 08657cb7b0
commit fb89971e73
1 changed files with 6 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/nn/state.py
+6 -11
View File
@@ -297,14 +297,9 @@ def gguf_load(tensor: Tensor) -> Tuple[Dict, Dict[str, Tensor]]:
kv_data, state_dict = gguf_load(gguf_tensor)
```
"""
if tensor.dtype != dtypes.uint8 or len(tensor.shape) != 1: raise ValueError("GGUF tensor must be 1d and of dtype uint8!")
pos, read_buffer, rb_start, kv_data, state_dict = 0, memoryview(bytes()), 0, {}, {}
def read_bytes(n: int):
nonlocal pos, read_buffer, rb_start
if rb_start + len(read_buffer) < pos + n: rb_start, read_buffer = pos, tensor[pos:(pos+max(n, 1000_000))].data()
return read_buffer[pos-rb_start:(pos:=pos+n)-rb_start]
def read_unpack(fmt: str, n: int): return struct.unpack(fmt, read_bytes(n))[0]
def read_str(): return str(read_bytes(read_uint64()), "utf-8")
reader, kv_data, state_dict = io.BufferedReader(TensorIO(tensor), 1000_000), {}, {}
def read_unpack(fmt: str, n: int): return struct.unpack(fmt, reader.read(n))[0]
def read_str(): return str(reader.read(read_uint64()), "utf-8")
def read_arr():
reader, n = readers[read_int32()], read_uint64()
return [ reader() for _ in range(n) ]
@@ -313,15 +308,15 @@ def gguf_load(tensor: Tensor) -> Tuple[Dict, Dict[str, Tensor]]:
(1,"b",1), (2,"H",2), (3,"h",2), (4,"I",4), (5,"i",4), (6,"f",4), (7,"?",1), (10,"Q",8), (11,"q",8), (12,"d",8) ] } }
read_uint32, read_int32, read_uint64, read_int64 = readers[4], readers[5], readers[10], readers[11]
magic, version, n_tensors, n_kv = read_bytes(4), read_int32(), read_int64(), read_int64()
magic, version, n_tensors, n_kv = reader.read(4), read_int32(), read_int64(), read_int64()
if magic != b"GGUF" or version not in [2, 3]: raise ValueError("Invalid GGUF format!")
for _ in range(n_kv):
k, typ = read_str(), read_int32()
kv_data[k] = readers[typ]()
t_infos = [ (read_str(), tuple(read_uint64() for _ in range(read_uint32())), read_int32(), read_uint64()) for _ in range(n_tensors) ]
alignment = kv_data.get("general.alignment", 32)
data_start = pos = pos + (alignment - pos % alignment if pos % alignment != 0 else 0)
alignment, pos = kv_data.get("general.alignment", 32), reader.tell()
data_start = pos + (alignment - pos % alignment if pos % alignment != 0 else 0)
for name, dims, typ, off in t_infos: state_dict[name] = ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims))