chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,809 @@
|
||||
# Copyright (c) 2026 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 collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.autograd import PyLayer
|
||||
from paddle.distributed.fleet.utils.tensor_fusion_helper import (
|
||||
align,
|
||||
alignment,
|
||||
get_current_device_type,
|
||||
)
|
||||
|
||||
# Global registry for fsdp_context
|
||||
_g_fsdp_context = None
|
||||
|
||||
|
||||
def register_fsdp_context(context):
|
||||
global _g_fsdp_context
|
||||
_g_fsdp_context = context
|
||||
|
||||
|
||||
def get_fsdp_context():
|
||||
return _g_fsdp_context
|
||||
|
||||
|
||||
class BufferState(Enum):
|
||||
# Buffer status for lazy double buffer mechanism
|
||||
#
|
||||
# State transitions:
|
||||
# FREED ──all_gather──> USING ──computation done──> READY ──release──> FREED
|
||||
# ^ │
|
||||
# │ (reuse) │
|
||||
# └────────────────────────────┘
|
||||
|
||||
FREED = 1 # Released, buffer data is sharded, tmp_buffer not allocated
|
||||
USING = 2 # Unsharded and actively in use
|
||||
READY = 3 # Unsharded, marked for lazy release, can be reused
|
||||
SYNCING = 4 # Communication in progress
|
||||
|
||||
|
||||
@dataclass
|
||||
class BufferGroup:
|
||||
params: list = field(default_factory=list)
|
||||
dtype: object = None
|
||||
trainable: bool = None
|
||||
fsdp_unit_id: int = None
|
||||
is_tie: bool = False
|
||||
is_expert_param: bool = False
|
||||
fsdp_group: object = None
|
||||
params_buffer: 'TensorFusionBuffer' = None
|
||||
grads_buffer: 'TensorFusionBuffer' = None
|
||||
params_use_sum: int = 0
|
||||
params_use_cnt: int = 0
|
||||
grads_use_sum: int = 0
|
||||
grads_use_cnt: int = 0
|
||||
|
||||
|
||||
def _dtensor_from_local(local_tensor, mesh, placements):
|
||||
global_dims = list(local_tensor.shape)
|
||||
for idx, placement in enumerate(placements):
|
||||
if placement.is_shard():
|
||||
global_dims[placement.get_dim()] = (
|
||||
global_dims[placement.get_dim()] * mesh.shape[idx]
|
||||
)
|
||||
place = paddle.framework._current_expected_place()
|
||||
place = paddle.framework._get_paddle_place(place)
|
||||
|
||||
return paddle.Tensor(
|
||||
local_tensor,
|
||||
dims=global_dims,
|
||||
process_mesh=mesh,
|
||||
placements=placements,
|
||||
place=place,
|
||||
)
|
||||
|
||||
|
||||
class TensorFusionBuffer:
|
||||
def __init__(
|
||||
self,
|
||||
group_id,
|
||||
params,
|
||||
fsdp_degree,
|
||||
dtype,
|
||||
is_params=False,
|
||||
main_grad_dtype=None,
|
||||
):
|
||||
# Calculate total buffer size needed (with padding)
|
||||
self.group_id = group_id
|
||||
self.fsdp_degree = fsdp_degree
|
||||
self.dtype = dtype
|
||||
self.main_grad_dtype = (
|
||||
main_grad_dtype if main_grad_dtype is not None else dtype
|
||||
)
|
||||
self.total_buffer_size = 0
|
||||
self.param_offsets = {}
|
||||
self.tmp_data_buffer = None
|
||||
self.comm_task = None
|
||||
self.trainable = params[0].trainable
|
||||
|
||||
for param in params:
|
||||
self.param_offsets[param.name] = self.total_buffer_size
|
||||
self.total_buffer_size += self.get_padded_size(param)
|
||||
|
||||
if is_params:
|
||||
# Create fused params_buffer
|
||||
# TODO(lizhenxing): Build full params_buffer on CPU and only move shards to GPU to minimize mem peaks
|
||||
self.data_buffer = paddle.zeros(
|
||||
shape=[self.total_buffer_size],
|
||||
dtype=dtype,
|
||||
)
|
||||
# Use BufferState enum instead of is_shard boolean, initial state is FREED (sharded)
|
||||
self.status = BufferState.FREED
|
||||
|
||||
for param in params:
|
||||
offset = self.param_offsets[param.name]
|
||||
stop_gradient = param.stop_gradient
|
||||
local_shape = param._local_shape
|
||||
param.stop_gradient = True
|
||||
param._local_value().flatten_()
|
||||
paddle.assign(
|
||||
param._local_value(),
|
||||
self.data_buffer._slice(
|
||||
offset,
|
||||
offset + param._numel(),
|
||||
),
|
||||
)
|
||||
|
||||
param._clear_data()
|
||||
param.stop_gradient = stop_gradient
|
||||
param._local_value().get_tensor()._set_dims(local_shape)
|
||||
paddle.device.cuda.empty_cache()
|
||||
|
||||
mesh = dist.auto_parallel.get_mesh()
|
||||
curr_global_rank = paddle.distributed.get_rank()
|
||||
if curr_global_rank in mesh.process_ids:
|
||||
total_nums = self.data_buffer.shape[0]
|
||||
num_of_pieces = mesh.shape[0]
|
||||
piece_len = (total_nums + num_of_pieces - 1) // num_of_pieces
|
||||
rank_relative = mesh.process_ids.index(curr_global_rank)
|
||||
start = rank_relative * piece_len
|
||||
end = min(start + piece_len, total_nums)
|
||||
self.data_buffer = paddle.slice(
|
||||
self.data_buffer, [0], [start], [end]
|
||||
).clone()
|
||||
|
||||
# Init params_buffer attr
|
||||
self.data_buffer.name = "fuse_params_" + str(group_id)
|
||||
self.data_buffer.stop_gradient = params[0].stop_gradient
|
||||
self.data_buffer.optimize_attr = params[0].optimize_attr
|
||||
else:
|
||||
# Create fused grads_buffer with shard
|
||||
self.data_buffer = paddle.zeros(
|
||||
shape=[self.total_buffer_size // self.fsdp_degree],
|
||||
dtype=self.main_grad_dtype,
|
||||
)
|
||||
|
||||
# Register get_main_grad method for each param, returns view_slice of grad_buffer
|
||||
for param in params:
|
||||
if param.trainable:
|
||||
param._fusion_buffer = self
|
||||
param._param_offsets = self.param_offsets
|
||||
|
||||
def get_grad_from_tmp_buf(param):
|
||||
tmp_buffer = param._fusion_buffer.get_tmp_buffer()
|
||||
offset = param._param_offsets[param.name]
|
||||
main_grad = paddle._C_ops.view_slice(
|
||||
tmp_buffer,
|
||||
offset,
|
||||
offset + param._numel(),
|
||||
)
|
||||
return main_grad
|
||||
|
||||
param.get_main_grad = get_grad_from_tmp_buf.__get__(param)
|
||||
|
||||
def get_padded_size(self, param):
|
||||
size = np.prod(param.shape)
|
||||
align_size = (
|
||||
alignment[get_current_device_type()]
|
||||
// align[param.dtype]
|
||||
* self.fsdp_degree
|
||||
)
|
||||
return ((size + align_size - 1) // align_size) * align_size
|
||||
|
||||
def get_tmp_buffer(self):
|
||||
# Reuse tmp_buffer if exists, else create
|
||||
if self.tmp_data_buffer is None:
|
||||
self.tmp_data_buffer = paddle.zeros(
|
||||
shape=[self.total_buffer_size], dtype=self.dtype
|
||||
)
|
||||
return self.tmp_data_buffer
|
||||
|
||||
def clear_tmp_buffer(self):
|
||||
if self.tmp_data_buffer is not None:
|
||||
self.tmp_data_buffer._clear_data()
|
||||
self.tmp_data_buffer = None
|
||||
# paddle.device.cuda.empty_cache()
|
||||
|
||||
|
||||
class FSDPBufferManager:
|
||||
def __init__(
|
||||
self, model, mesh, fsdp_unit_layers=None, moe_layers_name=None
|
||||
):
|
||||
self.model = model
|
||||
self._fsdp_group = mesh.get_group("dp")
|
||||
self.main_grad_dtype = paddle.float32
|
||||
|
||||
# Get EP group if "ep" dimension exists in mesh
|
||||
if "ep" in mesh.dim_names:
|
||||
self._ep_fsdp_group = mesh.get_group("ep")
|
||||
else:
|
||||
self._ep_fsdp_group = self._fsdp_group
|
||||
|
||||
topk = None
|
||||
if hasattr(self.model, 'config') and hasattr(
|
||||
self.model.config, 'num_experts_per_tok'
|
||||
):
|
||||
topk = self.model.config.num_experts_per_tok
|
||||
|
||||
# Layer types to wrap as FSDP sharding layers
|
||||
# Note: 'Qwen3VLTextDecoderLayer' is temporary; fleet models all use 'TransformerLayer'
|
||||
self.fsdp_unit_layers = fsdp_unit_layers or [
|
||||
'TransformerLayer',
|
||||
'Qwen3VLTextDecoderLayer',
|
||||
'Qwen3MoeDecoderLayer',
|
||||
]
|
||||
# Layer types to identify MoE expert layers
|
||||
self.moe_layers_name = moe_layers_name or [
|
||||
'StandardMLPExpert',
|
||||
]
|
||||
|
||||
# Get tie_param_name if using tie_weights
|
||||
self.tie_param_name = None
|
||||
# Note: need add get_input_embeddings in fleet modeling
|
||||
# if hasattr(self.model, "get_input_embeddings"):
|
||||
# self.tie_param_name = self.model.get_input_embeddings().weight.name
|
||||
|
||||
# Create buffer_groups
|
||||
grouped_params, group_is_expert = self._build_groups()
|
||||
self.buffer_groups = []
|
||||
self.param_to_buffer_id = {}
|
||||
|
||||
# Create params_buffer, grads_buffer with groups
|
||||
for gid, params in grouped_params.items():
|
||||
is_expert = group_is_expert.get(gid, False)
|
||||
# Use EP group for expert params, DP group for regular params
|
||||
fsdp_group = self._ep_fsdp_group if is_expert else self._fsdp_group
|
||||
|
||||
params_buffer = TensorFusionBuffer(
|
||||
gid,
|
||||
params,
|
||||
fsdp_group.nranks,
|
||||
params[0].dtype,
|
||||
is_params=True,
|
||||
)
|
||||
|
||||
if not params[0].stop_gradient:
|
||||
grads_buffer = TensorFusionBuffer(
|
||||
gid,
|
||||
params,
|
||||
fsdp_group.nranks,
|
||||
params[0].dtype,
|
||||
main_grad_dtype=self.main_grad_dtype,
|
||||
)
|
||||
else:
|
||||
grads_buffer = None
|
||||
|
||||
if is_expert:
|
||||
_params_use_sum = topk
|
||||
_grads_use_sum = topk
|
||||
else:
|
||||
_params_use_sum = len(params)
|
||||
_grads_use_sum = len(params)
|
||||
self.buffer_groups.append(
|
||||
BufferGroup(
|
||||
params=params,
|
||||
dtype=params[0].dtype,
|
||||
trainable=params[0].trainable,
|
||||
is_expert_param=is_expert,
|
||||
fsdp_group=fsdp_group,
|
||||
params_buffer=params_buffer,
|
||||
grads_buffer=grads_buffer,
|
||||
params_use_sum=_params_use_sum,
|
||||
params_use_cnt=0,
|
||||
grads_use_sum=_grads_use_sum,
|
||||
grads_use_cnt=0,
|
||||
)
|
||||
)
|
||||
|
||||
for param in params:
|
||||
self.param_to_buffer_id[param.name] = gid
|
||||
|
||||
def _build_groups(self):
|
||||
parameters = self.model.parameters()
|
||||
grouped_params = OrderedDict()
|
||||
group_is_expert = {}
|
||||
curr_gid = 0
|
||||
|
||||
param_to_unit_id = {}
|
||||
for unit_id, module in enumerate(self.model.modules()):
|
||||
if type(module).__name__ in self.fsdp_unit_layers:
|
||||
for param in module.parameters():
|
||||
param_to_unit_id[param.name] = unit_id
|
||||
if type(module).__name__ in self.moe_layers_name:
|
||||
for param in module.parameters():
|
||||
param.is_moe_param = True
|
||||
|
||||
temp_groups = []
|
||||
for param in parameters:
|
||||
name = param.name
|
||||
is_expert = getattr(param, "is_moe_param", False)
|
||||
if is_expert:
|
||||
continue
|
||||
is_tie = (
|
||||
self.tie_param_name is not None and name == self.tie_param_name
|
||||
)
|
||||
|
||||
param_attrs = {
|
||||
"dtype": param.dtype,
|
||||
"trainable": param.trainable,
|
||||
"fsdp_unit_id": param_to_unit_id.get(name),
|
||||
"is_tie": is_tie,
|
||||
"is_expert_param": is_expert,
|
||||
}
|
||||
|
||||
found_group = False
|
||||
for param_group in temp_groups:
|
||||
if (
|
||||
param_group.dtype == param_attrs["dtype"]
|
||||
and param_group.trainable == param_attrs["trainable"]
|
||||
and param_group.fsdp_unit_id == param_attrs["fsdp_unit_id"]
|
||||
and param_group.is_tie == param_attrs["is_tie"]
|
||||
and param_group.is_expert_param
|
||||
== param_attrs["is_expert_param"]
|
||||
):
|
||||
param_group.params.append(param)
|
||||
found_group = True
|
||||
break
|
||||
|
||||
# Create new group if no matching
|
||||
if not found_group:
|
||||
temp_groups.append(BufferGroup(params=[param], **param_attrs))
|
||||
|
||||
def group_sort_key(group):
|
||||
priority = 0 if group.is_tie else (1 if not group.trainable else 2)
|
||||
return (
|
||||
priority,
|
||||
group.fsdp_unit_id
|
||||
if group.fsdp_unit_id is not None
|
||||
else float('inf'),
|
||||
)
|
||||
|
||||
sorted_groups = sorted(temp_groups, key=group_sort_key)
|
||||
|
||||
# For each sorted parameter group, buffer them by execution order
|
||||
for param_group in sorted_groups:
|
||||
cur_params = param_group.params
|
||||
if len(cur_params) == 0:
|
||||
continue
|
||||
for p in cur_params:
|
||||
grouped_params.setdefault(curr_gid, []).append(p)
|
||||
group_is_expert[curr_gid] = param_group.is_expert_param
|
||||
curr_gid += 1
|
||||
|
||||
return grouped_params, group_is_expert
|
||||
|
||||
|
||||
class FSDPCommManager:
|
||||
def __init__(
|
||||
self,
|
||||
buffer_manager,
|
||||
enable_overlap=True,
|
||||
double_buffer_limit=2,
|
||||
):
|
||||
self.buffer_manager = buffer_manager
|
||||
self.enable_overlap = enable_overlap
|
||||
self.grad_reduce_queue = []
|
||||
|
||||
# for double buffer mechanism config
|
||||
self.double_buffer_limit = double_buffer_limit
|
||||
self.buffer_cnt_in_using = 0
|
||||
self._need_zero_grads = True
|
||||
|
||||
def _release_one_buffer_if_needed(self):
|
||||
# Release a buffer with the READY status if needed
|
||||
while self.buffer_cnt_in_using >= self.double_buffer_limit:
|
||||
found = False
|
||||
for group in self.buffer_manager.buffer_groups:
|
||||
if group.params_buffer.status == BufferState.READY:
|
||||
group.params_buffer.status = BufferState.FREED
|
||||
group.params_buffer.clear_tmp_buffer()
|
||||
self.buffer_cnt_in_using -= 1
|
||||
found = True
|
||||
break
|
||||
if not found:
|
||||
break
|
||||
|
||||
def _next_buffer_id(self, gid, is_backward):
|
||||
# Get next buffer id for prefetch
|
||||
if is_backward:
|
||||
next_gid = gid - 1
|
||||
# Search backward for trainable buffer_groups
|
||||
while (
|
||||
next_gid >= 0
|
||||
and not self.buffer_manager.buffer_groups[
|
||||
next_gid
|
||||
].params_buffer.trainable
|
||||
):
|
||||
next_gid -= 1
|
||||
return max(next_gid, 0)
|
||||
else:
|
||||
return min(gid + 1, len(self.buffer_manager.buffer_groups) - 1)
|
||||
|
||||
def all_gather_params(self, params, is_backward=False):
|
||||
if len(params) == 0:
|
||||
return
|
||||
for param in params:
|
||||
if hasattr(param, "is_moe_param"):
|
||||
continue
|
||||
gid = self.buffer_manager.param_to_buffer_id[param.name]
|
||||
group = self.buffer_manager.buffer_groups[gid]
|
||||
group.params_use_cnt += 1
|
||||
params_buffer = group.params_buffer
|
||||
# Use group-specific fsdp_group
|
||||
fsdp_group = group.fsdp_group or self.buffer_manager._fsdp_group
|
||||
|
||||
# Double buffer: reuse buffer if status is READY
|
||||
if params_buffer.status == BufferState.READY:
|
||||
# Reuse: READY -> USING, no need to all_gather again
|
||||
params_buffer.status = BufferState.USING
|
||||
|
||||
# Overlap prefetch comm
|
||||
if self.enable_overlap:
|
||||
next_gid = self._next_buffer_id(gid, is_backward)
|
||||
next_group = self.buffer_manager.buffer_groups[next_gid]
|
||||
next_params_buffer = next_group.params_buffer
|
||||
next_fsdp_group = (
|
||||
next_group.fsdp_group or self.buffer_manager._fsdp_group
|
||||
)
|
||||
if next_params_buffer.status == BufferState.FREED:
|
||||
# Check double_buffer_limit before prefetch
|
||||
self._release_one_buffer_if_needed()
|
||||
next_params_buffer.status = BufferState.SYNCING
|
||||
tmp_buffer_prefetch = next_params_buffer.get_tmp_buffer()
|
||||
next_params_buffer.comm_task = (
|
||||
paddle.distributed.all_gather(
|
||||
tmp_buffer_prefetch,
|
||||
next_params_buffer.data_buffer,
|
||||
group=next_fsdp_group,
|
||||
sync_op=False,
|
||||
)
|
||||
)
|
||||
self.buffer_cnt_in_using += 1
|
||||
|
||||
# Wait for async comm to complete: SYNCING -> USING
|
||||
if params_buffer.status == BufferState.SYNCING:
|
||||
params_buffer.status = BufferState.USING
|
||||
params_buffer.comm_task.wait()
|
||||
params_buffer.comm_task = None
|
||||
|
||||
tmp_buffer = params_buffer.get_tmp_buffer()
|
||||
# Do all_gather in sync: FREED -> USING
|
||||
if params_buffer.status == BufferState.FREED:
|
||||
fsdp_group.process_group.all_gather(
|
||||
params_buffer.data_buffer, tmp_buffer
|
||||
).wait()
|
||||
params_buffer.status = BufferState.USING
|
||||
self.buffer_cnt_in_using += 1
|
||||
|
||||
# Bind the unsharded param to the real param
|
||||
offset = params_buffer.param_offsets[param.name]
|
||||
tmp_param = paddle._C_ops.view_slice(
|
||||
tmp_buffer,
|
||||
offset,
|
||||
offset + param._numel(),
|
||||
)
|
||||
tmp_param.get_tensor()._set_dims(param.shape)
|
||||
tmp_param = _dtensor_from_local(
|
||||
tmp_param,
|
||||
param.process_mesh,
|
||||
param.placements,
|
||||
)
|
||||
param.get_tensor()._share_data_with(tmp_param.get_tensor())
|
||||
|
||||
def shard_params(self, params, is_backward=False):
|
||||
affected_gids = set()
|
||||
for param in params:
|
||||
if hasattr(param, "is_moe_param"):
|
||||
continue
|
||||
gid = self.buffer_manager.param_to_buffer_id.get(param.name)
|
||||
|
||||
group = self.buffer_manager.buffer_groups[gid]
|
||||
stop_gradient = param.stop_gradient
|
||||
local_shape = param._local_shape
|
||||
param._clear_data()
|
||||
param.stop_gradient = stop_gradient
|
||||
param._local_value().get_tensor()._set_dims(local_shape)
|
||||
|
||||
affected_gids.add(gid)
|
||||
|
||||
for gid in affected_gids:
|
||||
group = self.buffer_manager.buffer_groups[gid]
|
||||
if group.params_buffer.status == BufferState.USING:
|
||||
group.params_buffer.status = BufferState.READY
|
||||
|
||||
def reduce_scatter_grads(self, param):
|
||||
if self._need_zero_grads:
|
||||
self._need_zero_grads = False
|
||||
for group in self.buffer_manager.buffer_groups:
|
||||
if group.grads_buffer is not None:
|
||||
group.grads_buffer.data_buffer.zero_()
|
||||
gid = self.buffer_manager.param_to_buffer_id.get(param.name)
|
||||
group = self.buffer_manager.buffer_groups[gid]
|
||||
group.grads_use_cnt += 1
|
||||
fsdp_group = group.fsdp_group or self.buffer_manager._fsdp_group
|
||||
param.main_grad = None
|
||||
|
||||
if group.grads_use_cnt == group.grads_use_sum:
|
||||
group.grads_use_cnt = 0
|
||||
|
||||
# reduce_scatter from tmp_grad_buffer into grads_buffer
|
||||
grads_buffer = group.grads_buffer
|
||||
|
||||
# Grad queue mechanism: wait and release completed reduce_scatter async tasks
|
||||
self._wait_for_grad_comm()
|
||||
|
||||
tmp_buffer = grads_buffer.get_tmp_buffer()
|
||||
shard_size = grads_buffer.data_buffer.shape[0]
|
||||
grad_buffer_shard = tmp_buffer._slice(0, shard_size)
|
||||
if self.enable_overlap:
|
||||
# Comm grads async and check all comm_task before optimizer update
|
||||
grads_buffer.comm_task = paddle.distributed.reduce_scatter(
|
||||
grad_buffer_shard,
|
||||
tmp_buffer,
|
||||
op=paddle.distributed.ReduceOp.SUM,
|
||||
group=fsdp_group,
|
||||
sync_op=False,
|
||||
)
|
||||
|
||||
# Add async task to queue
|
||||
self.grad_reduce_queue.append(grads_buffer)
|
||||
else:
|
||||
paddle.distributed.reduce_scatter(
|
||||
grad_buffer_shard,
|
||||
tmp_buffer,
|
||||
op=paddle.distributed.ReduceOp.SUM,
|
||||
group=fsdp_group,
|
||||
sync_op=False,
|
||||
).wait()
|
||||
grads_buffer.data_buffer.add_(grad_buffer_shard)
|
||||
grads_buffer.clear_tmp_buffer()
|
||||
|
||||
def _wait_for_grad_comm(self, queue_limit=2):
|
||||
# Wait for async reduce_scatter tasks to complete and release resources
|
||||
# queue_limit: max queue size, default use 2, 0 means wait for all
|
||||
while len(self.grad_reduce_queue) > queue_limit:
|
||||
grads_buffer = self.grad_reduce_queue.pop(0)
|
||||
if grads_buffer.comm_task is not None:
|
||||
grads_buffer.comm_task.wait()
|
||||
grads_buffer.comm_task = None
|
||||
tmp_buffer = grads_buffer.get_tmp_buffer()
|
||||
shard_size = grads_buffer.data_buffer.shape[0]
|
||||
grad_buffer_shard = tmp_buffer._slice(0, shard_size)
|
||||
grads_buffer.data_buffer.add_(grad_buffer_shard)
|
||||
grads_buffer.clear_tmp_buffer()
|
||||
|
||||
def _finish_grads_sync(self):
|
||||
# Wait for all async reduce_scatter tasks, call before optimizer.step()
|
||||
self._wait_for_grad_comm(queue_limit=0)
|
||||
|
||||
def _reset_params_buffer_status(self):
|
||||
for group in self.buffer_manager.buffer_groups:
|
||||
params_buffer = group.params_buffer
|
||||
if params_buffer.status in (BufferState.READY, BufferState.USING):
|
||||
# Clear stale tmp_buffer to force re-all_gather with updated data_buffer
|
||||
params_buffer.clear_tmp_buffer()
|
||||
params_buffer.status = BufferState.FREED
|
||||
if self.buffer_cnt_in_using > 0:
|
||||
self.buffer_cnt_in_using -= 1
|
||||
|
||||
|
||||
class FusionBackwardHook(PyLayer):
|
||||
@staticmethod
|
||||
def forward(ctx, *inputs, layer, comm_manager, recursive=False):
|
||||
ctx.layer = layer
|
||||
ctx.comm_manager = comm_manager
|
||||
ctx.recursive = recursive
|
||||
return inputs if len(inputs) > 1 else inputs[0]
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args):
|
||||
trainable_params = []
|
||||
|
||||
for param in ctx.layer.parameters(include_sublayers=ctx.recursive):
|
||||
if param.trainable:
|
||||
trainable_params.append(param)
|
||||
|
||||
ctx.comm_manager.all_gather_params(trainable_params, is_backward=True)
|
||||
return args
|
||||
|
||||
|
||||
class FusionForwardHook(PyLayer):
|
||||
@staticmethod
|
||||
def forward(ctx, *inputs, layer, comm_manager, recursive=False):
|
||||
ctx.layer = layer
|
||||
ctx.comm_manager = comm_manager
|
||||
ctx.recursive = recursive
|
||||
return inputs
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args):
|
||||
ctx.comm_manager.shard_params(
|
||||
ctx.layer.parameters(include_sublayers=ctx.recursive),
|
||||
is_backward=True,
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class FullyShardFusion:
|
||||
def __init__(
|
||||
self, model, mesh, fsdp_unit_layers=None, moe_layers_name=None
|
||||
):
|
||||
self.model = model
|
||||
self.mesh = self._check_mesh(mesh)
|
||||
self._shard_all_params()
|
||||
self.buffer_manager = FSDPBufferManager(
|
||||
self.model, self.mesh, fsdp_unit_layers, moe_layers_name
|
||||
)
|
||||
self.comm_manager = FSDPCommManager(self.buffer_manager)
|
||||
self.register_tensor_fusion_hooks(self.model)
|
||||
register_fsdp_context(self)
|
||||
|
||||
def _check_mesh(self, mesh, pp_idx=0):
|
||||
if "pp" in mesh.dim_names:
|
||||
mesh = mesh.get_mesh_with_dim("pp", pp_idx)
|
||||
return mesh
|
||||
|
||||
def _shard_all_params(self):
|
||||
def shard_layer_param(layer):
|
||||
for param_name in list(layer._parameters.keys()):
|
||||
param = getattr(layer, param_name)
|
||||
if param is not None:
|
||||
param_placements = [
|
||||
dist.Replicate() for _ in range(len(self.mesh.shape))
|
||||
]
|
||||
if not param.is_dist():
|
||||
param = dist.shard_tensor(
|
||||
param, self.mesh, param_placements
|
||||
)
|
||||
setattr(layer, param_name, param)
|
||||
|
||||
for name, layer in self.model.named_sublayers(include_self=True):
|
||||
shard_layer_param(layer)
|
||||
|
||||
def comm_sync_and_reset_status(self):
|
||||
self.comm_manager._finish_grads_sync()
|
||||
self.comm_manager._reset_params_buffer_status()
|
||||
self.comm_manager._need_zero_grads = True
|
||||
# Reset main_grad for all trainable parameters
|
||||
for param in self.model.parameters():
|
||||
if param.trainable:
|
||||
param.main_grad = None
|
||||
|
||||
def register_tensor_fusion_hooks(self, model):
|
||||
def _pre_forward_hook(sublayers, recursive=False):
|
||||
comm_manager = self.comm_manager
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def all_gather_comm(*_):
|
||||
comm_manager.all_gather_params(
|
||||
sublayers.parameters(include_sublayers=recursive)
|
||||
)
|
||||
|
||||
return all_gather_comm
|
||||
|
||||
def _post_forward_hook(sublayers, recursive=False):
|
||||
comm_manager = self.comm_manager
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def shard_comm(*_):
|
||||
comm_manager.shard_params(
|
||||
sublayers.parameters(include_sublayers=recursive)
|
||||
)
|
||||
|
||||
return shard_comm
|
||||
|
||||
def _update_main_grad_hook(param):
|
||||
comm_manager = self.comm_manager
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def comm_hook(grad):
|
||||
if grad is not None and grad._is_initialized():
|
||||
# Share mem with grads_tmp_buffer
|
||||
_main_grad = param.get_main_grad()
|
||||
_main_grad.get_tensor()._set_dims(grad._local_shape)
|
||||
param.main_grad = _dtensor_from_local(
|
||||
_main_grad,
|
||||
grad.process_mesh,
|
||||
grad.placements,
|
||||
)
|
||||
param.main_grad._local_value().copy_(grad._local_value())
|
||||
grad._clear_data()
|
||||
comm_manager.shard_params([param], is_backward=True)
|
||||
comm_manager.reduce_scatter_grads(param)
|
||||
|
||||
return comm_hook
|
||||
|
||||
def _post_backward_hook(param):
|
||||
param.main_grad = None
|
||||
if hasattr(param, "get_main_grad"):
|
||||
param._register_grad_hook(_update_main_grad_hook(param))
|
||||
|
||||
for param in model.parameters():
|
||||
if param.trainable:
|
||||
_post_backward_hook(param)
|
||||
|
||||
def _register_recursive(layer):
|
||||
is_unit = (
|
||||
type(layer).__name__ in self.buffer_manager.fsdp_unit_layers
|
||||
)
|
||||
|
||||
if is_unit:
|
||||
# For FSDP Unit, register recursive hooks and stop recursion
|
||||
layer.register_forward_pre_hook(
|
||||
_pre_forward_hook(layer, recursive=True)
|
||||
)
|
||||
layer.register_forward_post_hook(
|
||||
_post_forward_hook(layer, recursive=True)
|
||||
)
|
||||
self._register_fusion_layer_hooks(layer, recursive=True)
|
||||
return
|
||||
|
||||
if layer.parameters(include_sublayers=False):
|
||||
layer.register_forward_pre_hook(
|
||||
_pre_forward_hook(layer, recursive=False)
|
||||
)
|
||||
layer.register_forward_post_hook(
|
||||
_post_forward_hook(layer, recursive=False)
|
||||
)
|
||||
self._register_fusion_layer_hooks(layer, recursive=False)
|
||||
|
||||
for child in layer.children():
|
||||
_register_recursive(child)
|
||||
|
||||
_register_recursive(model)
|
||||
|
||||
def _register_fusion_layer_hooks(self, layer, recursive=False):
|
||||
def _forward_post_hook(layer, inputs, outputs):
|
||||
if isinstance(outputs, dict):
|
||||
for key, value in outputs.items():
|
||||
if (
|
||||
isinstance(value, paddle.Tensor)
|
||||
and not value.stop_gradient
|
||||
):
|
||||
outputs[key] = FusionBackwardHook.apply(
|
||||
value,
|
||||
layer=layer,
|
||||
comm_manager=self.comm_manager,
|
||||
recursive=recursive,
|
||||
)
|
||||
return outputs
|
||||
elif isinstance(outputs, tuple):
|
||||
result = FusionBackwardHook.apply(
|
||||
*outputs,
|
||||
layer=layer,
|
||||
comm_manager=self.comm_manager,
|
||||
recursive=recursive,
|
||||
)
|
||||
if not isinstance(result, tuple):
|
||||
result = (result,)
|
||||
return result
|
||||
else:
|
||||
return FusionBackwardHook.apply(
|
||||
outputs,
|
||||
layer=layer,
|
||||
comm_manager=self.comm_manager,
|
||||
recursive=recursive,
|
||||
)
|
||||
|
||||
def _forward_pre_hook(layer, inputs):
|
||||
return FusionForwardHook.apply(
|
||||
*inputs,
|
||||
layer=layer,
|
||||
comm_manager=self.comm_manager,
|
||||
recursive=recursive,
|
||||
)
|
||||
|
||||
layer.register_forward_post_hook(_forward_post_hook)
|
||||
|
||||
# Register an additional hook for tie_weights shard_params
|
||||
for param in layer.parameters(include_sublayers=False):
|
||||
if param.name == self.comm_manager.buffer_manager.tie_param_name:
|
||||
layer.register_forward_pre_hook(_forward_pre_hook)
|
||||
Reference in New Issue
Block a user