Buckets:
| """Distributed current-policy rollout and globally normalized OPD updates.""" | |
| from __future__ import annotations | |
| import os | |
| import socket | |
| import traceback | |
| from dataclasses import dataclass | |
| from datetime import timedelta | |
| from typing import Any, Iterable, Sequence, TypeVar | |
| import torch | |
| import torch.distributed as dist | |
| T = TypeVar("T") | |
| def rank_shard(values: Sequence[T], *, rank: int, world_size: int) -> tuple[T, ...]: | |
| if world_size <= 0 or not 0 <= rank < world_size: | |
| raise ValueError("rank/world_size are invalid") | |
| return tuple(values[rank::world_size]) | |
| def restore_global_order( | |
| shards: Sequence[Sequence[T]], *, ordered_keys: Sequence[Any], key | |
| ) -> tuple[T, ...]: | |
| positions = {value: index for index, value in enumerate(ordered_keys)} | |
| flattened = tuple(item for shard in shards for item in shard) | |
| if len(flattened) != len(ordered_keys): | |
| raise ValueError("distributed shards do not cover the logical batch") | |
| try: | |
| restored = tuple(sorted(flattened, key=lambda item: positions[key(item)])) | |
| except KeyError as error: | |
| raise ValueError("distributed shard contains an unknown logical key") from error | |
| if tuple(key(item) for item in restored) != tuple(ordered_keys): | |
| raise ValueError("distributed shards duplicate or omit logical keys") | |
| return restored | |
| def cost_balanced_shards( | |
| values: Sequence[T], *, costs: Sequence[float], world_size: int | |
| ) -> tuple[tuple[T, ...], ...]: | |
| """Greedily balance deterministic variable-length work across local ranks.""" | |
| if world_size <= 0 or len(values) != len(costs): | |
| raise ValueError("balanced shards require aligned values/costs and positive world size") | |
| assignments: list[list[tuple[int, T]]] = [[] for _ in range(world_size)] | |
| loads = [0.0] * world_size | |
| ordered = sorted( | |
| enumerate(zip(values, costs, strict=True)), | |
| key=lambda item: (-float(item[1][1]), item[0]), | |
| ) | |
| for position, (value, cost) in ordered: | |
| rank = min(range(world_size), key=lambda index: (loads[index], index)) | |
| assignments[rank].append((position, value)) | |
| loads[rank] += max(float(cost), 0.0) | |
| return tuple( | |
| tuple(value for _position, value in sorted(assignment)) | |
| for assignment in assignments | |
| ) | |
| class OPDDistributedContext: | |
| rank: int | |
| local_rank: int | |
| world_size: int | |
| object_group: Any = None | |
| owns_process_group: bool = False | |
| def enabled(self) -> bool: | |
| return self.world_size > 1 | |
| def is_main(self) -> bool: | |
| return self.rank == 0 | |
| def initialize(cls, config) -> "OPDDistributedContext": | |
| runtime_world = int(os.environ.get("WORLD_SIZE", "1")) | |
| if runtime_world != config.world_size: | |
| raise ValueError( | |
| "runtime WORLD_SIZE differs from trainer.world_size: " | |
| f"{runtime_world} != {config.world_size}" | |
| ) | |
| rank = int(os.environ.get("RANK", "0")) | |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("OPD training requires scheduler-assigned CUDA devices") | |
| torch.cuda.set_device(local_rank) | |
| device = torch.device("cuda", local_rank) | |
| owns = False | |
| object_group = None | |
| if runtime_world > 1: | |
| os.environ["TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC"] = str( | |
| config.nccl_heartbeat_timeout_seconds | |
| ) | |
| timeout = timedelta(minutes=config.process_group_timeout_minutes) | |
| if not dist.is_initialized(): | |
| dist.init_process_group("nccl", device_id=device, timeout=timeout) | |
| owns = True | |
| rank = dist.get_rank() | |
| runtime_world = dist.get_world_size() | |
| object_group = dist.new_group(backend="gloo", timeout=timeout) | |
| identities = [None] * runtime_world | |
| identity = ( | |
| socket.gethostname(), | |
| local_rank, | |
| str(torch.cuda.get_device_properties(device).uuid), | |
| ) | |
| dist.all_gather_object(identities, identity, group=object_group) | |
| if len({(host, uuid) for host, _local, uuid in identities}) != runtime_world: | |
| raise RuntimeError(f"OPD ranks share a GPU: {identities}") | |
| if rank == 0: | |
| print( | |
| "OPD_DISTRIBUTED_READY " | |
| + str( | |
| { | |
| "world_size": runtime_world, | |
| "ranks": identities, | |
| "gradient_sync": "post_backward_sum", | |
| "torch": torch.__version__, | |
| } | |
| ), | |
| flush=True, | |
| ) | |
| return cls(rank, local_rank, runtime_world, object_group, owns) | |
| def barrier(self) -> None: | |
| if self.enabled: | |
| torch.cuda.synchronize() | |
| dist.barrier(group=self.object_group) | |
| def all_gather_object(self, value: Any) -> list[Any]: | |
| if not self.enabled: | |
| return [value] | |
| values = [None] * self.world_size | |
| dist.all_gather_object(values, value, group=self.object_group) | |
| return values | |
| def broadcast_object(self, value: Any, *, source: int = 0) -> Any: | |
| if not self.enabled: | |
| return value | |
| values = [value] | |
| dist.broadcast_object_list(values, src=source, group=self.object_group) | |
| return values[0] | |
| def gather_stage(self, value: Any, error: str | None, *, stage: str) -> list[Any]: | |
| payloads = self.all_gather_object({"value": value, "error": error}) | |
| errors = [payload["error"] for payload in payloads if payload["error"]] | |
| if errors: | |
| raise RuntimeError(f"distributed {stage} failed: " + "; ".join(errors)) | |
| return [payload["value"] for payload in payloads] | |
| def run_on_main(self, function, *, stage: str): | |
| """Broadcast rank-zero I/O success or failure before any rank proceeds.""" | |
| payload = None | |
| if self.is_main: | |
| try: | |
| payload = {"value": function(), "error": None} | |
| except Exception: | |
| payload = {"value": None, "error": traceback.format_exc()} | |
| payload = self.broadcast_object(payload) | |
| if payload["error"]: | |
| raise RuntimeError(f"distributed {stage} failed: {payload['error']}") | |
| return payload["value"] | |
| def all_true(self, value: bool, device: torch.device) -> bool: | |
| if not self.enabled: | |
| return value | |
| tensor = torch.tensor(int(value), dtype=torch.uint8, device=device) | |
| dist.all_reduce(tensor, op=dist.ReduceOp.MIN) | |
| return bool(tensor.item()) | |
| def broadcast_parameters( | |
| self, | |
| parameters: Iterable[torch.nn.Parameter], | |
| *, | |
| bucket_bytes: int = 32 * 1024 * 1024, | |
| ) -> None: | |
| """Make every trainable replica bit-identical to rank zero once.""" | |
| if not self.enabled: | |
| return | |
| parameters = tuple(parameters) | |
| schema = tuple( | |
| (tuple(parameter.shape), str(parameter.dtype)) | |
| for parameter in parameters | |
| ) | |
| schemas = self.all_gather_object(schema) | |
| if any(value != schemas[0] for value in schemas[1:]): | |
| raise RuntimeError("distributed OPD parameter schemas differ across ranks") | |
| bucket: list[torch.nn.Parameter] = [] | |
| size = 0 | |
| for parameter in (*parameters, None): | |
| nbytes = ( | |
| parameter.numel() * parameter.element_size() | |
| if parameter is not None | |
| else 0 | |
| ) | |
| boundary = bucket and ( | |
| parameter is None | |
| or size + nbytes > bucket_bytes | |
| or parameter.dtype != bucket[0].dtype | |
| or parameter.device != bucket[0].device | |
| ) | |
| if boundary: | |
| flat = torch.cat([value.detach().reshape(-1) for value in bucket]) | |
| dist.broadcast(flat, src=0) | |
| offset = 0 | |
| for value in bucket: | |
| count = value.numel() | |
| value.copy_(flat[offset : offset + count].view_as(value)) | |
| offset += count | |
| bucket, size = [], 0 | |
| if parameter is not None: | |
| bucket.append(parameter) | |
| size += nbytes | |
| def synchronize_gradients( | |
| self, | |
| parameters: Iterable[torch.nn.Parameter], | |
| *, | |
| bucket_bytes: int = 32 * 1024 * 1024, | |
| ) -> None: | |
| """SUM globally normalized local gradients once per logical update.""" | |
| if not self.enabled: | |
| return | |
| parameters = tuple(parameters) | |
| used = torch.tensor( | |
| [parameter.grad is not None for parameter in parameters], | |
| dtype=torch.uint8, | |
| device=torch.device("cuda", self.local_rank), | |
| ) | |
| dist.all_reduce(used, op=dist.ReduceOp.MAX) | |
| active = [] | |
| for parameter, globally_used in zip(parameters, used.tolist(), strict=True): | |
| if not globally_used: | |
| parameter.grad = None | |
| continue | |
| if parameter.grad is None: | |
| parameter.grad = torch.zeros_like(parameter) | |
| if parameter.grad.is_sparse: | |
| raise ValueError("distributed OPD requires dense gradients") | |
| active.append(parameter) | |
| bucket: list[torch.nn.Parameter] = [] | |
| size = 0 | |
| for parameter in (*active, None): | |
| nbytes = ( | |
| parameter.numel() * parameter.element_size() | |
| if parameter is not None | |
| else 0 | |
| ) | |
| boundary = bucket and ( | |
| parameter is None | |
| or size + nbytes > bucket_bytes | |
| or parameter.dtype != bucket[0].dtype | |
| or parameter.device != bucket[0].device | |
| ) | |
| if boundary: | |
| flat = torch.cat([value.grad.reshape(-1) for value in bucket]) | |
| dist.all_reduce(flat, op=dist.ReduceOp.SUM) | |
| offset = 0 | |
| for value in bucket: | |
| count = value.numel() | |
| value.grad.copy_(flat[offset : offset + count].view_as(value)) | |
| offset += count | |
| bucket, size = [], 0 | |
| if parameter is not None: | |
| bucket.append(parameter) | |
| size += nbytes | |
| def close(self) -> None: | |
| if self.object_group is not None and dist.is_initialized(): | |
| dist.destroy_process_group(self.object_group) | |
| if self.owns_process_group and dist.is_initialized(): | |
| dist.destroy_process_group() | |
| __all__ = [ | |
| "cost_balanced_shards", | |
| "OPDDistributedContext", | |
| "rank_shard", | |
| "restore_global_order", | |
| ] | |
Xet Storage Details
- Size:
- 11 kB
- Xet hash:
- 4fd17c8899963b84c4eaebeb6403c96401e0abd78388749f7c489ca0843d9a7c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.