wehub-resource-sync
60 lines
2.1 KiB
Python
60 lines
2.1 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 unittest
|
|
|
|
import paddle
|
|
from paddle.distributed import fleet
|
|
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
|
|
DygraphShardingOptimizerV2,
|
|
)
|
|
|
|
|
|
class TestClearParamStorage(unittest.TestCase):
|
|
def test_clear_param_storage(self):
|
|
class TestLayer(paddle.nn.Layer):
|
|
def __init__(self, dtype):
|
|
super().__init__()
|
|
self._w = self.create_parameter([2, 3], dtype=dtype)
|
|
self._b = self.create_parameter([2, 3], dtype=dtype)
|
|
self._w.color = {"color": "_w"}
|
|
self._b.color = {"color": "_b"}
|
|
|
|
@paddle.amp.debugging.check_layer_numerics
|
|
def forward(self, x):
|
|
return x * self._w + self._b
|
|
|
|
strategy = fleet.DistributedStrategy()
|
|
strategy.hybrid_configs = {
|
|
"dp_degree": 1,
|
|
"mp_degree": 1,
|
|
"pp_degree": 1,
|
|
"sharding_degree": 2,
|
|
}
|
|
fleet.init(is_collective=True, strategy=strategy)
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
dtype = 'float32'
|
|
model = TestLayer(dtype)
|
|
|
|
optimizer = paddle.optimizer.AdamW(parameters=model.parameters())
|
|
optimizer = DygraphShardingOptimizerV2(optimizer, hcg)
|
|
optimizer.clear_param_storage("_w")
|
|
optimizer.clear_param_storage("_b")
|
|
optimizer.clear_param_storage(None)
|
|
optimizer.reset_param_storage()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|