wehub-resource-sync
756 lines
24 KiB
Python
756 lines
24 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import copy
|
|
import os
|
|
import re
|
|
from collections import defaultdict
|
|
from dataclasses import replace
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
from safetensors.numpy import safe_open
|
|
|
|
import paddle
|
|
from paddle.distributed.fleet.utils.log_util import logger
|
|
|
|
from ..aoa.aoa_engine import (
|
|
postprocess_transpose,
|
|
)
|
|
from .metadata import (
|
|
LocalTensorIndex,
|
|
LocalTensorMetadata,
|
|
Metadata,
|
|
)
|
|
from .sharded_weight import (
|
|
ShardedWeight,
|
|
ShardedWeightDesc,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle.framework import core
|
|
|
|
|
|
def get_coordinator(mesh: np.array | list[list[int]], rank: int):
|
|
mesh = paddle.to_tensor(mesh)
|
|
rand_coordinator = (mesh == rank).nonzero()
|
|
assert rand_coordinator.shape[0] in (
|
|
0,
|
|
1,
|
|
), f"rand_coordinator.shape: {rand_coordinator.shape}"
|
|
return (
|
|
rand_coordinator[0].tolist() if rand_coordinator.shape[0] > 0 else None
|
|
)
|
|
|
|
|
|
# NOTE(zhangbo): Refer to the BalancedSplit function in the reshard_utils.cc file.
|
|
def balanced_split(total_nums, num_of_pieces):
|
|
has_remainder = total_nums % num_of_pieces != 0
|
|
result = [(total_nums + num_of_pieces - 1) // num_of_pieces] * num_of_pieces
|
|
if has_remainder:
|
|
last_value = result[-1]
|
|
result[-1] = last_value - (last_value * num_of_pieces - total_nums)
|
|
return result
|
|
|
|
|
|
def compute_local_shape_and_global_offset(
|
|
global_shape: list[int],
|
|
process_mesh: core.ProcessMesh,
|
|
placements: list[core.Placement],
|
|
) -> tuple[tuple[int], tuple[int]]:
|
|
from paddle.distributed.auto_parallel.placement_type import (
|
|
placemetns_to_dist_status,
|
|
)
|
|
|
|
mesh = np.array(process_mesh.process_ids).reshape(process_mesh.shape)
|
|
# deal with cross mesh case
|
|
if paddle.distributed.get_rank() not in mesh:
|
|
return (None, None)
|
|
rank_coordinator = get_coordinator(mesh, paddle.distributed.get_rank())
|
|
local_shape = copy.copy(global_shape)
|
|
global_offset = [0 for _ in global_shape]
|
|
|
|
dims_mapping, _ = placemetns_to_dist_status(placements, len(global_shape))
|
|
for tensor_dim, mesh_dims in enumerate(dims_mapping):
|
|
if len(mesh_dims) == 0:
|
|
continue
|
|
local_offset = [0] * len(global_shape)
|
|
for mesh_dim in mesh_dims:
|
|
chunk_idx = rank_coordinator[mesh_dim]
|
|
chunks = balanced_split(
|
|
local_shape[tensor_dim], process_mesh.shape[mesh_dim]
|
|
)
|
|
local_shape[tensor_dim] = chunks[chunk_idx]
|
|
local_offset[tensor_dim] = sum(chunks[:chunk_idx])
|
|
|
|
if global_offset[tensor_dim] <= local_offset[tensor_dim]:
|
|
global_offset[tensor_dim] = local_offset[tensor_dim]
|
|
else:
|
|
global_offset[tensor_dim] += local_offset[tensor_dim]
|
|
|
|
return tuple(local_shape), tuple(global_offset)
|
|
|
|
|
|
def flatten_state_dict(state_dict):
|
|
"""
|
|
Flatten the nested dict to a flat dict.
|
|
{"model": {"w0": xxx}} -> {model.w0: xxx}
|
|
"""
|
|
flatten_state_dict = {}
|
|
mapping = {}
|
|
|
|
def _flatten(key, value):
|
|
nonlocal _flatten
|
|
if isinstance(value, dict):
|
|
for k, v in value.items():
|
|
assert isinstance(k, str), f"The key should be str, but is {k}"
|
|
_flatten((*key, k), v)
|
|
elif isinstance(value, (paddle.Tensor, ShardedWeight)):
|
|
flatten_key_str = ".".join(key)
|
|
flatten_state_dict[flatten_key_str] = value
|
|
mapping[flatten_key_str] = key
|
|
else:
|
|
raise ValueError(
|
|
f"The value should be dict or paddle.Tensor, but is {value}"
|
|
)
|
|
|
|
_flatten((), state_dict)
|
|
del _flatten # force python gc of recursive closure
|
|
|
|
return flatten_state_dict, mapping
|
|
|
|
|
|
def unflatten_state_dict(flat_state_dict, mapping):
|
|
"""
|
|
Unflatten the flat dict to a nested dict.
|
|
{model.w0: xxx} -> {"model": {"w0": xxx}}
|
|
"""
|
|
state_dict = {}
|
|
for key, value in flat_state_dict.items():
|
|
key_tuple = mapping[key]
|
|
assert isinstance(key_tuple, tuple), (
|
|
f"The key should be tuple, but is {key_tuple}"
|
|
)
|
|
tmp = state_dict
|
|
for i in range(len(key_tuple) - 1):
|
|
key = key_tuple[i]
|
|
tmp = tmp.setdefault(key, {})
|
|
tmp[key_tuple[-1]] = value
|
|
|
|
return state_dict
|
|
|
|
|
|
def get_max_id(path):
|
|
numbers = [0]
|
|
pattern = re.compile(r"^(\d+)_(\d+)\.distcp$")
|
|
files = os.listdir(path)
|
|
for file in files:
|
|
match = pattern.match(file)
|
|
if match:
|
|
numbers.append(int(match.group(2)))
|
|
return max(numbers) if numbers else None
|
|
|
|
|
|
def check_unique_id(unique_id, process_group):
|
|
all_unique_id = []
|
|
paddle.distributed.all_gather_object(
|
|
all_unique_id, unique_id, process_group
|
|
)
|
|
for id in all_unique_id[1:]:
|
|
assert id == all_unique_id[0], f"id:{id} != all_unique_id[0]"
|
|
|
|
|
|
def ravel_index(indices, shape):
|
|
idx = 0
|
|
for i, dim in zip(indices, shape):
|
|
idx = idx * dim + i
|
|
return idx
|
|
|
|
|
|
def unravel_index(idx, shape):
|
|
indices = []
|
|
for dim in reversed(shape):
|
|
indices.append(idx % dim)
|
|
idx //= dim
|
|
return tuple(reversed(indices))
|
|
|
|
|
|
def minimal_nd_slice(shape, flat_start, flat_end):
|
|
start_idx = unravel_index(flat_start, shape)
|
|
end_idx = unravel_index(flat_end - 1, shape)
|
|
min_slices = []
|
|
for axis in range(len(shape)):
|
|
if axis == 0:
|
|
s = start_idx[axis]
|
|
e = end_idx[axis] + 1
|
|
else:
|
|
if start_idx[axis - 1] == end_idx[axis - 1]:
|
|
s = min(start_idx[axis], end_idx[axis])
|
|
e = max(start_idx[axis], end_idx[axis]) + 1
|
|
else:
|
|
s = 0
|
|
e = shape[axis]
|
|
min_slices.append((s, e))
|
|
return min_slices, start_idx, end_idx
|
|
|
|
|
|
def flat_range_in_min_slice(shape, min_slices, flat_start, flat_end):
|
|
min_starts = tuple(s[0] for s in min_slices)
|
|
min_flat_start = ravel_index(min_starts, shape)
|
|
return flat_start - min_flat_start, flat_end - min_flat_start
|
|
|
|
|
|
def is_sharded_state_dict(state_dict, use_dist=True, process_group=None):
|
|
values = list(state_dict.values())
|
|
is_all_sharded = all(isinstance(v, ShardedWeight) for v in values)
|
|
has_sharded = any(isinstance(v, ShardedWeight) for v in values)
|
|
|
|
if has_sharded and not is_all_sharded:
|
|
raise TypeError(
|
|
"All values must be ShardedWeight if any value is ShardedWeight."
|
|
)
|
|
|
|
if not use_dist:
|
|
return is_all_sharded
|
|
|
|
if is_all_sharded:
|
|
flag = 1
|
|
elif len(values) == 0:
|
|
flag = 0
|
|
else:
|
|
flag = -1
|
|
|
|
all_flags = []
|
|
paddle.distributed.all_gather_object(all_flags, flag, process_group)
|
|
|
|
assert all(f >= 0 for f in all_flags) or all(f <= 0 for f in all_flags), (
|
|
"Not support mixed type of ShardedWeight and non-ShardedWeight in the same state_dict!"
|
|
)
|
|
return all(f >= 0 for f in all_flags)
|
|
|
|
|
|
def get_overlap_region(desc_offset, desc_shape, shard_offset, shard_shape):
|
|
ndim = len(desc_offset)
|
|
overlap_offset = []
|
|
overlap_shape = []
|
|
desc_starts = []
|
|
shard_starts = []
|
|
for i in range(ndim):
|
|
desc_lo = desc_offset[i]
|
|
desc_hi = desc_offset[i] + desc_shape[i]
|
|
shard_lo = shard_offset[i]
|
|
shard_hi = shard_offset[i] + shard_shape[i]
|
|
# overlap
|
|
lo = max(desc_lo, shard_lo)
|
|
hi = min(desc_hi, shard_hi)
|
|
if lo >= hi:
|
|
return False, None, None, None, None
|
|
overlap_offset.append(lo)
|
|
overlap_shape.append(hi - lo)
|
|
desc_starts.append(lo - desc_lo)
|
|
shard_starts.append(lo - shard_lo)
|
|
return True, overlap_offset, overlap_shape, desc_starts, shard_starts
|
|
|
|
|
|
def assign_sharded_slice(
|
|
src_desc, src_shard, dst_desc, dst_shard, postprocess_list=None
|
|
):
|
|
src_has, _, overlap_shape, src_desc_starts, src_shard_starts = (
|
|
get_overlap_region(
|
|
src_desc.global_offset,
|
|
src_desc.local_shape,
|
|
src_shard.global_offset,
|
|
src_shard.local_shape,
|
|
)
|
|
)
|
|
|
|
dst_has, _, overlap_shape2, dst_desc_starts, dst_shard_starts = (
|
|
get_overlap_region(
|
|
dst_desc.global_offset,
|
|
dst_desc.local_shape,
|
|
dst_shard.global_offset,
|
|
dst_shard.local_shape,
|
|
)
|
|
)
|
|
|
|
assert src_has or dst_has, "no overlap!"
|
|
if overlap_shape != overlap_shape2:
|
|
assert postprocess_list is not None, (
|
|
"only post transpose operation could make overlap shape mismatch"
|
|
)
|
|
transposed_src_overlap_shape = postprocess_transpose(
|
|
overlap_shape, postprocess_list
|
|
)
|
|
|
|
assert transposed_src_overlap_shape == overlap_shape2, (
|
|
f"overlap shape mismatch: {transposed_src_overlap_shape} vs {overlap_shape2}"
|
|
)
|
|
axes = list(range(len(transposed_src_overlap_shape)))
|
|
|
|
src_tensor_slice = paddle.slice(
|
|
src_shard.local_tensor,
|
|
axes=axes,
|
|
starts=src_shard_starts,
|
|
ends=[s + o for s, o in zip(src_shard_starts, overlap_shape)],
|
|
)
|
|
|
|
dst_tensor_slice = paddle.slice(
|
|
dst_shard.local_tensor,
|
|
axes=axes,
|
|
starts=dst_shard_starts,
|
|
ends=[s + o for s, o in zip(dst_shard_starts, overlap_shape2)],
|
|
)
|
|
|
|
else:
|
|
axes = list(range(len(overlap_shape)))
|
|
|
|
src_tensor_slice = paddle.slice(
|
|
src_shard.local_tensor,
|
|
axes=axes,
|
|
starts=src_shard_starts,
|
|
ends=[s + o for s, o in zip(src_shard_starts, overlap_shape)],
|
|
)
|
|
|
|
dst_tensor_slice = paddle.slice(
|
|
dst_shard.local_tensor,
|
|
axes=axes,
|
|
starts=dst_shard_starts,
|
|
ends=[s + o for s, o in zip(dst_shard_starts, overlap_shape)],
|
|
)
|
|
|
|
if postprocess_list is not None:
|
|
for ps in postprocess_list:
|
|
is_list, result = is_list_string(ps)
|
|
if is_list:
|
|
src_tensor_slice = paddle.transpose(src_tensor_slice, result)
|
|
else:
|
|
if isinstance(ps, str):
|
|
src_tensor_slice = paddle.cast(src_tensor_slice, ps)
|
|
|
|
paddle.assign(src_tensor_slice, dst_tensor_slice)
|
|
|
|
|
|
def merge_shard_info_list(list_of_dicts):
|
|
merged = defaultdict(list)
|
|
for info in list_of_dicts:
|
|
for k, v in info.items():
|
|
merged[k].extend(v)
|
|
return dict(merged)
|
|
|
|
|
|
def build_shard_desc(val):
|
|
return ShardedWeightDesc(
|
|
key=val.key,
|
|
local_shape=tuple(val.local_shape),
|
|
global_shape=tuple(val.global_shape),
|
|
global_offset=tuple(val.global_offset),
|
|
dtype=str(val.local_tensor.dtype).split(".")[-1],
|
|
)
|
|
|
|
|
|
def is_list_string(s):
|
|
try:
|
|
result = ast.literal_eval(s)
|
|
return (True, result) if isinstance(result, list) else (False, None)
|
|
except:
|
|
return False, None
|
|
|
|
|
|
def write_to_file_if_empty(data, path):
|
|
lock_path = f"{path}.lock"
|
|
try:
|
|
fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
|
os.close(fd)
|
|
try:
|
|
if os.path.exists(path) and os.path.getsize(path) > 0:
|
|
logger.info(
|
|
f"Process {os.getpid()} found the metadata file already written."
|
|
)
|
|
return
|
|
paddle.save(data, path)
|
|
logger.info(
|
|
f"Process {os.getpid()} successfully wrote the metadata to the file."
|
|
)
|
|
finally:
|
|
if os.path.exists(lock_path):
|
|
os.remove(lock_path)
|
|
except FileExistsError:
|
|
logger.info(
|
|
f"Process {os.getpid()} could not acquire the lock; another process is writing or has written the metadata."
|
|
)
|
|
|
|
|
|
def build_global_state_shard_info(sharded_state_dict, process_group):
|
|
state_shard_info = defaultdict(list)
|
|
for key, val in sharded_state_dict.items():
|
|
desc = build_shard_desc(val)
|
|
state_shard_info[key].append(desc)
|
|
|
|
gathered_info = []
|
|
|
|
use_dist = True if paddle.distributed.get_world_size() > 1 else False
|
|
if use_dist:
|
|
paddle.distributed.all_gather_object(
|
|
gathered_info, dict(state_shard_info), process_group
|
|
)
|
|
else:
|
|
gathered_info = [dict(state_shard_info)]
|
|
|
|
return merge_shard_info_list(gathered_info)
|
|
|
|
|
|
def merge_state_dict_metadata(global_state_dict_metadata):
|
|
assert isinstance(global_state_dict_metadata, list), (
|
|
"The global_state_dict should be a list."
|
|
)
|
|
out = {}
|
|
for state_dict in global_state_dict_metadata:
|
|
for key, val in state_dict.items():
|
|
if key not in out:
|
|
out[key] = []
|
|
|
|
if isinstance(val, list):
|
|
for item in val:
|
|
if item not in out[key]:
|
|
out[key].append(item)
|
|
else:
|
|
if val not in out[key]:
|
|
out[key].append(val)
|
|
|
|
return out
|
|
|
|
|
|
def recover_shard_tensor_from_shards(sharded_weights: list, sw):
|
|
def _assign_slice(dst_tensor, dst_starts, dst_ends, src_tensor):
|
|
axes = list(range(len(dst_starts)))
|
|
view = paddle.slice(
|
|
dst_tensor, axes=axes, starts=dst_starts, ends=dst_ends
|
|
)
|
|
paddle.assign(src_tensor, output=view)
|
|
return dst_tensor
|
|
|
|
dims = len(sw.global_offset)
|
|
sw_glo_start = sw.global_offset
|
|
sw_glo_end = [sw.global_offset[i] + sw.local_shape[i] for i in range(dims)]
|
|
sw_shape = sw.local_shape
|
|
|
|
for s in sharded_weights:
|
|
s_glo_start = s.global_offset
|
|
s_glo_end = [s.global_offset[i] + s.local_shape[i] for i in range(dims)]
|
|
|
|
overlap = []
|
|
for i in range(dims):
|
|
ol_start = max(s_glo_start[i], sw_glo_start[i])
|
|
ol_end = min(s_glo_end[i], sw_glo_end[i])
|
|
if ol_start >= ol_end:
|
|
break
|
|
overlap.append((ol_start, ol_end))
|
|
else:
|
|
s_starts = [ol[0] - s_glo_start[i] for i, ol in enumerate(overlap)]
|
|
s_ends = [ol[1] - s_glo_start[i] for i, ol in enumerate(overlap)]
|
|
sw_starts = [
|
|
ol[0] - sw_glo_start[i] for i, ol in enumerate(overlap)
|
|
]
|
|
sw_ends = [ol[1] - sw_glo_start[i] for i, ol in enumerate(overlap)]
|
|
|
|
axes = list(range(len(s_starts)))
|
|
src = paddle.slice(
|
|
s.local_tensor, axes=axes, starts=s_starts, ends=s_ends
|
|
)
|
|
_assign_slice(sw.local_tensor, sw_starts, sw_ends, src)
|
|
|
|
return sw
|
|
|
|
|
|
def create_hf_ckpt_metadata(
|
|
ckpt_path: str,
|
|
process_group=None,
|
|
):
|
|
dtype_mapping = {
|
|
'U16': 'bfloat16',
|
|
'U8': 'uint8',
|
|
'I8': 'int8',
|
|
'I16': 'int16',
|
|
'BOOL': 'bool',
|
|
'F16': 'float16',
|
|
'F32': 'float32',
|
|
'F64': 'float64',
|
|
'BF16': 'bfloat16',
|
|
'I64': 'int64',
|
|
}
|
|
|
|
use_dist = paddle.distributed.get_world_size() > 1
|
|
cur_rank = paddle.distributed.get_rank() if use_dist else 0
|
|
|
|
accessible_files = os.listdir(ckpt_path)
|
|
safetensors_files = [
|
|
file for file in accessible_files if file.endswith(".safetensors")
|
|
]
|
|
if use_dist:
|
|
rank_visible_files = []
|
|
local_files = {cur_rank: safetensors_files}
|
|
paddle.distributed.all_gather_object(
|
|
rank_visible_files, local_files, process_group
|
|
)
|
|
rank_visible_files = {
|
|
rank: files for d in rank_visible_files for rank, files in d.items()
|
|
}
|
|
else:
|
|
rank_visible_files = {0: safetensors_files}
|
|
|
|
def assign_files(
|
|
rank_visible_files: dict[int, list[str]],
|
|
) -> dict[int, list[str]]:
|
|
all_files = set()
|
|
for files in rank_visible_files.values():
|
|
all_files.update(files)
|
|
all_files = list(all_files)
|
|
|
|
file2ranks = defaultdict(list)
|
|
for rank, files in rank_visible_files.items():
|
|
for f in files:
|
|
file2ranks[f].append(rank)
|
|
|
|
result = defaultdict(list)
|
|
|
|
all_files.sort(key=lambda f: (len(file2ranks[f]), f))
|
|
|
|
rank_load = dict.fromkeys(rank_visible_files, 0)
|
|
|
|
for f in all_files:
|
|
candidates = file2ranks[f]
|
|
min_rank = min(candidates, key=lambda r: (rank_load[r], r))
|
|
result[min_rank].append(f)
|
|
rank_load[min_rank] += 1
|
|
|
|
return {rank: result.get(rank, []) for rank in rank_visible_files}
|
|
|
|
rank2file = assign_files(rank_visible_files)
|
|
need_handle_files = rank2file[cur_rank]
|
|
|
|
local_state_dict_metadata = defaultdict(set)
|
|
local_storage_metadata = {}
|
|
for file_name in need_handle_files:
|
|
file_path = os.path.join(ckpt_path, file_name)
|
|
with safe_open(file_path, framework="np") as f:
|
|
for key in f.keys():
|
|
t_s = f.get_slice(key)
|
|
shape = tuple(t_s.get_shape())
|
|
dtype = t_s.get_dtype()
|
|
assert dtype in dtype_mapping, f"{dtype} is not supported yet."
|
|
dtype = dtype_mapping[dtype]
|
|
ltm = LocalTensorMetadata(
|
|
global_offset=(0,) * len(shape),
|
|
local_shape=shape,
|
|
dtype=dtype,
|
|
global_shape=shape,
|
|
is_flattened=False,
|
|
)
|
|
lti = LocalTensorIndex(
|
|
tensor_key=key,
|
|
global_offset=(0,) * len(shape),
|
|
is_flattened=False,
|
|
local_shape=shape,
|
|
)
|
|
local_state_dict_metadata[key].add(ltm)
|
|
local_storage_metadata[lti] = file_name
|
|
|
|
if use_dist:
|
|
global_state_dict_metadata = []
|
|
global_storage_metadata = []
|
|
paddle.distributed.all_gather_object(
|
|
global_state_dict_metadata,
|
|
dict(local_state_dict_metadata),
|
|
process_group,
|
|
)
|
|
paddle.distributed.all_gather_object(
|
|
global_storage_metadata, local_storage_metadata, process_group
|
|
)
|
|
else:
|
|
global_state_dict_metadata = [dict(local_state_dict_metadata)]
|
|
global_storage_metadata = [local_storage_metadata]
|
|
|
|
state_dict_metadata = defaultdict(set)
|
|
for md in global_state_dict_metadata:
|
|
for k, v in md.items():
|
|
state_dict_metadata[k].update(v)
|
|
state_dict_metadata = {k: list(v) for k, v in state_dict_metadata.items()}
|
|
|
|
storage_metadata = {}
|
|
for md in global_storage_metadata:
|
|
storage_metadata.update(md)
|
|
|
|
metadata = Metadata(
|
|
state_dict_metadata=state_dict_metadata,
|
|
storage_metadata=storage_metadata,
|
|
)
|
|
|
|
METADATA_FILE_NAME = "flex-ckpt.auto_generated.metadata"
|
|
write_to_file_if_empty(
|
|
metadata, os.path.join(ckpt_path, METADATA_FILE_NAME)
|
|
)
|
|
|
|
if use_dist:
|
|
paddle.distributed.barrier(process_group)
|
|
|
|
|
|
def get_target_tensor(target_state_dict, read_item):
|
|
use_dist = paddle.distributed.get_world_size() > 1
|
|
if any(isinstance(k, tuple) for k in target_state_dict):
|
|
key = (read_item.tensor_name, read_item.dst_global_offset)
|
|
else:
|
|
key = read_item.tensor_name
|
|
|
|
tensor = target_state_dict[key]
|
|
return tensor._local_value() if use_dist and tensor.is_dist() else tensor
|
|
|
|
|
|
def slice_tensor(tensor, slice_begin, slice_shape):
|
|
if not slice_shape:
|
|
assert not tensor.shape, (
|
|
"Only 0-dimensional tensor supports empty slice_shape."
|
|
)
|
|
return tensor
|
|
|
|
slice_end = [
|
|
start + length for start, length in zip(slice_begin, slice_shape)
|
|
]
|
|
axes = list(range(tensor.ndim))
|
|
return paddle.slice(tensor, axes=axes, starts=slice_begin, ends=slice_end)
|
|
|
|
|
|
def extract_tensor_metadata(val):
|
|
if isinstance(val, paddle.Tensor):
|
|
# Case1: not initialized means this tensor is placed in another mesh which do not contain this rank
|
|
if not val._is_initialized():
|
|
return None, None
|
|
if val.is_dist():
|
|
local_tensor = val._local_value()
|
|
# Note: The local_tensor must keep the same name with the original tensor. Otherwise, the StructuredToParameterName@@ mapping will be wrong.
|
|
local_tensor.name = val.name
|
|
# when val is scalar, the shape is []
|
|
(
|
|
local_shape,
|
|
global_offset,
|
|
) = (
|
|
compute_local_shape_and_global_offset(
|
|
val.shape,
|
|
val.process_mesh,
|
|
val.placements,
|
|
)
|
|
if len(val.shape) > 0
|
|
else ((), ())
|
|
)
|
|
global_shape = val.shape
|
|
if local_shape is None or global_offset is None:
|
|
return None, None
|
|
else:
|
|
local_shape = tuple(val.shape)
|
|
global_offset = (
|
|
tuple([0] * len(val.shape)) if len(val.shape) > 0 else ()
|
|
)
|
|
global_shape = local_shape
|
|
local_tensor = val
|
|
is_flattened = False
|
|
flattened_range = None
|
|
elif isinstance(val, ShardedWeight):
|
|
local_tensor = val.local_tensor
|
|
local_shape = val.local_shape
|
|
global_offset = val.global_offset
|
|
global_shape = val.global_shape
|
|
is_flattened = val.is_flattened
|
|
flattened_range = val.flattened_range
|
|
else:
|
|
raise ValueError(
|
|
f"The value of state_dict should be a paddle.Tensor, but got: {val}"
|
|
)
|
|
|
|
local_tensor_dtype = str(local_tensor.dtype).split('.')[1]
|
|
if flattened_range is not None:
|
|
flattened_range = (flattened_range.start, flattened_range.stop)
|
|
else:
|
|
flattened_range = None
|
|
local_tensor_metadata = LocalTensorMetadata(
|
|
tuple(global_offset),
|
|
tuple(local_shape),
|
|
local_tensor_dtype,
|
|
tuple(global_shape),
|
|
is_flattened,
|
|
flattened_range,
|
|
)
|
|
assert (local_tensor is None) == (local_tensor_metadata is None), (
|
|
"local_tensor and local_tensor_metadata must both be None or both not None!"
|
|
)
|
|
return local_tensor, local_tensor_metadata
|
|
|
|
|
|
def check_resumable_locally(
|
|
path, state_dict, metadata_manager, use_dist, process_group
|
|
):
|
|
local_load = True
|
|
rank = paddle.distributed.get_rank() if use_dist else 0
|
|
checkpoint_file = f"{rank}_0.distcp"
|
|
file_path = os.path.join(path, checkpoint_file)
|
|
|
|
if not os.path.isfile(file_path):
|
|
local_load = False
|
|
|
|
state_dict_metadata = {}
|
|
for key, value in state_dict.items():
|
|
_, local_tensor_metadata = extract_tensor_metadata(value)
|
|
if local_tensor_metadata is not None:
|
|
state_dict_metadata[key] = local_tensor_metadata
|
|
|
|
if local_load:
|
|
file_storage_info = metadata_manager.get_file_storage_info()
|
|
cur_file_storage = {
|
|
replace(index, replica_id=None)
|
|
for index in file_storage_info.get(checkpoint_file, [])
|
|
}
|
|
|
|
for key, local_tensor_metadata in state_dict_metadata.items():
|
|
local_tensor_index = LocalTensorIndex(
|
|
tensor_key=key,
|
|
global_offset=local_tensor_metadata.global_offset,
|
|
is_flattened=local_tensor_metadata.is_flattened,
|
|
flattened_range=local_tensor_metadata.flattened_range,
|
|
local_shape=local_tensor_metadata.local_shape,
|
|
replica_id=None,
|
|
)
|
|
if local_tensor_index not in cur_file_storage:
|
|
local_load = False
|
|
break
|
|
|
|
if use_dist:
|
|
global_local_loads = []
|
|
paddle.distributed.all_gather_object(
|
|
global_local_loads, local_load, process_group
|
|
)
|
|
return all(global_local_loads)
|
|
else:
|
|
return local_load
|
|
|
|
|
|
def need_transpose(postprocess_list):
|
|
if postprocess_list is None:
|
|
return False
|
|
|
|
for pp in postprocess_list:
|
|
if "[" in pp:
|
|
return True
|
|
else:
|
|
return False
|