Buckets:
| """Small process-group boundary for config-driven GRPO data parallelism.""" | |
| from __future__ import annotations | |
| import os | |
| import socket | |
| from dataclasses import dataclass | |
| from typing import Any, Iterable | |
| import torch | |
| import torch.distributed as dist | |
| import torch.multiprocessing as mp | |
| class DistributedContext: | |
| rank: int = 0 | |
| local_rank: int = 0 | |
| world_size: int = 1 | |
| owns_process_group: bool = False | |
| object_group: Any = None | |
| def enabled(self) -> bool: | |
| return self.world_size > 1 | |
| def is_main(self) -> bool: | |
| return self.rank == 0 | |
| def initialize(cls, expected_world_size: int) -> "DistributedContext": | |
| runtime_world_size = int(os.environ.get("WORLD_SIZE", "1")) | |
| if runtime_world_size != expected_world_size: | |
| raise ValueError( | |
| "runtime WORLD_SIZE differs from trainer.world_size: " | |
| f"{runtime_world_size} != {expected_world_size}" | |
| ) | |
| rank = int(os.environ.get("RANK", "0")) | |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) | |
| owns = False | |
| if runtime_world_size > 1 and not dist.is_initialized(): | |
| dist.init_process_group(backend="nccl") | |
| owns = True | |
| object_group = ( | |
| dist.new_group(backend="gloo") if runtime_world_size > 1 else None | |
| ) | |
| return cls(rank, local_rank, runtime_world_size, owns, object_group) | |
| def barrier(self) -> None: | |
| if self.enabled: | |
| dist.barrier(group=self.object_group) | |
| def all_gather_object(self, value: Any) -> list[Any]: | |
| if not self.enabled: | |
| return [value] | |
| gathered = [None] * self.world_size | |
| dist.all_gather_object(gathered, value, group=self.object_group) | |
| return gathered | |
| 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 main_call(self, function): | |
| """Publish rank-zero I/O errors before peers enter another stage.""" | |
| if not self.enabled: | |
| return function() | |
| value = error = None | |
| if self.is_main: | |
| try: | |
| value = function() | |
| except Exception as exception: | |
| error = f"{type(exception).__name__}: {exception}" | |
| value, error = self.broadcast_object((value, error)) | |
| if error: | |
| raise RuntimeError(error) | |
| return value | |
| def gather_stage( | |
| self, | |
| value: Any, | |
| error: str | None, | |
| *, | |
| stage: str, | |
| ) -> list[Any]: | |
| """Collect one stage result and surface every rank-local error.""" | |
| 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 synchronize_gradients( | |
| self, parameters: Iterable[torch.nn.Parameter] | |
| ) -> None: | |
| """Sum sharded globally-normalized gradients without changing semantics.""" | |
| if not self.enabled: | |
| return | |
| parameters = tuple(parameters) | |
| device = next( | |
| (parameter.device for parameter in parameters), | |
| torch.device("cuda", self.local_rank), | |
| ) | |
| used = torch.tensor( | |
| [parameter.grad is not None for parameter in parameters], | |
| device=device, | |
| dtype=torch.uint8, | |
| ) | |
| dist.all_reduce(used, op=dist.ReduceOp.MAX) | |
| 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) | |
| dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) | |
| 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() | |
| def assigned_group_indices(group_size: int, context: DistributedContext) -> tuple[int, ...]: | |
| return tuple(range(context.rank, group_size, context.world_size)) | |
| def _free_local_port() -> int: | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: | |
| listener.bind(("127.0.0.1", 0)) | |
| return int(listener.getsockname()[1]) | |
| def _spawn_worker(local_rank: int, config, port: int, worker) -> None: | |
| os.environ.update( | |
| MASTER_ADDR="127.0.0.1", | |
| MASTER_PORT=str(port), | |
| WORLD_SIZE=str(config.trainer.world_size), | |
| RANK=str(local_rank), | |
| LOCAL_RANK=str(local_rank), | |
| ) | |
| worker(config) | |
| def launch_configured_workers(config, worker): | |
| """Run directly or spawn the configured number of local CUDA workers.""" | |
| runtime_world_size = int(os.environ.get("WORLD_SIZE", "1")) | |
| if runtime_world_size > 1 or config.trainer.world_size == 1: | |
| return worker(config) | |
| available = torch.cuda.device_count() | |
| if available != config.trainer.world_size: | |
| raise RuntimeError( | |
| "trainer.world_size requires exactly that many scheduler-visible GPUs: " | |
| f"configured={config.trainer.world_size}, visible={available}" | |
| ) | |
| mp.spawn( | |
| _spawn_worker, | |
| args=(config, _free_local_port(), worker), | |
| nprocs=config.trainer.world_size, | |
| join=True, | |
| ) | |
| return None | |
| __all__ = [ | |
| "assigned_group_indices", | |
| "DistributedContext", | |
| "launch_configured_workers", | |
| ] | |
Xet Storage Details
- Size:
- 5.94 kB
- Xet hash:
- 7b9048f83a1c169518e13ab27640a71527a11d995b0b724688b5245372fd9a70
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.