wehub-resource-sync
722 lines
27 KiB
Python
722 lines
27 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.
|
|
|
|
import types
|
|
from abc import ABC, abstractmethod
|
|
from collections import defaultdict
|
|
from dataclasses import replace
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle.distributed.collective import Group
|
|
from paddle.distributed.fleet.utils.log_util import logger
|
|
|
|
from .resharder import ReadItem
|
|
from .utils import (
|
|
get_target_tensor,
|
|
slice_tensor,
|
|
)
|
|
|
|
GROUPED_BATCH_SIZE = 10
|
|
|
|
|
|
class CommunicatorFactory:
|
|
registry = {}
|
|
|
|
@classmethod
|
|
def register(cls, method, creator):
|
|
cls.registry[method] = creator
|
|
|
|
@classmethod
|
|
def create(cls, comm_method, **kwargs):
|
|
if comm_method not in cls.registry:
|
|
raise ValueError(
|
|
f"Unknown communication method '{comm_method}'. "
|
|
f"Available: {list(cls.registry.keys())}"
|
|
)
|
|
return cls.registry[comm_method](**kwargs)
|
|
|
|
|
|
class AbstractCommunicator(ABC):
|
|
@staticmethod
|
|
def schedule_read_items(
|
|
read_items: list[ReadItem],
|
|
) -> dict[str, list[ReadItem]]:
|
|
order_rules = lambda read_item: (
|
|
read_item.tensor_name,
|
|
read_item.src_rank,
|
|
read_item.src_global_offset,
|
|
read_item.dst_rank,
|
|
read_item.dst_local_offset,
|
|
read_item.dst_global_offset
|
|
if read_item.dst_global_offset is not None
|
|
else (),
|
|
read_item.src_local_offset,
|
|
read_item.slice_shape,
|
|
read_item.file_name,
|
|
read_item.dtype,
|
|
)
|
|
# Step 1: Group by tensor_name
|
|
tensor_groups = defaultdict(list)
|
|
for item in read_items:
|
|
tensor_groups[item.tensor_name].append(item)
|
|
|
|
scheduled_items = defaultdict(list)
|
|
|
|
# Step 2: For each tensor_name group, further group by all attributes except dst_rank
|
|
for tensor_name, items in tensor_groups.items():
|
|
grouped_items = defaultdict(list)
|
|
for item in items:
|
|
key = (
|
|
item.src_global_offset,
|
|
item.dst_global_offset,
|
|
item.src_rank,
|
|
item.dst_local_offset,
|
|
item.src_local_offset,
|
|
item.slice_shape,
|
|
item.file_name,
|
|
item.dtype,
|
|
)
|
|
grouped_items[key].append(item)
|
|
|
|
# Step 3: Combine items with the same key into a single ReadItem with all dst_ranks
|
|
for key, grouped_item in grouped_items.items():
|
|
combined_dst_rank = []
|
|
for item in grouped_item:
|
|
combined_dst_rank.extend(item.dst_rank)
|
|
combined_dst_rank = sorted(
|
|
set(combined_dst_rank)
|
|
) # Remove duplicates
|
|
|
|
# Create a new ReadItem with combined dst_ranks
|
|
scheduled_item = ReadItem(
|
|
tensor_name=tensor_name,
|
|
src_global_offset=key[0],
|
|
dst_global_offset=key[1],
|
|
dst_rank=tuple(combined_dst_rank),
|
|
src_rank=key[2],
|
|
dst_local_offset=key[3],
|
|
src_local_offset=key[4],
|
|
slice_shape=key[5],
|
|
file_name=key[6],
|
|
dtype=key[7],
|
|
)
|
|
scheduled_items[tensor_name].append(scheduled_item)
|
|
for key, items in scheduled_items.items():
|
|
scheduled_items[key] = sorted(items, key=order_rules)
|
|
|
|
return dict(sorted(scheduled_items.items()))
|
|
|
|
@staticmethod
|
|
def split_read_items(
|
|
read_items: list[ReadItem],
|
|
) -> (list[ReadItem], list[ReadItem]):
|
|
local_read_items = []
|
|
comm_read_items = []
|
|
|
|
for item in read_items:
|
|
assert len(item.dst_rank) == 1, (
|
|
"Before read_items is split, each ReadItem describes a communication task between one rank and another."
|
|
)
|
|
if item.src_rank == item.dst_rank[0]:
|
|
local_read_items.append(item)
|
|
else:
|
|
comm_read_items.append(item)
|
|
|
|
return local_read_items, comm_read_items
|
|
|
|
@staticmethod
|
|
def process_local_copy_tasks(
|
|
local_tasks, cur_rank, source_state_dict, target_state_dict
|
|
):
|
|
"""
|
|
Complete local copy tasks.
|
|
"""
|
|
logger.debug(
|
|
f"Rank {cur_rank} starting local copy for {len(local_tasks)} tasks."
|
|
)
|
|
for task in local_tasks:
|
|
if task.src_rank != cur_rank:
|
|
continue
|
|
|
|
src_tensor = source_state_dict[task.file_name][task.tensor_name]
|
|
dst_tensor = get_target_tensor(target_state_dict, task)
|
|
src_chunk_tensor = slice_tensor(
|
|
src_tensor, task.src_local_offset, task.slice_shape
|
|
)
|
|
|
|
dst_chunk_tensor = slice_tensor(
|
|
dst_tensor, task.dst_local_offset, task.slice_shape
|
|
)
|
|
if src_chunk_tensor.place == dst_chunk_tensor.place:
|
|
paddle.assign(src_chunk_tensor, dst_chunk_tensor)
|
|
logger.debug(f"Local copy (same device) for task {task}.")
|
|
else:
|
|
tmp = (
|
|
src_chunk_tensor.cuda()
|
|
if dst_chunk_tensor.place.is_gpu_place()
|
|
else src_chunk_tensor.cpu()
|
|
)
|
|
paddle.assign(tmp, dst_chunk_tensor)
|
|
del tmp
|
|
logger.debug(f"Local copy (cross device) for task {task}.")
|
|
|
|
@abstractmethod
|
|
def communicate(self, read_items, state, context):
|
|
pass
|
|
|
|
|
|
class BroadcastCommunicator(AbstractCommunicator):
|
|
"""
|
|
Communicator that uses broadcast operation for data transfer.
|
|
"""
|
|
|
|
def communicate(self, read_items, state, context):
|
|
cur_rank = context['rank']
|
|
process_group = context['process_group']
|
|
|
|
source_state_dict = state['source_state_dict']
|
|
target_state_dict = state['target_state_dict']
|
|
|
|
local_read_items, comm_read_items = (
|
|
BroadcastCommunicator.split_read_items(read_items)
|
|
)
|
|
|
|
logger.info(f"Generated {len(comm_read_items)} communication tasks.")
|
|
logger.info(f"Generated {len(local_read_items)} local tasks.")
|
|
|
|
BroadcastCommunicator.process_local_copy_tasks(
|
|
local_read_items,
|
|
cur_rank,
|
|
source_state_dict,
|
|
target_state_dict,
|
|
)
|
|
|
|
logger.info(
|
|
f"Rank {cur_rank} finished local copy and entered communication phase."
|
|
)
|
|
|
|
comm_tasks = BroadcastCommunicator.schedule_read_items(comm_read_items)
|
|
cnt = 0
|
|
total_task_len = len(comm_tasks)
|
|
for tensor_name, read_items in comm_tasks.items():
|
|
cnt += 1
|
|
if cnt % 500 == 0 or cnt == total_task_len:
|
|
logger.info(
|
|
f"{cnt}/{total_task_len} tasks have been sent/received successfully!"
|
|
)
|
|
|
|
source_tensors = {}
|
|
destination_tensors = {}
|
|
for item in read_items:
|
|
logger.debug(f"Beginning to send/recv task {item}.")
|
|
if item.src_rank == cur_rank:
|
|
src_tensor = source_state_dict[item.file_name][
|
|
item.tensor_name
|
|
]
|
|
if not src_tensor.place.is_gpu_place():
|
|
src_tensor = src_tensor.cuda()
|
|
source_tensors[(tensor_name, item.file_name)] = src_tensor
|
|
elif cur_rank in item.dst_rank:
|
|
dst_tensor = get_target_tensor(target_state_dict, item)
|
|
if not dst_tensor.place.is_gpu_place():
|
|
gpu_dst_tensor = dst_tensor.cuda()
|
|
gpu_dst_tensor.need_cross_device_copy = True
|
|
gpu_dst_tensor.target_tensor = dst_tensor
|
|
destination_tensors[
|
|
(tensor_name, cur_rank, item.dst_global_offset)
|
|
] = gpu_dst_tensor
|
|
else:
|
|
gpu_dst_tensor = dst_tensor
|
|
gpu_dst_tensor.target_tensor = dst_tensor
|
|
destination_tensors[
|
|
(tensor_name, cur_rank, item.dst_global_offset)
|
|
] = dst_tensor
|
|
|
|
for item in read_items:
|
|
logger.debug(f"Beginning to send/recv task {item}.")
|
|
if item.src_rank == cur_rank:
|
|
src_tensor = source_tensors[(tensor_name, item.file_name)]
|
|
src_chunk_tensor = slice_tensor(
|
|
src_tensor, item.src_local_offset, item.slice_shape
|
|
)
|
|
buffer_tensor = src_chunk_tensor.contiguous()
|
|
elif cur_rank in item.dst_rank:
|
|
dst_tensor = destination_tensors[
|
|
(tensor_name, cur_rank, item.dst_global_offset)
|
|
]
|
|
dst_chunk_tensor = slice_tensor(
|
|
dst_tensor, item.dst_local_offset, item.slice_shape
|
|
)
|
|
buffer_tensor = paddle.zeros_like(dst_chunk_tensor)
|
|
paddle.assign(dst_chunk_tensor, buffer_tensor)
|
|
|
|
else:
|
|
buffer_tensor = paddle.zeros(item.slice_shape, item.dtype)
|
|
paddle.distributed.broadcast(
|
|
buffer_tensor, src=item.src_rank, group=process_group
|
|
)
|
|
if cur_rank in item.dst_rank:
|
|
paddle.assign(buffer_tensor, dst_chunk_tensor)
|
|
del buffer_tensor
|
|
|
|
for dst_tensor in destination_tensors.values():
|
|
if getattr(dst_tensor, 'need_cross_device_copy', False):
|
|
target_tensor = dst_tensor.target_tensor
|
|
delattr(dst_tensor, "target_tensor")
|
|
target_tensor.copy_(dst_tensor)
|
|
else:
|
|
target_tensor = dst_tensor.target_tensor
|
|
delattr(dst_tensor, "target_tensor")
|
|
paddle.assign(dst_tensor, target_tensor)
|
|
del dst_tensor
|
|
|
|
del source_tensors
|
|
paddle.distributed.barrier(process_group)
|
|
|
|
logger.info("All communication tasks completed.")
|
|
|
|
|
|
class MultiGroupBroadcastCommunicator(AbstractCommunicator):
|
|
"""
|
|
Communicator that uses broadcast for data transfer across multiple communication groups.
|
|
"""
|
|
|
|
def __init__(self, worker_groups):
|
|
if worker_groups is None:
|
|
raise ValueError(
|
|
"worker_groups must be specified when using multi_group_broadcast."
|
|
)
|
|
self.worker_groups = worker_groups
|
|
|
|
@staticmethod
|
|
def schedule_read_items(
|
|
comm_read_items: list[ReadItem],
|
|
worker_groups: list[Group],
|
|
) -> list[list[ReadItem]]:
|
|
group_members = {}
|
|
name_to_groups = {}
|
|
read_items = []
|
|
|
|
order_rules = lambda read_item: (
|
|
read_item.tensor_name,
|
|
read_item.src_rank,
|
|
read_item.src_global_offset,
|
|
read_item.dst_rank,
|
|
read_item.dst_local_offset,
|
|
read_item.dst_global_offset
|
|
if read_item.dst_global_offset is not None
|
|
else (),
|
|
read_item.src_local_offset,
|
|
read_item.slice_shape,
|
|
read_item.file_name,
|
|
read_item.dtype,
|
|
)
|
|
|
|
def _find_min_group(need_ranks, group_members, name_to_groups):
|
|
min_group = None
|
|
min_size = None
|
|
for name, ranks in group_members.items():
|
|
if need_ranks <= ranks:
|
|
if (min_size is None) or (len(ranks) < min_size):
|
|
min_size = len(ranks)
|
|
min_group = name_to_groups[name]
|
|
assert min_group is not None, f"No group found for {need_ranks}!"
|
|
return min_group
|
|
|
|
for group in worker_groups:
|
|
if len(group.ranks) <= 1:
|
|
continue
|
|
group_members[group.name] = set(group.ranks)
|
|
name_to_groups[group.name] = group
|
|
|
|
for read_item in comm_read_items:
|
|
need_ranks = need_ranks = {*read_item.dst_rank, read_item.src_rank}
|
|
group = _find_min_group(
|
|
need_ranks,
|
|
group_members,
|
|
name_to_groups,
|
|
)
|
|
read_items.append(replace(read_item, comm_group=group))
|
|
|
|
read_items = sorted(read_items, key=order_rules)
|
|
|
|
def _build_group_conflict(group_members: dict[str, set]):
|
|
member_to_groups = defaultdict(set)
|
|
for g, members in group_members.items():
|
|
for m in members:
|
|
member_to_groups[m].add(g)
|
|
group_conflict = defaultdict(set)
|
|
for group_set in member_to_groups.values():
|
|
for g1 in group_set:
|
|
for g2 in group_set:
|
|
if g1 != g2:
|
|
group_conflict[g1].add(g2)
|
|
return group_conflict
|
|
|
|
def _dsatur_coloring(group_conflict: dict[str, set]) -> dict[str, int]:
|
|
import heapq
|
|
|
|
all_groups = sorted(group_conflict.keys())
|
|
sorted_conflict = {g: sorted(group_conflict[g]) for g in all_groups}
|
|
|
|
color_map = {}
|
|
neighbor_colors = {g: set() for g in all_groups}
|
|
uncolored = set(all_groups)
|
|
|
|
degree = {g: len(sorted_conflict[g]) for g in all_groups}
|
|
|
|
heap = []
|
|
for g in all_groups:
|
|
heapq.heappush(heap, (0, -degree[g], g))
|
|
saturation = dict.fromkeys(all_groups, 0)
|
|
|
|
while uncolored:
|
|
while True:
|
|
_, _, node = heapq.heappop(heap)
|
|
if node in uncolored:
|
|
break
|
|
used = neighbor_colors[node]
|
|
color = 0
|
|
while color in used:
|
|
color += 1
|
|
color_map[node] = color
|
|
uncolored.remove(node)
|
|
for neighbor in sorted_conflict[node]:
|
|
if neighbor in uncolored:
|
|
if color not in neighbor_colors[neighbor]:
|
|
neighbor_colors[neighbor].add(color)
|
|
saturation[neighbor] += 1
|
|
heapq.heappush(
|
|
heap,
|
|
(
|
|
-saturation[neighbor],
|
|
-degree[neighbor],
|
|
neighbor,
|
|
),
|
|
)
|
|
return color_map
|
|
|
|
def _assign_batches(tasks, group_color_map):
|
|
batches = defaultdict(list)
|
|
for t in tasks:
|
|
g = t.comm_group.name
|
|
batches[group_color_map[g]].append(t)
|
|
return [
|
|
sorted(batches[c], key=order_rules) for c in sorted(batches)
|
|
]
|
|
|
|
group_conflict = _build_group_conflict(group_members)
|
|
group_color_map = _dsatur_coloring(group_conflict)
|
|
results = _assign_batches(read_items, group_color_map)
|
|
return results
|
|
|
|
def communicate(self, read_items, state, context):
|
|
cur_rank = context['rank']
|
|
process_group = context['process_group']
|
|
worker_groups = self.worker_groups
|
|
|
|
source_state_dict = state['source_state_dict']
|
|
target_state_dict = state['target_state_dict']
|
|
|
|
local_read_items, comm_read_items = (
|
|
MultiGroupBroadcastCommunicator.split_read_items(read_items)
|
|
)
|
|
|
|
logger.info(f"Generated {len(comm_read_items)} communication tasks.")
|
|
logger.info(f"Generated {len(local_read_items)} local tasks.")
|
|
|
|
MultiGroupBroadcastCommunicator.process_local_copy_tasks(
|
|
local_read_items,
|
|
cur_rank,
|
|
source_state_dict,
|
|
target_state_dict,
|
|
)
|
|
results = MultiGroupBroadcastCommunicator.schedule_read_items(
|
|
comm_read_items, worker_groups
|
|
)
|
|
logger.info(
|
|
f"Communication task scheduling completed, {len(results)} batches in total."
|
|
)
|
|
for read_items in results:
|
|
source_tensors = {}
|
|
destination_tensors = {}
|
|
for item in read_items:
|
|
tensor_name = item.tensor_name
|
|
if item.src_rank == cur_rank:
|
|
src_tensor = source_state_dict[item.file_name][tensor_name]
|
|
if not src_tensor.place.is_gpu_place():
|
|
src_tensor = src_tensor.cuda()
|
|
source_tensors[(tensor_name, item.file_name)] = src_tensor
|
|
elif cur_rank in item.dst_rank:
|
|
dst_tensor = get_target_tensor(target_state_dict, item)
|
|
if not dst_tensor.place.is_gpu_place():
|
|
gpu_dst_tensor = dst_tensor.cuda()
|
|
gpu_dst_tensor.need_cross_device_copy = True
|
|
gpu_dst_tensor.target_tensor = dst_tensor
|
|
destination_tensors[
|
|
(tensor_name, cur_rank, item.dst_global_offset)
|
|
] = gpu_dst_tensor
|
|
else:
|
|
gpu_dst_tensor = dst_tensor
|
|
gpu_dst_tensor.target_tensor = dst_tensor
|
|
destination_tensors[
|
|
(tensor_name, cur_rank, item.dst_global_offset)
|
|
] = dst_tensor
|
|
|
|
for item in read_items:
|
|
logger.debug(f"Beginning to send/recv task {item}.")
|
|
tensor_name = item.tensor_name
|
|
if item.src_rank == cur_rank:
|
|
src_tensor = source_tensors[(tensor_name, item.file_name)]
|
|
src_chunk_tensor = slice_tensor(
|
|
src_tensor, item.src_local_offset, item.slice_shape
|
|
)
|
|
buffer_tensor = src_chunk_tensor.contiguous()
|
|
elif cur_rank in item.dst_rank:
|
|
dst_tensor = destination_tensors[
|
|
(tensor_name, cur_rank, item.dst_global_offset)
|
|
]
|
|
dst_chunk_tensor = slice_tensor(
|
|
dst_tensor, item.dst_local_offset, item.slice_shape
|
|
)
|
|
buffer_tensor = paddle.zeros_like(dst_chunk_tensor)
|
|
paddle.assign(dst_chunk_tensor, buffer_tensor)
|
|
|
|
elif cur_rank in item.comm_group.ranks:
|
|
buffer_tensor = paddle.zeros(item.slice_shape, item.dtype)
|
|
else:
|
|
buffer_tensor = None
|
|
|
|
if cur_rank in item.comm_group.ranks:
|
|
paddle.distributed.broadcast(
|
|
buffer_tensor, src=item.src_rank, group=item.comm_group
|
|
)
|
|
|
|
if cur_rank in item.dst_rank:
|
|
paddle.assign(buffer_tensor, dst_chunk_tensor)
|
|
del buffer_tensor
|
|
|
|
for dst_tensor in destination_tensors.values():
|
|
if getattr(dst_tensor, 'need_cross_device_copy', False):
|
|
target_tensor = dst_tensor.target_tensor
|
|
delattr(dst_tensor, "target_tensor")
|
|
target_tensor.copy_(dst_tensor)
|
|
else:
|
|
target_tensor = dst_tensor.target_tensor
|
|
delattr(dst_tensor, "target_tensor")
|
|
paddle.assign(dst_tensor, target_tensor)
|
|
del dst_tensor
|
|
|
|
del source_tensors
|
|
|
|
paddle.distributed.barrier(process_group)
|
|
logger.info("All communication tasks completed.")
|
|
|
|
|
|
class SendRecvCommunicator(AbstractCommunicator):
|
|
"""
|
|
Communicator that uses send/recv operations for data transfer.
|
|
|
|
The process is broken down into batches to manage memory and communication overhead.
|
|
"""
|
|
|
|
def __init__(self, use_group):
|
|
self.use_group = use_group
|
|
|
|
@staticmethod
|
|
def schedule_read_items(
|
|
read_items: list[ReadItem],
|
|
) -> dict[str, list[ReadItem]]:
|
|
order_rules = lambda read_item: (
|
|
read_item.tensor_name,
|
|
read_item.src_rank,
|
|
read_item.src_global_offset,
|
|
read_item.dst_rank,
|
|
read_item.dst_local_offset,
|
|
read_item.dst_global_offset
|
|
if read_item.dst_global_offset is not None
|
|
else (),
|
|
read_item.src_local_offset,
|
|
read_item.slice_shape,
|
|
read_item.file_name,
|
|
read_item.dtype,
|
|
)
|
|
|
|
tensor_groups = defaultdict(list)
|
|
for item in read_items:
|
|
tensor_groups[item.tensor_name].append(item)
|
|
|
|
return dict(sorted(tensor_groups.items()))
|
|
|
|
def communicate(self, read_items, state, context):
|
|
comm_tasks = SendRecvCommunicator.schedule_read_items(read_items)
|
|
cur_rank = context['rank']
|
|
process_group = context['process_group']
|
|
|
|
source_state_dict = state['source_state_dict']
|
|
target_state_dict = state['target_state_dict']
|
|
|
|
total_items = sum(len(items) for items in comm_tasks.values())
|
|
processed_items = 0
|
|
|
|
for batch_data in self._process_batches(
|
|
comm_tasks, cur_rank, source_state_dict
|
|
):
|
|
received_slices = {}
|
|
self._execute_p2p_ops(
|
|
batch_data, cur_rank, use_group=self.use_group
|
|
)
|
|
|
|
for item, tensor in batch_data.source_slices.items():
|
|
if item not in batch_data.local_copy_tasks:
|
|
tensor._clear()
|
|
|
|
received_slices.update(batch_data.target_slices)
|
|
|
|
processed_items += len(batch_data.read_items)
|
|
progress = processed_items / total_items * 100
|
|
logger.info(
|
|
f"Batch communication completed. Progress: {processed_items}/{total_items} ({progress:.1f}%)."
|
|
)
|
|
|
|
self._assign_received_data(received_slices, target_state_dict)
|
|
|
|
for received_slice in received_slices.values():
|
|
received_slice._clear()
|
|
|
|
del received_slices
|
|
|
|
if self.use_group:
|
|
paddle.distributed.barrier(process_group)
|
|
logger.info("All communication tasks completed successfully.")
|
|
|
|
def _process_batches(self, comm_tasks, cur_rank, source_state_dict):
|
|
total_items = sum(len(items) for items in comm_tasks.values())
|
|
item_count = 0
|
|
|
|
batch_read_items = []
|
|
batch_source_slices = {}
|
|
batch_target_slices = {}
|
|
batch_local_copy_tasks = set()
|
|
|
|
for tensor_name, read_items in comm_tasks.items():
|
|
tensors_to_clear = set()
|
|
for item in read_items:
|
|
item_count += 1
|
|
batch_read_items.append(item)
|
|
if cur_rank == item.src_rank:
|
|
src_tensor = source_state_dict[item.file_name][
|
|
item.tensor_name
|
|
]
|
|
src_slice = (
|
|
slice_tensor(
|
|
src_tensor, item.src_local_offset, item.slice_shape
|
|
)
|
|
.cuda()
|
|
.clone()
|
|
)
|
|
batch_source_slices[item] = src_slice
|
|
tensors_to_clear.add(src_tensor)
|
|
if cur_rank in item.dst_rank:
|
|
if cur_rank == item.src_rank:
|
|
batch_local_copy_tasks.add(item)
|
|
batch_target_slices[item] = batch_source_slices[item]
|
|
else:
|
|
dst_slice = paddle.zeros(
|
|
item.slice_shape, dtype=item.dtype
|
|
)
|
|
batch_target_slices[item] = dst_slice
|
|
|
|
if ((item_count % GROUPED_BATCH_SIZE) == 0) or (
|
|
item_count == total_items
|
|
):
|
|
batch_data = types.SimpleNamespace(
|
|
read_items=batch_read_items,
|
|
source_slices=batch_source_slices,
|
|
target_slices=batch_target_slices,
|
|
local_copy_tasks=batch_local_copy_tasks,
|
|
)
|
|
yield batch_data
|
|
batch_read_items = []
|
|
batch_source_slices = {}
|
|
batch_target_slices = {}
|
|
batch_local_copy_tasks = set()
|
|
|
|
for tensor in tensors_to_clear:
|
|
tensor._clear_to_zero_allocation()
|
|
|
|
def _execute_p2p_ops(self, batch_data, cur_rank, use_group):
|
|
p2p_ops = []
|
|
for item in batch_data.read_items:
|
|
if item.src_rank == cur_rank:
|
|
for rank in item.dst_rank:
|
|
if rank != cur_rank:
|
|
send_tensor = batch_data.source_slices[item]
|
|
if use_group:
|
|
p2p_ops.append(
|
|
dist.P2POp(dist.isend, send_tensor, rank)
|
|
)
|
|
else:
|
|
dist.send(send_tensor, rank)
|
|
|
|
if cur_rank in item.dst_rank and item.src_rank != cur_rank:
|
|
recv_tensor = batch_data.target_slices[item]
|
|
if use_group:
|
|
p2p_ops.append(
|
|
dist.P2POp(dist.irecv, recv_tensor, item.src_rank)
|
|
)
|
|
else:
|
|
dist.recv(recv_tensor, item.src_rank)
|
|
|
|
if use_group and p2p_ops:
|
|
logger.info(
|
|
f"Starting batched send/recv for {len(p2p_ops)} P2P operations."
|
|
)
|
|
reqs = dist.batch_isend_irecv(p2p_ops)
|
|
for req in reqs:
|
|
req.wait()
|
|
logger.info("Batched send/recv finished.")
|
|
|
|
def _assign_received_data(self, received_slices, target_state_dict):
|
|
for item, received_slice in received_slices.items():
|
|
dest_tensor = get_target_tensor(target_state_dict, item)
|
|
if not dest_tensor._is_initialized():
|
|
buffer = paddle.zeros_like(dest_tensor)
|
|
buffer._share_buffer_to(dest_tensor)
|
|
|
|
dest_slice = slice_tensor(
|
|
dest_tensor, item.dst_local_offset, item.slice_shape
|
|
)
|
|
|
|
if dest_slice.place != received_slice.place:
|
|
received_slice = received_slice.to(dest_slice.place)
|
|
|
|
paddle.assign(received_slice, dest_slice)
|
|
|
|
|
|
CommunicatorFactory.register(
|
|
"multi_group_broadcast",
|
|
lambda worker_groups: MultiGroupBroadcastCommunicator(worker_groups),
|
|
)
|
|
CommunicatorFactory.register(
|
|
"send_recv", lambda **kwargs: SendRecvCommunicator(use_group=False)
|
|
)
|
|
CommunicatorFactory.register(
|
|
"grouped_send_recv", lambda **kwargs: SendRecvCommunicator(use_group=True)
|
|
)
|
|
CommunicatorFactory.register(
|
|
"broadcast", lambda **kwargs: BroadcastCommunicator()
|
|
)
|