wehub-resource-sync
200 lines
7.5 KiB
Python
200 lines
7.5 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 copy
|
|
import os
|
|
from types import MethodType
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle.autograd import PyLayer
|
|
|
|
from .auto_dp_utils import in_auto_dp_mode
|
|
from .fully_shard_fusion import FullyShardFusion
|
|
|
|
|
|
def shard_accumulators(parameters_and_grads, optimizer, target_block):
|
|
if getattr(optimizer, "_has_sharded_accumulators", False):
|
|
return
|
|
optimizer._has_sharded_accumulators = True
|
|
for param, _ in parameters_and_grads:
|
|
optimizer._create_accumulators(
|
|
target_block,
|
|
[param],
|
|
)
|
|
target_name = param.name
|
|
if param.name in optimizer._master_weights.keys():
|
|
master_weight = optimizer._master_weights[param.name]
|
|
target_name = master_weight.name
|
|
for key in optimizer._accumulators.keys():
|
|
accumulator = optimizer._accumulators[key][target_name]
|
|
if accumulator.is_dist():
|
|
continue
|
|
origin_accumulator_name = accumulator.name
|
|
|
|
if 'beta' not in key:
|
|
placements = copy.deepcopy(param.placements)
|
|
else:
|
|
placements = [
|
|
dist.Replicate()
|
|
for _ in range(len(param.process_mesh.shape))
|
|
]
|
|
optimizer._accumulators[key][target_name] = dist.shard_tensor(
|
|
accumulator,
|
|
mesh=param.process_mesh,
|
|
placements=placements,
|
|
)
|
|
optimizer._accumulators[key][
|
|
target_name
|
|
].name = origin_accumulator_name
|
|
|
|
def _finish_update_impl(self, block, parameters_and_grads):
|
|
if not isinstance(parameters_and_grads, list):
|
|
parameters_and_grads = parameters_and_grads['params']
|
|
for param, _ in parameters_and_grads:
|
|
param.main_grad = None
|
|
|
|
optimizer._finish_update = MethodType(_finish_update_impl, optimizer)
|
|
|
|
|
|
class FullyShardAuto:
|
|
def __init__(
|
|
self,
|
|
model,
|
|
mesh,
|
|
enable_tensor_fusion_and_overlap=True,
|
|
fsdp_unit_layers=None,
|
|
moe_layers_name=None,
|
|
):
|
|
if enable_tensor_fusion_and_overlap:
|
|
FullyShardFusion(model, mesh, fsdp_unit_layers, moe_layers_name)
|
|
else:
|
|
self.model = model
|
|
self.mesh = mesh
|
|
# use first dims as sharding axis
|
|
self._shard_fn = dist.ShardingStage3(0, self.mesh)
|
|
for param in self.model.parameters():
|
|
param._need_shard_auto = True
|
|
self._shard_fn._shard_parameter(param)
|
|
if not in_auto_dp_mode():
|
|
self._shard_fn._register_hook_for_param_grad(param)
|
|
if in_auto_dp_mode():
|
|
self._register_comm_hook(self.model)
|
|
os.environ["skip_sharding3_output_reshard"] = "1"
|
|
|
|
def _register_comm_hook(self, model):
|
|
def _pre_forward_hook(sublayers):
|
|
@paddle.autograd.no_grad()
|
|
def gather_comm(*_):
|
|
dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
|
|
for key, param in sublayers._parameters.items():
|
|
if param.placements[dp_axis] != dist.Replicate():
|
|
new_placements = copy.deepcopy(param.placements)
|
|
new_placements[dp_axis] = dist.Replicate()
|
|
replicate_param = dist.reshard(
|
|
param, param.process_mesh, new_placements
|
|
)
|
|
param.get_tensor()._share_data_with(
|
|
replicate_param.get_tensor()
|
|
)
|
|
|
|
return gather_comm
|
|
|
|
def _post_forward_hook(sublayers):
|
|
@paddle.autograd.no_grad()
|
|
def shard_comm(*_):
|
|
dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
|
|
for key, param in sublayers._parameters.items():
|
|
if (
|
|
param.trainable
|
|
and param.placements[dp_axis] == dist.Replicate()
|
|
):
|
|
new_placements = copy.deepcopy(param.placements)
|
|
new_placements[dp_axis] = dist.Shard(dp_axis)
|
|
shard_param = dist.reshard(
|
|
param, param.process_mesh, new_placements
|
|
)
|
|
param.get_tensor()._share_data_with(
|
|
shard_param.get_tensor()
|
|
)
|
|
|
|
return shard_comm
|
|
|
|
def _post_backward_hook(param):
|
|
def shard_comm(grad):
|
|
dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
|
|
if param.placements[dp_axis] == dist.Replicate():
|
|
new_placements = copy.deepcopy(param.placements)
|
|
new_placements[dp_axis] = dist.Shard(dp_axis)
|
|
shard_param = dist.reshard(
|
|
param, param.process_mesh, new_placements
|
|
)
|
|
param.get_tensor()._share_data_with(
|
|
shard_param.get_tensor()
|
|
)
|
|
return grad
|
|
|
|
param.register_hook(shard_comm)
|
|
|
|
# register forward hooks
|
|
for name, sublayers in model.named_sublayers(include_self=True):
|
|
sublayers.register_forward_pre_hook(_pre_forward_hook(sublayers))
|
|
sublayers.register_forward_post_hook(_post_forward_hook(sublayers))
|
|
|
|
# register backward hooks
|
|
for param in model.parameters():
|
|
if param.trainable:
|
|
_post_backward_hook(param)
|
|
|
|
# register layer hooks for param sync in tie weights
|
|
self._register_layer_hooks(model)
|
|
|
|
def _register_layer_hooks(self, layer, name="last_layer"):
|
|
def _forward_post_hook(layer, inputs, outputs):
|
|
return LayerHook.apply(
|
|
outputs,
|
|
layer=layer,
|
|
)
|
|
|
|
if layer.parameters(include_sublayers=False):
|
|
layer.register_forward_post_hook(_forward_post_hook)
|
|
for name, sub_layer in layer.named_children():
|
|
self._register_layer_hooks(sub_layer, name)
|
|
|
|
|
|
class LayerHook(PyLayer):
|
|
@staticmethod
|
|
def forward(ctx, inputs, layer):
|
|
ctx.layer = layer
|
|
return inputs
|
|
|
|
@staticmethod
|
|
def backward(ctx, *args):
|
|
layer = ctx.layer
|
|
dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
|
|
for param in layer.parameters(include_sublayers=False):
|
|
if (
|
|
param.trainable
|
|
and param.placements[dp_axis] != dist.Replicate()
|
|
):
|
|
new_placements = copy.deepcopy(param.placements)
|
|
new_placements[dp_axis] = dist.Replicate()
|
|
replicate_param = dist.reshard(
|
|
param, param.process_mesh, new_placements
|
|
)
|
|
param.get_tensor()._share_data_with(
|
|
replicate_param.get_tensor()
|
|
)
|
|
return args
|