diff --git a/test/unit/test_function.py b/test/unit/test_function.py index e333d8f9a6..242189bc80 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -557,6 +557,15 @@ class TestFunctionTuple(unittest.TestCase): def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0] with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize()) + def test_realize_inside_function(self): + constants = [] + @function(precompile=True, allow_implicit=True) + def f(a:Tensor): + constants.append(Tensor([2], device=a.device).realize()) + return a + constants[-1] + np.testing.assert_equal(f(Tensor([1], device="CPU").realize()).numpy(), [3]) + self.assertTrue(constants[0].uop.is_realized) + def test_custom_kernel_write_only_persistent_output_is_implicit(self): # a write-only custom_kernel output that is a realized buffer must be captured def write(C:UOp, A:UOp) -> UOp: diff --git a/test/unit/test_llm_amd.py b/test/unit/test_llm_amd.py index 04b92c6dd8..6aa552a431 100644 --- a/test/unit/test_llm_amd.py +++ b/test/unit/test_llm_amd.py @@ -1,10 +1,9 @@ import unittest import numpy as np -from tinygrad import Tensor, dtypes -from tinygrad.llm.kernels import cached_attention +from tinygrad import Tensor, dtypes, nn +from tinygrad.llm.kernels import Linear, cached_attention from tinygrad.llm.kernels.amd import q8_quantize from tinygrad.llm.gguf import ggml_data_to_tensor -from tinygrad.llm.model import Linear class TestQ8Quantize(unittest.TestCase): def test_values_and_scales(self): @@ -23,10 +22,9 @@ class TestQ8Quantize(unittest.TestCase): raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:] decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256) linear = Linear(256, 1, bias=False) - offset = linear.set_quantized(decoded) - assert offset is not None - linear._raw_offset_uop = offset.realize().uop + nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False) self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item())) + self.assertIsNotNone(linear._raw_offset) def test_attention_uses_physical_cache_length(self): if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required") diff --git a/tinygrad/llm/kernels/__init__.py b/tinygrad/llm/kernels/__init__.py index 5989545f11..48c1b0093c 100644 --- a/tinygrad/llm/kernels/__init__.py +++ b/tinygrad/llm/kernels/__init__.py @@ -10,7 +10,7 @@ class Linear(nn.Linear): def __init__(self, in_features:int, out_features:int, bias=True): super().__init__(in_features, out_features, bias) self.in_features, self.out_features = in_features, out_features - self._raw_offset_uop:UOp|None = None + self._raw_offset:Tensor|None = None def set_quantized(self, decoded:Tensor) -> Tensor|None: packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in ((13, 176), (14, 210), (23, 136))} raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None) @@ -21,8 +21,10 @@ class Linear(nn.Linear): if self.ggml_type == 23 and str(self.weight.device).startswith("AMD"): from tinygrad.llm.kernels.amd import iq4_half_lut iq4_half_lut(str(self.weight.device)) - return Tensor([raw_offset // 4], dtype=dtypes.uint64, device=self.weight.device) + self._raw_offset = Tensor([raw_offset//4], dtype=dtypes.uint64, device=self.weight.device).realize() + return self._raw_offset def __call__(self, x:Tensor) -> Tensor: + if self.ggml_type is None and str(self.weight.device).startswith("AMD"): self.set_quantized(self.weight) if self.ggml_type in (13, 14, 23) and str(self.weight.device).startswith("AMD"): from tinygrad.llm.kernels.amd import q8_linear return q8_linear(self, x) diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index fc6102eef1..7c869c1f8c 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -395,10 +395,10 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, raw_offset:UOp return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma") def q8_linear(layer:Linear, x:Tensor) -> Tensor: - assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS) and layer._raw_offset_uop is not None + assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS) and layer._raw_offset is not None tokens = int(x.numel()) // layer.in_features out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop - raw, offset = layer.weight.uop.buf_uop, layer._raw_offset_uop + raw, offset = layer.weight.uop.buf_uop, layer._raw_offset.uop out_features, in_features = layer.out_features, layer.in_features use_wmma = tokens % 16 == 0 and layer.out_features % 16 == 0 def run(fxn:Callable[..., UOp], *srcs:UOp) -> Tensor: diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index bbe8b72d2d..aaf1735d1e 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -1,21 +1,11 @@ from __future__ import annotations import functools, itertools, pathlib from dataclasses import dataclass, replace -from typing import cast from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context, dtypes from tinygrad.llm.kernels import Linear, cached_attention, gated_delta_prefill from tinygrad.llm.gguf import gguf_load from tinygrad.uop.ops import resolve -def load_state_dict(model:Transformer, state_dict:dict[str, Tensor]): - nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) - layers = cast(dict[str, Linear], nn.state.get_state_dict(model, tensor_type=Linear)).values() - packed:list[tuple[Linear, Tensor]] = [] - for layer in layers: - if str(layer.weight.device).startswith("AMD") and (offset:=layer.set_quantized(layer.weight)) is not None: packed.append((layer, offset)) - if packed: Tensor.realize(*(offset for _,offset in packed)) - for layer,offset in packed: layer._raw_offset_uop = offset.uop - @functools.cache def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor: freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim)) @@ -408,7 +398,7 @@ class Transformer: qkv_bias='blk.0.attn_q.bias' in state_dict, expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict) model = Transformer(config) - load_state_dict(model, state_dict) # NOTE: rope_freqs.weight (32,) is unused + nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) # NOTE: rope_freqs.weight (32,) is unused # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster if realize: for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous()) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 71b4fa8f9d..e81e8fdedd 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -4,7 +4,7 @@ import time, functools, sys, inspect, pathlib, hashlib, weakref from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING if TYPE_CHECKING: import numpy from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst -from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey +from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey, Context from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike from tinygrad.mixin.rand import RandMixin @@ -16,6 +16,16 @@ from tinygrad.callify import transform_to_call # *** all in scope Tensors are here. this gets relevant UOps *** all_tensors: dict[weakref.ref[Tensor], None] = {} +def _realize_tensors(tensors:tuple[Tensor, ...], do_update_stats=True): + for x in tensors: + if x.uop.op is Ops.COPY and x.uop.src[0].op is Ops.BUFFER and x.uop.src[0].device == "PYTHON" and \ + isinstance(x.device, str) and x.nbytes() <= 8: + out = UOp.new_buffer(x.device, x.uop.max_numel(), x.dtype) + with Context(ALLOW_DEVICE_USAGE=1): + cast(Buffer, out.buffer).ensure_allocated().copy_from(cast(Buffer, x.uop.src[0].buffer).ensure_allocated()) + x.uop = out.reshape(x.shape) + to_realize = [x for x in tensors if not x.uop.is_virtual and not x.uop.has_buffer_identity()] + if to_realize: run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats) def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None: with cpu_profile(TracingKey(name), "TINY"): # get tensors in scope @@ -190,9 +200,7 @@ class Tensor(RandMixin): @disable_gc() def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor: """Triggers the computation needed to create these Tensor(s).""" - to_realize = [x for x in (self,)+lst if not x.uop.is_virtual and not x.uop.has_buffer_identity()] - if len(to_realize): - run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats) + _realize_tensors((self,)+lst, do_update_stats) return self def replace(self, x:Tensor) -> Tensor: