|
| 1 | +import contextlib |
| 2 | +from argparse import Namespace |
| 3 | +from collections import deque |
| 4 | +from collections.abc import Iterator, Sequence |
| 5 | +from concurrent.futures import Future |
| 6 | + |
| 7 | +import torch |
| 8 | +import torch.multiprocessing |
| 9 | +from sglang.srt.utils import MultiprocessingSerializer |
| 10 | + |
| 11 | +from miles.backends.sglang_utils.sglang_api_client import SGLangApiClient |
| 12 | +from miles.backends.training_utils.parallel import ParallelState |
| 13 | +from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement |
| 14 | +from miles.backends.training_utils.weight_update.protocol import WeightTransferProtocol |
| 15 | +from miles.backends.training_utils.weight_update.protocols.utils.data_replica import ( |
| 16 | + OwnedBuckets, |
| 17 | + get_data_replica_rank_and_size, |
| 18 | +) |
| 19 | +from miles.backends.training_utils.weight_update.protocols.utils.flattened_bucket import flatten_weight_bucket |
| 20 | +from miles.backends.training_utils.weight_update.rollout_cell_updater import create_rollout_cell_updaters |
| 21 | +from miles.backends.training_utils.weight_update.session import check_weight_sync_results |
| 22 | +from miles.utils import async_utils |
| 23 | + |
| 24 | + |
| 25 | +# buckets in flight across the host, enough to keep the engines loading while the next one is copied out; arbitrary |
| 26 | +# otherwise. Each replica rank gets its share, so /dev/shm holds about this many buckets whatever the rank count. |
| 27 | +_MAX_BUCKETS_IN_FLIGHT_PER_HOST = 4 |
| 28 | + |
| 29 | + |
| 30 | +class UpdateWeightFromSharedMemory(WeightTransferProtocol): |
| 31 | + """ |
| 32 | + Host weight sync through shared memory, for engines on the trainer's host. The ranks holding the |
| 33 | + gathered buckets split them round-robin, copy each flattened bucket into a shared-memory tensor |
| 34 | + once the collective gathers are done, and hand the engines its handle through |
| 35 | + /update_weights_from_tensor; each engine maps it in place, with no file written or parsed. |
| 36 | + A bucket is published as soon as it is on the host, so the copies overlap the engine loads. About |
| 37 | + `_MAX_BUCKETS_IN_FLIGHT_PER_HOST` host copies are alive at once across the host, so host memory stays a |
| 38 | + few buckets at any model size. |
| 39 | + """ |
| 40 | + |
| 41 | + def __init__(self, args: Namespace) -> None: |
| 42 | + super().__init__(args) |
| 43 | + # each published host copy with its engine calls; a copy lives until every engine has loaded it |
| 44 | + self._in_flight: deque[tuple[torch.Tensor, list[Future]]] = deque() |
| 45 | + |
| 46 | + def connect( |
| 47 | + self, |
| 48 | + rollout_engines: Sequence[SGLangApiClient], |
| 49 | + engine_gpu_counts: Sequence[int] | None, |
| 50 | + engine_gpu_offsets: Sequence[int] | None, |
| 51 | + engine_cell_ids: Sequence[str], |
| 52 | + parallel_state: ParallelState, |
| 53 | + placement: WeightUpdatePlacement, |
| 54 | + selector: str, |
| 55 | + ) -> None: |
| 56 | + self.rollout_engines = rollout_engines |
| 57 | + self.cell_updaters_of_cell_id = create_rollout_cell_updaters(self.args, self.rollout_engines, engine_cell_ids) |
| 58 | + self.group_name = "miles-shm" |
| 59 | + self.is_sender = True |
| 60 | + self._selector = selector |
| 61 | + self._engine_tp_sizes = engine_gpu_counts |
| 62 | + replica_rank, replica_size = get_data_replica_rank_and_size(parallel_state, placement) |
| 63 | + self._owned_buckets = OwnedBuckets(replica_rank=replica_rank, replica_size=replica_size) |
| 64 | + self._max_in_flight = max(1, _MAX_BUCKETS_IN_FLIGHT_PER_HOST // replica_size) |
| 65 | + |
| 66 | + def begin_sync(self, weight_version: int, iter_buckets) -> bool: |
| 67 | + self._owned_buckets.reset() |
| 68 | + return True |
| 69 | + |
| 70 | + def send_bucket(self, bucket: list[tuple[str, torch.Tensor]]) -> None: |
| 71 | + self._owned_buckets.offer(bucket) |
| 72 | + |
| 73 | + def after_base_weights(self) -> None: |
| 74 | + # copied out only after the collective gathers finish, so the ranks' host copies run in parallel |
| 75 | + cell_updaters = list(self.cell_updaters_of_cell_id.values()) |
| 76 | + for _, bucket in self._owned_buckets.take(): |
| 77 | + if len(self._in_flight) == self._max_in_flight: |
| 78 | + self._wait_for_oldest_bucket() |
| 79 | + host_tensor, metadata = _copy_to_host_bucket(bucket) |
| 80 | + with _file_system_sharing(): |
| 81 | + futures = self._publish(cell_updaters, {"flattened_tensor": host_tensor, "metadata": metadata}) |
| 82 | + self._in_flight.append((host_tensor, futures)) |
| 83 | + |
| 84 | + def _publish(self, cell_updaters: list, bucket_payload: dict) -> list[Future]: |
| 85 | + return [ |
| 86 | + updater.submit_client_call( |
| 87 | + "update_weights_from_tensor", |
| 88 | + # one serialization per receiving process: each takes its own reference to the /dev/shm file, |
| 89 | + # which is unlinked once the references are released, so a shared payload fails a later reader |
| 90 | + serialized_named_tensors=[ |
| 91 | + MultiprocessingSerializer.serialize(bucket_payload, output_str=True) for _ in range(tp_size) |
| 92 | + ], |
| 93 | + load_format="flattened_bucket", |
| 94 | + selector=self._selector, |
| 95 | + ) |
| 96 | + for updater, tp_size in zip(cell_updaters, self._engine_tp_sizes, strict=True) |
| 97 | + # a payload no engine opens keeps its reference, so the file would outlive the sync |
| 98 | + if not updater.is_errored |
| 99 | + ] |
| 100 | + |
| 101 | + def finalize(self, weight_version: int) -> None: |
| 102 | + while self._in_flight: |
| 103 | + self._wait_for_oldest_bucket() |
| 104 | + |
| 105 | + def _wait_for_oldest_bucket(self) -> None: |
| 106 | + _, futures = self._in_flight.popleft() |
| 107 | + check_weight_sync_results(async_utils.wait_futures(futures), is_lora=False) |
| 108 | + |
| 109 | + |
| 110 | +def _copy_to_host_bucket(bucket: list[tuple[str, torch.Tensor]]) -> tuple[torch.Tensor, list]: |
| 111 | + """The bucket as one shared host byte tensor in the engine's flattened layout, copied tensor by tensor. |
| 112 | +
|
| 113 | + Flattening on the device would view each tensor as bytes, which TPU tiles with a padded trailing dimension |
| 114 | + (64x for bf16), so the layout comes from meta tensors and each tensor is copied into its slice. |
| 115 | + """ |
| 116 | + metadata = flatten_weight_bucket([(name, tensor.to("meta")) for name, tensor in bucket]).get_metadata() |
| 117 | + with _file_system_sharing(): |
| 118 | + host_tensor = torch.empty(metadata[-1].end_idx, dtype=torch.uint8).share_memory_() |
| 119 | + for (_, tensor), meta in zip(bucket, metadata, strict=True): |
| 120 | + host_tensor[meta.start_idx : meta.end_idx].view(tensor.dtype).view(tensor.shape).copy_(tensor) |
| 121 | + return host_tensor, metadata |
| 122 | + |
| 123 | + |
| 124 | +@contextlib.contextmanager |
| 125 | +def _file_system_sharing() -> Iterator[None]: |
| 126 | + # the payload must name the /dev/shm file, which an unrelated process (the engine) can open; the strategy is |
| 127 | + # process-wide, so it is set only while allocating and serializing |
| 128 | + previous = torch.multiprocessing.get_sharing_strategy() |
| 129 | + torch.multiprocessing.set_sharing_strategy("file_system") |
| 130 | + try: |
| 131 | + yield |
| 132 | + finally: |
| 133 | + torch.multiprocessing.set_sharing_strategy(previous) |
0 commit comments