import asyncio import contextlib import gc import random import time from contextlib import suppress import numpy as np import psutil import pytest import torch import hivemind import hivemind.averaging.averager from hivemind import MPFuture, get_logger from hivemind.averaging.allreduce import AveragingMode from hivemind.averaging.key_manager import GroupKeyManager from hivemind.averaging.load_balancing import load_balance_peers from hivemind.p2p import PeerID from hivemind.proto.runtime_pb2 import CompressionType from test_utils.dht_swarms import launch_dht_instances logger = get_logger(__name__) @contextlib.contextmanager def cleanup_children(): yield gc.collect() # Call .__del__() for removed objects children = psutil.Process().children(recursive=True) if children: logger.info(f"Cleaning up {len(children)} leftover child processes") for child in children: with suppress(psutil.NoSuchProcess): child.terminate() psutil.wait_procs(children, timeout=1) # Broken code or killing of child processes may leave the MPFuture backend corrupted MPFuture.reset_backend() @pytest.mark.forked @pytest.mark.asyncio async def test_key_manager(): dht = hivemind.DHT(start=True) key_manager = GroupKeyManager( dht, prefix="test_averaging", initial_group_bits="10110", target_group_size=2, ) alice = dht.peer_id bob = PeerID(b"bob") t = hivemind.get_dht_time() key = key_manager.current_key await key_manager.declare_averager(key, alice, expiration_time=t + 60) await key_manager.declare_averager(key, bob, expiration_time=t + 61) q1 = await key_manager.get_averagers(key, only_active=True) await key_manager.declare_averager(key, alice, expiration_time=t + 66) q2 = await key_manager.get_averagers(key, only_active=True) await key_manager.declare_averager(key, bob, expiration_time=t + 61, looking_for_group=False) q3 = await key_manager.get_averagers(key, only_active=True) q4 = await key_manager.get_averagers(key, only_active=False) q5 = await key_manager.get_averagers("nonexistent_key.0b0101", only_active=False) assert len(q1) == 2 and (alice, t + 60) in q1 and (bob, t + 61) in q1 assert len(q2) == 2 and (alice, t + 66) in q2 and (bob, t + 61) in q2 assert len(q3) == 1 and (alice, t + 66) in q3 assert len(q4) == 2 and (alice, t + 66) in q4 and (bob, t + 61) in q2 assert len(q5) == 0 dht.shutdown() def _test_allreduce_once(n_clients, n_aux): n_peers = 4 modes = ( [AveragingMode.CLIENT] * n_clients + [AveragingMode.AUX] * n_aux + [AveragingMode.NODE] * (n_peers - n_clients - n_aux) ) random.shuffle(modes) tensors1 = [torch.randn(123), torch.zeros(3)] tensors2 = [torch.rand(123), torch.ones(3)] tensors3 = [-torch.rand(123), torch.arange(3).to(torch.float32)] tensors4 = [torch.randn(123) ** 3, torch.arange(3).to(torch.float32) / 2] peer_tensors = [tensors1, tensors2, tensors3, tensors4] reference = [ sum(tensors[i] for tensors, mode in zip(peer_tensors, modes) if mode != AveragingMode.AUX) / max(1, n_peers - n_aux) for i in range(len(tensors1)) ] dht_instances = launch_dht_instances(len(peer_tensors)) averagers = [ hivemind.averaging.DecentralizedAverager( tensors, dht=dht, target_group_size=4, averaging_expiration=15, prefix="mygroup", client_mode=mode == AveragingMode.CLIENT, auxiliary=mode == AveragingMode.AUX, start=True, ) for tensors, dht, mode in zip(peer_tensors, dht_instances, modes) ] futures = [] for averager in averagers: futures.append(averager.step(wait=False)) for future in futures: result = future.result() for averager in averagers: assert averager.endpoint in result for averager in averagers: if averager.mode != AveragingMode.AUX: with averager.get_tensors() as averaged_tensors: for ref, our in zip(reference, averaged_tensors): assert torch.allclose(ref, our, atol=1e-6) for instance in averagers + dht_instances: instance.shutdown() @pytest.mark.forked @pytest.mark.parametrize("n_clients", [0, 1, 2]) @pytest.mark.parametrize("n_aux", [0, 1, 2]) def test_allreduce_once(n_clients, n_aux): _test_allreduce_once(n_clients, n_aux) @pytest.mark.forked @pytest.mark.parametrize("n_clients, n_aux", [(0, 4), (1, 3), (0, 3)]) def test_allreduce_once_edge_cases(n_clients, n_aux): _test_allreduce_once(n_clients, n_aux) @pytest.mark.forked def test_allreduce_weighted(n_client_mode_peers: int = 2): n_peers = 4 client_modes = [True] * n_client_mode_peers + [False] * (n_peers - n_client_mode_peers) random.shuffle(client_modes) tensors1 = [torch.randn(123), torch.zeros(3)] tensors2 = [torch.rand(123), torch.ones(3)] tensors3 = [-torch.rand(123), torch.arange(3).to(torch.float32)] tensors4 = [torch.randn(123) ** 3, torch.arange(3).to(torch.float32) / 2] dht_instances = launch_dht_instances(4) averagers = [ hivemind.averaging.DecentralizedAverager( tensors, dht=dht, target_group_size=4, averaging_expiration=15, prefix="mygroup", client_mode=client_mode, start=True, ) for tensors, dht, client_mode in zip([tensors1, tensors2, tensors3, tensors4], dht_instances, client_modes) ] weights = list(map(float, np.random.rand(len(averagers)) * 10 + 0.01)) reference = [ (tensors1[i] * weights[0] + tensors2[i] * weights[1] + tensors3[i] * weights[2] + tensors4[i] * weights[3]) / sum(weights) for i in range(len(tensors1)) ] futures = [] for averager, weight in zip(averagers, weights): futures.append(averager.step(weight=weight, wait=False)) for future in futures: future.result() for future, averager in zip(futures, averagers): with averager.get_tensors() as averaged_tensors: for ref, our in zip(reference, averaged_tensors): assert torch.allclose(ref, our, atol=1e-6) for instance in averagers + dht_instances: instance.shutdown() @pytest.mark.forked def test_allreduce_compression(): """this test ensures that compression works correctly when multiple tensors have different compression types""" tensors1 = [torch.linspace(0, 500, 1000) ** 0.5, torch.randn(1000)] tensors2 = [torch.linspace(300, 800, 1000) ** 0.5, torch.randn(1000)] results = {} FLOAT16, UINT8 = CompressionType.FLOAT16, CompressionType.UNIFORM_8BIT for compression_type_pair in [(FLOAT16, FLOAT16), (FLOAT16, UINT8), (UINT8, FLOAT16), (UINT8, UINT8)]: dht_instances = launch_dht_instances(2) averager1 = hivemind.averaging.DecentralizedAverager( [x.clone() for x in tensors1], dht=dht_instances[0], compression_type=compression_type_pair, client_mode=True, target_group_size=2, prefix="mygroup", start=True, ) averager2 = hivemind.averaging.DecentralizedAverager( [x.clone() for x in tensors2], dht=dht_instances[1], compression_type=compression_type_pair, target_group_size=2, prefix="mygroup", start=True, ) for future in averager1.step(wait=False), averager2.step(wait=False): future.result() with averager1.get_tensors() as averaged_tensors: results[compression_type_pair] = averaged_tensors for instance in [averager1, averager2] + dht_instances: instance.shutdown() assert torch.allclose(results[UINT8, FLOAT16][0], results[UINT8, UINT8][0]) assert torch.allclose(results[UINT8, FLOAT16][1], results[FLOAT16, FLOAT16][1]) assert torch.allclose(results[UINT8, UINT8][1], results[FLOAT16, UINT8][1]) assert torch.allclose(results[FLOAT16, UINT8][0], results[FLOAT16, FLOAT16][0]) assert not torch.allclose(results[UINT8, FLOAT16][1], results[UINT8, UINT8][1]) assert not torch.allclose(results[UINT8, FLOAT16][0], results[FLOAT16, FLOAT16][0]) assert not torch.allclose(results[UINT8, UINT8][0], results[FLOAT16, UINT8][0]) assert not torch.allclose(results[FLOAT16, UINT8][1], results[FLOAT16, FLOAT16][1]) reference = [(tensors1[i] + tensors2[i]) / 2 for i in range(len(tensors1))] for i in range(2): assert 0 < torch.mean(torch.square(results[FLOAT16, FLOAT16][i] - reference[i])).item() <= 1e-5 assert 1e-5 < torch.mean(torch.square(results[UINT8, UINT8][i] - reference[i])).item() <= 1e-2 def compute_mean_std(averagers, unbiased=True): results = [] for averager in averagers: with averager.get_tensors() as tensors: results.append([tensor.clone() for tensor in tensors]) results_stacked_per_tensor = list(map(torch.stack, zip(*results))) means = [stack.mean(dim=0) for stack in results_stacked_per_tensor] stds = [stack.std(dim=0, unbiased=unbiased) for stack in results_stacked_per_tensor] return means, stds @pytest.mark.forked def test_allreduce_grid(): dht_instances = launch_dht_instances(8) averagers = [ hivemind.averaging.DecentralizedAverager( averaged_tensors=[torch.randn(3)], dht=dht, target_group_size=2, prefix="mygroup", initial_group_bits=bin(i // 2)[2:].rjust(2, "0"), start=True, ) for i, dht in enumerate(dht_instances) ] [means0], [stds0] = compute_mean_std(averagers) assert not torch.allclose(stds0, torch.zeros_like(stds0)) prev_means, prev_stds = means0, stds0 for i in range(5): step_futures = [averager.step(wait=False) for averager in averagers] groups = [future.result() for future in step_futures] [means], [stds] = compute_mean_std(averagers) assert torch.allclose(means, prev_means, atol=1e-6, rtol=0) assert all(len(group) == 2 for group in groups) if i <= 2: assert torch.all(torch.le(stds, prev_stds)) else: assert torch.allclose(stds, torch.zeros_like(stds), atol=1e-6, rtol=0) for averager in averagers + dht_instances: averager.shutdown() @pytest.mark.forked def test_allgather(): dht_instances = launch_dht_instances(8) averagers = [ hivemind.averaging.DecentralizedAverager( [torch.ones(1)], dht=dht, target_group_size=4, averaging_expiration=15, prefix="mygroup", initial_group_bits="000", start=True, ) for dht in dht_instances ] futures = [] for i, averager in enumerate(averagers): futures.append(averager.step(wait=False, gather=dict(batch_size=123 + i, foo="bar"))) gathered_data = [future.result() for future in futures] gathered_data_reprs = [ repr(sorted({peer_id.to_base58(): data for peer_id, data in result.items()})) for result in gathered_data ] assert len(set(gathered_data_reprs)) == 2 reference_metadata = { averager.endpoint: dict(batch_size=123 + i, foo="bar") for i, averager in enumerate(averagers) } for future in futures: gathered = future.result() assert len(gathered) == 4 for endpoint in gathered: assert gathered[endpoint] == reference_metadata[endpoint] for averager in averagers + dht_instances: averager.shutdown() def get_cost(vector_size, partitions, bandwidths): return max( (vector_size - partitions[i] + (len(partitions) - 1) * partitions[i]) / max(bandwidths[i], 1e-9) for i in range(len(partitions)) ) def check_optimality(vector_size, bandwidths, ref_partitions): partitions = list(load_balance_peers(vector_size, bandwidths)) assert get_cost(vector_size, partitions, bandwidths) <= get_cost(vector_size, ref_partitions, bandwidths) @pytest.mark.forked def test_load_balancing(): check_optimality(60, np.array([0.25, 0.25, 0.25, 0.25]), [15, 15, 15, 15]) check_optimality(1024, np.array([0.3, 0.5, 0.9]), [0, 255, 769]) check_optimality(60, np.array([0.44, 0.33, 0.22]), [42, 18, 0]) check_optimality(60, np.array([0.55, 0.44, 0.40]), [35, 16, 9]) check_optimality(1024 * 1024, np.array([0.3, 0.5, 0.9, 0.6]), [0, 169327, 602629, 276620]) check_optimality(1024 * 1024, np.array([0.0, 0.5, 0.0, 0.6]), [0, 428963, 0, 619613]) assert load_balance_peers(60, np.array([0.55, 0.44, 0.40]), min_size=10) == (41, 19, 0) assert load_balance_peers(60, np.array([0.32, 0.55, 0.44]), min_size=10) == (0, 40, 20) assert load_balance_peers(2, np.array([0.55, 0.20, 0.44]), min_size=10) == (1, 0, 1) assert load_balance_peers(1, np.array([0.55, 0.20, 0.44]), min_size=10) == (1, 0, 0) assert load_balance_peers(100, (None, None)) == (50, 50) assert load_balance_peers(100, (None, None, None, None, None)) == (20, 20, 20, 20, 20) assert load_balance_peers(100, (0, 0, 0, None, None)) == (0, 0, 0, 50, 50) with pytest.raises(AssertionError): load_balance_peers(100, (0, 0, 0)) for i in range(10): vector_size = np.random.randint(1, 1024 ** 3) num_peers = np.random.randint(1, 256) scale = 1e-9 + np.random.rand() * 1e5 bandwidths = np.random.rand(num_peers) * scale + 1e-6 min_size = np.random.choice([0, np.random.randint(0, vector_size // 10)]) assignment = load_balance_peers(vector_size, bandwidths, min_size) assert np.sum(assignment) == vector_size assert np.min(assignment) >= 0 @pytest.mark.forked def test_too_few_peers(): dht_instances = launch_dht_instances(4) averagers = [ hivemind.averaging.DecentralizedAverager( averaged_tensors=[torch.randn(3)], dht=dht, target_group_size=2, averaging_expiration=1, request_timeout=0.5, prefix="mygroup", initial_group_bits=bin(i)[2:].rjust(3, "0"), start=True, ) for i, dht in enumerate(dht_instances) ] step_futures = [averager.step(wait=False) for averager in averagers] for future in step_futures: assert len(future.result()) == 2 for averager in averagers + dht_instances: averager.shutdown() @pytest.mark.skip( reason="The current implementation of elasticity (multi-stage averaging for the case when " "num_peers > ~3 * target_group_size) is incorrect (TODO @justheuristic)" ) @pytest.mark.forked def test_overcrowded(num_peers=16): dht_instances = launch_dht_instances(num_peers) averagers = [ hivemind.averaging.DecentralizedAverager( averaged_tensors=[torch.randn(3)], dht=dht, target_group_size=2, averaging_expiration=1, request_timeout=0.5, prefix="mygroup", initial_group_bits="", start=True, ) for dht in dht_instances ] for _ in range(5): step_futures = [averager.step(wait=False, timeout=5) for averager in averagers] assert sum(len(future.result() or []) == 2 for future in step_futures) >= len(averagers) - 1 for averager in averagers + dht_instances: averager.shutdown() @pytest.mark.forked def test_load_state_from_peers(): num_calls = 0 super_metadata = dict(x=123) super_tensors = (torch.randn(3), torch.randint(0, 5, (3,))) class TestAverager(hivemind.averaging.DecentralizedAverager): def get_current_state(self): """ Get current state and send it to a peer. executed in the host process. Meant to be overriden. :returns: a tuple of (serializable_small_metadata, sequence of torch tensors) """ nonlocal num_calls, super_metadata, super_tensors num_calls += 1 return super_metadata, super_tensors dht_instances = launch_dht_instances(2) averager1 = TestAverager( [torch.randn(3), torch.rand(5)], dht=dht_instances[0], start=True, prefix="demo-run", target_group_size=2, ) dht_instances[1].get("demo-run.all_averagers") averager2 = TestAverager( [torch.randn(3), torch.rand(5)], dht=dht_instances[1], start=True, prefix="demo-run", target_group_size=2, ) assert num_calls == 0 got_metadata, got_tensors = averager2.load_state_from_peers() assert num_calls == 1 assert got_metadata == super_metadata assert all(map(torch.allclose, got_tensors, super_tensors)) super_metadata["y"] = 123 super_tensors[1][2] = 9 assert num_calls == 1 assert got_metadata != super_metadata assert not all(map(torch.allclose, got_tensors, super_tensors)) got_metadata, got_tensors = averager2.load_state_from_peers() assert num_calls == 2 assert got_metadata == super_metadata assert all(map(torch.allclose, got_tensors, super_tensors)) averager1.allow_state_sharing = False assert averager2.load_state_from_peers() is None averager1.allow_state_sharing = True got_metadata, got_tensors = averager2.load_state_from_peers() assert num_calls == 3 assert got_metadata == super_metadata for instance in [averager1, averager2] + dht_instances: instance.shutdown() @pytest.mark.forked def test_getset_bits(): dht = hivemind.DHT(start=True) averager = hivemind.averaging.DecentralizedAverager( [torch.randn(3)], dht=dht, start=True, prefix="test_prefix", target_group_size=2, ) averager.set_group_bits("00101011101010") assert averager.get_group_bits() == "00101011101010" @pytest.mark.forked def test_training_averager(n_steps: int = 10, n_dims: int = 16): torch.manual_seed(42) dht_instances = launch_dht_instances(2) common_kwargs = { "start": True, "prefix": "demo-run", "target_group_size": 2, } x1 = torch.randn(n_dims, requires_grad=True) opt1 = torch.optim.Adam([x1], lr=0.05) averager1 = hivemind.averaging.TrainingAverager( opt1, average_gradients=True, average_parameters=True, average_opt_statistics=["exp_avg_sq"], dht=dht_instances[0], **common_kwargs ) x2 = torch.randn(n_dims, requires_grad=True) opt2 = torch.optim.Adam([x2], lr=0.05) averager2 = hivemind.averaging.TrainingAverager( opt2, average_gradients=True, average_parameters=True, average_opt_statistics=["exp_avg_sq"], dht=dht_instances[1], **common_kwargs ) a = torch.ones(n_dims) for i in range(n_steps): opt1.zero_grad() opt2.zero_grad() (x1 - a).pow(2).sum().backward() (x2 - a).pow(2).sum().backward() opt1.step() opt2.step() with torch.no_grad(): x_avg = 0.5 * (x1 + x2) grad_avg = 0.5 * (x1.grad + x2.grad) stats_avg = 0.5 * (opt1.state[x1]["exp_avg_sq"] + opt2.state[x2]["exp_avg_sq"]) # we set wait=False in order to prevent deadlock, when averager1 locks and waits for averager2 f1 = averager1.step(wait=False) f2 = averager2.step(wait=False) f1.result() f2.result() assert torch.allclose(x1, x_avg) assert torch.allclose(x2, x_avg) assert torch.allclose(x1.grad, grad_avg) assert torch.allclose(x2.grad, grad_avg) assert torch.allclose(opt1.state[x1]["exp_avg_sq"], stats_avg) assert torch.allclose(opt2.state[x2]["exp_avg_sq"], stats_avg) for instance in [averager1, averager2] + dht_instances: instance.shutdown() if __name__ == '__main__': with cleanup_children(): print(f"BEGAN test_key_manager()") loop = asyncio.new_event_loop() loop.run_until_complete(test_key_manager()) print(f"PASSED test_key_manager()") del loop for n_clients in [0, 1, 2]: for n_aux in [0, 1, 2]: with cleanup_children(): print(f"BEGAN _test_allreduce_once({n_clients}, {n_aux})") _test_allreduce_once(n_clients, n_aux) print(f"PASSED _test_allreduce_once({n_clients}, {n_aux})") for n_clients, n_aux in [(0, 4), (1, 3), (0, 3)]: with cleanup_children(): print(f"BEGAN _test_allreduce_once({n_clients}, {n_aux})") _test_allreduce_once(n_clients, n_aux) print(f"PASSED _test_allreduce_once({n_clients}, {n_aux})") with cleanup_children(): print(f"BEGAN _test_allreduce_weighted") test_allreduce_weighted() print(f"PASSED _test_allreduce_weighted") with cleanup_children(): print(f"BEGAN _test_allreduce_grid") test_allreduce_grid() print(f"PASSED _test_allreduce_grid") with cleanup_children(): print("BEGAN _test_allreduce_compression") test_allreduce_compression() print(f"PASSED _test_allreduce_compression") with cleanup_children(): test_allgather() print(f"PASSED _test_allreduce_allgather") with cleanup_children(): test_load_balancing() print(f"PASSED _test_loadbalancing") with cleanup_children(): test_too_few_peers() print(f"PASSED _test_too_few") with cleanup_children(): test_load_state_from_peers() print(f"PASSED _test_load_state") with cleanup_children(): test_getset_bits() print(f"PASSED _test_getset") with cleanup_children(): test_training_averager() print(f"PASSED _test_training") print("DONE!")