mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-08-24 01:33:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

The smallest promotion

Browse Source
This commit is contained in:
firestar5683
2026-08-10 22:53:46 -05:00
parent d938c9daa6
commit 6bd48c28e8
582 changed files with 64651 additions and 17100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad_repo/test/unit/test_multitensor.py
+57 -18
View File
@@ -2,7 +2,7 @@ import unittest, numpy as np
from tinygrad import Tensor, Variable, Context, Device, TinyJit, GlobalCounters, dtypes, UOp, nn, getenv
from tinygrad.nn.state import get_parameters, get_state_dict
from tinygrad.uop.ops import Ops
from test.helpers import not_support_multi_device, needs_second_gpu, slow
from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count, KernelCountException
from hypothesis import given, strategies as strat, settings
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -60,8 +60,15 @@ class TestMultiTensor(unittest.TestCase):
def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]])
def test_alu_deviceless_const(self):
s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0))).numpy(), [2, 3, 4, 5])
np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float))).numpy(), [2, 3, 4, 5])
np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
def test_add_rank_expand_shard(self):
# a sharded src keeps its own rank under implicit broadcast, its shard axis right-aligns into the output
a = Tensor([1.,2.,3.,4.]).shard(devices_2, 0)
b = Tensor([[10.,20.,30.,40.]]).shard(devices_2, None)
self.assertEqual((a+b).uop.axis, 1)
np.testing.assert_equal((a+b).numpy(), [[11.,22.,33.,44.]])
def test_shard_reduce(self):
self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6)
@@ -72,14 +79,19 @@ class TestMultiTensor(unittest.TestCase):
with self.assertRaises(RuntimeError):
X.shard_(devices_3, 0)
def test_shard_reshape_cross_boundary(self):
X = Tensor.ones(5, 4).contiguous().realize().shard(devices_2, 1)
with self.assertRaises(RuntimeError): X.reshape(10, 2).uop.axis
def test_tensor_from_multi(self):
X = Tensor([1, 2], dtype=dtypes.int).shard_(devices_2, 0)
Y = Tensor(X.uop)
self.assertEqual(Y.device, devices_2)
np.testing.assert_equal(X.numpy(), Y.numpy())
with self.assertRaises(AssertionError):
_ = Tensor(X.uop, dtype=dtypes.float)
Z = Tensor(X.uop, dtype=dtypes.float)
self.assertEqual(Z.dtype, dtypes.float)
np.testing.assert_equal(Z.numpy(), [1.0, 2.0])
def test_sharded_arange(self):
sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0)
@@ -115,6 +127,27 @@ class TestMultiTensor(unittest.TestCase):
with Context(RING=use_ring):
np.testing.assert_equal(t.shard(devices_2, axis=axis).sum().item(), 10)
def test_allreduce_cast_half(self, assign=False, kernel_count=8):
devices = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
a_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
b_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
a = a_src.shard(devices, axis=0).realize()
b = b_src.shard(devices, axis=0).realize()
# assigning creates a copy of the output before allreduce
if assign:
tst = Tensor.empty_like(b)
tst.assign(a + b)
else:
tst = a + b
tst = tst.float().sum(0)
GlobalCounters.reset()
with Context(ALLREDUCE_CAST=1, RING=0, ALL2ALL=0):
tst.realize()
assert_kernel_count(kernel_count)
np.testing.assert_allclose(tst.numpy(), (a_src.numpy()+b_src.numpy()).sum(0))
def test_allreduce_cast_half_assign(self): self.test_allreduce_cast_half(assign=True, kernel_count=10)
def test_multiple_to_single_device(self):
kernel_counts = {}
for ring in (0, 2):
@@ -207,7 +240,7 @@ class TestMultiTensor(unittest.TestCase):
out.numpy()
def test_backprop_conv(self):
with Tensor.train():
with Context(TRAINING=1):
conv = nn.Conv2d(3, 16, 3)
for p in get_parameters(conv): p.shard_(devices_2)
optim = nn.optim.Adam(get_parameters(conv))
@@ -351,6 +384,12 @@ class TestMultiTensor(unittest.TestCase):
np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
assert jf.captured is not None
def test_symbolic_broadcast_copy(self):
rows = Variable("rows", 1, 4).bind(3)
out = Tensor.ones(rows, 8).to(devices_2).realize()
self.assertEqual(out.shape, (rows, 8))
np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8)))
def test_multitensor_jit_in_list(self):
# test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
@TinyJit
@@ -511,7 +550,7 @@ class TestMultiTensor(unittest.TestCase):
def test_full_like_on_shard_axis(self): self.test_full_like_on_shard(0)
def test_dropout_on_shard(self):
with Tensor.train():
with Context(TRAINING=1):
X = Tensor.ones(256).to(devices_2)
output = X.dropout(0.5).numpy()
unique, counts = np.unique(output, return_counts=True)
@@ -519,7 +558,7 @@ class TestMultiTensor(unittest.TestCase):
assert 96 < counts[0] < 160, counts[0]
def test_dropout_on_shard_axis(self):
with Tensor.train():
with Context(TRAINING=1):
X = Tensor.ones(512).shard(devices_2, axis=0)
output = X.dropout(0.5).numpy()
unique, counts = np.unique(output, return_counts=True)
@@ -550,7 +589,7 @@ class TestMultiTensor(unittest.TestCase):
zeros = Tensor.zeros(3).realize()
b = a.to(devices_2)*zeros.to(devices_2)
sched = b.schedule_linear().src
self.assertEqual(len(sched), 0)
if len(sched) != 0: raise KernelCountException(0, len(sched))
self.assertListEqual(b.tolist(), [0, 0, 0])
@unittest.skipIf(not_support_multi_device(), "no multi")
@@ -563,7 +602,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
t = Tensor.arange(64).reshape(8, 8).clone().realize()
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
with self.assertRaises(AssertionError):
with self.assertRaises(RuntimeError):
# sharded axis shrink on non-device boundry is not allowed
a = t.shrink(((0, 3), (0, 8))).contiguous()
a.schedule_linear()
@@ -585,7 +624,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
t = Tensor.arange(64).reshape(8, 8).clone().realize()
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
for i in range(4):
for i in range(2):
print(f"{i=}")
a = t.shrink(((0+2*i,2+2*i),None))
b = Tensor(t.numpy()[0+2*i:2+2*i])
@@ -601,8 +640,8 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
np.testing.assert_allclose((a+a).numpy(), (b+b).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_equal((a+1).numpy(), (b+1).numpy())
np.testing.assert_equal((1+a).numpy(), (1+b).numpy())
np.testing.assert_allclose((a.where(a+a, a)).numpy(), (b.where(b+b, b)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose((a.where(1, 0)).numpy(), (b.where(1, 0)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose((a.bool().where(a+a, a)).numpy(), (b.bool().where(b+b, b)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose((a.bool().where(1, 0)).numpy(), (b.bool().where(1, 0)).numpy(), rtol=1e-7, atol=1e-3)
# reduce
np.testing.assert_allclose(a.max().numpy(), b.max().numpy(), rtol=1e-7, atol=1e-3)
@@ -664,7 +703,7 @@ class TestBatchNorm(unittest.TestCase):
def setUp(self): pass
def test_unsynced_backprop_conv_bn(self):
with Tensor.train():
with Context(TRAINING=1):
from extra.lr_scheduler import OneCycleLR
convs = [nn.Conv2d(3, 16, 3), nn.Conv2d(3, 16, 3)]
@@ -709,7 +748,7 @@ class TestBatchNorm(unittest.TestCase):
bn_ts.append(bni)
return bn_ts[0].cat(*bn_ts[1:])
with Tensor.train():
with Context(TRAINING=1):
conv = nn.Conv2d(3, 16, 3)
bn = BatchNorm(16)
@@ -731,7 +770,7 @@ class TestBatchNorm(unittest.TestCase):
from examples.hlb_cifar10 import UnsyncedBatchNorm
GPUS = (d1, d2)
with Tensor.train():
with Context(TRAINING=1):
conv = nn.Conv2d(3, 16, 3)
bn = UnsyncedBatchNorm(16, num_devices=len(GPUS))
@@ -756,7 +795,7 @@ class TestBatchNorm(unittest.TestCase):
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0)
with Tensor.train(is_training):
with Context(TRAINING=is_training):
bns = []
for _ in range(len(devices)):
bn = nn.BatchNorm2d(8)
@@ -777,7 +816,7 @@ class TestBatchNorm(unittest.TestCase):
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
x = Tensor.ones(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0)
with Tensor.train():
with Context(TRAINING=1):
synced_bn = BatchNorm2d(8)
unsynced_bn = UnsyncedBatchNorm(8, num_devices=len(devices))