chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,721 @@
|
||||
# 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()
|
||||
)
|
||||
Reference in New Issue
Block a user