File size: 5,136 Bytes
6011e08 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from datetime import timedelta
from typing import Any
import torch
import torch.distributed as dist
from torch.distributed.distributed_c10d import (
Backend,
PrefixStore,
Store,
_new_process_group_helper,
_world,
default_pg_timeout,
rendezvous,
)
GLOO_GROUP = None
def init_gloo_group():
"""Initialize Gloo group for distributed communication."""
global GLOO_GROUP
if GLOO_GROUP is None:
GLOO_GROUP = dist.new_group(backend="gloo")
return GLOO_GROUP
def get_gloo_group():
"""Get the Gloo group for distributed communication."""
global GLOO_GROUP
if GLOO_GROUP is None:
raise RuntimeError("Gloo group has not been initialized. Call _init_gloo_group() first.")
return GLOO_GROUP
# Copy from pytorch to allow creating multiple main groups.
# https://github.com/pytorch/pytorch/blob/main/torch/distributed/distributed_c10d.py
def init_process_group(
backend: str | Backend = None,
init_method: str | None = None,
timeout: timedelta | None = None,
world_size: int = -1,
rank: int = -1,
store: Store | None = None,
group_name: str = None,
pg_options: Any | None = None,
):
assert (store is None) or (init_method is None), "Cannot specify both init_method and store."
if store is not None:
assert world_size > 0, "world_size must be positive if using store"
assert rank >= 0, "rank must be non-negative if using store"
elif init_method is None:
init_method = "env://"
if backend:
backend = Backend(backend)
else:
backend = Backend("undefined")
if timeout is None:
timeout = default_pg_timeout
# backward compatible API
if store is None:
rendezvous_iterator = rendezvous(init_method, rank, world_size, timeout=timeout)
store, rank, world_size = next(rendezvous_iterator)
store.set_timeout(timeout)
# Use a PrefixStore to avoid accidental overrides of keys used by
# different systems (e.g. RPC) in case the store is multi-tenant.
store = PrefixStore(group_name, store)
# NOTE: The pg_options parameter was renamed into backend_options in PyTorch 2.6.0
# https://github.com/pytorch/pytorch/commit/a0c7029a75628cd5fa8df83c0de0ea98ee7fd844
# We need to determine the appropriate parameter name based on PyTorch version
pg_options_param_name = "backend_options" if str(torch.__version__) >= "2.6" else "pg_options"
pg, _ = _new_process_group_helper(
world_size,
rank,
[],
backend,
store,
group_name=group_name,
**{pg_options_param_name: pg_options},
timeout=timeout,
)
_world.pg_group_ranks[pg] = {i: i for i in range(world_size)}
return pg
def distributed_masked_whiten(
values: torch.Tensor,
mask: torch.Tensor,
process_group: dist.ProcessGroup | None = None,
shift_mean: bool = True,
epsilon: float = 1e-8,
):
"""
Performs whitening on a tensor using global statistics from all participating GPUs.
It calculates the global mean and variance across all ranks in the default
process group (the WORLD) and uses these global statistics to normalize the
local data on each rank.
Args:
values (torch.Tensor): The local tensor of values to whiten.
mask (torch.Tensor): The local mask corresponding to the values.
process_group: The process group for all_reduce.
If None, uses the default world group.
shift_mean (bool): If True, the output is zero-mean. Defaults to True.
epsilon (float): A small value for numerical stability.
Returns:
torch.Tensor: The locally whitened tensor using global statistics.
"""
# Calculate local intermediate statistics
local_sum = (values * mask).sum()
local_sum_sq = ((values**2) * mask).sum()
local_mask_sum = mask.sum()
stats_tensor = torch.tensor(
[local_sum, local_sum_sq, local_mask_sum],
device=values.device,
dtype=torch.float32,
)
# Aggregate via all_reduce within the DP group
dist.all_reduce(stats_tensor, group=process_group)
# Calculate global stats from aggregated results
global_sum, global_sum_sq, global_mask_sum = stats_tensor
if global_mask_sum.item() == 0:
raise ValueError("The global mask sum across all participating GPUs is zero.")
global_mean = global_sum / global_mask_sum
global_mean_sq = global_sum_sq / global_mask_sum
global_var = global_mean_sq - global_mean**2
# Bessel's correction for unbiased estimate
if global_mask_sum.item() >= 2:
bessel_correction = global_mask_sum / (global_mask_sum - 1)
global_var = global_var * bessel_correction
# Whiten local data using global stats
whitened_values = (values - global_mean) * torch.rsqrt(global_var + epsilon)
if not shift_mean:
whitened_values += global_mean
return whitened_values
|