Anastasiia Filippova
347 lines
11 KiB
Python
347 lines
11 KiB
Python
# Copyright © 2024 Apple Inc.
|
|
|
|
import mlx.core as mx
|
|
import mlx.optimizers as optim
|
|
import mlx_distributed_tests
|
|
import mlx_tests
|
|
from mlx.nn.utils import average_gradients, fsdp_apply_gradients
|
|
|
|
|
|
class TestNCCLDistributed(mlx_distributed_tests.MLXDistributedCommonTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
_ = mx.distributed.init(strict=True, backend="nccl")
|
|
cls.atol = 1e-4
|
|
cls.rtol = 1e-4
|
|
|
|
def test_sum_scatter(self):
|
|
|
|
world = mx.distributed.init()
|
|
|
|
dtypes = [
|
|
(mx.float32, 1e-6),
|
|
(mx.float16, 5e-3),
|
|
(mx.bfloat16, 1e-1),
|
|
]
|
|
sizes = [
|
|
(8,),
|
|
(64,),
|
|
(1024,),
|
|
(1024, 1024),
|
|
]
|
|
key = mx.random.key(world.rank())
|
|
|
|
for dt, rtol in dtypes:
|
|
for sh in sizes:
|
|
x = (mx.random.uniform(shape=sh, key=key) * 10).astype(dt) # shape=sh
|
|
|
|
# Sum scatter
|
|
y = mx.distributed.sum_scatter(x) # shape=sh/world.size()
|
|
z = mx.distributed.all_sum(x) # shape=sh
|
|
chunk = sh[0] // world.size()
|
|
start = world.rank() * chunk
|
|
stop = start + chunk
|
|
z_ref = z[start:stop]
|
|
|
|
maxrelerror = (y - z_ref).abs()
|
|
if rtol > 0:
|
|
maxrelerror /= z_ref.abs()
|
|
maxrelerror = maxrelerror.max()
|
|
self.assertLessEqual(maxrelerror, rtol)
|
|
|
|
def test_groups(self):
|
|
world = mx.distributed.init()
|
|
self.assertEqual(world.size(), 8)
|
|
self.assertTrue(0 <= world.rank() < 8)
|
|
|
|
world2 = mx.distributed.init()
|
|
self.assertEqual(world.size(), world2.size())
|
|
self.assertEqual(world.rank(), world2.rank())
|
|
|
|
sub = world.split(world.rank() % 2)
|
|
self.assertEqual(sub.size(), 4)
|
|
self.assertEqual(sub.rank(), world.rank() // 2)
|
|
|
|
sub = world.split(world.rank() // 2)
|
|
self.assertEqual(sub.size(), 2)
|
|
|
|
def test_all_reduce_split(self):
|
|
world = mx.distributed.init()
|
|
dtypes = [
|
|
(mx.float32, 1e-6),
|
|
(mx.float16, 5e-3),
|
|
(mx.bfloat16, 1e-1),
|
|
]
|
|
sizes = [
|
|
(7,),
|
|
(10,),
|
|
(1024,),
|
|
(1024, 1024),
|
|
]
|
|
key = mx.random.key(0)
|
|
group = world.split(world.rank() % 2)
|
|
|
|
for dt, rtol in dtypes:
|
|
for sh in sizes:
|
|
x = (
|
|
mx.random.uniform(shape=(group.size(),) + sh, key=key) * 10
|
|
).astype(dt)
|
|
|
|
# All sum
|
|
y = mx.distributed.all_sum(x[group.rank()], group=group)
|
|
z = x.sum(0)
|
|
maxrelerror = (y - z).abs()
|
|
if rtol > 0:
|
|
maxrelerror /= z.abs()
|
|
maxrelerror = maxrelerror.max()
|
|
self.assertLessEqual(maxrelerror, rtol)
|
|
|
|
# All max
|
|
y = mx.distributed.all_max(x[group.rank()], group=group)
|
|
z = x.max(0)
|
|
self.assertTrue(mx.all(y == z))
|
|
|
|
# All min
|
|
y = mx.distributed.all_min(x[group.rank()], group=group)
|
|
z = x.min(0)
|
|
self.assertTrue(mx.all(y == z))
|
|
|
|
def test_all_gather_split(self):
|
|
world = mx.distributed.init()
|
|
dtypes = [mx.float32, mx.float16, mx.bfloat16]
|
|
sub = world.split(world.rank() % 2)
|
|
for dt in dtypes:
|
|
x = mx.ones((2, 2, 4), dtype=dt)
|
|
y = mx.distributed.all_gather(x, group=sub)
|
|
self.assertEqual(y.shape, (sub.size() * 2, 2, 4))
|
|
self.assertTrue(mx.all(y == 1))
|
|
|
|
def test_fsdp_apply_gradients(self):
|
|
world = mx.distributed.init()
|
|
N = world.size()
|
|
|
|
params = {
|
|
"w1": mx.ones((N * 10, 8)),
|
|
"w2": mx.ones((N * 20,)),
|
|
}
|
|
grads = {
|
|
"w1": mx.ones((N * 10, 8)) * 0.1,
|
|
"w2": mx.ones((N * 20,)) * 0.1,
|
|
}
|
|
|
|
optimizer = optim.SGD(learning_rate=0.1)
|
|
updated_params_fsdp = fsdp_apply_gradients(grads, params, optimizer)
|
|
mx.eval(updated_params_fsdp)
|
|
|
|
self.assertEqual(updated_params_fsdp["w1"].shape, (N * 10, 8))
|
|
self.assertEqual(updated_params_fsdp["w2"].shape, (N * 20,))
|
|
|
|
self.assertTrue(
|
|
mx.allclose(
|
|
updated_params_fsdp["w1"], mx.ones((N * 10, 8)) * 0.99, atol=1e-6
|
|
)
|
|
)
|
|
self.assertTrue(
|
|
mx.allclose(updated_params_fsdp["w2"], mx.ones((N * 20,)) * 0.99, atol=1e-6)
|
|
)
|
|
|
|
grads = {
|
|
"w1": mx.ones((N * 10, 8)) * 10.0,
|
|
"w2": mx.ones((N * 20,)) * 10.0,
|
|
}
|
|
|
|
new_params_clipped, grad_norm = fsdp_apply_gradients(
|
|
grads, params, optimizer, max_norm=1.0
|
|
)
|
|
mx.eval(new_params_clipped, grad_norm)
|
|
|
|
self.assertIsNotNone(grad_norm)
|
|
expected_norm = mx.sqrt((N * 10 * 8 + N * 20) * 100.0)
|
|
self.assertTrue(mx.allclose(grad_norm, expected_norm, atol=1e-4, rtol=1e-4))
|
|
self.assertEqual(new_params_clipped["w1"].shape, (N * 10, 8))
|
|
self.assertEqual(new_params_clipped["w2"].shape, (N * 20,))
|
|
|
|
scale = 1.0 / expected_norm
|
|
expected_update = 1.0 - 0.1 * 10.0 * scale
|
|
self.assertTrue(
|
|
mx.allclose(
|
|
new_params_clipped["w1"],
|
|
mx.ones((N * 10, 8)) * expected_update,
|
|
atol=1e-4,
|
|
rtol=1e-4,
|
|
)
|
|
)
|
|
self.assertTrue(
|
|
mx.allclose(
|
|
new_params_clipped["w2"],
|
|
mx.ones((N * 20,)) * expected_update,
|
|
atol=1e-4,
|
|
rtol=1e-4,
|
|
)
|
|
)
|
|
params = {"w": mx.ones((N * 4,))}
|
|
grads = {"w": mx.ones((N * 4,)) * 0.5}
|
|
|
|
optimizer_fsdp = optim.SGD(learning_rate=0.1)
|
|
updated_params_fsdp = fsdp_apply_gradients(grads, params, optimizer_fsdp)
|
|
|
|
optimizer_ddp = optim.SGD(learning_rate=0.1)
|
|
avg_grads = average_gradients(grads)
|
|
updated_params_ddp = optimizer_ddp.apply_gradients(avg_grads, params)
|
|
mx.eval(updated_params_ddp, updated_params_fsdp)
|
|
|
|
self.assertTrue(
|
|
mx.allclose(
|
|
updated_params_fsdp["w"], updated_params_ddp["w"], atol=1e-6, rtol=1e-4
|
|
),
|
|
)
|
|
|
|
def test_fsdp_ddp_apply_gradients(self):
|
|
world = mx.distributed.init()
|
|
N = world.size()
|
|
S = 4
|
|
fsdp_group = world.split(world.rank() // S)
|
|
dp_group = world.split(world.rank() % S)
|
|
|
|
self.assertEqual(fsdp_group.size(), S)
|
|
self.assertEqual(dp_group.size(), N // S)
|
|
|
|
params = {
|
|
"w1": mx.ones((S * 10, 8)),
|
|
"w2": mx.ones((S * 20,)),
|
|
}
|
|
grads = {
|
|
"w1": mx.ones((S * 10, 8)) * 0.1,
|
|
"w2": mx.ones((S * 20,)) * 0.1,
|
|
}
|
|
|
|
optimizer = optim.SGD(learning_rate=0.1)
|
|
updated = fsdp_apply_gradients(
|
|
grads,
|
|
params,
|
|
optimizer,
|
|
fsdp_group=fsdp_group,
|
|
dp_group=dp_group,
|
|
)
|
|
mx.eval(updated)
|
|
|
|
self.assertEqual(updated["w1"].shape, (S * 10, 8))
|
|
self.assertEqual(updated["w2"].shape, (S * 20,))
|
|
|
|
self.assertTrue(
|
|
mx.allclose(updated["w1"], mx.ones((S * 10, 8)) * 0.99, atol=1e-6)
|
|
)
|
|
self.assertTrue(
|
|
mx.allclose(updated["w2"], mx.ones((S * 20,)) * 0.99, atol=1e-6)
|
|
)
|
|
|
|
grads_big = {
|
|
"w1": mx.ones((S * 10, 8)) * 10.0,
|
|
"w2": mx.ones((S * 20,)) * 10.0,
|
|
}
|
|
|
|
optimizer2 = optim.SGD(learning_rate=0.1)
|
|
clipped, grad_norm = fsdp_apply_gradients(
|
|
grads_big,
|
|
params,
|
|
optimizer2,
|
|
fsdp_group=fsdp_group,
|
|
dp_group=dp_group,
|
|
max_norm=1.0,
|
|
)
|
|
mx.eval(clipped, grad_norm)
|
|
|
|
self.assertIsNotNone(grad_norm)
|
|
expected_norm = mx.sqrt((S * 10 * 8 + S * 20) * 100.0)
|
|
self.assertTrue(mx.allclose(grad_norm, expected_norm, atol=1e-4, rtol=1e-4))
|
|
self.assertEqual(clipped["w1"].shape, (S * 10, 8))
|
|
self.assertEqual(clipped["w2"].shape, (S * 20,))
|
|
|
|
scale = 1.0 / expected_norm
|
|
expected_update = 1.0 - 0.1 * 10.0 * scale
|
|
self.assertTrue(
|
|
mx.allclose(
|
|
clipped["w1"],
|
|
mx.ones((S * 10, 8)) * expected_update,
|
|
atol=1e-4,
|
|
rtol=1e-4,
|
|
)
|
|
)
|
|
self.assertTrue(
|
|
mx.allclose(
|
|
clipped["w2"],
|
|
mx.ones((S * 20,)) * expected_update,
|
|
atol=1e-4,
|
|
rtol=1e-4,
|
|
)
|
|
)
|
|
|
|
params_eq = {"w": mx.ones((S * 4,))}
|
|
grads_eq = {"w": mx.ones((S * 4,)) * 0.5}
|
|
|
|
optimizer_hybrid = optim.SGD(learning_rate=0.1)
|
|
updated_hybrid = fsdp_apply_gradients(
|
|
grads_eq,
|
|
params_eq,
|
|
optimizer_hybrid,
|
|
fsdp_group=fsdp_group,
|
|
dp_group=dp_group,
|
|
)
|
|
|
|
optimizer_ddp = optim.SGD(learning_rate=0.1)
|
|
avg_grads = average_gradients(grads_eq)
|
|
updated_ddp = optimizer_ddp.apply_gradients(avg_grads, params_eq)
|
|
mx.eval(updated_hybrid, updated_ddp)
|
|
|
|
self.assertTrue(
|
|
mx.allclose(updated_hybrid["w"], updated_ddp["w"], atol=1e-6, rtol=1e-4),
|
|
)
|
|
|
|
def test_fsdp_peak_memory(self):
|
|
world = mx.distributed.init()
|
|
N = world.size()
|
|
mx.random.seed(42)
|
|
params = {
|
|
"w1": mx.random.normal((N * 1024, 1024)),
|
|
"w2": mx.random.normal((N * 2048, 512)),
|
|
}
|
|
grads = {
|
|
"w1": mx.random.normal((N * 1024, 1024)),
|
|
"w2": mx.random.normal((N * 2048, 512)),
|
|
}
|
|
mx.eval(params, grads)
|
|
optimizer_ddp = optim.Adam(learning_rate=0.01)
|
|
optimizer_fsdp = optim.Adam(learning_rate=0.01)
|
|
|
|
def pseudo_step_ddp(grads, params, optimizer):
|
|
grads = average_gradients(grads)
|
|
grads, grad_norm = optim.clip_grad_norm(grads, max_norm=1.0)
|
|
params = optimizer.apply_gradients(grads, params)
|
|
return grad_norm, params
|
|
|
|
def pseudo_step_fsdp(grads, params, optimizer):
|
|
params, grad_norm = fsdp_apply_gradients(
|
|
grads, params, optimizer, max_norm=1.0
|
|
)
|
|
return grad_norm, params
|
|
|
|
mx.reset_peak_memory()
|
|
|
|
for i in range(10):
|
|
grad_norm, params = pseudo_step_ddp(grads, params, optimizer_ddp)
|
|
mx.eval(grad_norm, params)
|
|
|
|
ddp_peak_memory = mx.get_peak_memory()
|
|
mx.reset_peak_memory()
|
|
|
|
for i in range(10):
|
|
grad_norm, params = pseudo_step_fsdp(grads, params, optimizer_fsdp)
|
|
mx.eval(grad_norm, params)
|
|
|
|
fsdp_peak_memory = mx.get_peak_memory()
|
|
self.assertTrue(fsdp_peak_memory < ddp_peak_memory)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mlx_tests.MLXTestRunner()
|