wehub-resource-sync
241 lines
11 KiB
Python
241 lines
11 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
import re
|
|
import torch
|
|
import types
|
|
from typing import List, Tuple, Union
|
|
from dataclasses import dataclass
|
|
from .constants import (FP32_WEIGHT_KEY, PARAM, VOCAB_TENSOR, CAT_DIM, PARAM_N_SUB_PARAMS, SUB_PARAM_SHAPE,
|
|
EP_IS_EXPERT_PARAM, EP_NUM_EXPERTS, DS_AUTOTP_UC_META)
|
|
|
|
|
|
@dataclass
|
|
class SubparamShape:
|
|
patterns: List[str]
|
|
shape: Tuple[Union[Tuple[int], int]]
|
|
partition_dim: int
|
|
|
|
|
|
def _get_param_uc_restore_meta(param):
|
|
"""Return the restore-facing view of AutoTP UC metadata for a parameter.
|
|
|
|
AutoTP parameter metadata intentionally serves two separate consumers:
|
|
- restore-time fields at the top level, consumed here by UC loading
|
|
- conversion-time fields under `conversion`, consumed by
|
|
`collect_autotp_universal_checkpoint_info()` in `layers.py`
|
|
"""
|
|
return getattr(param, DS_AUTOTP_UC_META, None)
|
|
|
|
|
|
def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank, tp_world_size):
|
|
meta = _get_param_uc_restore_meta(current_param)
|
|
if not meta:
|
|
return None
|
|
|
|
partition_dim = meta.get('partition_dim')
|
|
logical_shape = meta.get('logical_shape')
|
|
sub_param_shape = meta.get('sub_param_shape')
|
|
sub_param_sizes = meta.get('sub_param_sizes')
|
|
replicated = meta.get('replicated', False)
|
|
|
|
if replicated:
|
|
assert partition_dim is None
|
|
slice_tensor = full_hp_param
|
|
return slice_tensor.flatten()
|
|
|
|
if partition_dim is None:
|
|
return None
|
|
|
|
if logical_shape is None:
|
|
return None
|
|
|
|
full_view = full_hp_param.view(logical_shape)
|
|
|
|
if sub_param_shape is not None:
|
|
if hasattr(sub_param_shape, "shape") and hasattr(sub_param_shape, "partition_dim"):
|
|
shape_spec = sub_param_shape.shape
|
|
partition_dim = sub_param_shape.partition_dim
|
|
else:
|
|
shape_spec = sub_param_shape
|
|
|
|
sub_dim_sizes = shape_spec[partition_dim]
|
|
if not isinstance(sub_dim_sizes, tuple):
|
|
sub_dim_sizes = (sub_dim_sizes, )
|
|
|
|
offset = 0
|
|
merged_chunks = []
|
|
for sub_dim_size in sub_dim_sizes:
|
|
sub_slice = full_view.narrow(partition_dim, offset, sub_dim_size) \
|
|
.chunk(tp_world_size, dim=partition_dim)[tp_rank]
|
|
merged_chunks.append(sub_slice)
|
|
offset += sub_dim_size
|
|
|
|
slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
|
|
return slice_tensor.flatten()
|
|
|
|
if sub_param_sizes is not None:
|
|
if not isinstance(sub_param_sizes, (tuple, list)):
|
|
sub_param_sizes = (sub_param_sizes, )
|
|
|
|
offset = 0
|
|
merged_chunks = []
|
|
for sub_dim_size in sub_param_sizes:
|
|
sub_slice = full_view.narrow(partition_dim, offset, sub_dim_size) \
|
|
.chunk(tp_world_size, dim=partition_dim)[tp_rank]
|
|
merged_chunks.append(sub_slice)
|
|
offset += sub_dim_size
|
|
|
|
slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
|
|
return slice_tensor.flatten()
|
|
|
|
slice_tensor = full_view.chunk(tp_world_size, dim=partition_dim)[tp_rank]
|
|
return slice_tensor.flatten()
|
|
|
|
|
|
def load_hp_checkpoint_state(self, folder, tp_rank, tp_world_size, ep_rank=0, ep_size=1):
|
|
hp_mapping = self._hp_mapping
|
|
hp_mapping.optim_fragment = {}
|
|
|
|
hp_keys = []
|
|
for file in os.listdir(folder):
|
|
# We expect files named something like "exp_avg.pt", "exp_avg_sq.pt", "fp32.pt"
|
|
pattern = r'(.+).pt'
|
|
match = re.search(pattern, file)
|
|
if match:
|
|
hp_keys.append(match.group(1))
|
|
|
|
step = None
|
|
for key in hp_keys:
|
|
ckpt_file = os.path.join(folder, f"{key}.pt")
|
|
ckpt_dict = torch.load(ckpt_file, weights_only=False)
|
|
|
|
if key == "step":
|
|
step = ckpt_dict
|
|
continue
|
|
|
|
full_hp_param = ckpt_dict[PARAM]
|
|
|
|
# EP-aware slicing for expert parameters saved in universal format.
|
|
# Must happen BEFORE shape-match check so that after slicing,
|
|
# full_hp_param.shape == self.shape triggers tp_rank=0, tp_world_size=1.
|
|
is_expert_param = ckpt_dict.get(EP_IS_EXPERT_PARAM, False)
|
|
if is_expert_param and ep_size > 1:
|
|
ep_num_experts = ckpt_dict.get(EP_NUM_EXPERTS)
|
|
assert ep_num_experts is not None, \
|
|
f"Expert param in {ckpt_file} missing '{EP_NUM_EXPERTS}' metadata"
|
|
assert full_hp_param.shape[0] == ep_num_experts, \
|
|
f"Expert param dim 0 ({full_hp_param.shape[0]}) != {EP_NUM_EXPERTS} ({ep_num_experts})"
|
|
assert ep_num_experts % ep_size == 0, \
|
|
f"num_experts ({ep_num_experts}) not divisible by ep_size ({ep_size})"
|
|
num_local = ep_num_experts // ep_size
|
|
ep_start = ep_rank * num_local
|
|
ep_end = ep_start + num_local
|
|
full_hp_param = full_hp_param[ep_start:ep_end]
|
|
|
|
# need to deal with slices that were averaged.
|
|
# the opposite of averaging here becomes an exact copy of the first slice
|
|
# I thought of 2 ways:
|
|
# implementation a. find a way for a client to pass a dict with patterns
|
|
# if any(re.search(pattern, folder) for pattern in WEIGHTS_TO_AVERAGE_PATTERNS):
|
|
# tp_rank = 0
|
|
# tp_world_size = 1
|
|
# the other approach is to assume that the saved data is correct and if full_hp_param.shape ==
|
|
# self.shape that means we automatically copy?
|
|
# implementation b.
|
|
# this version requires no additional data passed from the client
|
|
# if the shapes already match it must be slices that were averaged - so we just hack around those
|
|
if full_hp_param.shape == self.shape:
|
|
tp_rank = 0
|
|
tp_world_size = 1
|
|
|
|
# special case for word_embeddings weights which get padded differently depending on TP degree.
|
|
# the converter to universal currently strips the original padding completely so the saved
|
|
# weight is padding-free and we just need to add new padding depending on the target TP
|
|
# degree
|
|
is_vocab_tensor = ckpt_dict.get(VOCAB_TENSOR, False) and not is_expert_param
|
|
if is_vocab_tensor:
|
|
# In the absence of data passed from the user wrt new padded vocab specific to tp degree
|
|
# we can again derive that data by reverse engineering the target shapes like so:
|
|
padded_target_vocab_size = self.shape[0] * tp_world_size
|
|
assert padded_target_vocab_size >= full_hp_param.shape[0], \
|
|
f'Vocab tensor padded size {padded_target_vocab_size} < loaded universal size {full_hp_param.shape[0]}'
|
|
if padded_target_vocab_size > full_hp_param.shape[0]:
|
|
padding_size = padded_target_vocab_size - full_hp_param.shape[0]
|
|
full_hp_param = torch.nn.functional.pad(full_hp_param, (0, 0, 0, padding_size), "constant", 0)
|
|
|
|
autotp_tp_hp_slice = _resolve_autotp_partition(self, ckpt_dict, full_hp_param, tp_rank, tp_world_size)
|
|
if autotp_tp_hp_slice is not None:
|
|
tp_hp_slice = autotp_tp_hp_slice
|
|
else:
|
|
full_param_numel = full_hp_param.numel()
|
|
tp_slice_numel = self.numel()
|
|
assert full_param_numel == tp_world_size * tp_slice_numel, \
|
|
f'Loading {ckpt_file} full param numel {full_param_numel} != tensor slice numel {tp_slice_numel} * tp_world_size {tp_world_size}'
|
|
|
|
# print(f"{full_hp_param.shape=} {full_param_numel=} {folder=}")
|
|
# print(f"{dst_tensor.shape=} {dst_tensor.numel()=}{folder=}")
|
|
|
|
sub_param_shape = ckpt_dict.get(SUB_PARAM_SHAPE, None)
|
|
# since when we do many to 1 on tp we cat sometimes on dim=0 and other times on dim=1 we have to do exactly the same in reverse
|
|
# special case is when a single parameter is effectively a container for multiple sub parameters
|
|
# (more details at PARAM_N_SUB_PARAMS definition)
|
|
chunk_dim = ckpt_dict.get(CAT_DIM, 0)
|
|
n_sub_params = ckpt_dict.get(PARAM_N_SUB_PARAMS, 1)
|
|
if sub_param_shape:
|
|
partition_dim = sub_param_shape.partition_dim
|
|
sub_dim_sizes = sub_param_shape.shape[partition_dim]
|
|
if not isinstance(sub_dim_sizes, tuple):
|
|
sub_dim_sizes = (sub_dim_sizes, )
|
|
|
|
partition_shape = [sum(d) if isinstance(d, tuple) else d for d in sub_param_shape.shape]
|
|
full_hp_param = full_hp_param.view(partition_shape)
|
|
|
|
offset = 0
|
|
merged_chunks = []
|
|
for sub_dim_size in sub_dim_sizes:
|
|
sub_params_tp_slice = full_hp_param.narrow(partition_dim,
|
|
offset, sub_dim_size).chunk(tp_world_size,
|
|
dim=partition_dim)[tp_rank]
|
|
merged_chunks.append(sub_params_tp_slice)
|
|
offset += sub_dim_size
|
|
tp_hp_slice = torch.cat(merged_chunks, dim=partition_dim)
|
|
|
|
elif n_sub_params > 1:
|
|
sub_params = full_hp_param.chunk(n_sub_params, dim=chunk_dim)
|
|
sub_params_tp_slice = [p.chunk(tp_world_size, dim=chunk_dim)[tp_rank] for p in sub_params]
|
|
tp_hp_slice = torch.cat(sub_params_tp_slice, dim=chunk_dim)
|
|
else:
|
|
# this performs the opposite of cat when merging TP slices
|
|
tp_hp_slice = full_hp_param.chunk(tp_world_size, chunk_dim)[tp_rank]
|
|
|
|
tp_hp_slice = tp_hp_slice.flatten()
|
|
|
|
lp_frag_address = hp_mapping.lp_fragment_address
|
|
tp_hp_fragment = tp_hp_slice.narrow(0, lp_frag_address.start, lp_frag_address.numel)
|
|
|
|
# print(f"{key} SHAPE: {tp_hp_slice.shape=}")
|
|
# print(f"{key} SHAPE: {dst_tensor.shape=}")
|
|
# print(f"{key} SHAPE: {tp_hp_fragment.shape=}")
|
|
|
|
if key == FP32_WEIGHT_KEY:
|
|
dst_tensor = hp_mapping.get_hp_fragment()
|
|
assert dst_tensor.numel() == lp_frag_address.numel, \
|
|
f'Load checkpoint {key} dst numel {dst_tensor.numel()} != src numel {lp_frag_address.numel}'
|
|
dst_tensor.data.copy_(tp_hp_fragment.data)
|
|
else:
|
|
assert tp_hp_fragment.numel() == lp_frag_address.numel, \
|
|
f'Load checkpoint {key} dst numel {tp_hp_fragment.numel()} != src numel {lp_frag_address.numel}'
|
|
|
|
hp_mapping.optim_fragment[key] = tp_hp_fragment.clone().detach()
|
|
|
|
return step
|
|
|
|
|
|
def enable_universal_checkpoint(param_list):
|
|
for param in param_list:
|
|
param.load_hp_checkpoint_state = types.MethodType(load_hp_checkpoint_state, param)
|