diff --git a/test/test_multitensor.py b/test/test_multitensor.py index ff21906016..2de7052e66 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -75,7 +75,7 @@ class TestMultiTensor(unittest.TestCase): ei.run() assert names[-2] == names[-1], "function was relinearized" - @unittest.skip("this doesn't fold because from_sharded calls contiguous on all lbs") + @unittest.skip("this doesn't fold because shard_ calls contiguous on all lbs") def test_sharded_memory(self): # Buffer may be stuck in track_cross_buffer for x in (d0, d1, d2, d3, d4): Device[x].synchronize() diff --git a/tinygrad/multi.py b/tinygrad/multi.py index ab521e9978..8e65f46ac9 100644 --- a/tinygrad/multi.py +++ b/tinygrad/multi.py @@ -1,6 +1,6 @@ from __future__ import annotations import functools, itertools, operator -from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv, ceildiv +from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv from tinygrad.dtype import DType from tinygrad.ops import Ops, MathTrait, UOp, sint @@ -63,18 +63,6 @@ class MultiLazyBuffer(MathTrait): def __repr__(self): return f"" - @staticmethod - def from_sharded(lb:UOp, devices:tuple[str, ...], axis:int|None): - if axis is not None: - if not isinstance(total:=lb.shape[axis], int): raise RuntimeError(f"cannot shard symbolic shape {lb.shape=}, {axis=}") - sz = ceildiv(total, len(devices)) - splits = tuple([max(0, min(sz, total - sz*i)) for i in range(len(devices))]) - bounds = tuple(itertools.pairwise(itertools.accumulate(splits, initial=0))) - lbs = [lb] * len(devices) - sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(to_sharded(lbs, axis, bounds) if axis is not None and bounds is not None else lbs, devices)] - # NOTE: this contiguous is making it impossible for the scheduler to do late const folding - return MultiLazyBuffer([lb.contiguous(allow_buffer_view=False) for lb in sharded_lbs], axis) - def copy_to_device(self, device:str) -> UOp: if self.axis is None: # if we already have a copy on the device, return that diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 53ebb9b86c..71fe8b5693 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -179,7 +179,7 @@ class Tensor(SimpleMathTrait): # data might be on a different device if isinstance(device, str): self.lazydata:Union[UOp, MultiLazyBuffer] = data if data.device == device else data.copy_to_device(device) # if device is a tuple, we should have/construct a MultiLazyBuffer - elif isinstance(data, UOp): self.lazydata = MultiLazyBuffer.from_sharded(data, device, None) + elif isinstance(data, UOp): self.lazydata = Tensor(data).shard(device).lazydata else: assert data.device == device, f"MultiLazyBuffer device mismatch, {data.device} != {device}" self.lazydata = data @@ -405,8 +405,16 @@ class Tensor(SimpleMathTrait): """ assert isinstance(self.lazydata, UOp), "can't shard a MultiLazyBuffer" devices = tuple(Device.canonicalize(x) for x in devices) - if axis is not None: axis = self._resolve_dim(axis) - return Tensor(MultiLazyBuffer.from_sharded(self.lazydata, devices, axis), device=devices, requires_grad=self.requires_grad) + if axis is None: lbs = [self.lazydata] * len(devices) + else: + axis = self._resolve_dim(axis) + sz = ceildiv(self.shape[axis], len(devices)) + sizes = [max(0, min(sz, self.shape[axis] - sz*i)) for i in range(len(devices))] + lbs = [cast(UOp, t.lazydata) for t in self.split(sizes, axis)] + sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(lbs, devices)] + # NOTE: this contiguous is making it impossible for the scheduler to do late const folding + mlb = MultiLazyBuffer([lb.contiguous(allow_buffer_view=False) for lb in sharded_lbs], axis) + return Tensor(mlb, device=devices, requires_grad=self.requires_grad) def shard_(self, devices:tuple[str, ...], axis:Optional[int]=None): """