William Lin
7d61590213
Ray sets CUDA_VISIBLE_DEVICES of each GPU actor to the GPU of that actor. The import of fastvideo initializes CUDA, so the executor cannot change the device list later. On a node with more than one GPU, each worker after the first worker failed with "invalid device ordinal". On NVIDIA GPUs, each worker actor now keeps the device list of its raylet (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1). Each worker uses the position of its GPU in that list as its local rank. The safetensors loader now reads each file on node-group rank 0, which is the rank that sends the broadcast.
319 lines
13 KiB
Python
319 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py
|
|
"""Utilities for downloading and initializing model weights."""
|
|
import glob
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from collections.abc import Generator
|
|
from pathlib import Path
|
|
|
|
import filelock
|
|
import huggingface_hub.constants
|
|
import torch
|
|
import torch.distributed as dist
|
|
from safetensors.torch import safe_open
|
|
from tqdm.auto import tqdm
|
|
|
|
import fastvideo.distributed.parallel_state as parallel_state
|
|
from fastvideo.logger import init_logger
|
|
|
|
SAFETENSORS_TO_TORCH_DTYPE = {
|
|
'F16': torch.float16,
|
|
'BF16': torch.bfloat16,
|
|
'F32': torch.float32,
|
|
'F64': torch.float64,
|
|
'I8': torch.int8,
|
|
'I16': torch.int16,
|
|
'I32': torch.int32,
|
|
'I64': torch.int64,
|
|
'U8': torch.uint8,
|
|
'BOOL': torch.bool,
|
|
'F8_E4M3': torch.float8_e4m3fn,
|
|
'F8_E5M2': torch.float8_e5m2,
|
|
}
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# use system-level temp directory for file locks, so that multiple users
|
|
# can share the same lock without error.
|
|
# lock files in the temp directory will be automatically deleted when the
|
|
# system reboots, so users will not complain about annoying lock files
|
|
temp_dir = tempfile.gettempdir()
|
|
|
|
|
|
def enable_hf_transfer() -> None:
|
|
"""automatically activates hf_transfer
|
|
"""
|
|
if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
|
|
try:
|
|
# enable hf hub transfer if available
|
|
import hf_transfer # type: ignore # noqa
|
|
huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
enable_hf_transfer()
|
|
|
|
|
|
class DisabledTqdm(tqdm):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs, disable=True)
|
|
|
|
|
|
def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
|
|
lock_dir = cache_dir or temp_dir
|
|
model_name_or_path = str(model_name_or_path)
|
|
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
|
|
model_name = model_name_or_path.replace("/", "-")
|
|
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
|
# add hash to avoid conflict with old users' lock files
|
|
lock_file_name = hash_name + model_name + ".lock"
|
|
# mode 0o666 is required for the filelock to be shared across users
|
|
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
|
return lock
|
|
|
|
|
|
# For models like Mistral-7B-v0.3, there are both sharded
|
|
# safetensors files and a consolidated safetensors file.
|
|
# Passing both of these to the weight loader functionality breaks.
|
|
# So, we use the index_file to
|
|
# look up which safetensors files should be used.
|
|
def filter_duplicate_safetensors_files(hf_weights_files: list[str], hf_folder: str, index_file: str) -> list[str]:
|
|
# model.safetensors.index.json is a mapping from keys in the
|
|
# torch state_dict to safetensors file holding that weight.
|
|
index_file_name = os.path.join(hf_folder, index_file)
|
|
if not os.path.isfile(index_file_name):
|
|
return hf_weights_files
|
|
|
|
# Iterate through the weight_map (weight_name: safetensors files)
|
|
# to identify weights that we should use.
|
|
with open(index_file_name) as f:
|
|
weight_map = json.load(f)["weight_map"]
|
|
weight_files_in_index = set()
|
|
for weight_name in weight_map:
|
|
weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
|
|
# Filter out any fields that are not found in the index file.
|
|
hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
|
|
return hf_weights_files
|
|
|
|
|
|
SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json"
|
|
|
|
|
|
def resolve_safetensors_files(model_path: str) -> list[str]:
|
|
"""Discover safetensors files in a model directory."""
|
|
files = sorted(glob.glob(os.path.join(model_path, "*.safetensors")))
|
|
if not files:
|
|
raise FileNotFoundError(f"No .safetensors files found in {model_path}")
|
|
index_file = os.path.join(model_path, SAFE_WEIGHTS_INDEX_NAME)
|
|
if os.path.exists(index_file):
|
|
files = filter_duplicate_safetensors_files(files, model_path, SAFE_WEIGHTS_INDEX_NAME)
|
|
return files
|
|
|
|
|
|
def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]:
|
|
"""
|
|
Exclude files that are not needed for inference.
|
|
|
|
See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
|
|
"""
|
|
blacklist = [
|
|
"training_args.bin",
|
|
"optimizer.bin",
|
|
"optimizer.pt",
|
|
"scheduler.pt",
|
|
"scaler.pt",
|
|
]
|
|
hf_weights_files = [f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist)]
|
|
return hf_weights_files
|
|
|
|
|
|
# explicitly use pure text format, with a newline at the end
|
|
# this makes it impossible to see the animation in the progress bar
|
|
# but will avoid messing up with ray or multiprocessing, which wraps
|
|
# each line of output with some prefix.
|
|
_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501
|
|
|
|
|
|
def _get_initialized_node_group():
|
|
# Single-rank / pytest paths may not have initialized the node group.
|
|
# Fall through to the old get_local_torch_device() behavior when uninitialized.
|
|
if torch.distributed.is_initialized() and parallel_state._NODE is not None:
|
|
return parallel_state.get_node_group()
|
|
return None
|
|
|
|
|
|
def _is_node_leader(node_group) -> bool:
|
|
"""Return True on the rank that reads the checkpoint files for its node.
|
|
|
|
The node group broadcasts from its rank 0, so that rank reads. The local
|
|
rank is the CUDA device ordinal, and a node can have no rank on device 0
|
|
(for example, Ray workers on GPUs 2 and 3 of a 4-GPU node).
|
|
"""
|
|
if node_group is not None:
|
|
return node_group.rank_in_group == 0
|
|
return int(os.environ.get("LOCAL_RANK", 0)) == 0
|
|
|
|
|
|
def safetensors_weights_iterator(hf_weights_files: list[str],
|
|
to_cpu: bool = False,
|
|
broadcast: bool = True,
|
|
async_broadcast: bool = False) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
"""Iterate over the weights in the model safetensor files.
|
|
Args:
|
|
hf_weights_files: List of safetensor files to load.
|
|
to_cpu: Whether to load the weights to CPU. If False, will load to the GPU device bound to the current
|
|
process.
|
|
broadcast: Whether node-group rank 0 should read GPU weights and broadcast them to the other ranks of the node.
|
|
async_broadcast: Whether to overlap loading from disk and broadcasting to other ranks. If True,
|
|
must iterate over all the weights before use. Only used when broadcast is True and to_cpu is False.
|
|
"""
|
|
node_group = _get_initialized_node_group()
|
|
is_leader = _is_node_leader(node_group)
|
|
device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu"
|
|
enable_tqdm = not torch.distributed.is_initialized() or is_leader
|
|
if to_cpu or not broadcast or node_group is None:
|
|
async_broadcast = False
|
|
|
|
handles = []
|
|
for st_file in tqdm(
|
|
hf_weights_files,
|
|
desc="Loading safetensors checkpoint shards",
|
|
disable=not enable_tqdm,
|
|
bar_format=_BAR_FORMAT,
|
|
):
|
|
with safe_open(st_file, framework="pt", device=device) as f:
|
|
for name in f.keys(): # noqa: SIM118
|
|
if to_cpu:
|
|
param = f.get_tensor(name)
|
|
elif broadcast and node_group is not None:
|
|
if is_leader:
|
|
param = f.get_tensor(name)
|
|
else:
|
|
sl = f.get_slice(name)
|
|
shape = sl.get_shape()
|
|
dtype = SAFETENSORS_TO_TORCH_DTYPE[sl.get_dtype()]
|
|
param = torch.empty(shape, device=device, dtype=dtype)
|
|
# broadcast to local ranks
|
|
# TODO(Wenxuan): scatter instead of broadcast
|
|
if node_group.world_size > 1:
|
|
group = node_group.device_group
|
|
if async_broadcast:
|
|
handle = dist.broadcast(param,
|
|
src=dist.get_global_rank(group, 0),
|
|
async_op=True,
|
|
group=group)
|
|
handles.append(handle)
|
|
else:
|
|
dist.broadcast(param, src=dist.get_global_rank(group, 0), group=group)
|
|
else:
|
|
param = f.get_tensor(name)
|
|
yield name, param
|
|
|
|
if async_broadcast:
|
|
for handle in handles:
|
|
handle.wait()
|
|
handles.clear()
|
|
|
|
|
|
def pt_weights_iterator(hf_weights_files: list[str],
|
|
to_cpu: bool = False,
|
|
broadcast: bool = True) -> Generator[tuple[str, torch.Tensor], None, None]:
|
|
"""Iterate over the weights in the model bin/pt files.
|
|
|
|
Args:
|
|
hf_weights_files: List of bin/pt files to load.
|
|
to_cpu: Whether to load the weights to CPU.
|
|
broadcast: Accepted for API symmetry. PT weights are loaded through
|
|
torch.load and do not use the safetensors broadcast path.
|
|
"""
|
|
node_group = _get_initialized_node_group()
|
|
device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu"
|
|
enable_tqdm = not torch.distributed.is_initialized() or _is_node_leader(node_group)
|
|
for bin_file in tqdm(
|
|
hf_weights_files,
|
|
desc="Loading pt checkpoint shards",
|
|
disable=not enable_tqdm,
|
|
bar_format=_BAR_FORMAT,
|
|
):
|
|
state = torch.load(bin_file, map_location=device, weights_only=True)
|
|
yield from state.items()
|
|
del state
|
|
|
|
|
|
def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
|
|
"""Default weight loader."""
|
|
try:
|
|
if param.numel() == 1 and loaded_weight.numel() == 1:
|
|
# Sometimes scalar values aren't considered tensors with shapes
|
|
# so if both param and loaded_weight are a scalar,
|
|
# "broadcast" instead of copy
|
|
param.data.fill_(loaded_weight.item())
|
|
else:
|
|
assert param.size() == loaded_weight.size(), (f"Attempted to load weight ({loaded_weight.size()}) "
|
|
f"into parameter ({param.size()})")
|
|
|
|
param.data.copy_(loaded_weight)
|
|
except Exception:
|
|
# NOTE: This exception is added for the purpose of setting breakpoint to
|
|
# debug weight loading issues.
|
|
raise
|
|
|
|
|
|
def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
|
|
"""Remap the name of FP8 k/v_scale parameters.
|
|
|
|
This function handles the remapping of FP8 k/v_scale parameter names.
|
|
It detects if the given name ends with a suffix and attempts to remap
|
|
it to the expected name format in the model. If the remapped name is not
|
|
found in the params_dict, a warning is printed and None is returned.
|
|
|
|
Args:
|
|
name (str): The original loaded checkpoint parameter name.
|
|
params_dict (dict): Dictionary containing the model's named parameters.
|
|
|
|
Returns:
|
|
str: The remapped parameter name if successful, or the original name
|
|
if no remapping is needed.
|
|
None: If the remapped name is not found in params_dict.
|
|
"""
|
|
if name.endswith(".kv_scale"):
|
|
logger.warning_once("DEPRECATED. Found kv_scale in the checkpoint. "
|
|
"This format is deprecated in favor of separate k_scale and "
|
|
"v_scale tensors and will be removed in a future release. "
|
|
"Functionally, we will remap kv_scale to k_scale and duplicate "
|
|
"k_scale to v_scale")
|
|
# NOTE: we remap the deprecated kv_scale to k_scale
|
|
remapped_name = name.replace(".kv_scale", ".attn.k_scale")
|
|
if remapped_name not in params_dict:
|
|
logger.warning_once(f"Found kv_scale in the checkpoint (e.g. {name}), "
|
|
"but not found the expected name in the model "
|
|
f"(e.g. {remapped_name}). kv_scale is "
|
|
"not loaded.")
|
|
return None
|
|
return remapped_name
|
|
|
|
possible_scale_names = [".k_scale", ".v_scale"]
|
|
modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"]
|
|
for scale_name in possible_scale_names:
|
|
if name.endswith(scale_name):
|
|
if any(mo_scale_name in name for mo_scale_name in modelopt_scale_names):
|
|
remapped_name = name.replace(f".self_attn.{scale_name[1]}_proj{scale_name}",
|
|
f".self_attn.attn{scale_name}")
|
|
else:
|
|
remapped_name = name.replace(scale_name, f".attn{scale_name}")
|
|
if remapped_name not in params_dict:
|
|
logger.warning_once(f"Found {scale_name} in the checkpoint (e.g. {name}), "
|
|
"but not found the expected name in the model "
|
|
f"(e.g. {remapped_name}). {scale_name} is "
|
|
"not loaded.")
|
|
return None
|
|
return remapped_name
|
|
|
|
# If there were no matches, return the untouched param name
|
|
return name
|