diff --git a/examples/other_mnist/beautiful_mnist_torch.py b/examples/other_mnist/beautiful_mnist_torch.py index 8e0b7dd64d27a..5684e5d172a08 100644 --- a/examples/other_mnist/beautiful_mnist_torch.py +++ b/examples/other_mnist/beautiful_mnist_torch.py @@ -28,6 +28,7 @@ def forward(self, x): if __name__ == "__main__": if getenv("TINY_BACKEND"): import tinygrad.nn.torch # noqa: F401 + import extra.torch_backend.compile # noqa: F401 device = torch.device("tiny") else: device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu")) @@ -43,7 +44,7 @@ def forward(self, x): optimizer = optim.Adam(model.parameters(), 1e-3) loss_fn = nn.CrossEntropyLoss() - #@torch.compile + @torch.compile def step(samples): X,Y = X_train[samples], Y_train[samples] out = model(X) diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index 85466aa8253be..fde7982bef0a6 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -55,22 +55,36 @@ def device_count(self): return getenv("GPUS", 1) # TODO: device count in tiny? torch.utils.generate_methods_for_privateuse1_backend() aten = torch.ops.aten +# a tiny tensor has no Storage at all, but dynamo fakifies its inputs through untyped_storage(). hand it a meta storage, +# torch's own "a storage with no data", which is what MetaConverter rebuilds from it anyway. dynamo compares storage +# identity to decide if two tensors alias, so it is cached on the tinygrad Tensor: t, t.detach() and nn.Parameter(t) +# are different torch tensors over one tinygrad Tensor and do alias +_untyped_storage = torch.Tensor.untyped_storage +def untyped_storage(self:torch.Tensor): + # only a tiny tensor is storage-less by our doing: sparse and nested ones are torch's, and keep torch's own error + if torch._C._has_storage(self) or self.device.type != "tiny": return _untyped_storage(self) + if getattr(x:=unwrap(self), "_meta_storage", None) is None: x._meta_storage = torch.UntypedStorage(self.nbytes, device="meta") + return x._meta_storage +torch.Tensor.untyped_storage = untyped_storage + # track view relationships for in place operations def canonical_base(view: Tensor): return getattr(view, "_view_base", view) def derived_views(base: Tensor): return [t for tref in getattr(base, "_views", set()) if (t:=tref()) is not None] def unwrap_args(args, kwargs): return [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args], {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()} +# record that `view` is `base` with `ops` applied, so a later in place write to either goes through the base +def register_view(base: Tensor, view: Tensor, ops) -> Tensor: + view._view_base, view._view_ops = base, ops + base._views = getattr(base, "_views", set()) + base._views.add(weakref.ref(view)) + return view + def wrap_view_op(fn): @functools.wraps(fn) def _wrap(*args, **kwargs): args, kwargs = unwrap_args(args, kwargs) ret = fn(*args, **kwargs) - base = canonical_base(args[0]) - ret._view_base = base - base._views = getattr(base, "_views", set()) - base._views.add(weakref.ref(ret)) - ret._view_ops = _get_view_ops(args[0]) + [(fn, args[1:], kwargs)] - return wrap(ret) + return wrap(register_view(canonical_base(args[0]), ret, _get_view_ops(args[0]) + [(fn, args[1:], kwargs)])) return _wrap # NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index @@ -349,6 +363,8 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic from torch._decomp import get_decompositions decomps = [ aten.native_layer_norm_backward, + # functionalizing a training BatchNorm under torch.compile gives this instead of aten.native_batch_norm + aten._native_batch_norm_legit_functional, aten.native_group_norm_backward, aten.linalg_cross, aten.addmm, diff --git a/extra/torch_backend/compile.py b/extra/torch_backend/compile.py new file mode 100644 index 0000000000000..7cd923bf860b1 --- /dev/null +++ b/extra/torch_backend/compile.py @@ -0,0 +1,52 @@ +# a dynamo backend that lowers the AOT-autograd FX graph into a TinyJit, so a compiled function replays kernels +# instead of dispatching op-by-op. importing this makes a bare torch.compile use it +import torch +from torch._dynamo.backends.registry import register_backend, set_default_backend +from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_func +from tinygrad import Tensor, TinyJit +from extra.torch_backend.backend import wrap, unwrap, canonical_base, register_view, _get_view_ops + +# TinyJit is static shape, and a dim dynamo marked dynamic arrives as a SymInt with no Tensor to bind. NOTE: process wide +torch._dynamo.config.automatic_dynamic_shapes = False + +def _jit_input(x:torch.Tensor) -> Tensor: + # a JIT input has to be a real buffer: .tiny() brings other devices over, .clone() gives a deviceless const a device + t = unwrap(x.tiny()) + return t if t.device is not None else t.replace(t.clone(t._torch_device)) + +def _copy_out(outs:list[Tensor|None]) -> list[Tensor|None]: + # copy every output, recording a returned view as a view of its base's copy. sharing one buffer instead would hand the + # next graph two JIT inputs on it, which jit.py rejects. one realize for all: per-output is a graph rewrite each + ret:list[Tensor|None] = [None if t is None else t.clone() for t in outs] + copies = {t: c for t, c in zip(outs, ret) if t is not None and t is canonical_base(t)} + for t, c in zip(outs, ret): + if t is None or t is (b:=canonical_base(t)): continue + if (base_copy:=copies.get(b)) is not None: register_view(base_copy, c, _get_view_ops(t)) + if (real:=[t for t in ret if t is not None]): Tensor.realize(*real) + return ret + +def _tiny_compiler(gm:torch.fx.GraphModule, sample_inputs): + # dynamo emits op-free graphs at training graph breaks and TinyJit cannot capture one, so run it eagerly + if not any(n.op in ("call_function", "call_method", "call_module") for n in gm.graph.nodes): + return make_boxed_func(torch.compiler.disable(gm)) + # the JIT hands back the buffers it captured: copy inside so a passthrough output is written on this call, copy outside + # so a retained output survives the next. a backward graph has a None output per input that wanted no grad + @TinyJit + def jitted(*args:Tensor): return _copy_out([x if x is None else unwrap(x) for x in gm(*[wrap(a) for a in args])]) + # this runs under an active dynamo, which would otherwise trace the backend itself + @torch.compiler.disable + def run(*args:torch.Tensor): + return [t if t is None else wrap(t) for t in _copy_out(jitted(*map(_jit_input, args)))] + return make_boxed_func(run) + +# mode= and options= are inductor knobs torch hands to any named backend, drop them like torch's own aot_autograd does +@register_backend +def tiny(gm:torch.fx.GraphModule, sample_inputs, **_ignored): + # a non-Tensor input is a SymInt for a dynamic dim. NOTE: never format a tiny tensor in here, repr() of one runs a + # float64 reduction that METAL cannot compile + if not all(isinstance(x, torch.Tensor) for x in sample_inputs): + raise RuntimeError(f"the tiny backend needs static shapes, got {[type(x).__name__ for x in sample_inputs]}") + return aot_module_simplified(gm, sample_inputs, fw_compiler=_tiny_compiler) + +# torch.compile defaults to inductor, which cannot run a tiny tensor. NOTE: process wide, like the config above +set_default_backend("tiny") diff --git a/extra/torch_backend/torch_tests.py b/extra/torch_backend/torch_tests.py index cd8505c094371..cafde424dc7d2 100644 --- a/extra/torch_backend/torch_tests.py +++ b/extra/torch_backend/torch_tests.py @@ -3,8 +3,9 @@ import torch import extra.torch_backend.backend -from torch.testing._internal.common_utils import TestCase, is_privateuse1_backend_available -assert is_privateuse1_backend_available() and torch._C._get_privateuse1_backend_name() == "tiny" +from torch.testing._internal.common_utils import TestCase +# NOTE: this is what torch's own is_privateuse1_backend_available did, inlined: it went private in torch 2.13 +assert torch.tiny.is_available() and torch._C._get_privateuse1_backend_name() == "tiny" from torch.testing._internal.common_device_type import ops, onlyOn, instantiate_device_type_tests from torch.testing._internal.common_methods_invocations import unary_ufuncs, binary_ufuncs, reduction_ops, shape_funcs diff --git a/pyproject.toml b/pyproject.toml index a039dcbffbd02..3fbb1755e99f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,7 @@ linting = [ # ] testing_minimal = [ "numpy", - "torch==2.9.1", + "torch==2.13.0", "pytest", "pytest-xdist", "pytest-split",