wehub-resource-sync
278 lines
11 KiB
Python
278 lines
11 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.
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
|
|
|
|
class TestSemiAutoParallelShardOptimizerAPI:
|
|
def __init__(self):
|
|
self._backend = os.getenv("backend")
|
|
self._seed = eval(os.getenv("seed"))
|
|
self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
|
|
self._ckpt_path = os.getenv("ckpt_path")
|
|
|
|
def check_tensor_eq(self, a, b, rtol=1e-05, atol=0, verbose=True):
|
|
np.testing.assert_allclose(a, b, rtol=rtol, atol=atol, verbose=verbose)
|
|
|
|
def get_single_card_rst(self):
|
|
paddle.seed(self._seed)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
batch = paddle.rand(shape=[10, 10])
|
|
opt = paddle.optimizer.AdamW(parameters=linear.parameters())
|
|
for _ in range(5):
|
|
loss = linear(batch)
|
|
loss.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
self.weight = linear.weight.numpy()
|
|
self.bias = linear.bias.numpy()
|
|
|
|
def shard_layer_fn(self, layer_name, layer, process_mesh):
|
|
layer.weight = dist.shard_tensor(
|
|
layer.weight, process_mesh, [dist.Shard(1)]
|
|
)
|
|
layer.bias = dist.shard_tensor(
|
|
layer.bias, process_mesh, [dist.Shard(0)]
|
|
)
|
|
|
|
def test_opt(self, opt):
|
|
for key in opt._accumulators.keys():
|
|
for k, v in opt._accumulators[key].items():
|
|
assert opt._accumulators[key][k].is_dist()
|
|
if 'moment' in key:
|
|
assert (
|
|
opt._accumulators[key][k].shape[-1]
|
|
== opt._accumulators[key][k]._local_shape[-1] * 2
|
|
)
|
|
else:
|
|
assert opt._accumulators[key][k].shape == [1]
|
|
assert opt._accumulators[key][k]._local_shape == [1]
|
|
|
|
def test_shard_optimizer_mp(self):
|
|
paddle.seed(self._seed)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
|
|
batch = paddle.rand(shape=[10, 10])
|
|
opt = paddle.optimizer.AdamW(parameters=linear.parameters())
|
|
opt = dist.shard_optimizer(opt)
|
|
for _ in range(5):
|
|
loss = linear(batch)
|
|
loss.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
self.test_opt(opt)
|
|
self.check_tensor_eq(self.weight, linear.weight.numpy())
|
|
self.check_tensor_eq(self.bias, linear.bias.numpy())
|
|
# save load
|
|
ckpt_state_dict = opt.state_dict()
|
|
ckpt_state_dict_keys = list(ckpt_state_dict.keys())
|
|
dist.save_state_dict(ckpt_state_dict, self._ckpt_path)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
|
|
new_opt = paddle.optimizer.AdamW(parameters=linear.parameters())
|
|
new_opt = dist.shard_optimizer(new_opt)
|
|
new_state_dict = new_opt.state_dict()
|
|
new_state_dict = {
|
|
ckpt_state_dict_keys[i]: v
|
|
for i, (k, v) in enumerate(new_state_dict.items())
|
|
}
|
|
dist.load_state_dict(new_state_dict, self._ckpt_path)
|
|
assert len(new_state_dict) > 0, "load_state_dict fail"
|
|
for k, v in new_state_dict.items():
|
|
assert k in ckpt_state_dict
|
|
if k in ["master_weights", "LR_Scheduler"]:
|
|
continue
|
|
self.check_tensor_eq(v, ckpt_state_dict[k])
|
|
|
|
def test_shard_optimizer_from_non_shard_layer(self):
|
|
paddle.seed(self._seed)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
batch = paddle.rand(shape=[10, 10])
|
|
opt = paddle.optimizer.AdamW(parameters=linear.parameters())
|
|
opt = dist.shard_optimizer(opt)
|
|
for _ in range(5):
|
|
loss = linear(batch)
|
|
loss.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
self.check_tensor_eq(self.weight, linear.weight.numpy())
|
|
self.check_tensor_eq(self.bias, linear.bias.numpy())
|
|
# save load
|
|
ckpt_state_dict = opt.state_dict()
|
|
ckpt_state_dict_keys = list(ckpt_state_dict.keys())
|
|
ckpt_path = os.path.join(
|
|
self._ckpt_path, "test_shard_optimizer_from_non_shard_layer"
|
|
)
|
|
dist.save_state_dict(ckpt_state_dict, ckpt_path)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
new_opt = paddle.optimizer.AdamW(parameters=linear.parameters())
|
|
new_opt = dist.shard_optimizer(new_opt)
|
|
new_state_dict = new_opt.state_dict()
|
|
new_state_dict = {
|
|
ckpt_state_dict_keys[i]: v
|
|
for i, (k, v) in enumerate(new_state_dict.items())
|
|
}
|
|
dist.load_state_dict(new_state_dict, ckpt_path)
|
|
assert len(new_state_dict) > 0, "load_state_dict fail"
|
|
for k, v in new_state_dict.items():
|
|
assert k in ckpt_state_dict
|
|
if k in ["master_weights", "LR_Scheduler"]:
|
|
continue
|
|
self.check_tensor_eq(v, ckpt_state_dict[k])
|
|
|
|
def shard_opt_fn(self, accumulator_name, param, accumulator):
|
|
if param.is_dist():
|
|
if 'beta' not in accumulator_name:
|
|
placements = param.placements
|
|
else:
|
|
placements = [
|
|
dist.Replicate()
|
|
for _ in range(len(param.process_mesh.shape))
|
|
]
|
|
return dist.shard_tensor(
|
|
accumulator, param.process_mesh, placements
|
|
)
|
|
return accumulator
|
|
|
|
def test_shard_optimizer_shard_fn(self):
|
|
paddle.seed(self._seed)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
|
|
batch = paddle.rand(shape=[10, 10])
|
|
opt = paddle.optimizer.AdamW(parameters=linear.parameters())
|
|
opt = dist.shard_optimizer(opt, self.shard_opt_fn)
|
|
loss = linear(batch)
|
|
loss.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
self.test_opt(opt)
|
|
|
|
def test_shard_optimizer_master_params(self):
|
|
paddle.seed(self._seed)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
batch = paddle.rand(shape=[10, 10], dtype="float16")
|
|
linear = paddle.amp.decorate(linear, level="O2", dtype="float16")
|
|
dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
|
|
opt = paddle.optimizer.AdamW(
|
|
parameters=linear.parameters(), multi_precision=True
|
|
)
|
|
opt = dist.shard_optimizer(opt)
|
|
loss = linear(batch)
|
|
loss.backward()
|
|
opt.step()
|
|
self.test_opt(opt)
|
|
for k, v in opt._master_weights.items():
|
|
assert v.dtype == paddle.float32
|
|
assert v.is_dist()
|
|
assert v.shape[-1] == v._local_shape[-1] * 2
|
|
|
|
# save load
|
|
ckpt_state_dict = opt.state_dict()
|
|
ckpt_path = os.path.join(
|
|
self._ckpt_path, "test_shard_optimizer_master_params"
|
|
)
|
|
dist.save_state_dict(ckpt_state_dict, ckpt_path)
|
|
paddle.distributed.barrier()
|
|
expected_local_state_dict = {}
|
|
expected_local_state_dict.setdefault("master_weights", {})
|
|
need_load_state_dict = {}
|
|
need_load_state_dict.setdefault("master_weights", {})
|
|
for k, v in ckpt_state_dict.items():
|
|
if k == "LR_Scheduler":
|
|
continue
|
|
elif k == "master_weights":
|
|
assert isinstance(v, dict), v
|
|
for mk, mv in v.items():
|
|
expected_local_state_dict[k][mk] = mv._local_value().clone()
|
|
need_load_state_dict[k][mk] = paddle.zeros_like(mv)
|
|
else:
|
|
expected_local_state_dict[k] = v._local_value().clone()
|
|
need_load_state_dict[k] = paddle.zeros_like(v)
|
|
opt.set_state_dict(need_load_state_dict)
|
|
after_set_state_dict = opt.state_dict()
|
|
for k, v in after_set_state_dict.items():
|
|
if k == "master_weights":
|
|
assert isinstance(v, dict), v
|
|
for mk, mv in v.items():
|
|
assert mv.numpy().sum() == 0.0, (
|
|
f"state_dict {k} in master_weights is not zero"
|
|
)
|
|
assert need_load_state_dict[k][mk].numpy().sum() == 0.0, (
|
|
f"state_dict {k} in master_weights is not zero"
|
|
)
|
|
else:
|
|
assert v.numpy().sum() == 0.0, f"state_dict {k} is not zero"
|
|
assert k in need_load_state_dict, f"state_dict {k} is not found"
|
|
assert need_load_state_dict[k].numpy().sum() == 0.0, (
|
|
f"state_dict {k} is not zero"
|
|
)
|
|
dist.load_state_dict(need_load_state_dict, ckpt_path)
|
|
opt.set_state_dict(need_load_state_dict)
|
|
new_state_dict = opt.state_dict()
|
|
assert "master_weights" in new_state_dict, new_state_dict
|
|
for k, v in new_state_dict.items():
|
|
assert k in expected_local_state_dict
|
|
if k == "master_weights":
|
|
for mk, mv in v.items():
|
|
np.testing.assert_equal(
|
|
mv._local_value().numpy(),
|
|
expected_local_state_dict[k][mk].numpy(),
|
|
)
|
|
else:
|
|
np.testing.assert_equal(
|
|
v._local_value().numpy(),
|
|
expected_local_state_dict[k].numpy(),
|
|
)
|
|
|
|
def test_shard_optimizer_params_group(self):
|
|
paddle.seed(self._seed)
|
|
linear = paddle.nn.Linear(10, 10)
|
|
dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
|
|
batch = paddle.rand(shape=[10, 10])
|
|
linear.weight.optimize_attr = {'lr': 1}
|
|
linear.bias.optimize_attr = {'lr': 1}
|
|
params_group = [{'params': linear.weight}, {'params': linear.bias}]
|
|
opt = paddle.optimizer.AdamW(parameters=params_group)
|
|
opt = dist.shard_optimizer(opt)
|
|
loss = linear(batch)
|
|
loss.backward()
|
|
opt.step()
|
|
opt.clear_grad()
|
|
self.test_opt(opt)
|
|
|
|
def run_test_case(self):
|
|
if self._backend == "cpu":
|
|
paddle.set_device("cpu")
|
|
elif self._backend == "gpu":
|
|
paddle.set_device("gpu:" + str(dist.get_rank()))
|
|
else:
|
|
raise ValueError("Only support cpu or gpu backend.")
|
|
|
|
self.get_single_card_rst()
|
|
self.test_shard_optimizer_params_group()
|
|
self.test_shard_optimizer_shard_fn()
|
|
if self._backend == "gpu":
|
|
self.test_shard_optimizer_master_params()
|
|
self.test_shard_optimizer_mp()
|
|
self.test_shard_optimizer_from_non_shard_layer()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
TestSemiAutoParallelShardOptimizerAPI().run_test_case()
|