Buckets:
| import unittest | |
| from unittest.mock import patch | |
| import torch | |
| from opd.distributed import ( | |
| OPDDistributedContext, | |
| cost_balanced_shards, | |
| rank_shard, | |
| restore_global_order, | |
| ) | |
| class OPDDistributedTest(unittest.TestCase): | |
| def test_broadcast_parameters_copies_rank_zero_flat_buckets_back(self): | |
| parameters = [ | |
| torch.nn.Parameter(torch.tensor([1.0, 2.0])), | |
| torch.nn.Parameter(torch.tensor([3.0])), | |
| ] | |
| context = OPDDistributedContext(rank=1, local_rank=0, world_size=2) | |
| def rank_zero_payload(flat, src): | |
| self.assertEqual(src, 0) | |
| flat.copy_(torch.tensor([7.0, 8.0, 9.0])) | |
| schema = tuple((tuple(value.shape), str(value.dtype)) for value in parameters) | |
| with ( | |
| patch.object( | |
| OPDDistributedContext, | |
| "all_gather_object", | |
| return_value=[schema, schema], | |
| ), | |
| patch("torch.distributed.broadcast", side_effect=rank_zero_payload), | |
| ): | |
| context.broadcast_parameters(parameters) | |
| torch.testing.assert_close(parameters[0], torch.tensor([7.0, 8.0])) | |
| torch.testing.assert_close(parameters[1], torch.tensor([9.0])) | |
| def test_broadcast_parameters_rejects_rank_schema_mismatch_first(self): | |
| parameter = torch.nn.Parameter(torch.tensor([1.0, 2.0])) | |
| context = OPDDistributedContext(rank=1, local_rank=0, world_size=2) | |
| mismatched = [ | |
| (((2,), "torch.float32"),), | |
| (((3,), "torch.float32"),), | |
| ] | |
| with ( | |
| patch.object( | |
| OPDDistributedContext, | |
| "all_gather_object", | |
| return_value=mismatched, | |
| ), | |
| patch("torch.distributed.broadcast") as broadcast, | |
| self.assertRaisesRegex(RuntimeError, "parameter schemas differ"), | |
| ): | |
| context.broadcast_parameters([parameter]) | |
| broadcast.assert_not_called() | |
| def test_rank_shards_restore_one_global_logical_order(self): | |
| values = tuple(range(10)) | |
| shards = [rank_shard(values, rank=rank, world_size=4) for rank in range(4)] | |
| self.assertEqual(shards, [(0, 4, 8), (1, 5, 9), (2, 6), (3, 7)]) | |
| self.assertEqual( | |
| restore_global_order(shards, ordered_keys=values, key=lambda value: value), | |
| values, | |
| ) | |
| def test_restore_rejects_duplicate_or_missing_members(self): | |
| with self.assertRaisesRegex(ValueError, "duplicate or omit"): | |
| restore_global_order( | |
| ((0, 2), (1, 2)), | |
| ordered_keys=(0, 1, 2, 3), | |
| key=lambda value: value, | |
| ) | |
| def test_cost_balancing_preserves_members_and_reduces_tail_load(self): | |
| values = tuple("abcdefgh") | |
| costs = (9, 8, 7, 6, 5, 4, 3, 2) | |
| shards = cost_balanced_shards(values, costs=costs, world_size=4) | |
| by_value = dict(zip(values, costs, strict=True)) | |
| loads = [sum(by_value[value] for value in shard) for shard in shards] | |
| self.assertEqual(set().union(*(set(shard) for shard in shards)), set(values)) | |
| self.assertLessEqual(max(loads) - min(loads), 1) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 3.23 kB
- Xet hash:
- 47fef855bcb90e794d5bdfe61c55e487caa65e23b69889598555e9d198a2f368
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.