wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
paddlepaddle--paddle/test/flex_checkpoint/test_aoa_engine.py
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

742 lines
23 KiB
Python
Raw Permalink Blame History

# 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,
ShardedWeightDesc,
ShardMappingEntry,
)
class TestAOAEngine(unittest.TestCase):
def test_aoa_spilt_merge(self):
# ------------------------------------------------------
# 1. Define source tensor shards (s0 and s1).
# Each is a (2,2) tensor, fully covering its global shape.
#
# s0 (2,2): s1 (2,2):
# +----+----+ +----+----+
# | | | | | |
# +----+----+ +----+----+
# | | | | | |
# +----+----+ +----+----+
s0 = ShardedWeightDesc(
key="s0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
)
s1 = ShardedWeightDesc(
key="s1",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
)
# ------------------------------------------------------
# 2. Define destination tensor shards (d0 and d1).
# Both are (1,4) tensors, i.e., a single row with 4 columns.
#
# d0 (1,4): d1 (1,4):
# +--+--+--+--+ +--+--+--+--+
# | | | | | | | | | |
# +--+--+--+--+ +--+--+--+--+
d0 = ShardedWeightDesc(
key="d0",
local_shape=(1, 4),
global_shape=(1, 4),
global_offset=(0, 0),
)
d1 = ShardedWeightDesc(
key="d1",
local_shape=(1, 4),
global_shape=(1, 4),
global_offset=(0, 0),
)
# ------------------------------------------------------
# 3. Record the shard info for sources and destinations
source_state_shard_info = {
"s0": [s0],
"s1": [s1],
}
destination_state_shard_info = {
"d0": [d0],
"d1": [d1],
}
# ------------------------------------------------------
# 4. AOA statements define axis mapping for concatenation and splitting:
# - "s" is formed by concatenating s0 and s1 along axis 1 (columns).
# - d0 and d1 are obtained by splitting "s" along axis 0 (rows).
aoa_statements = [
"s0, s1 -> s, axis = 1 \n",
"s -> d0, d1, axis = 0 \n",
]
# ------------------------------------------------------
# 5. Create the AOAEngine with this configuration
aoa_engine = AOAEngine(
aoa_config={"aoa_statements": aoa_statements},
source_state_shard_info=source_state_shard_info,
destination_state_shard_info=destination_state_shard_info,
)
queries = []
answers = []
# ======================================================
# Query 1: Find source for the first half of d0 (columns 0-1)
# d0 shard: key="d0", local_shape=(1,2), global_shape=(1,4), global_offset=(0,0)
# Covers d0[:, 0:2]
#
# d0 (1,4):
# +------+------+------+------+
# |(0,0) |(0,1) | | |
# +------+------+------+------+
#
# This region is mapped from s0, row 0, columns 0-1
query = ShardedWeightDesc(
key="d0",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 0),
)
src_sharded_weight_desc = ShardedWeightDesc(
key="s0",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(0, 0),
)
shard_mapping_entry = ShardMappingEntry(
target_slice=query,
source_slice=src_sharded_weight_desc,
postprocess_list=None,
)
answer = [shard_mapping_entry]
queries.append(query)
answers.append(answer)
# ======================================================
# Query 2: Find source for the second half of d1 (columns 2-3)
# d1 shard: key="d1", local_shape=(1,2), global_shape=(1,4), global_offset=(0,2)
# Covers d1[:, 2:4]
#
# d1 (1,4):
# +------+------+------+------+
# | | |(0,2)|(0,3)|
# +------+------+------+------+
#
# This region is mapped from s1, row 1, columns 0-1
query = ShardedWeightDesc(
key="d1",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 2),
)
src_sharded_weight_desc = ShardedWeightDesc(
key="s1",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0),
)
shard_mapping_entry = ShardMappingEntry(
target_slice=query,
source_slice=src_sharded_weight_desc,
postprocess_list=None,
)
answer = [shard_mapping_entry]
queries.append(query)
answers.append(answer)
# ======================================================
# Query 3: Find sources for the entire d1 (full row)
# d1 shard: key="d1", local_shape=(1,4), global_shape=(1,4), global_offset=(0,0)
# Layout: covers all columns
#
# d1 (1,4):
# +------+------+------+------+
# | s0 | s0 | s1 | s1 |
# |(0,0) |(0,1) |(0,2) |(0,3) |
# +------+------+------+------+
# The first two columns come from s0, the last two from s1.
#
# Source slices:
# s0, local_shape=(1,2), global_shape=(2,2), global_offset=(1,0)
# +----+----+
# |(1,0)|(1,1)| <- used for d1 (0,0)-(0,1)
# +----+----+
#
# s1, local_shape=(1,2), global_shape=(2,2), global_offset=(1,0)
# +----+----+
# |(1,0)|(1,1)| <- used for d1 (0,2)-(0,3)
# +----+----+
#
# The answer consists of two mapping entries:
# 1. d1[:, 0:2] <-- s0[1, :]
# 2. d1[:, 2:4] <-- s1[1, :]
query = ShardedWeightDesc(
key="d1",
local_shape=(1, 4),
global_shape=(1, 4),
global_offset=(0, 0),
)
# d1[:, 0:2] <--- s0[1, :]
src_sharded_weight_desc0 = ShardedWeightDesc(
key="s0",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0), # row 1, columns 0:2
)
dst_sharded_weight_desc0 = ShardedWeightDesc(
key="d1",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 0),
)
# Visual mapping:
# d1 (0,0)-(0,1) <--- s0 (1,0)-(1,1)
# +------+------+------+------+
# |==s0==|==s0==| | |
# +------+------+------+------+
src_sharded_weight_desc1 = ShardedWeightDesc(
key="s1",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0),
)
dst_sharded_weight_desc1 = ShardedWeightDesc(
key="d1",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 2),
)
# Visual mapping:
# d1 (0,2)-(0,3) <--- s1 (1,0)-(1,1)
# +------+------+------+------+
# | | |==s1==|==s1==|
# +------+------+------+------+
shard_mapping_entry0 = ShardMappingEntry(
target_slice=dst_sharded_weight_desc0,
source_slice=src_sharded_weight_desc0,
postprocess_list=None,
)
shard_mapping_entry1 = ShardMappingEntry(
target_slice=dst_sharded_weight_desc1,
source_slice=src_sharded_weight_desc1,
postprocess_list=None,
)
answer = [shard_mapping_entry0, shard_mapping_entry1]
queries.append(query)
answers.append(answer)
# Visual answer summary:
# d1 (row 0):
# +------+------+------+------+
# |==s0==|==s0==|==s1==|==s1==|
# +------+------+------+------+
# ^ ^ ^ ^
# | | | |
# |______| |______|
# from s0 from s1
# ------------------------------------------------------
# ======================================================
# Query 4: for optimizer state
query = ShardedWeightDesc(
key="d1.moment1_0",
local_shape=(1, 4),
global_shape=(1, 4),
global_offset=(0, 0),
)
# d1[:, 0:2] <--- s0[1, :]
src_sharded_weight_desc0 = ShardedWeightDesc(
key="s0.moment1_0",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0), # row 1, columns 0:2
)
dst_sharded_weight_desc0 = ShardedWeightDesc(
key="d1.moment1_0",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 0),
)
src_sharded_weight_desc1 = ShardedWeightDesc(
key="s1.moment1_0",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0),
)
dst_sharded_weight_desc1 = ShardedWeightDesc(
key="d1.moment1_0",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 2),
)
shard_mapping_entry0 = ShardMappingEntry(
target_slice=dst_sharded_weight_desc0,
source_slice=src_sharded_weight_desc0,
postprocess_list=None,
)
shard_mapping_entry1 = ShardMappingEntry(
target_slice=dst_sharded_weight_desc1,
source_slice=src_sharded_weight_desc1,
postprocess_list=None,
)
answer = [shard_mapping_entry0, shard_mapping_entry1]
queries.append(query)
answers.append(answer)
# ======================================================
# Query 5: for optimizer state
query = ShardedWeightDesc(
key="d1.w_0",
local_shape=(1, 4),
global_shape=(1, 4),
global_offset=(0, 0),
)
# d1[:, 0:2] <--- s0[1, :]
src_sharded_weight_desc0 = ShardedWeightDesc(
key="s0.w_0",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0), # row 1, columns 0:2
)
dst_sharded_weight_desc0 = ShardedWeightDesc(
key="d1.w_0",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 0),
)
src_sharded_weight_desc1 = ShardedWeightDesc(
key="s1.w_0",
local_shape=(1, 2),
global_shape=(2, 2),
global_offset=(1, 0),
)
dst_sharded_weight_desc1 = ShardedWeightDesc(
key="d1.w_0",
local_shape=(1, 2),
global_shape=(1, 4),
global_offset=(0, 2),
)
shard_mapping_entry0 = ShardMappingEntry(
target_slice=dst_sharded_weight_desc0,
source_slice=src_sharded_weight_desc0,
postprocess_list=None,
)
shard_mapping_entry1 = ShardMappingEntry(
target_slice=dst_sharded_weight_desc1,
source_slice=src_sharded_weight_desc1,
postprocess_list=None,
)
answer = [shard_mapping_entry0, shard_mapping_entry1]
queries.append(query)
answers.append(answer)
# 6. Run the queries and check results
for idx in range(len(queries)):
query = queries[idx]
answer = answers[idx]
result = aoa_engine.find_shard_sources(query)
self.assertEqual(result, answer)
def test_aoa_cast(self):
"""Test AOA cast primitive for dtype conversion."""
s0 = ShardedWeightDesc(
key="s0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
dtype="int32",
)
d0 = ShardedWeightDesc(
key="d0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
dtype="float32",
)
source_state_shard_info = {
"s0": [s0],
}
destination_state_shard_info = {
"d0": [d0],
}
aoa_statements = [
's0 -> d0, dtype="float32" \n',
]
aoa_engine = AOAEngine(
aoa_config={"aoa_statements": aoa_statements},
source_state_shard_info=source_state_shard_info,
destination_state_shard_info=destination_state_shard_info,
)
query = ShardedWeightDesc(
key="d0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
dtype="float32",
)
src_sharded_weight_desc = ShardedWeightDesc(
key="s0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
dtype="int32",
)
shard_mapping_entry = ShardMappingEntry(
target_slice=query,
source_slice=src_sharded_weight_desc,
postprocess_list=['float32'],
)
answer = [shard_mapping_entry]
result = aoa_engine.find_shard_sources(query)
self.assertEqual(result, answer)
def test_aoa_add(self):
"""Test AOA add primitive for adding new keys that don't exist in source."""
d0 = ShardedWeightDesc(
key="d0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
dtype="float32",
)
source_state_shard_info = {}
destination_state_shard_info = {
"d0": [d0],
}
aoa_statements = [
"_ -> d0 \n",
]
aoa_engine = AOAEngine(
aoa_config={"aoa_statements": aoa_statements},
source_state_shard_info=source_state_shard_info,
destination_state_shard_info=destination_state_shard_info,
)
query = ShardedWeightDesc(
key="d0",
local_shape=(2, 2),
global_shape=(2, 2),
global_offset=(0, 0),
dtype="float32",
)
answer = []
result = aoa_engine.find_shard_sources(query)
self.assertEqual(result, answer)
def test_mixed_aoa_statements(self):
# test fused_ffn and transposed,rename,test_get_var_mapping_chain_macro
s0 = ShardedWeightDesc(
key="layers.0.gate_up_fused_proj.weight",
local_shape=(2, 4),
global_shape=(2, 8),
global_offset=(0, 0),
dtype="float32",
)
s1 = ShardedWeightDesc(
key="layers.0.gate_up_fused_proj.weight",
local_shape=(2, 4),
global_shape=(2, 8),
global_offset=(0, 4),
dtype="float32",
)
d0 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 0),
dtype="float32",
)
d1 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 2),
dtype="float32",
)
d2 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 4),
dtype="float32",
)
d3 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 6),
dtype="float32",
)
source_state_shard_info = {
"layers.0.gate_up_fused_proj.weight": [s0, s1],
}
destination_state_shard_info = {
"new_name_layers.0.gate_up_fused_proj.weight": [d0, d1, d2, d3],
}
# find temp_var -> dst
aoa_statements = [
"layers.0.gate_up_fused_proj.weight -> temp_var, fused_ffn \n",
"temp_var^T -> new_name_layers.0.gate_up_fused_proj.weight \n",
]
aoa_engine = AOAEngine(
aoa_config={"aoa_statements": aoa_statements},
source_state_shard_info=source_state_shard_info,
destination_state_shard_info=destination_state_shard_info,
)
# new_name_up_proj_0
query = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 1),
global_shape=(2, 8),
global_offset=(0, 0),
dtype="float32",
)
target_slice_1 = ShardedWeightDesc(
key='new_name_layers.0.gate_up_fused_proj.weight',
local_shape=(1, 1),
global_shape=(2, 8),
global_offset=(0, 0),
dtype='float32',
)
target_slice_2 = ShardedWeightDesc(
key='new_name_layers.0.gate_up_fused_proj.weight',
local_shape=(1, 1),
global_shape=(2, 8),
global_offset=(1, 0),
dtype='float32',
)
src_slice_1 = ShardedWeightDesc(
key='layers.0.gate_up_fused_proj.weight',
local_shape=(1, 1),
global_shape=(2, 8),
global_offset=(0, 0),
dtype='float32',
)
src_slice_2 = ShardedWeightDesc(
key='layers.0.gate_up_fused_proj.weight',
local_shape=(1, 1),
global_shape=(2, 8),
global_offset=(0, 2),
dtype='float32',
)
shard_mapping_entry_1 = ShardMappingEntry(
target_slice=target_slice_1,
source_slice=src_slice_1,
postprocess_list=['[1, 0]'],
)
shard_mapping_entry_2 = ShardMappingEntry(
target_slice=target_slice_2,
source_slice=src_slice_2,
postprocess_list=['[1, 0]'],
)
answer = [shard_mapping_entry_1, shard_mapping_entry_2]
result = aoa_engine.find_shard_sources(query)
self.assertEqual(result, answer)
s0 = ShardedWeightDesc(
key="layers.0.gate_up_fused_proj.weight",
local_shape=(4, 2),
global_shape=(8, 2),
global_offset=(0, 0),
dtype="float32",
)
s1 = ShardedWeightDesc(
key="layers.0.gate_up_fused_proj.weight",
local_shape=(4, 2),
global_shape=(8, 2),
global_offset=(4, 0),
dtype="float32",
)
d0 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 0),
dtype="float32",
)
d1 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 2),
dtype="float32",
)
d2 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 4),
dtype="float32",
)
d3 = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 2),
global_shape=(2, 8),
global_offset=(0, 6),
dtype="float32",
)
source_state_shard_info = {
"layers.0.gate_up_fused_proj.weight": [s0, s1],
}
destination_state_shard_info = {
"new_name_layers.0.gate_up_fused_proj.weight": [d0, d1, d2, d3],
}
# find temp_var -> src
aoa_statements = [
"layers.0.gate_up_fused_proj.weight^T -> temp_var \n",
"temp_var -> new_name_layers.0.gate_up_fused_proj.weight,fused_ffn\n",
]
aoa_engine = AOAEngine(
aoa_config={"aoa_statements": aoa_statements},
source_state_shard_info=source_state_shard_info,
destination_state_shard_info=destination_state_shard_info,
)
# new_name_up_proj_0
query = ShardedWeightDesc(
key="new_name_layers.0.gate_up_fused_proj.weight",
local_shape=(2, 1),
global_shape=(2, 8),
global_offset=(0, 0),
dtype="float32",
)
src_slice_1 = ShardedWeightDesc(
key='layers.0.gate_up_fused_proj.weight',
local_shape=(1, 2),
global_shape=(8, 2),
global_offset=(0, 0),
dtype='float32',
)
shard_mapping_entry_1 = ShardMappingEntry(
target_slice=query,
source_slice=src_slice_1,
postprocess_list=['[1, 0]'],
)
answer = [shard_mapping_entry_1]
result = aoa_engine.find_shard_sources(query)
self.assertEqual(result, answer)
def test_aoa_transpose_reverse(self):
s0 = ShardedWeightDesc(
key="s0",
local_shape=(1, 2, 4),
global_shape=(1, 2, 4),
global_offset=(0, 0, 0),
dtype="float32",
)
d0 = ShardedWeightDesc(
key="d0",
local_shape=(4, 1, 2),
global_shape=(4, 1, 2),
global_offset=(0, 0, 0),
dtype="float32",
)
aoa_statements = [
"s0 -> d0, permute= '[2, 0, 1]' \n",
]
aoa_config = {
"aoa_statements": aoa_statements,
"aoa_config_reverse": True,
}
source_state_shard_info = {
"d0": [d0],
}
destination_state_shard_info = {
"s0": [s0],
}
aoa_engine = AOAEngine(
aoa_config=aoa_config,
source_state_shard_info=source_state_shard_info,
destination_state_shard_info=destination_state_shard_info,
)
query = ShardedWeightDesc(
key="s0",
local_shape=(1, 2, 4),
global_shape=(1, 2, 4),
global_offset=(0, 0, 0),
dtype="float32",
)
answer = [
ShardMappingEntry(
target_slice=s0,
source_slice=d0,
postprocess_list=['[1, 2, 0]'],
)
]
result = aoa_engine.find_shard_sources(query)
self.assertEqual(result, answer)
if __name__ == '__main__':
unittest.main()
Reference in New Issue View Git Blame Copy Permalink