File size: 5,139 Bytes
0839907 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | # 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 fastgen.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.
"""
@functools.wraps(func)
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_all(local_all: bool, device: torch.device) -> bool:
"""Synchronize local all across distributed ranks.
Args:
local_all: all() in each rank
device: Device for tensor operations
Returns:
global_all
"""
global_all = torch.tensor([local_all], dtype=torch.uint8, device=device)
if world_size() > 1:
# MIN reduction: global_all is True only if all ranks have all samples in second stage
torch.distributed.all_reduce(global_all, op=torch.distributed.ReduceOp.MIN)
return global_all.to(torch.bool).item()
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()
def move_module_to_device(
module: torch.nn.Module,
device: torch.device,
dtype: Optional[torch.dtype] = None,
name: str = "module",
) -> None:
"""
Move a module to the target device and precision, handling meta tensors.
When using FSDP meta initialization, non-rank-0 processes have meta tensors
that need to be materialized and synchronized from rank 0.
Args:
module: The module to move
device: Target device
dtype: Target dtype (optional)
name: Name of the module for logging
"""
# Check if ANY rank has meta tensors (need collective check for broadcast)
is_meta = any(p.device.type == "meta" for p in module.parameters())
any_meta = sync_any(is_meta, device=device) if world_size() > 1 else is_meta
if any_meta:
logger.info(f"{name}: some ranks have meta tensors, materializing and broadcasting from rank 0")
if is_meta:
# This rank has meta tensors - materialize them
module.to_empty(device=device)
if dtype is not None:
module.to(dtype=dtype)
else:
# Rank 0 has real weights - just move to device
module.to(dtype=dtype, device=device)
# Broadcast weights from rank 0 to all other ranks (collective operation)
if world_size() > 1:
for param in module.parameters():
dist.broadcast(param.data, src=0)
for buffer in module.buffers():
dist.broadcast(buffer.data, src=0)
else:
# No meta tensors on any rank - simple move
module.to(dtype=dtype, device=device)
synchronize()
|