Spaces:
Sleeping
Sleeping
| """ | |
| Lightweight inference runtime (replaces Megatron/DeepSpeed training stack in release builds). | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import os | |
| from typing import Any, Dict, Optional | |
| import torch | |
| import torch.distributed as dist | |
| _GLOBAL_ARGS: Optional[Any] = None | |
| def set_global_args(args: Any) -> None: | |
| global _GLOBAL_ARGS | |
| _GLOBAL_ARGS = args | |
| def get_args() -> Any: | |
| if _GLOBAL_ARGS is None: | |
| raise RuntimeError("Inference args not initialized. Call set_global_args() first.") | |
| return _GLOBAL_ARGS | |
| def print_rank_0(msg: str) -> None: | |
| if not dist.is_initialized() or dist.get_rank() == 0: | |
| print(msg, flush=True) | |
| class _NoOpTimer: | |
| def start(self, *args, **kwargs): | |
| return self | |
| def stop(self, *args, **kwargs): | |
| return self | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, *args): | |
| return False | |
| def get_timers(): | |
| return _NoOpTimer() | |
| def get_model_parallel_world_size() -> int: | |
| return 1 | |
| def get_tensor_model_parallel_rank() -> int: | |
| return 0 | |
| def get_data_parallel_world_size() -> int: | |
| # When sequence parallelism is on, the data-parallel axis is the complement | |
| # of the SP axis (each SP group works on ONE sample together). | |
| from .ulysses import sequence_parallel_is_enabled, get_dp_world_size | |
| if sequence_parallel_is_enabled(): | |
| return get_dp_world_size() | |
| if dist.is_initialized(): | |
| return dist.get_world_size() | |
| return 1 | |
| def get_data_parallel_rank() -> int: | |
| from .ulysses import sequence_parallel_is_enabled, get_dp_rank | |
| if sequence_parallel_is_enabled(): | |
| return get_dp_rank() | |
| if dist.is_initialized(): | |
| return dist.get_rank() | |
| return 0 | |
| def broadcast_data(keys, data: Dict[str, Any], dtype: torch.dtype): | |
| if get_model_parallel_world_size() <= 1: | |
| return data | |
| return data | |
| class MPU: | |
| get_model_parallel_world_size = staticmethod(get_model_parallel_world_size) | |
| get_tensor_model_parallel_rank = staticmethod(get_tensor_model_parallel_rank) | |
| get_data_parallel_world_size = staticmethod(get_data_parallel_world_size) | |
| get_data_parallel_rank = staticmethod(get_data_parallel_rank) | |
| broadcast_data = staticmethod(broadcast_data) | |
| mpu = MPU() | |
| def init_method_normal(std: float): | |
| def init_(tensor): | |
| return torch.nn.init.normal_(tensor, mean=0.0, std=std) | |
| return init_ | |
| def scaled_init_method_normal(std: float, num_layers: int): | |
| scale = std / math.sqrt(2.0 * num_layers) | |
| def init_(tensor): | |
| return torch.nn.init.normal_(tensor, mean=0.0, std=scale) | |
| return init_ | |
| def init_distributed(backend: str = "nccl") -> tuple[int, int, int]: | |
| if dist.is_initialized(): | |
| rank = dist.get_rank() | |
| world_size = dist.get_world_size() | |
| else: | |
| # When launched with plain `python` (single process, no torchrun), the | |
| # standard env:// variables are absent. Fill in single-process defaults | |
| # so init_process_group does not hang waiting for a rendezvous. | |
| os.environ.setdefault("MASTER_ADDR", "127.0.0.1") | |
| os.environ.setdefault("MASTER_PORT", "29500") | |
| os.environ.setdefault("RANK", "0") | |
| os.environ.setdefault("WORLD_SIZE", "1") | |
| os.environ.setdefault("LOCAL_RANK", "0") | |
| dist.init_process_group(backend=backend, init_method="env://") | |
| rank = dist.get_rank() | |
| world_size = dist.get_world_size() | |
| local_rank = int(os.environ.get("LOCAL_RANK", rank % max(torch.cuda.device_count(), 1))) | |
| torch.cuda.set_device(local_rank) | |
| return rank, local_rank, world_size | |