Balance disk and shm bucket ownership by bytes instead of round-robin

Each bucket goes to the replica rank with the fewest bytes so far, as
Megatron's fully parallel save and PyTorch DCP's dedup_save_plans balance
replicated writes; round-robin by count let an embedding-sized bucket skew
one rank's share.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: HJSang

miles/backends/training_utils/weight_update/protocols/utils/data_replica.py

@@ -21,21 +21,24 @@ def get_data_replica_rank_and_size(parallel_state: ParallelState, placement: Wei
 
 
 class OwnedBuckets:
-    """The buckets one rank publishes, copied to the host: the ranks holding identical gathered data take them
-    round-robin."""
+    """The buckets one rank publishes, copied to the host: each bucket goes to the replica rank with the fewest
+    bytes so far, so the ranks holding identical gathered data share the writes evenly."""
 
     def __init__(self, *, replica_rank: int, replica_size: int) -> None:
         self._replica_rank = replica_rank
         self._replica_size = replica_size
-        self._next_index = 0
-        self._owned: list[tuple[int, list[tuple[str, torch.Tensor]]]] = []
+        self.reset()
 
     def reset(self) -> None:
         self._next_index = 0
-        self._owned = []
+        self._bytes_per_rank = [0] * self._replica_size
+        self._owned: list[tuple[int, list[tuple[str, torch.Tensor]]]] = []
 
     def offer(self, bucket: list[tuple[str, torch.Tensor]]) -> None:
-        if self._next_index % self._replica_size == self._replica_rank:
+        # every rank sees the same buckets in the same order, so every rank picks the same owner
+        owner = min(range(self._replica_size), key=lambda rank: (self._bytes_per_rank[rank], rank))
+        self._bytes_per_rank[owner] += sum(tensor.nbytes for _, tensor in bucket)
+        if owner == self._replica_rank:
             # kept on the device until the stream ends, the owned buckets would add model/replica_size to each
             # rank's peak device memory. Non-blocking, so the copy overlaps the next gather; into unpinned host
             # tensors, since `.to("cpu", non_blocking=True)` pins memory that CUDA's host allocator keeps cached

tests/fast/backends/training_utils/weight_update/protocols/test_data_replica.py

@@ -0,0 +1,17 @@
+import torch
+
+from miles.backends.training_utils.weight_update.protocols.utils.data_replica import OwnedBuckets
+
+
+def _owned_names(replica_rank: int, buckets: list) -> list[str]:
+    owned = OwnedBuckets(replica_rank=replica_rank, replica_size=2)
+    for bucket in buckets:
+        owned.offer(bucket)
+    return [name for _, bucket in owned.take() for name, _ in bucket]
+
+
+def test_one_large_bucket_is_balanced_by_bytes_not_by_count():
+    """Round-robin would give rank 0 the large bucket and half the small ones; bytes give it the large one alone."""
+    buckets = [[("embed", torch.zeros(300))]] + [[(f"layer{i}", torch.zeros(100))] for i in range(3)]
+    assert _owned_names(0, buckets) == ["embed"]
+    assert _owned_names(1, buckets) == ["layer0", "layer1", "layer2"]

tests/fast/backends/training_utils/weight_update/protocols/test_shm.py

@@ -111,14 +111,18 @@ def test_another_process_rebuilds_every_tp_rank_bucket():
     protocol.finalize(weight_version=1)
 
 
-def test_a_replica_rank_publishes_only_its_own_buckets():
-    protocol = _protocol(replica_rank=1, replica_size=2, tp_sizes=[1])
-    _publish(protocol)
-
-    (updater,) = protocol.cell_updaters_of_cell_id.values()
-    (payload,) = updater.calls[0][1]["serialized_named_tensors"]
-    assert len(updater.calls) == 1
-    assert [name for name, _, _ in _rebuild_in_another_process([payload])[0]] == ["layers.1.w"]
+def test_the_replica_ranks_publish_every_bucket_exactly_once():
+    published = []
+    for replica_rank in range(2):
+        protocol = _protocol(replica_rank=replica_rank, replica_size=2, tp_sizes=[1])
+        _publish(protocol)
+        (updater,) = protocol.cell_updaters_of_cell_id.values()
+        payloads = [kwargs["serialized_named_tensors"][0] for _, kwargs in updater.calls]
+        assert payloads, f"replica rank {replica_rank} published nothing"
+        published += [[name for name, _, _ in bucket] for bucket in _rebuild_in_another_process(payloads)]
+        protocol.finalize(weight_version=1)
+
+    assert sorted(published) == sorted([name for name, _ in bucket] for bucket in _BUCKETS)
 
 
 def test_a_sync_keeps_at_most_the_window_of_host_copies_alive(monkeypatch):