This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,419 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user