Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import functools | |
| import os | |
| from typing import Optional, Callable | |
| import torch | |
| import torch.distributed as dist | |
| import lipforcing.utils.logging_utils as logger | |
| def world_size(): | |
| """Get the world size.""" | |
| if dist.is_initialized() and torch.cuda.is_available(): | |
| return dist.get_world_size() | |
| return 1 | |
| def get_rank(group: Optional[dist.ProcessGroup] = None) -> int: | |
| """Get the rank (GPU device) of the worker. | |
| Returns: | |
| rank (int): The rank of the worker. | |
| """ | |
| rank = 0 | |
| if dist.is_available() and dist.is_initialized(): | |
| rank = dist.get_rank(group) | |
| return rank | |
| def is_rank0() -> bool: | |
| """Return True if this is rank 0 (the primary loading rank).""" | |
| return get_rank() == 0 | |
| def synchronize(): | |
| """ | |
| Synchronize all devices. | |
| This method checks if the current running environment | |
| is distributed with a world-size greater than 1. | |
| If so, we use `dist.barrier` to synchronize | |
| all processes. | |
| """ | |
| if not dist.is_available(): | |
| return | |
| if not dist.is_initialized(): | |
| return | |
| world_size = dist.get_world_size() | |
| if world_size == 1: | |
| return | |
| logger.debug(f"Synchronizing all devices with world size {world_size}") | |
| dist.barrier(device_ids=[int(os.environ.get("LOCAL_RANK", "0"))]) | |
| logger.debug(f"Synchronized all devices with world size {world_size}") | |
| def rank0_only(func: Callable) -> Callable: | |
| """Apply this function only to the master GPU. | |
| Example usage: | |
| @rank0_only | |
| def func(x): | |
| return x + 1 | |
| Args: | |
| func (Callable): any function. | |
| Returns: | |
| (Callable): A function wrapper executing the function only on the master GPU. | |
| """ | |
| def wrapper(*args, **kwargs): | |
| if is_rank0(): | |
| return func(*args, **kwargs) | |
| else: | |
| return None | |
| return wrapper | |
| def clean_up(): | |
| if dist.is_available() and dist.is_initialized(): | |
| try: | |
| logger.info("Distributed clean up.") | |
| dist.destroy_process_group() | |
| except ValueError as e: | |
| logger.error(f"Error destroying default process group: {e}") | |
| def sync_any(local_any: bool, device: torch.device) -> bool: | |
| """Synchronize local any across distributed ranks. | |
| Args: | |
| local_any: any() in each rank | |
| device: Device for tensor operations | |
| Returns: | |
| global_any | |
| """ | |
| global_any = torch.tensor([local_any], dtype=torch.uint8, device=device) | |
| if world_size() > 1: | |
| # MAX reduction: global_any is True if any rank has any samples in second stage | |
| torch.distributed.all_reduce(global_any, op=torch.distributed.ReduceOp.MAX) | |
| return global_any.to(torch.bool).item() | |