wehub-resource-sync
420 lines
14 KiB
Python
420 lines
14 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import gc
|
|
import hashlib
|
|
import numpy as np
|
|
import os
|
|
import pickle
|
|
import time
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn.functional as F
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from datasets.utils.filelock import FileLock
|
|
from datetime import timedelta
|
|
from modelscope.hub.utils.utils import get_cache_dir
|
|
from transformers.utils import is_torch_cuda_available, is_torch_mps_available, is_torch_npu_available
|
|
from typing import Any, Mapping, Optional, Union
|
|
|
|
from .env import get_dist_setting, get_node_setting, is_dist, is_local_master, is_master, is_mp
|
|
|
|
|
|
def nanstd(tensor: torch.Tensor, dim: Optional[Union[int, tuple]] = None, keepdim: bool = False) -> torch.Tensor:
|
|
"""Compute the standard deviation of a tensor, ignoring NaNs.
|
|
|
|
Refer: trl/trainer/utils.py
|
|
|
|
Args:
|
|
tensor (`torch.Tensor`):
|
|
Input tensor.
|
|
dim (`int` or `tuple[int, ...]`, *optional*):
|
|
Dimension to reduce. Defaults to all dimensions.
|
|
keepdim (`bool`, *optional*, defaults to `False`):
|
|
Whether to keep reduced dimensions.
|
|
|
|
Returns:
|
|
`torch.Tensor`:
|
|
Standard deviation of the tensor, ignoring NaNs.
|
|
"""
|
|
mean = torch.nanmean(tensor, dim=dim, keepdim=True)
|
|
variance = torch.nanmean((tensor - mean)**2, dim=dim, keepdim=True)
|
|
count = torch.sum(~torch.isnan(tensor), dim=dim, keepdim=True)
|
|
correction = count / (count - 1)
|
|
correction = torch.where(count > 1, correction, torch.full_like(correction, float('nan')))
|
|
variance *= correction # Bessel's correction
|
|
std = torch.sqrt(variance)
|
|
if keepdim:
|
|
return std
|
|
if dim is None:
|
|
return std.squeeze()
|
|
if isinstance(dim, int):
|
|
return std.squeeze(dim)
|
|
dims = [(d if d >= 0 else d + std.ndim) for d in dim]
|
|
for d in sorted(dims, reverse=True):
|
|
std = std.squeeze(d)
|
|
return std
|
|
|
|
|
|
def _find_local_mac() -> str:
|
|
mac = uuid.getnode()
|
|
mac_address = ':'.join(('%012x' % mac)[i:i + 2] for i in range(0, 12, 2))
|
|
return mac_address
|
|
|
|
|
|
def synchronize(device: Union[torch.device, str, int, None] = None):
|
|
if is_torch_npu_available():
|
|
torch.npu.synchronize(device)
|
|
elif is_torch_cuda_available():
|
|
torch.cuda.synchronize(device)
|
|
else:
|
|
torch.cuda.synchronize(device)
|
|
|
|
|
|
def time_synchronize() -> float:
|
|
synchronize()
|
|
return time.perf_counter() # second
|
|
|
|
|
|
_DISABLE_USE_BARRIER = False
|
|
|
|
|
|
@contextmanager
|
|
def disable_safe_ddp_context_use_barrier():
|
|
global _DISABLE_USE_BARRIER
|
|
_DISABLE_USE_BARRIER = True
|
|
try:
|
|
yield
|
|
finally:
|
|
_DISABLE_USE_BARRIER = False
|
|
|
|
|
|
@contextmanager
|
|
def safe_ddp_context(hash_id: Optional[str], use_barrier: bool = True):
|
|
if _DISABLE_USE_BARRIER:
|
|
use_barrier = False
|
|
if use_barrier and dist.is_initialized():
|
|
if is_dist():
|
|
if not is_master():
|
|
dist.barrier()
|
|
if not is_local_master():
|
|
# Compatible with multi-machine scenarios,
|
|
# where each machine uses different storage hardware.
|
|
dist.barrier()
|
|
yield
|
|
if is_dist():
|
|
if is_master():
|
|
dist.barrier()
|
|
if is_local_master():
|
|
dist.barrier()
|
|
elif hash_id is not None:
|
|
lock_dir = os.path.join(get_cache_dir(), 'lockers')
|
|
os.makedirs(lock_dir, exist_ok=True)
|
|
file_path = hashlib.sha256(hash_id.encode('utf-8')).hexdigest() + '.lock'
|
|
file_path = os.path.join(lock_dir, file_path)
|
|
with FileLock(file_path):
|
|
yield
|
|
else:
|
|
yield
|
|
|
|
|
|
def get_device(local_rank: Optional[Union[str, int]] = None) -> str:
|
|
if local_rank is None:
|
|
local_rank = max(0, get_dist_setting()[1])
|
|
local_rank = str(local_rank)
|
|
if is_torch_npu_available():
|
|
device = 'npu:{}'.format(local_rank)
|
|
elif is_torch_mps_available():
|
|
device = 'mps:{}'.format(local_rank)
|
|
elif is_torch_cuda_available():
|
|
device = 'cuda:{}'.format(local_rank)
|
|
else:
|
|
device = 'cpu'
|
|
|
|
return device
|
|
|
|
|
|
def get_current_device():
|
|
if is_torch_npu_available():
|
|
current_device = torch.npu.current_device()
|
|
elif is_torch_cuda_available():
|
|
current_device = torch.cuda.current_device()
|
|
elif is_torch_mps_available():
|
|
current_device = 'mps'
|
|
else:
|
|
current_device = 'cpu'
|
|
return current_device
|
|
|
|
|
|
def get_torch_device():
|
|
if is_torch_cuda_available():
|
|
return torch.cuda
|
|
elif is_torch_npu_available():
|
|
return torch.npu
|
|
elif is_torch_mps_available():
|
|
return torch.mps
|
|
else:
|
|
return torch.cpu
|
|
|
|
|
|
def set_device(local_rank: Optional[Union[str, int]] = None):
|
|
if local_rank is None:
|
|
local_rank = max(0, get_dist_setting()[1])
|
|
if is_torch_npu_available():
|
|
torch.npu.set_device(local_rank)
|
|
elif is_torch_cuda_available():
|
|
torch.cuda.set_device(local_rank)
|
|
|
|
|
|
def get_device_count() -> int:
|
|
if is_torch_npu_available():
|
|
return torch.npu.device_count()
|
|
elif is_torch_cuda_available():
|
|
return torch.cuda.device_count()
|
|
else:
|
|
return 0
|
|
|
|
|
|
def is_torch_rocm() -> bool:
|
|
"""True on an AMD ROCm/HIP torch build.
|
|
|
|
PyTorch on ROCm exposes devices through the ``torch.cuda.*`` API (hipify), so
|
|
``is_torch_cuda_available()`` is True on ROCm too; the distinguishing signal is
|
|
a non-None ``torch.version.hip``.
|
|
"""
|
|
return is_torch_cuda_available() and getattr(torch.version, 'hip', None) is not None
|
|
|
|
|
|
def get_physical_device_count() -> int:
|
|
"""Number of physical accelerators, ignoring ``*_VISIBLE_DEVICES`` masks.
|
|
|
|
Unlike :func:`get_device_count` (which honors CUDA/HIP_VISIBLE_DEVICES via the
|
|
runtime), this queries the driver/topology layer, so a process that inherited a
|
|
restricted mask can still discover the full device set. It still respects the
|
|
container's device exposure (``--gpus`` / ``NVIDIA_VISIBLE_DEVICES`` / passed
|
|
``/dev/dri`` + ``/dev/kfd``). Falls back to the mask-dependent runtime count if
|
|
the driver query is unavailable.
|
|
"""
|
|
import glob
|
|
|
|
if is_torch_rocm():
|
|
# ROCm: the kfd topology lists one node per device; GPU nodes have
|
|
# ``simd_count > 0`` (CPU/host nodes have 0). Reading sysfs neither
|
|
# initializes HIP nor is affected by the visibility mask.
|
|
count = 0
|
|
for prop in glob.glob('/sys/class/kfd/kfd/topology/nodes/*/properties'):
|
|
try:
|
|
with open(prop) as f:
|
|
for line in f:
|
|
if line.startswith('simd_count'):
|
|
if int(line.split()[1]) > 0:
|
|
count += 1
|
|
break
|
|
except Exception:
|
|
continue
|
|
if count > 0:
|
|
return count
|
|
elif is_torch_cuda_available():
|
|
# NVIDIA: NVML/procfs/nvidia-smi enumerate at the driver level and are not
|
|
# filtered by CUDA_VISIBLE_DEVICES (only by container-level exposure).
|
|
try:
|
|
import pynvml
|
|
pynvml.nvmlInit()
|
|
try:
|
|
count = pynvml.nvmlDeviceGetCount()
|
|
finally:
|
|
pynvml.nvmlShutdown()
|
|
if count > 0:
|
|
return count
|
|
except Exception:
|
|
pass
|
|
count = len(glob.glob('/proc/driver/nvidia/gpus/*'))
|
|
if count > 0:
|
|
return count
|
|
try:
|
|
import subprocess
|
|
out = subprocess.check_output(['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'],
|
|
stderr=subprocess.DEVNULL,
|
|
timeout=5)
|
|
count = len([line for line in out.decode().splitlines() if line.strip()])
|
|
if count > 0:
|
|
return count
|
|
except Exception:
|
|
pass
|
|
|
|
return get_device_count()
|
|
|
|
|
|
def empty_cache():
|
|
if is_torch_npu_available():
|
|
torch.npu.empty_cache()
|
|
elif is_torch_mps_available():
|
|
torch.mps.empty_cache()
|
|
elif is_torch_cuda_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def ipc_collect():
|
|
if is_torch_cuda_available():
|
|
torch.cuda.ipc_collect()
|
|
elif is_torch_npu_available():
|
|
torch.npu.ipc_collect()
|
|
|
|
|
|
def gc_collect() -> None:
|
|
gc.collect()
|
|
empty_cache()
|
|
|
|
|
|
def get_last_valid_indices(attention_mask: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Get the last valid (non-padding) token position indices for each sample.
|
|
|
|
This function correctly handles sequences with different padding directions (left/right/none)
|
|
within the same batch by computing the last valid index for each sequence individually.
|
|
|
|
Args:
|
|
attention_mask: Attention mask [batch_size, seq_len] where 1=valid, 0=padding
|
|
|
|
Returns:
|
|
torch.Tensor: Indices of last valid positions [batch_size]
|
|
|
|
Examples:
|
|
>>> # Right padding
|
|
>>> attention_mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 0]])
|
|
>>> get_last_valid_indices(attention_mask)
|
|
tensor([2, 3])
|
|
|
|
>>> # Left padding
|
|
>>> attention_mask = torch.tensor([[0, 0, 1, 1, 1], [0, 1, 1, 1, 1]])
|
|
>>> get_last_valid_indices(attention_mask)
|
|
tensor([4, 4])
|
|
"""
|
|
seq_len = attention_mask.shape[1]
|
|
|
|
# Flip the mask horizontally to bring the last elements to the front.
|
|
# `argmax` will then find the index of the first '1', which corresponds to the last valid token.
|
|
last_valid_indices = torch.fliplr(attention_mask).argmax(dim=1)
|
|
|
|
# Convert the index from the right-to-left frame to the original left-to-right frame.
|
|
indices = seq_len - 1 - last_valid_indices
|
|
|
|
return indices
|
|
|
|
|
|
class Serializer:
|
|
|
|
@staticmethod
|
|
def to_tensor(obj):
|
|
res = pickle.dumps(obj)
|
|
res = np.array([len(res)], dtype=np.int64).tobytes() + res
|
|
res = np.frombuffer(res, dtype=np.uint8).copy()
|
|
res = torch.from_numpy(res)
|
|
return res
|
|
|
|
@staticmethod
|
|
def from_tensor(obj):
|
|
if isinstance(obj, torch.Tensor):
|
|
obj = obj.cpu().numpy()
|
|
res = obj.tobytes()
|
|
buffer_size = np.frombuffer(res[:8], dtype=np.int64)[0]
|
|
res = res[8:]
|
|
return pickle.loads(res[:buffer_size])
|
|
|
|
|
|
def set_default_ddp_config():
|
|
# It runs normally with Python as well.
|
|
rank, local_rank, _, _ = get_dist_setting()
|
|
if rank == -1 or local_rank == -1:
|
|
os.environ['NPROC_PER_NODE'] = '1'
|
|
os.environ['RANK'] = '0'
|
|
os.environ['LOCAL_RANK'] = '0'
|
|
os.environ['WORLD_SIZE'] = '1'
|
|
os.environ['LOCAL_WORLD_SIZE'] = '1'
|
|
os.environ['MASTER_ADDR'] = '127.0.0.1'
|
|
os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '29500')
|
|
|
|
|
|
def init_process_group(backend: Optional[str] = None, timeout: int = 18000000):
|
|
if dist.is_initialized():
|
|
return
|
|
set_device()
|
|
if backend is None:
|
|
if is_torch_npu_available():
|
|
backend = 'hccl'
|
|
elif torch.cuda.is_available():
|
|
backend = 'nccl'
|
|
else:
|
|
backend = 'gloo'
|
|
timeout = timedelta(seconds=timeout)
|
|
dist.init_process_group(backend=backend, timeout=timeout)
|
|
|
|
|
|
def check_shared_disk(error, cache_dir: Optional[str] = None):
|
|
nnodes = get_node_setting()[1]
|
|
if nnodes <= 1:
|
|
return True
|
|
assert dist.is_initialized()
|
|
if cache_dir is None:
|
|
cache_dir = os.path.join(get_cache_dir(), 'tmp')
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
tmp_path = os.path.join(cache_dir, 'check_shared_disk.tmp')
|
|
is_shared_disk = True
|
|
|
|
try:
|
|
with safe_ddp_context(None, True):
|
|
if is_master():
|
|
with open(tmp_path, 'w'):
|
|
pass
|
|
if not os.path.exists(tmp_path):
|
|
is_shared_disk = False
|
|
shared_state = [None] * dist.get_world_size()
|
|
dist.all_gather_object(shared_state, is_shared_disk)
|
|
finally:
|
|
if is_master() and os.path.exists(tmp_path):
|
|
os.remove(tmp_path)
|
|
if not all(shared_state):
|
|
raise error
|
|
|
|
|
|
def to_float_dtype(data: Any, dtype: torch.dtype) -> Any:
|
|
"""Change the float inputs to a dtype"""
|
|
if isinstance(data, Mapping):
|
|
return type(data)({k: to_float_dtype(v, dtype) for k, v in data.items()})
|
|
elif isinstance(data, (tuple, list)):
|
|
return type(data)(to_float_dtype(v, dtype) for v in data)
|
|
elif isinstance(data, torch.Tensor) and torch.is_floating_point(data):
|
|
return data.to(dtype=dtype)
|
|
else:
|
|
return data
|
|
|
|
|
|
def to_device(data: Any, device: Union[str, torch.device, int], non_blocking: bool = False) -> Any:
|
|
"""Move inputs to a device"""
|
|
if isinstance(data, Mapping):
|
|
return type(data)({k: to_device(v, device, non_blocking) for k, v in data.items()})
|
|
elif isinstance(data, (tuple, list)):
|
|
return type(data)(to_device(v, device, non_blocking) for v in data)
|
|
elif isinstance(data, torch.Tensor):
|
|
return data.to(device=device, non_blocking=non_blocking)
|
|
else:
|
|
return data
|
|
|
|
|
|
def get_generative_reranker_logits(lm_head_weight, tokenizer, hidden_states):
|
|
positive_token = os.environ.get('GENERATIVE_RERANKER_POSITIVE_TOKEN', 'yes')
|
|
negative_token = os.environ.get('GENERATIVE_RERANKER_NEGATIVE_TOKEN', 'no')
|
|
positive_token_id = tokenizer.convert_tokens_to_ids(positive_token)
|
|
negative_token_id = tokenizer.convert_tokens_to_ids(negative_token)
|
|
weight = lm_head_weight[[positive_token_id, negative_token_id]]
|
|
logits = F.linear(hidden_states, weight)
|
|
return logits[..., 0:1] - logits[..., 1:2]
|
|
|
|
|
|
def get_max_reserved_memory() -> float:
|
|
devices = list(range(get_device_count())) if is_mp() else [None]
|
|
try:
|
|
mems = [get_torch_device().max_memory_reserved(device=device) for device in devices]
|
|
except AttributeError:
|
|
return 0 # fix mps
|
|
return sum(mems) / 1024**3
|