File size: 5,720 Bytes
d5147ba | 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 | """Distributed helpers for multi-GPU evaluation."""
from __future__ import annotations
import os
import tempfile
import time
from pathlib import Path
from typing import TypeVar
import torch.distributed as dist
from utils.device import get_dist_backend, set_device
T = TypeVar("T")
def init_eval_distributed() -> tuple[int, int, bool]:
"""Initialize HCCL if launched via torchrun. Returns (rank, world_size, is_main)."""
if "LOCAL_RANK" not in os.environ:
return 0, 1, True
if not dist.is_initialized():
dist.init_process_group(backend=get_dist_backend())
rank = dist.get_rank()
world_size = dist.get_world_size()
set_device(int(os.environ["LOCAL_RANK"]))
return rank, world_size, rank == 0
def prep_done_flag() -> Path:
"""Shared-filesystem flag for multi-node; /tmp is only OK for single-node."""
port = os.environ.get("MASTER_PORT", "0")
base = os.environ.get("EVAL_PREP_FLAG_DIR", "").strip()
if base:
path = Path(base)
path.mkdir(parents=True, exist_ok=True)
return path / f"loom_eval_{port}_prep.done"
return Path(tempfile.gettempdir()) / f"loom_eval_{port}_prep.done"
def clear_prep_done_flag() -> None:
flag = prep_done_flag()
if flag.exists():
flag.unlink()
def signal_prep_done() -> None:
prep_done_flag().touch()
def wait_for_prep_done(rank: int, world_size: int) -> None:
"""Let rank 0 load datasets without holding HCCL collectives open."""
if world_size <= 1 or rank == 0:
return
flag = prep_done_flag()
while not flag.exists():
time.sleep(0.5)
def shard_by_rank(items: list[T], rank: int, world_size: int) -> list[tuple[int, T]]:
return [(i, item) for i, item in enumerate(items) if i % world_size == rank]
def shard_by_rank_balanced(
items: list[T],
rank: int,
world_size: int,
weight_fn=None,
) -> list[tuple[int, T]]:
"""Greedy load-balanced sharding by estimated workload (default: prompt char length)."""
if world_size <= 1:
return list(enumerate(items))
weights = [float(weight_fn(item)) if weight_fn is not None else 1.0 for item in items]
order = sorted(range(len(items)), key=lambda i: weights[i], reverse=True)
rank_loads = [0.0] * world_size
rank_indices: list[list[int]] = [[] for _ in range(world_size)]
for idx in order:
target = min(range(world_size), key=lambda r: rank_loads[r])
rank_indices[target].append(idx)
rank_loads[target] += weights[idx]
return [(idx, items[idx]) for idx in rank_indices[rank]]
def gather_shard_balance_stats(
prompts: list[str],
rank: int,
world_size: int,
*,
char_weights: list[float] | None = None,
token_weights: list[float] | None = None,
) -> list[dict[str, float | int]]:
"""Per-rank sample count and estimated workload for modulo vs balanced sharding."""
if char_weights is None:
char_weights = [float(len(p)) for p in prompts]
if token_weights is None:
token_weights = char_weights
def _balanced_indices(weights: list[float]) -> list[list[int]]:
order = sorted(range(len(prompts)), key=lambda i: weights[i], reverse=True)
rank_loads = [0.0] * world_size
rank_indices: list[list[int]] = [[] for _ in range(world_size)]
for idx in order:
target = min(range(world_size), key=lambda r: rank_loads[r])
rank_indices[target].append(idx)
rank_loads[target] += weights[idx]
return rank_indices
def _stats(indices: list[int], weights: list[float]) -> dict[str, float | int]:
ws = [weights[i] for i in indices]
return {
"count": len(indices),
"weight_sum": int(sum(ws)),
"weight_max": int(max(ws) if ws else 0),
"weight_avg": round(sum(ws) / len(ws), 1) if ws else 0.0,
}
modulo_indices = [i for i in range(len(prompts)) if i % world_size == rank]
char_indices = _balanced_indices(char_weights)
token_indices = _balanced_indices(token_weights)
local = {
"modulo": _stats(modulo_indices, char_weights),
"balanced_chars": _stats(char_indices[rank], char_weights),
"balanced_tokens": _stats(token_indices[rank], token_weights),
# Legacy alias for callers expecting "balanced"
"balanced": _stats(token_indices[rank], token_weights),
}
if not dist.is_initialized() or world_size == 1:
return [local]
gathered: list[dict | None] = [None] * world_size
dist.all_gather_object(gathered, local)
return gathered # type: ignore[return-value]
def gather_indexed_results(local_results: list[tuple[int, str]], total_size: int) -> list[str]:
"""All-gather (idx, text) pairs from every rank and rebuild the full ordered list."""
if not dist.is_initialized() or dist.get_world_size() == 1:
outputs = [""] * total_size
for idx, text in local_results:
outputs[idx] = text
return outputs
gathered: list[list[tuple[int, str]] | None] = [None] * dist.get_world_size()
dist.all_gather_object(gathered, local_results)
outputs = [""] * total_size
for rank_items in gathered:
for idx, text in rank_items:
outputs[idx] = text
return outputs
def broadcast_from_main(obj: T | None, rank: int) -> T:
if not dist.is_initialized() or dist.get_world_size() == 1:
return obj # type: ignore[return-value]
container: list[T | None] = [obj] if rank == 0 else [None]
dist.broadcast_object_list(container, src=0)
return container[0] # type: ignore[return-value]
def barrier() -> None:
if dist.is_initialized():
dist.barrier()
|