339 lines
13 KiB
Python
339 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 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 local rank 0 should read GPU weights and broadcast them to the other local ranks.
|
|
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()
|
|
local_rank = node_group.local_rank if node_group is not None else int(
|
|
os.environ.get("LOCAL_RANK", 0))
|
|
device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu"
|
|
enable_tqdm = not torch.distributed.is_initialized() or local_rank == 0
|
|
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 local_rank == 0:
|
|
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()
|
|
local_rank = node_group.local_rank if node_group is not None else int(
|
|
os.environ.get("LOCAL_RANK", 0))
|
|
device = str(parallel_state.get_local_torch_device()) if not to_cpu else "cpu"
|
|
enable_tqdm = not torch.distributed.is_initialized() or local_rank == 0
|
|
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
|