wehub-resource-sync
850 lines
30 KiB
Python
850 lines
30 KiB
Python
# Copyright (c) 2025 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 abc
|
|
import math
|
|
from collections import defaultdict
|
|
from copy import deepcopy
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
)
|
|
|
|
import paddle
|
|
|
|
from ..aoa.aoa_engine import SUPPORTED_DTYPES, AOAEngine
|
|
from .resharder import (
|
|
ReadItem,
|
|
)
|
|
from .sharded_weight import (
|
|
ShardedWeight,
|
|
ShardedWeightDesc,
|
|
)
|
|
from .utils import (
|
|
assign_sharded_slice,
|
|
build_shard_desc,
|
|
merge_shard_info_list,
|
|
recover_shard_tensor_from_shards,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Generator, Iterable
|
|
|
|
from paddle.distributed.collective import Group
|
|
|
|
from .sharded_weight import ShardedStateDict
|
|
|
|
|
|
INTERNAL_PADDING_TENSOR_NAME = "__internal_padding_tensor_name__"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExtendReadItem(ReadItem):
|
|
target_tensor_names: tuple[str] | None = None
|
|
global_shape: tuple[int] | None = None
|
|
|
|
|
|
class BaseAssembler(abc.ABC):
|
|
"""
|
|
Abstract base class for assembling full parameters from sharded states.
|
|
|
|
This class encapsulates the common logic for:
|
|
1. Analyzing source and destination tensor mappings (AOA).
|
|
2. Creating a plan to read/communicate necessary tensor shards.
|
|
3. Assembling final tensors once all their source shards are available.
|
|
4. Managing memory by cleaning up consumed shards.
|
|
|
|
Subclasses must implement the `run` method, which defines the specific
|
|
distributed communication strategy to fetch the tensor shards.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sharded_state_dict: ShardedStateDict,
|
|
aoa_config: dict[str, list[str]] | None = None,
|
|
num_splits: int = 1,
|
|
idx: int = 0,
|
|
):
|
|
self.sharded_state_dict = sharded_state_dict
|
|
self.aoa_config = aoa_config or {}
|
|
self.num_splits = num_splits
|
|
self.idx = idx
|
|
|
|
self.cur_rank: int = paddle.distributed.get_rank()
|
|
self.world_size: int = paddle.distributed.get_world_size()
|
|
self.use_dist: bool = self.world_size > 1
|
|
|
|
self.filtered_sharded_state_dict = {}
|
|
self.aoa_engine = None
|
|
self.destination_sharded_weight_desc: dict[str, ShardedWeightDesc] = {}
|
|
self.destination_sharded_mappings = {}
|
|
|
|
self.source_to_target_names: dict[str, set[str]] = defaultdict(set)
|
|
self.source_consumers: dict[str, set[str]] = {}
|
|
self.ref_map: dict[str, set] = {}
|
|
self.read_items: list[ExtendReadItem] = []
|
|
|
|
self.sharded_desc_to_tensor: dict[ShardedWeightDesc, paddle.Tensor] = {}
|
|
|
|
def _prepare_metainfo(self, source_state_shard_info):
|
|
"""Builds destination descriptions and mappings using AOAEngine."""
|
|
self.aoa_engine = AOAEngine(
|
|
aoa_config=self.aoa_config,
|
|
source_state_shard_info=source_state_shard_info,
|
|
destination_state_shard_info=None,
|
|
)
|
|
|
|
output_vars = self.split_output_vars()
|
|
|
|
for k, v in output_vars.items():
|
|
dtype = self.infer_real_dtype(v)
|
|
self.destination_sharded_weight_desc[k] = ShardedWeightDesc(
|
|
key=k,
|
|
local_shape=v.shape,
|
|
global_shape=v.shape,
|
|
global_offset=(0,) * len(v.shape),
|
|
dtype=dtype,
|
|
)
|
|
|
|
for k, desc in self.destination_sharded_weight_desc.items():
|
|
self.destination_sharded_mappings[k] = (
|
|
self.aoa_engine.find_shard_sources(desc)
|
|
)
|
|
|
|
for tgt_name, mapping in self.destination_sharded_mappings.items():
|
|
for m in mapping:
|
|
self.source_to_target_names[m.source_slice.key].add(tgt_name)
|
|
|
|
self.filtered_sharded_state_dict = {
|
|
k: v
|
|
for k, v in self.sharded_state_dict.items()
|
|
if k in self.source_to_target_names
|
|
}
|
|
|
|
self.source_consumers = deepcopy(self.source_to_target_names)
|
|
|
|
def split_output_vars(self):
|
|
data_dict = self.aoa_engine.output_vars
|
|
if self.num_splits < 1:
|
|
raise ValueError('num_splits must be >= 1')
|
|
if self.idx < 0 or self.idx >= self.num_splits:
|
|
raise IndexError(f'idx must be in [0,{self.num_splits - 1}]')
|
|
|
|
sorted_keys = sorted(data_dict.keys())
|
|
total = len(sorted_keys)
|
|
base = total // self.num_splits
|
|
extra = total % self.num_splits
|
|
|
|
if self.idx < extra:
|
|
start = self.idx * (base + 1)
|
|
end = start + (base + 1)
|
|
else:
|
|
start = extra * (base + 1) + (self.idx - extra) * base
|
|
end = start + base
|
|
|
|
selected_keys = sorted_keys[start:end]
|
|
return {k: data_dict[k] for k in selected_keys}
|
|
|
|
def _assemble_and_yield_ready_tensors(
|
|
self, ready_tensor_names: list[str]
|
|
) -> Iterable[tuple[str, paddle.Tensor]]:
|
|
"""
|
|
Assembles, yields, and cleans up tensors whose dependencies are all met.
|
|
This logic is shared across different communication strategies.
|
|
"""
|
|
if not ready_tensor_names:
|
|
return
|
|
|
|
for name in ready_tensor_names:
|
|
target_desc = self.destination_sharded_weight_desc[name]
|
|
local_tensor = paddle.empty(
|
|
target_desc.local_shape, dtype=target_desc.dtype
|
|
)
|
|
cur_sharded_tensor = ShardedWeight(
|
|
key=target_desc.key,
|
|
local_tensor=local_tensor,
|
|
local_shape=target_desc.local_shape,
|
|
global_shape=target_desc.global_shape,
|
|
global_offset=target_desc.global_offset,
|
|
)
|
|
|
|
for mapping in self.destination_sharded_mappings[name]:
|
|
src_desc = mapping.source_slice
|
|
dst_desc = mapping.target_slice
|
|
|
|
src_shard_template = ShardedWeight(
|
|
key=src_desc.key,
|
|
local_tensor=paddle.zeros(
|
|
src_desc.local_shape, dtype=src_desc.dtype
|
|
),
|
|
local_shape=src_desc.local_shape,
|
|
global_shape=src_desc.global_shape,
|
|
global_offset=src_desc.global_offset,
|
|
)
|
|
|
|
received_shards = []
|
|
for desc, tensor in self.sharded_desc_to_tensor.items():
|
|
if desc.key == src_desc.key:
|
|
received_shards.append(
|
|
ShardedWeight(
|
|
key=desc.key,
|
|
local_tensor=tensor,
|
|
local_shape=desc.local_shape,
|
|
global_shape=desc.global_shape,
|
|
global_offset=desc.global_offset,
|
|
)
|
|
)
|
|
|
|
recover_shard_tensor_from_shards(
|
|
received_shards, src_shard_template
|
|
)
|
|
|
|
assign_sharded_slice(
|
|
src_desc=src_desc,
|
|
src_shard=src_shard_template,
|
|
dst_desc=dst_desc,
|
|
dst_shard=cur_sharded_tensor,
|
|
postprocess_list=mapping.postprocess_list,
|
|
)
|
|
src_shard_template.local_tensor._clear()
|
|
|
|
yield name, cur_sharded_tensor.local_tensor
|
|
|
|
need_clear_source_names = self._update_consumer_counts(
|
|
ready_tensor_names
|
|
)
|
|
|
|
self._cleanup_consumed_shards(need_clear_source_names)
|
|
|
|
def _update_consumer_counts(
|
|
self, ready_tensor_names: list[str]
|
|
) -> list[str]:
|
|
"""Decrement consumer counts and return source names that can be cleared."""
|
|
need_clear_source_names = []
|
|
del_keys = []
|
|
for source_name, target_names in self.source_consumers.items():
|
|
target_names.difference_update(ready_tensor_names)
|
|
if not target_names:
|
|
del_keys.append(source_name)
|
|
need_clear_source_names.append(source_name)
|
|
|
|
for k in del_keys:
|
|
del self.source_consumers[k]
|
|
|
|
return need_clear_source_names
|
|
|
|
def dedup_read_items(self, global_read_items):
|
|
group = defaultdict(list)
|
|
for item in global_read_items:
|
|
key = (item.tensor_name, item.src_global_offset, item.slice_shape)
|
|
group[key].append(item)
|
|
result = []
|
|
for key, items in group.items():
|
|
min_item = min(items, key=lambda x: x.src_rank)
|
|
result.append(min_item)
|
|
return result
|
|
|
|
def _cleanup_consumed_shards(self, source_names_to_clear: list[str]):
|
|
"""Delete cached tensors corresponding to the given source names."""
|
|
if not source_names_to_clear:
|
|
return
|
|
|
|
to_delete_descs = []
|
|
for desc, tensor in self.sharded_desc_to_tensor.items():
|
|
if desc.key in source_names_to_clear:
|
|
tensor._clear()
|
|
to_delete_descs.append(desc)
|
|
|
|
for desc in to_delete_descs:
|
|
del self.sharded_desc_to_tensor[desc]
|
|
|
|
@abc.abstractmethod
|
|
def prepare(self):
|
|
"""Subclasses must implement this to build their specific read plan."""
|
|
raise NotImplementedError
|
|
|
|
@abc.abstractmethod
|
|
def run(self) -> Generator[tuple[str, paddle.Tensor], None, None]:
|
|
"""
|
|
The main entry point. Subclasses must implement their communication
|
|
loop and yield final tensors.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abc.abstractmethod
|
|
def all_gather_fn(self, info, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
def infer_real_dtype(self, desc) -> str:
|
|
found_dtypes = []
|
|
for slice_ref in desc.slices:
|
|
key, sl_src, sl_dst, pp_list = slice_ref
|
|
if pp_list is None or len(pp_list) == 0:
|
|
continue
|
|
last_supported = None
|
|
for item in reversed(pp_list):
|
|
if item in SUPPORTED_DTYPES:
|
|
last_supported = item
|
|
break
|
|
if last_supported:
|
|
found_dtypes.append(last_supported)
|
|
if not found_dtypes:
|
|
return desc.dtype
|
|
|
|
dtype_set = set(found_dtypes)
|
|
if len(dtype_set) > 1:
|
|
raise ValueError(
|
|
f"Found multiple different dtypes from slices: {dtype_set}"
|
|
)
|
|
return found_dtypes[0]
|
|
|
|
def build_global_state_shard_info(self, **all_gather_args):
|
|
state_shard_info = defaultdict(list)
|
|
for key, val in self.sharded_state_dict.items():
|
|
desc = build_shard_desc(val)
|
|
state_shard_info[key].append(desc)
|
|
|
|
use_dist = True if paddle.distributed.get_world_size() > 1 else False
|
|
|
|
if use_dist:
|
|
gathered_info = self.all_gather_fn(
|
|
dict(state_shard_info), **all_gather_args
|
|
)
|
|
else:
|
|
gathered_info = [dict(state_shard_info)]
|
|
|
|
return merge_shard_info_list(gathered_info)
|
|
|
|
def get_read_items(
|
|
self,
|
|
all_gather_args=None,
|
|
):
|
|
current_rank = paddle.distributed.get_rank()
|
|
rank_vfile = f"{current_rank}.vdistcp"
|
|
|
|
local_read_plan = []
|
|
for tensor_name, shard_info in self.filtered_sharded_state_dict.items():
|
|
common_attrs = {
|
|
"tensor_name": tensor_name,
|
|
"src_rank": current_rank,
|
|
"src_global_offset": tuple(shard_info.global_offset),
|
|
"dst_global_offset": tuple(shard_info.global_offset),
|
|
"src_local_offset": (0,) * len(shard_info.local_shape),
|
|
"dst_local_offset": (0,) * len(shard_info.local_shape),
|
|
"slice_shape": tuple(shard_info.local_shape),
|
|
"global_shape": tuple(shard_info.global_shape),
|
|
"target_tensor_names": tuple(
|
|
self.source_to_target_names[tensor_name]
|
|
),
|
|
"file_name": rank_vfile,
|
|
"dtype": str(shard_info.local_tensor.dtype).split(".")[1],
|
|
"dst_rank": None,
|
|
"comm_group": None,
|
|
}
|
|
local_read_plan.append(ExtendReadItem(**common_attrs))
|
|
gathered_plans_per_rank = self.all_gather_fn(
|
|
local_read_plan, **(all_gather_args or {})
|
|
)
|
|
|
|
global_read_plan = [
|
|
item for plan in gathered_plans_per_rank for item in plan
|
|
]
|
|
|
|
return self.dedup_read_items(global_read_plan)
|
|
|
|
def group_read_items_by_tensor_name(self, global_read_items):
|
|
groups = defaultdict(list)
|
|
for item in global_read_items:
|
|
groups[item.tensor_name].append(item)
|
|
return groups
|
|
|
|
def sort_groups_for_early_release(self, groups, source_to_target_names):
|
|
def count_fn(name):
|
|
return len(source_to_target_names.get(name, []))
|
|
|
|
sorted_items = sorted(groups.items(), key=lambda x: -count_fn(x[0]))
|
|
return dict(sorted_items)
|
|
|
|
def build_reference_map(self, groups: dict[str, set[ExtendReadItem]]):
|
|
ref_map = defaultdict(set)
|
|
for _, items in groups.items():
|
|
for item in items:
|
|
for tgt in item.target_tensor_names:
|
|
ref_map[tgt].add(item)
|
|
return ref_map
|
|
|
|
def _build_read_plan(self, all_gather_args):
|
|
"""Creates an optimized, sorted list of read operations."""
|
|
read_items = self.get_read_items(
|
|
all_gather_args=all_gather_args,
|
|
)
|
|
grouped = self.group_read_items_by_tensor_name(read_items)
|
|
grouped = self.sort_groups_for_early_release(
|
|
grouped, self.source_to_target_names
|
|
)
|
|
self.ref_map = self.build_reference_map(grouped)
|
|
|
|
self.read_items = [
|
|
item for _, items in grouped.items() for item in items
|
|
]
|
|
|
|
def __iter__(self):
|
|
return self.run()
|
|
|
|
|
|
class SingleCommGroupFullParamAssembler(BaseAssembler):
|
|
"""
|
|
Implements the assembly logic from the original full_param function.
|
|
This version handles both single-card and distributed scenarios.
|
|
In the distributed case, it uses a broadcast-based communication strategy.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sharded_state_dict: ShardedStateDict,
|
|
aoa_config: dict[str, list[str]] | None = None,
|
|
process_group: Group | None = None,
|
|
num_splits: int = 1,
|
|
idx: int = 0,
|
|
):
|
|
super().__init__(sharded_state_dict, aoa_config, num_splits, idx)
|
|
self.process_group = process_group
|
|
|
|
def all_gather_fn(self, info, **kwargs):
|
|
process_group = kwargs.get('process_group', self.process_group)
|
|
gathered_info = []
|
|
paddle.distributed.all_gather_object(gathered_info, info, process_group)
|
|
return gathered_info
|
|
|
|
def is_identity_mapping(self, shard_mappings):
|
|
if len(shard_mappings) != 1:
|
|
return False
|
|
mapping = shard_mappings[0]
|
|
src = mapping.source_slice
|
|
dst = mapping.target_slice
|
|
return (
|
|
src.key == dst.key
|
|
and src.local_shape == dst.local_shape
|
|
and src.global_shape == dst.global_shape
|
|
and src.global_offset == dst.global_offset
|
|
and src.dtype == dst.dtype
|
|
and mapping.postprocess_list is None
|
|
)
|
|
|
|
def prepare(self):
|
|
"""Prepare metadata and build the read plan."""
|
|
source_state_shard_info = self.build_global_state_shard_info(
|
|
process_group=self.process_group
|
|
)
|
|
|
|
self._prepare_metainfo(source_state_shard_info)
|
|
|
|
if self.use_dist:
|
|
self._build_read_plan(
|
|
all_gather_args={"process_group": self.process_group}
|
|
)
|
|
|
|
def run(self) -> Generator[tuple[str, paddle.Tensor], None, None]:
|
|
"""Main execution generator."""
|
|
self.prepare()
|
|
if not self.use_dist:
|
|
yield from self._run_single_card()
|
|
else:
|
|
yield from self._run_distributed()
|
|
|
|
def _run_single_card(
|
|
self,
|
|
) -> Generator[tuple[str, paddle.Tensor], None, None]:
|
|
"""Simple assembly path for a single GPU."""
|
|
for k, v in self.filtered_sharded_state_dict.items():
|
|
assert v.local_shape == v.global_shape, (
|
|
"Single card params must not be sharded.But now the key is {k}, the local_shape is {v.local_shape}, the global_shape is {v.global_shape}."
|
|
)
|
|
|
|
for k, shard_mappings in self.destination_sharded_mappings.items():
|
|
if self.is_identity_mapping(shard_mappings):
|
|
src_key = shard_mappings[0].source_slice.key
|
|
yield (
|
|
k,
|
|
self.filtered_sharded_state_dict[
|
|
src_key
|
|
].local_tensor.clone(),
|
|
)
|
|
else:
|
|
desc = self.destination_sharded_weight_desc[k]
|
|
cur_sharded_tensor = ShardedWeight(
|
|
key=desc.key,
|
|
local_tensor=paddle.empty(
|
|
desc.local_shape, dtype=desc.dtype
|
|
),
|
|
local_shape=desc.local_shape,
|
|
global_shape=desc.global_shape,
|
|
global_offset=desc.global_offset,
|
|
)
|
|
for mapping in shard_mappings:
|
|
source_tensor = self.filtered_sharded_state_dict[
|
|
mapping.source_slice.key
|
|
]
|
|
assign_sharded_slice(
|
|
src_desc=mapping.source_slice,
|
|
src_shard=source_tensor,
|
|
dst_desc=mapping.target_slice,
|
|
dst_shard=cur_sharded_tensor,
|
|
postprocess_list=mapping.postprocess_list,
|
|
)
|
|
yield k, cur_sharded_tensor.local_tensor
|
|
|
|
def _run_distributed(
|
|
self,
|
|
) -> Generator[tuple[str, paddle.Tensor], None, None]:
|
|
"""Distributed assembly using broadcast and packed buffers."""
|
|
for item in self.read_items:
|
|
cur_src_rank = item.src_rank
|
|
|
|
if self.cur_rank == cur_src_rank:
|
|
local_tensor = self.filtered_sharded_state_dict[
|
|
item.tensor_name
|
|
].local_tensor.clone()
|
|
else:
|
|
local_tensor = paddle.empty(item.slice_shape, dtype=item.dtype)
|
|
|
|
on_cpu = local_tensor.place.is_cpu_place()
|
|
if on_cpu:
|
|
local_tensor = local_tensor.cuda()
|
|
paddle.distributed.broadcast(
|
|
local_tensor, src=cur_src_rank, group=self.process_group
|
|
)
|
|
if on_cpu:
|
|
local_tensor = local_tensor.cpu()
|
|
|
|
shard_desc = ShardedWeightDesc(
|
|
key=item.tensor_name,
|
|
local_shape=item.slice_shape,
|
|
global_shape=item.global_shape,
|
|
global_offset=item.src_global_offset,
|
|
dtype=item.dtype,
|
|
)
|
|
self.sharded_desc_to_tensor[shard_desc] = local_tensor
|
|
|
|
ready_tensor_names = []
|
|
for name in item.target_tensor_names:
|
|
self.ref_map[name].remove(item)
|
|
if len(self.ref_map[name]) == 0:
|
|
ready_tensor_names.append(name)
|
|
del self.ref_map[name]
|
|
|
|
yield from self._assemble_and_yield_ready_tensors(
|
|
ready_tensor_names
|
|
)
|
|
|
|
|
|
class OperationType(Enum):
|
|
GLOBAL_BROADCAST = 1
|
|
BROADCAST_ALLGATHER = 2
|
|
|
|
|
|
class HVCommGroupFullParamAssembler(BaseAssembler):
|
|
"""
|
|
Implements the assembly logic using a 2D-mesh communication strategy.
|
|
|
|
This strategy involves a broadcast along the vertical axis of the process
|
|
mesh, followed by an all-gather along the horizontal axis.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sharded_state_dict: ShardedStateDict,
|
|
horizontal_group: Group,
|
|
vertical_group: Group,
|
|
aoa_config: dict[str, list[str]] | None = None,
|
|
num_splits: int = 1,
|
|
idx: int = 0,
|
|
memory_growth_threshold: int = 8 * (2**30), # 8GB
|
|
):
|
|
super().__init__(sharded_state_dict, aoa_config, num_splits, idx)
|
|
self.h_group = horizontal_group
|
|
self.v_group = vertical_group
|
|
self.using_1d_comm_group = (
|
|
self.v_group is None or self.v_group.nranks == 1
|
|
)
|
|
|
|
self.topology: list[list[int]] = []
|
|
self.vertical_ranks: list[set[int]] = []
|
|
self.horizontal_index: dict[int, int] = {}
|
|
self.vertical_index: dict[int, int] = {}
|
|
self.cur_horizontal_index: int = -1
|
|
self.memory_growth_threshold = memory_growth_threshold
|
|
|
|
def all_gather_fn(self, info, **kwargs):
|
|
h_group = kwargs.get('h_group', self.h_group)
|
|
v_group = kwargs.get('v_group', self.v_group)
|
|
|
|
h_obj_list = []
|
|
paddle.distributed.all_gather_object(h_obj_list, info, h_group)
|
|
|
|
v_obj_list = []
|
|
if not self.using_1d_comm_group:
|
|
paddle.distributed.all_gather_object(
|
|
v_obj_list, h_obj_list, v_group
|
|
)
|
|
else:
|
|
v_obj_list = [h_obj_list]
|
|
|
|
gathered_info = [x for sublist in v_obj_list for x in sublist]
|
|
return gathered_info
|
|
|
|
def prepare(self):
|
|
"""Build topology, prepare metadata, and build the read plan."""
|
|
assert self.use_dist, (
|
|
"FullParamAssembler only supports distributed training."
|
|
)
|
|
self._build_topology()
|
|
|
|
source_state_shard_info = self.build_global_state_shard_info(
|
|
h_group=self.h_group, v_group=self.v_group
|
|
)
|
|
self._prepare_metainfo(source_state_shard_info)
|
|
self._build_read_plan(
|
|
all_gather_args={'h_group': self.h_group, 'v_group': self.v_group}
|
|
)
|
|
|
|
def _build_topology(self):
|
|
h_ranks = []
|
|
paddle.distributed.all_gather_object(
|
|
h_ranks, self.cur_rank, self.h_group
|
|
)
|
|
if not self.using_1d_comm_group:
|
|
paddle.distributed.all_gather_object(
|
|
self.topology, h_ranks, self.v_group
|
|
)
|
|
else:
|
|
self.topology = [h_ranks]
|
|
self.vertical_ranks = [set(col) for col in zip(*self.topology)]
|
|
self.horizontal_index = {
|
|
rank: i
|
|
for i, ranks in enumerate(self.vertical_ranks)
|
|
for rank in ranks
|
|
}
|
|
self.vertical_index = {
|
|
rank: i for i, row in enumerate(self.topology) for rank in row
|
|
}
|
|
self.cur_horizontal_index = self.horizontal_index[self.cur_rank]
|
|
|
|
def run(self) -> Generator[tuple[str, paddle.Tensor], None, None]:
|
|
"""Main execution generator using 2D-mesh communication."""
|
|
self.prepare()
|
|
|
|
while len(self.read_items) > 0:
|
|
ready_tensor_names = self._process_one_batch()
|
|
|
|
yield from self._assemble_and_yield_ready_tensors(
|
|
ready_tensor_names
|
|
)
|
|
|
|
def get_batch_read_items(self):
|
|
read_items = self.read_items
|
|
vertical_ranks = self.vertical_ranks
|
|
horizontal_index = self.horizontal_index
|
|
|
|
bathch_read_items = [None] * len(vertical_ranks)
|
|
read_item_index = [None] * len(vertical_ranks)
|
|
cnt = 0
|
|
cur_shape = None
|
|
cur_dtype = None
|
|
for i, item in enumerate(read_items):
|
|
src_rank = item.src_rank
|
|
h_index = horizontal_index[src_rank]
|
|
if bathch_read_items[h_index] is None and cnt == 0:
|
|
bathch_read_items[h_index] = item
|
|
read_item_index[h_index] = i
|
|
cnt += 1
|
|
cur_dtype = item.dtype
|
|
cur_shape = item.slice_shape
|
|
element_size = paddle.core.size_of_dtype(
|
|
getattr(paddle, cur_dtype)
|
|
)
|
|
memory_growth = (
|
|
element_size * math.prod(cur_shape) * len(vertical_ranks)
|
|
)
|
|
if memory_growth > self.memory_growth_threshold:
|
|
return (
|
|
bathch_read_items,
|
|
read_item_index,
|
|
OperationType.GLOBAL_BROADCAST,
|
|
)
|
|
if cnt == len(vertical_ranks):
|
|
return (
|
|
bathch_read_items,
|
|
read_item_index,
|
|
OperationType.GLOBAL_BROADCAST,
|
|
)
|
|
|
|
if bathch_read_items[h_index] is None and cnt != 0:
|
|
if item.slice_shape == cur_shape and item.dtype == cur_dtype:
|
|
bathch_read_items[h_index] = item
|
|
read_item_index[h_index] = i
|
|
cnt += 1
|
|
if cnt == len(vertical_ranks):
|
|
return (
|
|
bathch_read_items,
|
|
read_item_index,
|
|
OperationType.BROADCAST_ALLGATHER,
|
|
)
|
|
|
|
assert cur_shape is not None
|
|
assert cur_dtype is not None
|
|
|
|
for i, item in enumerate(bathch_read_items):
|
|
if item is None:
|
|
src_rank = min(vertical_ranks[i])
|
|
common_attrs = {
|
|
"tensor_name": INTERNAL_PADDING_TENSOR_NAME,
|
|
"src_rank": src_rank,
|
|
"src_global_offset": (0,) * len(cur_shape),
|
|
"dst_global_offset": (0,) * len(cur_shape),
|
|
"src_local_offset": (0,) * len(cur_shape),
|
|
"dst_local_offset": (0,) * len(cur_shape),
|
|
"slice_shape": cur_shape,
|
|
"global_shape": cur_shape,
|
|
"target_tensor_names": None,
|
|
"file_name": "padding_vfile",
|
|
"dtype": cur_dtype,
|
|
"comm_group": None,
|
|
}
|
|
|
|
padding_read_item = ExtendReadItem(
|
|
dst_rank=None, **common_attrs
|
|
)
|
|
bathch_read_items[i] = padding_read_item
|
|
|
|
return (
|
|
bathch_read_items,
|
|
read_item_index,
|
|
OperationType.BROADCAST_ALLGATHER,
|
|
)
|
|
|
|
def _process_one_batch(self) -> list[str]:
|
|
"""Performs V-Broadcast + H-AllGather for one batch of items."""
|
|
|
|
batch_items, batch_indices, op_type = self.get_batch_read_items()
|
|
|
|
if op_type == OperationType.BROADCAST_ALLGATHER:
|
|
read_item = batch_items[self.cur_horizontal_index]
|
|
else:
|
|
values = [x for x in batch_items if x is not None]
|
|
if len(values) == 1:
|
|
read_item = values[0]
|
|
else:
|
|
raise ValueError(
|
|
"When the comm op is GLOBAL_BROADCAST, read_items should be of length 1!"
|
|
)
|
|
batch_items = [read_item]
|
|
|
|
if self.cur_rank == read_item.src_rank:
|
|
buffer = (
|
|
paddle.empty(read_item.slice_shape, read_item.dtype)
|
|
if read_item.tensor_name == INTERNAL_PADDING_TENSOR_NAME
|
|
else self.filtered_sharded_state_dict[
|
|
read_item.tensor_name
|
|
].local_tensor.clone()
|
|
)
|
|
else:
|
|
buffer = paddle.empty(read_item.slice_shape, dtype=read_item.dtype)
|
|
|
|
if op_type == OperationType.BROADCAST_ALLGATHER:
|
|
if not self.using_1d_comm_group:
|
|
paddle.distributed.broadcast(
|
|
buffer, src=read_item.src_rank, group=self.v_group
|
|
)
|
|
tensor_list = []
|
|
paddle.distributed.all_gather(
|
|
tensor_list, buffer, group=self.h_group
|
|
)
|
|
else:
|
|
src_rank = read_item.src_rank
|
|
v_ranks = sorted(
|
|
self.vertical_ranks[self.horizontal_index[src_rank]]
|
|
)
|
|
if self.cur_rank in v_ranks:
|
|
if not self.using_1d_comm_group:
|
|
paddle.distributed.broadcast(
|
|
buffer, src=src_rank, group=self.v_group
|
|
)
|
|
src_rank = v_ranks[self.vertical_index[self.cur_rank]]
|
|
paddle.distributed.broadcast(
|
|
buffer, src=src_rank, group=self.h_group
|
|
)
|
|
tensor_list = [buffer]
|
|
|
|
for idx, item in enumerate(batch_items):
|
|
if item.tensor_name != INTERNAL_PADDING_TENSOR_NAME:
|
|
shard_desc = ShardedWeightDesc(
|
|
key=item.tensor_name,
|
|
local_shape=item.slice_shape,
|
|
global_shape=item.global_shape,
|
|
global_offset=item.src_global_offset,
|
|
dtype=item.dtype,
|
|
)
|
|
self.sharded_desc_to_tensor[shard_desc] = tensor_list[idx]
|
|
|
|
ready_tensor_names = []
|
|
for item in batch_items:
|
|
if item.target_tensor_names:
|
|
for name in item.target_tensor_names:
|
|
self.ref_map[name].remove(item)
|
|
if not self.ref_map[name]:
|
|
ready_tensor_names.append(name)
|
|
del self.ref_map[name]
|
|
|
|
for index in sorted(
|
|
[i for i in batch_indices if i is not None], reverse=True
|
|
):
|
|
del self.read_items[index]
|
|
|
|
return ready_tensor_names
|
|
|
|
|
|
@paddle.no_grad()
|
|
def full_param(
|
|
sharded_state_dict: ShardedStateDict,
|
|
aoa_config: dict[str, list[str]] | None = None,
|
|
**kwargs,
|
|
):
|
|
h_group = kwargs.pop("h_group", None)
|
|
v_group = kwargs.pop("v_group", None)
|
|
process_group = kwargs.pop("process_group", None)
|
|
num_splits = kwargs.pop("num_splits", 1)
|
|
memory_growth_threshold = kwargs.pop("memory_growth_threshold", 8 * (2**30))
|
|
idx = kwargs.pop("shard_idx", 0)
|
|
assert (h_group and v_group) or not (h_group or v_group), (
|
|
"Both horizontal and vertical groups must be provided when using FullParamAssembler."
|
|
)
|
|
if h_group and v_group:
|
|
return HVCommGroupFullParamAssembler(
|
|
sharded_state_dict,
|
|
h_group,
|
|
v_group,
|
|
aoa_config,
|
|
num_splits,
|
|
idx,
|
|
memory_growth_threshold,
|
|
)
|
|
else:
|
|
return SingleCommGroupFullParamAssembler(
|
|
sharded_state_dict, aoa_config, process_group
|
|
)
|