wehub-resource-sync
109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
# coding: utf-8
|
|
|
|
import os
|
|
import os.path as osp
|
|
import uuid
|
|
from typing import Any, Optional
|
|
|
|
import torch
|
|
from safetensors.torch import save_file as save_safetensors, save_model as save_safetensors_model
|
|
|
|
from .basic import get_global_rank
|
|
from .fs import copy, is_hdfs_path, mkdir
|
|
from .logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_local_dir = None
|
|
|
|
|
|
def get_local_dir():
|
|
"""
|
|
Get a local directory for temporary storage for this process.
|
|
"""
|
|
global _local_dir
|
|
if _local_dir is None:
|
|
_local_dir = os.path.join("local_save", "rank_" + str(get_global_rank()) + "_" + str(uuid.uuid4()))
|
|
mkdir(_local_dir)
|
|
return _local_dir
|
|
|
|
|
|
def set_local_dir(dirname):
|
|
"""
|
|
Set a local directory for temporary storage for this process.
|
|
"""
|
|
global _local_dir
|
|
if dirname is None:
|
|
return
|
|
_local_dir = os.path.join(dirname, str(uuid.uuid4()))
|
|
mkdir(_local_dir)
|
|
|
|
|
|
def get_local_path(path: str) -> str:
|
|
"""
|
|
Get a local path for storing the file.
|
|
If the path is already a local path, directly return.
|
|
"""
|
|
if is_hdfs_path(path):
|
|
path = os.path.join(get_local_dir(), os.path.basename(path))
|
|
else:
|
|
mkdir(os.path.dirname(path))
|
|
return path
|
|
|
|
|
|
def convert_dtype(states: Any, dtype: Optional[torch.dtype] = None):
|
|
"""
|
|
Recursively convert the state_dict to device and dtype.
|
|
"""
|
|
if dtype is None:
|
|
return states
|
|
if torch.is_tensor(states):
|
|
return states.to("cpu", dtype)
|
|
if isinstance(states, dict):
|
|
return {k: convert_dtype(v, dtype) for k, v in states.items()}
|
|
if isinstance(states, list):
|
|
return [convert_dtype(v, dtype) for v in states]
|
|
return states
|
|
|
|
|
|
def save(data: Any, path: str, blocking: bool = True, local_dir: Optional[str] = None):
|
|
"""
|
|
Safely save data to a local or HDFS path.
|
|
"""
|
|
if not is_hdfs_path(path):
|
|
if path.endswith(".safetensors"):
|
|
if isinstance(data, torch.nn.Module):
|
|
save_safetensors_model(data, path)
|
|
else:
|
|
save_safetensors(data, path)
|
|
else:
|
|
torch.save(data, path)
|
|
|
|
logger.info(f"Early saved to local path: {path}")
|
|
return
|
|
|
|
if local_dir is None:
|
|
local_dir = get_local_dir()
|
|
|
|
local_path = osp.join(local_dir, osp.basename(path))
|
|
if path.endswith(".safetensors"):
|
|
if isinstance(data, torch.nn.Module):
|
|
save_safetensors_model(data, local_path)
|
|
else:
|
|
save_safetensors(data, local_path)
|
|
else:
|
|
torch.save(data, local_path)
|
|
logger.info(f"Saved to local path: {local_path}")
|
|
|
|
copy(local_path, path, blocking=blocking)
|
|
logger.info(f"Copy {local_path} to HDFS or Local path: {path} done.")
|
|
|
|
|
|
def dummy_indexes_searchsorted(packed_text_indexes: torch.LongTensor, ce_loss_indexes: torch.LongTensor) -> torch.LongTensor:
|
|
"""
|
|
Find dummy indexes via searchsorted over sorted packed text indexes.
|
|
"""
|
|
sorted_vals, sorted_pos = torch.sort(packed_text_indexes)
|
|
loc = torch.searchsorted(sorted_vals, ce_loss_indexes)
|
|
return sorted_pos[loc]
|