206403950e
* add offload opt util (#10607) * fix offload optimizer (#10608) * fix sharding reshard bug (#10613) * reset ema (#10790) * Move fused_quanted_ops and token_dispatcher_utils to FleetY (#10803) * move setup_fp8.py (#10820) * Add tokens_zip_unique_add_subbatch and merge_subbatch_cast ops (#10822) * add tokens_zip_unique_add_subbatch and merge_subbatch_cast ops * enhance ut * enhance ut again * remove duplicate codes * enhance ut again * Sharding reshard supports mismatch parameter name (#10479) * Sharding Reshard Supports N -> 1 (#10532) * cherry-pick https://github.com/PaddlePaddle/PaddleNLP/pull/10531 * support tpdp-ep sharding reshard (#10568) * fix typo * revert sharding_first
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
# Copyright (c) 2023 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 paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
|
|
DygraphShardingOptimizer,
|
|
)
|
|
|
|
from ....transformers.model_utils import unwrap_optimizer
|
|
from .common import is_sharding_opt
|
|
|
|
|
|
def shard(node_model_state, model, optimizer):
|
|
cur_rank = max(node_model_state.group.rank, 0)
|
|
unwrapped_optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
|
|
if unwrapped_optimizer is not None:
|
|
optimizer = unwrapped_optimizer
|
|
assert not is_sharding_opt(optimizer)
|
|
param2rank = optimizer._param2rank
|
|
|
|
def filter_func(key):
|
|
names = key
|
|
param_name = names[1]
|
|
assert param_name in param2rank
|
|
dst_rank = param2rank[param_name]
|
|
return dst_rank == cur_rank
|
|
|
|
else:
|
|
assert not is_sharding_opt(optimizer)
|
|
filter_func = lambda key: True
|
|
|
|
node_model_state.reshard(filter_func)
|
|
return node_model_state
|
|
|
|
|
|
def restore(node_model_state, model, optimizer):
|
|
node_model_state.drop_rank()
|
|
return node_model_state
|