wehub-resource-sync
104 lines
3.5 KiB
Python
104 lines
3.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 unittest
|
|
|
|
from paddle.distributed.flex_checkpoint.aoa.aoa_engine import AOAEngine
|
|
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
|
|
ShardedWeightDesc,
|
|
)
|
|
|
|
|
|
def create_shard_info(keys, shape=(4, 4), dtype="float32"):
|
|
info = {}
|
|
for k in keys:
|
|
desc = ShardedWeightDesc(
|
|
key=k,
|
|
local_shape=shape,
|
|
global_shape=shape,
|
|
global_offset=(0, 0),
|
|
dtype=dtype,
|
|
)
|
|
info[k] = [desc]
|
|
return info
|
|
|
|
|
|
class TestMacroLayerOffsetError(unittest.TestCase):
|
|
def setUp(self):
|
|
self.source_keys = [f"model.layers.{i}.weight" for i in range(10)]
|
|
self.dest_keys = [f"model.layers.{i}.weight_out" for i in range(10)] + [
|
|
f"model.layers.{i}.weight_out2" for i in range(10)
|
|
]
|
|
|
|
self.src_info = create_shard_info(self.source_keys)
|
|
self.dst_info = create_shard_info(self.dest_keys)
|
|
|
|
def test_macro_error_chain(self):
|
|
"""
|
|
The statement contains fused_qkv_old and is missing a comma, expecting to trigger the assertion and print the chain.
|
|
"""
|
|
aoa_config = {
|
|
"aoa_statements": [
|
|
"model.layers.$LAYER_ID.weight^T -> model.layers.$LAYER_ID.weight_out, axis=0 fused_qkv_old, num_heads=20,num_key_value_groups=4",
|
|
],
|
|
"enable_traceback": True,
|
|
}
|
|
|
|
with self.assertRaises(AssertionError):
|
|
AOAEngine(
|
|
aoa_config=aoa_config,
|
|
source_state_shard_info=self.src_info,
|
|
destination_state_shard_info=self.dst_info,
|
|
)
|
|
|
|
def test_no_error_should_be_raised(self):
|
|
# No error should be raised
|
|
source_keys = ["model.layers.0.weight"]
|
|
dest_keys = ["model.layers.0.weight_out"]
|
|
src_info = create_shard_info(source_keys)
|
|
dst_info = create_shard_info(dest_keys)
|
|
aoa_config = {
|
|
"aoa_statements": [
|
|
"model.layers.0.weight^T -> model.layers.0.weight_out",
|
|
],
|
|
"enable_traceback": True,
|
|
}
|
|
AOAEngine(
|
|
aoa_config=aoa_config,
|
|
source_state_shard_info=src_info,
|
|
destination_state_shard_info=dst_info,
|
|
)
|
|
|
|
def test_shape_propagation_error_chain(self):
|
|
"""
|
|
when split/concat, only support one attr named `axis`, but got multiple attrs.
|
|
"""
|
|
aoa_config = {
|
|
"aoa_statements": [
|
|
"model.layers.0.weight -> model.layers.0.weight_out,model.layers.0.weight_out2,axis=0,axis=1",
|
|
],
|
|
"enable_traceback": True,
|
|
}
|
|
|
|
with self.assertRaises(ValueError):
|
|
AOAEngine(
|
|
aoa_config=aoa_config,
|
|
source_state_shard_info=self.src_info,
|
|
destination_state_shard_info=self.dst_info,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|