chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,944 @@
|
||||
# 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.
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from ..dcp.sharded_weight import ShardedWeightDesc
|
||||
from .lexer import Lexer
|
||||
from .parser import Parser
|
||||
from .traceback import AOATraceback
|
||||
|
||||
_ShardInfo = dict[str, list[ShardedWeightDesc]]
|
||||
|
||||
# SliceRef := (key, src_slice, dst_slice, postprocess_list)
|
||||
SliceRef = tuple[str, tuple[slice, ...], tuple[slice, ...], list[str] | None]
|
||||
SUPPORTED_DTYPES = ['float16', 'float32', 'bfloat16']
|
||||
|
||||
|
||||
class TensorDesc:
|
||||
def __init__(
|
||||
self,
|
||||
slices: list[SliceRef],
|
||||
shape: tuple[int],
|
||||
in_degree: int = 0,
|
||||
out_degree: int = 0,
|
||||
dtype: str | None = None,
|
||||
):
|
||||
self.slices = slices
|
||||
self.shape = shape
|
||||
self.in_degree = in_degree
|
||||
self.out_degree = out_degree
|
||||
self.dtype = dtype
|
||||
|
||||
def __repr__(self):
|
||||
s = []
|
||||
for key, sl_src, sl_dst, pp_list in self.slices:
|
||||
s.append(
|
||||
f"{key}{sl_src} -> self{sl_dst}, postprocess_list={pp_list}"
|
||||
)
|
||||
return f"Tensor(shape={self.shape}, slices={s}, in_degree={self.in_degree}, out_degree={self.out_degree}, dtype={self.dtype})"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShardMappingEntry:
|
||||
target_slice: ShardedWeightDesc
|
||||
source_slice: ShardedWeightDesc
|
||||
postprocess_list: list[str] | None = None
|
||||
|
||||
|
||||
ShardMapping = list[ShardMappingEntry]
|
||||
|
||||
OPTIMIZER_STATE_NAME = [
|
||||
".w_0",
|
||||
".moment1_0",
|
||||
".moment2_0",
|
||||
".beta1_pow_acc_0",
|
||||
".beta2_pow_acc_0",
|
||||
]
|
||||
|
||||
|
||||
def split_optimizer_state_key(key: str) -> tuple[str, str]:
|
||||
for opt_state_name in OPTIMIZER_STATE_NAME:
|
||||
if key.endswith(opt_state_name):
|
||||
return key[: -len(opt_state_name)], opt_state_name
|
||||
return key, None
|
||||
|
||||
|
||||
class AOAShardInfoContext:
|
||||
def __init__(
|
||||
self,
|
||||
source_state_shard_info: _ShardInfo,
|
||||
destination_state_shard_info: _ShardInfo,
|
||||
aoa_config_reverse: bool = False,
|
||||
) -> None:
|
||||
self.source_state_shard_info = source_state_shard_info
|
||||
self.destination_state_shard_info = destination_state_shard_info
|
||||
self.aoa_config_reverse = aoa_config_reverse
|
||||
self.left_var_to_right_var_mapping = {}
|
||||
self.right_var_from_left_var_mapping = {}
|
||||
self.src_state_keys = set()
|
||||
self.dst_state_keys = set()
|
||||
self.init_src_state_keys()
|
||||
self.init_dst_state_keys()
|
||||
|
||||
def init_src_state_keys(self):
|
||||
for k in self.source_state_shard_info.keys():
|
||||
model_state_key, _ = split_optimizer_state_key(k)
|
||||
self.src_state_keys.add(model_state_key)
|
||||
|
||||
def init_dst_state_keys(self):
|
||||
if self.destination_state_shard_info is None:
|
||||
return
|
||||
for k in self.destination_state_shard_info.keys():
|
||||
model_state_key, _ = split_optimizer_state_key(k)
|
||||
self.dst_state_keys.add(model_state_key)
|
||||
|
||||
def get_all_dst_state_keys(self):
|
||||
return self.dst_state_keys
|
||||
|
||||
def get_all_src_state_keys(self):
|
||||
return self.src_state_keys
|
||||
|
||||
def get_num_hidden_layers(
|
||||
self,
|
||||
name_with_layer_id: str,
|
||||
layer_id_macro_tag: str,
|
||||
) -> int:
|
||||
if layer_id_macro_tag not in name_with_layer_id:
|
||||
raise ValueError(
|
||||
f"layer_id_macro_tag '{layer_id_macro_tag}' not in name_with_layer_id '{name_with_layer_id}'"
|
||||
)
|
||||
prefix, suffix = name_with_layer_id.split(layer_id_macro_tag, 1)
|
||||
pattern = re.compile(rf"{re.escape(prefix)}(\d+){re.escape(suffix)}")
|
||||
match_layer_id = set()
|
||||
for key in self.get_all_src_state_keys():
|
||||
match = pattern.fullmatch(key)
|
||||
if match:
|
||||
layer_num = int(match.group(1))
|
||||
match_layer_id.add(layer_num)
|
||||
return match_layer_id
|
||||
|
||||
def get_src_state_shard_num(self, src_state_key: str) -> int:
|
||||
model_state_key, opt_state_name = split_optimizer_state_key(
|
||||
src_state_key
|
||||
)
|
||||
|
||||
assert opt_state_name is None, (
|
||||
"AOA notions apply only to the model state, but are automatically propagated to the optimizer state.Now the src_state_key is {src_state_key}, which is a optimizer state key."
|
||||
)
|
||||
reverse = True
|
||||
if self.aoa_config_reverse:
|
||||
reverse = False
|
||||
# Only need to parse the model state key for optimizer state shard num, because the optimizer state slice info is completely consistent with the model state slice info.
|
||||
resolved_model_state_key = self.resolve_mapping_chain(
|
||||
model_state_key, reverse=reverse
|
||||
)
|
||||
|
||||
state_keys = [
|
||||
resolved_model_state_key,
|
||||
f"{resolved_model_state_key}.w_0",
|
||||
f"{resolved_model_state_key}.moment1_0",
|
||||
f"{resolved_model_state_key}.moment2_0",
|
||||
]
|
||||
|
||||
shard_nums = {
|
||||
len(
|
||||
{
|
||||
shard_info.global_offset
|
||||
for shard_info in self.source_state_shard_info[key]
|
||||
}
|
||||
)
|
||||
for key in state_keys
|
||||
if key in self.source_state_shard_info
|
||||
}
|
||||
|
||||
if not shard_nums:
|
||||
logger.warning(
|
||||
f"No shard information found for any of the keys: {state_keys}, return 1."
|
||||
)
|
||||
return 1
|
||||
if len(shard_nums) > 1:
|
||||
raise AssertionError(
|
||||
f"Inconsistent shard numbers among keys in source_sharded_state_dict for the key {src_state_key}: shard_nums={shard_nums}."
|
||||
)
|
||||
return shard_nums.pop()
|
||||
|
||||
def get_dst_state_shard_num(self, dst_state_key: str) -> int:
|
||||
if self.destination_state_shard_info is None:
|
||||
# Default `dst_state_shard_num=1` if `destination_state_shard_info` is missing.
|
||||
return 1
|
||||
model_state_key, opt_state_name = split_optimizer_state_key(
|
||||
dst_state_key
|
||||
)
|
||||
|
||||
assert opt_state_name is None, (
|
||||
"AOA notions apply only to the model state, but are automatically propagated to the optimizer state.Now the dst_state_key is {dst_state_key}, which is a optimizer state key."
|
||||
)
|
||||
reverse = False
|
||||
if self.aoa_config_reverse:
|
||||
reverse = True
|
||||
# Only need to parse the model state key for optimizer state shard num, because the optimizer state slice info is completely consistent with the model state slice info.
|
||||
resolved_model_state_key = self.resolve_mapping_chain(
|
||||
model_state_key, reverse=reverse
|
||||
)
|
||||
|
||||
state_keys = [
|
||||
resolved_model_state_key,
|
||||
f"{resolved_model_state_key}.w_0",
|
||||
f"{resolved_model_state_key}.moment1_0",
|
||||
f"{resolved_model_state_key}.moment2_0",
|
||||
]
|
||||
|
||||
shard_nums = {
|
||||
len(
|
||||
{
|
||||
shard_info.global_offset
|
||||
for shard_info in self.destination_state_shard_info[key]
|
||||
}
|
||||
)
|
||||
for key in state_keys
|
||||
if key in self.destination_state_shard_info
|
||||
}
|
||||
|
||||
if not shard_nums:
|
||||
logger.warning(
|
||||
f"No shard information found for any of the keys: {state_keys}, return 1."
|
||||
)
|
||||
return 1
|
||||
if len(shard_nums) > 1:
|
||||
raise AssertionError(
|
||||
f"Inconsistent shard numbers among keys in destination_state_shard_info for the key {dst_state_key}: shard_nums={shard_nums}."
|
||||
)
|
||||
return shard_nums.pop()
|
||||
|
||||
def resolve_mapping_chain(self, key: str, reverse: bool = False) -> str:
|
||||
"""
|
||||
Recursively resolve the mapping chain, find the final leaf node
|
||||
|
||||
Args:
|
||||
key: The key to be resolved
|
||||
reverse: False use left_var_to_right_var_mapping，True use right_var_from_left_var_mapping
|
||||
|
||||
For example:
|
||||
- reverse=False: temp_var -> dst_key
|
||||
- reverse=True: temp_var -> src_key
|
||||
"""
|
||||
visited = set() # avoid infinite loop
|
||||
current_key = key
|
||||
|
||||
if reverse:
|
||||
mapping_dict = self.right_var_from_left_var_mapping
|
||||
else:
|
||||
mapping_dict = self.left_var_to_right_var_mapping
|
||||
|
||||
while current_key in mapping_dict:
|
||||
assert current_key not in visited, (
|
||||
f"Infinite loop detected in resolve_mapping_chain, which means the start key is not src_key or the end key is not dst_key, the aoa_config is error. current_key={current_key}, the loop is: {'->'.join(visited)}->{current_key}"
|
||||
)
|
||||
visited.add(current_key)
|
||||
if reverse and current_key in self.get_all_src_state_keys():
|
||||
break
|
||||
elif not reverse and current_key in self.get_all_dst_state_keys():
|
||||
break
|
||||
|
||||
mapped_vars = mapping_dict[current_key]
|
||||
if mapped_vars and len(mapped_vars) > 0:
|
||||
assert len(mapped_vars) == 1, (
|
||||
f"Reference chain resolution failed: "
|
||||
f"Unable to determine which leaf node the intermediate node '{key}' is directly associated with, "
|
||||
f"because a many-to-one mapping was found in the mapping relationship. "
|
||||
f"The many-to-one mapping is {current_key} : {mapped_vars}."
|
||||
)
|
||||
current_key = mapped_vars[0]
|
||||
else:
|
||||
break
|
||||
|
||||
return current_key
|
||||
|
||||
|
||||
class AOAEngine:
|
||||
def __init__(
|
||||
self,
|
||||
aoa_config: dict[str, list[str]],
|
||||
source_state_shard_info: _ShardInfo,
|
||||
destination_state_shard_info: _ShardInfo,
|
||||
):
|
||||
self.aoa_config = aoa_config
|
||||
self.source_state_shard_info = source_state_shard_info
|
||||
self.destination_state_shard_info = destination_state_shard_info
|
||||
self.aoa_config_reverse = self.aoa_config.get(
|
||||
"aoa_config_reverse", False
|
||||
)
|
||||
enable_traceback = self.aoa_config.get("enable_traceback", True)
|
||||
self.traceback = AOATraceback() if enable_traceback else None
|
||||
self.context = AOAShardInfoContext(
|
||||
source_state_shard_info,
|
||||
destination_state_shard_info,
|
||||
self.aoa_config_reverse,
|
||||
)
|
||||
self.lexer = Lexer(self.context, traceback=self.traceback)
|
||||
tokens = self.lexer.all_tokens(
|
||||
self.aoa_config.get("aoa_statements", [])
|
||||
)
|
||||
self.parser = Parser(tokens)
|
||||
self.statements = self.parser.parse_program()
|
||||
|
||||
if self.traceback and getattr(self.lexer, "final_expressions", None):
|
||||
final_exprs = self.lexer.final_expressions
|
||||
if len(final_exprs) == len(self.statements):
|
||||
for expr, stmt in zip(final_exprs, self.statements):
|
||||
self.traceback.record_children(
|
||||
expr, [repr(stmt)], macro_name="parser"
|
||||
)
|
||||
|
||||
if self.aoa_config_reverse:
|
||||
self.statements = list(reversed(self.statements))
|
||||
self.input_vars = self.build_input_vars()
|
||||
self.output_vars = {}
|
||||
self.intermediate_vars = {}
|
||||
self.need_remove_input_vars = set()
|
||||
self.need_add_output_vars = set()
|
||||
|
||||
self.shape_propagation()
|
||||
|
||||
def make_input_tensor(
|
||||
self, key: str, shape: tuple[int], dtype: str
|
||||
) -> TensorDesc:
|
||||
base_slice = tuple([slice(0, s) for s in shape])
|
||||
return TensorDesc(
|
||||
[(key, base_slice, base_slice, None)],
|
||||
shape,
|
||||
in_degree=0,
|
||||
out_degree=0,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def build_input_vars(self):
|
||||
input_vars = {}
|
||||
dtype = None
|
||||
for key, shards in sorted(self.source_state_shard_info.items()):
|
||||
global_shape = shards[0].global_shape
|
||||
model_state_key, opt_state_name = split_optimizer_state_key(key)
|
||||
if opt_state_name is None:
|
||||
dtype = shards[0].dtype
|
||||
if model_state_key in input_vars.keys() or opt_state_name in [
|
||||
".beta1_pow_acc_0",
|
||||
".beta2_pow_acc_0",
|
||||
]:
|
||||
continue
|
||||
input_vars[model_state_key] = self.make_input_tensor(
|
||||
model_state_key, global_shape, dtype
|
||||
)
|
||||
return input_vars
|
||||
|
||||
def split(
|
||||
self, tensor: TensorDesc, axis: int, sizes: list[int]
|
||||
) -> list[TensorDesc]:
|
||||
results = []
|
||||
start = 0
|
||||
tensor.out_degree += len(sizes)
|
||||
dtype = tensor.dtype
|
||||
for sz in sizes:
|
||||
sub_dst_slice = [slice(None)] * len(tensor.shape)
|
||||
sub_dst_slice[axis] = slice(0, sz)
|
||||
sub_slices = []
|
||||
for aidx, src_sl, dst_sl, pp_list in tensor.slices:
|
||||
if pp_list is not None:
|
||||
src_sl = postprocess_transpose(list(src_sl), pp_list)
|
||||
|
||||
dst_start = (
|
||||
dst_sl[axis].start if dst_sl[axis].start is not None else 0
|
||||
)
|
||||
dst_stop = (
|
||||
dst_sl[axis].stop
|
||||
if dst_sl[axis].stop is not None
|
||||
else tensor.shape[axis]
|
||||
)
|
||||
inter_begin = max(start, dst_start)
|
||||
inter_end = min(start + sz, dst_stop)
|
||||
if inter_begin < inter_end:
|
||||
src_axis_start = (
|
||||
src_sl[axis].start
|
||||
if src_sl[axis].start is not None
|
||||
else 0
|
||||
)
|
||||
sub_src_sl = list(src_sl)
|
||||
sub_dst_sl = list(dst_sl)
|
||||
offset = inter_begin - dst_start
|
||||
length = inter_end - inter_begin
|
||||
sub_src_sl[axis] = slice(
|
||||
src_axis_start + offset,
|
||||
src_axis_start + offset + length,
|
||||
)
|
||||
sub_dst_sl[axis] = slice(
|
||||
inter_begin - start, inter_begin - start + length
|
||||
)
|
||||
if pp_list is not None:
|
||||
sub_src_sl = postprocess_transpose(
|
||||
list(sub_src_sl), pp_list, reverse=True
|
||||
)
|
||||
sub_slices.append(
|
||||
(
|
||||
aidx,
|
||||
tuple(sub_src_sl),
|
||||
tuple(sub_dst_sl),
|
||||
pp_list.copy(),
|
||||
)
|
||||
)
|
||||
else:
|
||||
sub_slices.append(
|
||||
(aidx, tuple(sub_src_sl), tuple(sub_dst_sl), None)
|
||||
)
|
||||
new_shape = list(tensor.shape)
|
||||
new_shape[axis] = sz
|
||||
results.append(
|
||||
TensorDesc(
|
||||
sub_slices,
|
||||
tuple(new_shape),
|
||||
in_degree=1,
|
||||
out_degree=0,
|
||||
dtype=dtype,
|
||||
)
|
||||
)
|
||||
start += sz
|
||||
return results
|
||||
|
||||
def concat(self, tensors: list[TensorDesc], axis: int) -> TensorDesc:
|
||||
slices = []
|
||||
assert len(tensors) >= 1, (
|
||||
"When concatenating multiple tensors, there should be at least one!"
|
||||
)
|
||||
shape = list(tensors[0].shape)
|
||||
ndim = len(shape)
|
||||
assert 0 <= axis < ndim, (
|
||||
f"when concat, the axis {axis} is out of range for tensors "
|
||||
f"with shape {shape} (valid range: {0} to {ndim - 1})."
|
||||
)
|
||||
shape[axis] = sum(t.shape[axis] for t in tensors)
|
||||
dtype = tensors[0].dtype
|
||||
assert all(t.dtype == dtype for t in tensors), (
|
||||
f"All tensors must have the same dtype when concatenating multiple tensors!But the tensors {tensors} have different dtypes: {[t.dtype for t in tensors]}."
|
||||
)
|
||||
curr = 0
|
||||
for t in tensors:
|
||||
t.out_degree += 1
|
||||
for aidx, src_sl, dst_sl, pp_list in t.slices:
|
||||
new_dst_sl = list(dst_sl)
|
||||
dst_start = (
|
||||
dst_sl[axis].start if dst_sl[axis].start is not None else 0
|
||||
)
|
||||
dst_stop = (
|
||||
dst_sl[axis].stop
|
||||
if dst_sl[axis].stop is not None
|
||||
else t.shape[axis]
|
||||
)
|
||||
length = dst_stop - dst_start
|
||||
new_dst_sl[axis] = slice(
|
||||
dst_start + curr, dst_start + curr + length
|
||||
)
|
||||
if pp_list is not None:
|
||||
slices.append(
|
||||
(aidx, src_sl, tuple(new_dst_sl), pp_list.copy())
|
||||
)
|
||||
else:
|
||||
slices.append((aidx, src_sl, tuple(new_dst_sl), None))
|
||||
curr += t.shape[axis]
|
||||
return TensorDesc(
|
||||
slices,
|
||||
tuple(shape),
|
||||
in_degree=len(tensors),
|
||||
out_degree=0,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def transpose(self, tensor: TensorDesc, permutation: str) -> TensorDesc:
|
||||
slices = []
|
||||
tensor.out_degree += 1
|
||||
tensor_shape = transpose_list(
|
||||
tensor.shape, ast.literal_eval(permutation)
|
||||
)
|
||||
dtype = tensor.dtype
|
||||
for aidx, src_sl, dst_sl, pp_list in tensor.slices:
|
||||
trans_dst_sl = transpose_list(dst_sl, ast.literal_eval(permutation))
|
||||
if pp_list is not None:
|
||||
new_pp_list = pp_list.copy()
|
||||
new_pp_list.append(permutation)
|
||||
slices.append((aidx, src_sl, trans_dst_sl, new_pp_list))
|
||||
else:
|
||||
slices.append((aidx, src_sl, trans_dst_sl, [permutation]))
|
||||
return TensorDesc(
|
||||
slices, tensor_shape, in_degree=1, out_degree=0, dtype=dtype
|
||||
)
|
||||
|
||||
def cast(self, tensor: TensorDesc, dtype: str) -> TensorDesc:
|
||||
slices = []
|
||||
tensor.out_degree += 1
|
||||
for aidx, src_sl, dst_sl, pp_list in tensor.slices:
|
||||
if pp_list is not None:
|
||||
new_pp_list = pp_list.copy()
|
||||
new_pp_list.append(dtype)
|
||||
slices.append((aidx, src_sl, dst_sl, new_pp_list))
|
||||
else:
|
||||
slices.append((aidx, src_sl, dst_sl, [dtype]))
|
||||
# For the cast operation, post_process is required. Therefore, the returned
|
||||
# Tensor's dtype here is the same as the input tensor's dtype, rather than the casted dtype.
|
||||
return TensorDesc(
|
||||
slices, tensor.shape, in_degree=1, out_degree=0, dtype=tensor.dtype
|
||||
)
|
||||
|
||||
def identity(self, tensor: TensorDesc) -> TensorDesc:
|
||||
tensor.out_degree += 1
|
||||
return TensorDesc(
|
||||
tensor.slices,
|
||||
tensor.shape,
|
||||
in_degree=1,
|
||||
out_degree=0,
|
||||
dtype=tensor.dtype,
|
||||
)
|
||||
|
||||
def shape_propagation(self):
|
||||
def _get_var_ref(var):
|
||||
if var.name in self.intermediate_vars:
|
||||
return self.intermediate_vars[var.name]
|
||||
elif var.name in self.input_vars:
|
||||
return self.input_vars[var.name]
|
||||
else:
|
||||
raise ValueError(f"{var.name} should be assigned before!")
|
||||
|
||||
for stmt in self.statements:
|
||||
stmt_repr = repr(stmt)
|
||||
left_vars = stmt.left_vars
|
||||
right_vars = stmt.right_vars
|
||||
if self.aoa_config_reverse:
|
||||
left_vars, right_vars = right_vars, left_vars
|
||||
attrs = stmt.attrs
|
||||
|
||||
try:
|
||||
if len(left_vars) > 1 or len(right_vars) > 1:
|
||||
if not (len(attrs) == 1 and attrs[0].key == "axis"):
|
||||
raise ValueError(
|
||||
f"When split/concat, only support one attr named `axis`, but got {attrs}."
|
||||
)
|
||||
axis = attrs[0].value
|
||||
|
||||
if len(left_vars) == 1:
|
||||
in_name = left_vars[0].name
|
||||
in_ref = _get_var_ref(left_vars[0])
|
||||
ndim = len(in_ref.shape)
|
||||
assert 0 <= axis < ndim, (
|
||||
f"when split, the axis {axis} is out of range for tensor {in_name} "
|
||||
f"with shape {in_ref.shape} (valid range: {0} to {ndim - 1})."
|
||||
)
|
||||
assert in_ref.shape[axis] % len(right_vars) == 0, (
|
||||
f"when split, the shape of the input tensor {in_name} is {in_ref.shape}, the axis is {axis}, the number of right_vars is {len(right_vars)}, but the shape of the input tensor {in_name} is not divisible by the number of right_vars."
|
||||
)
|
||||
sizes = [
|
||||
in_ref.shape[axis] // len(right_vars)
|
||||
for var in right_vars
|
||||
]
|
||||
result = self.split(in_ref, axis, sizes)
|
||||
for out_var, out_ref in zip(right_vars, result):
|
||||
self.intermediate_vars[out_var.name] = out_ref
|
||||
if (
|
||||
out_var.name
|
||||
in self.context.get_all_dst_state_keys()
|
||||
):
|
||||
self.output_vars[out_var.name] = out_ref
|
||||
|
||||
elif len(right_vars) == 1:
|
||||
left_refs = [_get_var_ref(var) for var in left_vars]
|
||||
result = self.concat(left_refs, axis)
|
||||
out_name = right_vars[0].name
|
||||
self.intermediate_vars[out_name] = result
|
||||
if out_name in self.context.get_all_dst_state_keys():
|
||||
self.output_vars[out_name] = result
|
||||
|
||||
else:
|
||||
raise SyntaxError(
|
||||
f'Unexpected split/concat statement: {stmt}'
|
||||
)
|
||||
|
||||
elif len(left_vars) == 1 and len(right_vars) == 1:
|
||||
lvar, rvar = left_vars[0], right_vars[0]
|
||||
if rvar.name == "_":
|
||||
self.need_remove_input_vars.add(lvar.name)
|
||||
elif lvar.name == "_":
|
||||
self.need_add_output_vars.add(rvar.name)
|
||||
else:
|
||||
if len(attrs) > 0:
|
||||
assert len(attrs) == 1 or (
|
||||
len(attrs) == 2
|
||||
and {attr.key for attr in attrs}
|
||||
== {"src_dtype", "dst_dtype"}
|
||||
), (
|
||||
"Only support:\n"
|
||||
" - One operator, OR\n"
|
||||
" - Two operators with keys {'src_dtype', 'dst_dtype'}."
|
||||
)
|
||||
attr = attrs[0]
|
||||
in_ref = _get_var_ref(lvar)
|
||||
if attr.key == "permute":
|
||||
if attr.value == "[]":
|
||||
ndim = len(in_ref.shape)
|
||||
perm = str(list(range(ndim - 1, -1, -1)))
|
||||
else:
|
||||
perm = attr.value
|
||||
if self.aoa_config_reverse:
|
||||
perm = str(
|
||||
invert_permutation(
|
||||
ast.literal_eval(perm)
|
||||
)
|
||||
)
|
||||
result = self.transpose(in_ref, perm)
|
||||
elif attr.key == "dtype":
|
||||
assert not self.aoa_config_reverse, (
|
||||
"When `aoa_config_reverse=True`, the dtype must be specified as "
|
||||
"'src_dtype=...,dst_dtype=...'. Formats like 'dtype=xxx' are not supported."
|
||||
)
|
||||
assert attr.value in SUPPORTED_DTYPES, (
|
||||
f"Unsupported cast dtype: {attr.value}"
|
||||
)
|
||||
result = self.cast(in_ref, attr.value)
|
||||
elif (
|
||||
attrs[0].key == "src_dtype"
|
||||
and attrs[1].key == "dst_dtype"
|
||||
):
|
||||
src_dtype, dst_dtype = (
|
||||
attrs[0].value,
|
||||
attrs[1].value,
|
||||
)
|
||||
assert src_dtype in SUPPORTED_DTYPES, (
|
||||
f"Unsupported cast dtype: {src_dtype}"
|
||||
)
|
||||
assert dst_dtype in SUPPORTED_DTYPES, (
|
||||
f"Unsupported cast dtype: {dst_dtype}"
|
||||
)
|
||||
if self.aoa_config_reverse:
|
||||
src_dtype, dst_dtype = dst_dtype, src_dtype
|
||||
result = self.cast(in_ref, dst_dtype)
|
||||
elif attr.key == "axis":
|
||||
result = in_ref
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported attribute: {attr}"
|
||||
)
|
||||
|
||||
self.intermediate_vars[rvar.name] = result
|
||||
if (
|
||||
rvar.name
|
||||
in self.context.get_all_dst_state_keys()
|
||||
):
|
||||
self.output_vars[rvar.name] = result
|
||||
else:
|
||||
# rename operation
|
||||
in_ref = _get_var_ref(lvar)
|
||||
result = self.identity(in_ref)
|
||||
self.intermediate_vars[rvar.name] = result
|
||||
if (
|
||||
rvar.name
|
||||
in self.context.get_all_dst_state_keys()
|
||||
):
|
||||
self.output_vars[rvar.name] = result
|
||||
else:
|
||||
raise SyntaxError(f'Unexpected statement: {stmt}')
|
||||
except (
|
||||
AssertionError,
|
||||
ValueError,
|
||||
KeyError,
|
||||
SyntaxError,
|
||||
RuntimeError,
|
||||
) as e:
|
||||
if self.traceback:
|
||||
chain = self.traceback.build_chain(stmt_repr)
|
||||
self.traceback.add_error(
|
||||
error_message=str(e),
|
||||
stage="shape_propagation",
|
||||
chain=chain,
|
||||
error_type=type(e).__name__,
|
||||
)
|
||||
self.traceback.print()
|
||||
raise
|
||||
if self.destination_state_shard_info is not None:
|
||||
for name in self.destination_state_shard_info:
|
||||
model_state_key, _ = split_optimizer_state_key(name)
|
||||
if model_state_key not in self.output_vars:
|
||||
if model_state_key in self.need_add_output_vars:
|
||||
self.output_vars[model_state_key] = None
|
||||
else:
|
||||
assert model_state_key in self.input_vars, (
|
||||
f"{model_state_key} needs to be loaded, "
|
||||
f"but not found in checkpoint. "
|
||||
f"If the key exists in the current model but not in the loaded checkpoint, please use the add primitive in aoa_statements: "
|
||||
f"_ -> {model_state_key}, and {model_state_key} will be randomly initialized."
|
||||
)
|
||||
self.output_vars[model_state_key] = self.input_vars[
|
||||
model_state_key
|
||||
]
|
||||
else:
|
||||
# When destination_state_shard_info is not provided, the AOAEngine automatically derives it
|
||||
# from source_state_shard_info and aha_statements. In this case, all destination_states
|
||||
# remain unsharded (not partitioned).
|
||||
for name, ref_t in self.input_vars.items():
|
||||
if (
|
||||
name not in self.output_vars
|
||||
and ref_t.out_degree == 0
|
||||
and name not in self.need_remove_input_vars
|
||||
):
|
||||
self.output_vars[name] = self.identity(ref_t)
|
||||
for name, ref_t in self.intermediate_vars.items():
|
||||
if name not in self.output_vars and ref_t.out_degree == 0:
|
||||
self.output_vars[name] = self.identity(ref_t)
|
||||
|
||||
def find_source_slices(
|
||||
self, key: str, local_slice: tuple[slice, ...]
|
||||
) -> list[SliceRef]:
|
||||
assert key in self.output_vars, (
|
||||
f"The key {key} is not in the output_vars (which is built during load_state_dict)."
|
||||
)
|
||||
tensor = self.output_vars[key]
|
||||
if tensor is None:
|
||||
return []
|
||||
results = []
|
||||
assert len(local_slice) == len(tensor.shape), (
|
||||
f"For the key {key}, the target_tensor has {len(local_slice)} dimensions, "
|
||||
f"but the tensor in output_vars has {len(tensor.shape)} dimensions (shape={tensor.shape}). "
|
||||
)
|
||||
ndim = len(tensor.shape)
|
||||
|
||||
def slice_intersect(a: slice, b: slice):
|
||||
start = max(a.start, b.start)
|
||||
stop = min(a.stop, b.stop)
|
||||
if start >= stop:
|
||||
return None
|
||||
return slice(start, stop, 1)
|
||||
|
||||
for src_key, sl_src, sl_dst, pp_list in tensor.slices:
|
||||
intersection = []
|
||||
for i in range(ndim):
|
||||
inter = slice_intersect(local_slice[i], sl_dst[i])
|
||||
if inter is None:
|
||||
break
|
||||
intersection.append(inter)
|
||||
else:
|
||||
# Compute corresponding src_slice for the intersection
|
||||
if pp_list is not None:
|
||||
sl_src = postprocess_transpose(list(sl_src), pp_list)
|
||||
src_slice = []
|
||||
for i in range(ndim):
|
||||
dst = sl_dst[i]
|
||||
src = sl_src[i]
|
||||
dst_start = dst.start
|
||||
src_start = src.start
|
||||
inter_start, inter_stop = (
|
||||
intersection[i].start,
|
||||
intersection[i].stop,
|
||||
)
|
||||
offset = inter_start - dst_start
|
||||
src_inter_start = src_start + offset
|
||||
src_inter_stop = src_inter_start + (
|
||||
inter_stop - inter_start
|
||||
)
|
||||
src_slice.append(slice(src_inter_start, src_inter_stop, 1))
|
||||
if pp_list is not None:
|
||||
src_slice = postprocess_transpose(
|
||||
list(src_slice), pp_list, reverse=True
|
||||
)
|
||||
results.append(
|
||||
(
|
||||
src_key,
|
||||
tuple(src_slice),
|
||||
tuple(intersection),
|
||||
pp_list.copy(),
|
||||
),
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
(src_key, tuple(src_slice), tuple(intersection), None)
|
||||
)
|
||||
return results
|
||||
|
||||
def find_shard_sources(
|
||||
self,
|
||||
target: ShardedWeightDesc,
|
||||
) -> ShardMapping:
|
||||
target_key, opt_state_name = split_optimizer_state_key(target.key)
|
||||
target_local_shape = target.local_shape
|
||||
target_global_offset = target.global_offset
|
||||
target_global_shape = target.global_shape
|
||||
|
||||
if opt_state_name in [".beta1_pow_acc_0", ".beta2_pow_acc_0"]:
|
||||
assert target_key in self.output_vars, (
|
||||
f"The key {target_key} is not in the output_vars (which is built during load_state_dict)."
|
||||
)
|
||||
tensor = self.output_vars[target_key]
|
||||
target_local_shape = tensor.shape
|
||||
target_global_offset = (0,) * len(target_local_shape)
|
||||
target_global_shape = target_local_shape
|
||||
|
||||
slices = tuple(
|
||||
slice(offset, offset + size, 1)
|
||||
for offset, size in zip(target_global_offset, target_local_shape)
|
||||
)
|
||||
|
||||
results = self.find_source_slices(target_key, slices)
|
||||
|
||||
shard_mappings = []
|
||||
|
||||
target_key = (
|
||||
target_key + opt_state_name
|
||||
if opt_state_name is not None
|
||||
else target_key
|
||||
)
|
||||
|
||||
src_keys = {
|
||||
result[0]
|
||||
for result in results
|
||||
if result[0] not in self.need_remove_input_vars
|
||||
}
|
||||
if opt_state_name in [".beta1_pow_acc_0", ".beta2_pow_acc_0"]:
|
||||
if len(src_keys) == 0:
|
||||
return shard_mappings
|
||||
elif len(src_keys) > 1:
|
||||
logger.warning(
|
||||
f"{target_key} has multiple sources: {src_keys} (e.g., .beta1_pow_acc_0). Returning one arbitrarily."
|
||||
)
|
||||
src_key = next(iter(src_keys))
|
||||
else:
|
||||
src_key = next(iter(src_keys))
|
||||
return [
|
||||
ShardMappingEntry(
|
||||
target,
|
||||
ShardedWeightDesc(
|
||||
src_key + opt_state_name,
|
||||
target.local_shape,
|
||||
target.global_shape,
|
||||
target.global_offset,
|
||||
target.dtype,
|
||||
),
|
||||
None,
|
||||
)
|
||||
]
|
||||
|
||||
for src_key, src_slices, local_slices, pp_list in results:
|
||||
src_var = self.input_vars[src_key]
|
||||
target_model_state_key, target_opt_state_name = (
|
||||
split_optimizer_state_key(target.key)
|
||||
)
|
||||
if target_opt_state_name is None:
|
||||
if src_var.dtype != target.dtype:
|
||||
assert pp_list is not None and target.dtype in str(
|
||||
pp_list
|
||||
), (
|
||||
"Direct assignment of Tensors with different types is prohibited in AOA. "
|
||||
f"If you want to achieve this functionality, please use the cast semantics provided by AOA. "
|
||||
f"Now the src_var.dtype is {src_var.dtype}, the target.dtype is {target.dtype}, the pp_list is {pp_list}."
|
||||
f"The src_key is {src_key}, the target_key is {target.key}."
|
||||
)
|
||||
else:
|
||||
src_var.dtype = target.dtype
|
||||
|
||||
src_global_shape = src_var.shape
|
||||
|
||||
src_local_shape = tuple(slc.stop - slc.start for slc in src_slices)
|
||||
src_global_offset = tuple(slc.start for slc in src_slices)
|
||||
|
||||
tgt_local_shape = tuple(
|
||||
slc.stop - slc.start for slc in local_slices
|
||||
)
|
||||
tgt_global_offset = tuple(slc.start for slc in local_slices)
|
||||
|
||||
new_src_key = (
|
||||
src_key + opt_state_name
|
||||
if opt_state_name is not None
|
||||
else src_key
|
||||
)
|
||||
|
||||
source_sharded_weight = ShardedWeightDesc(
|
||||
new_src_key,
|
||||
src_local_shape,
|
||||
tuple(src_global_shape),
|
||||
src_global_offset,
|
||||
src_var.dtype,
|
||||
)
|
||||
target_sharded_weight = ShardedWeightDesc(
|
||||
target_key,
|
||||
tgt_local_shape,
|
||||
tuple(target_global_shape),
|
||||
tgt_global_offset,
|
||||
target.dtype,
|
||||
)
|
||||
|
||||
if src_key in self.need_remove_input_vars:
|
||||
mapping_entry = ShardMappingEntry(
|
||||
target_sharded_weight,
|
||||
source_sharded_weight,
|
||||
[],
|
||||
)
|
||||
continue
|
||||
|
||||
shard_mappings.append(
|
||||
ShardMappingEntry(
|
||||
target_sharded_weight,
|
||||
source_sharded_weight,
|
||||
pp_list,
|
||||
)
|
||||
)
|
||||
|
||||
return shard_mappings
|
||||
|
||||
|
||||
def postprocess_transpose(
|
||||
li: list[tuple[slice, ...]] | tuple[tuple[slice, ...]],
|
||||
postprocess_list: list[str],
|
||||
reverse: bool = False,
|
||||
) -> list[tuple[slice, ...]] | tuple[tuple[slice, ...]]:
|
||||
result = li
|
||||
if reverse:
|
||||
for pp in list(reversed(postprocess_list)):
|
||||
if pp.startswith("["):
|
||||
reversed_transpose = np.argsort(ast.literal_eval(pp)).tolist()
|
||||
result = transpose_list(result, reversed_transpose)
|
||||
else:
|
||||
for pp in postprocess_list:
|
||||
if pp.startswith("["):
|
||||
result = transpose_list(result, ast.literal_eval(pp))
|
||||
return result
|
||||
|
||||
|
||||
def transpose_list(
|
||||
li: list[tuple[slice, ...]] | tuple[tuple[slice, ...]],
|
||||
permutation: list[int],
|
||||
) -> list[tuple[slice, ...]] | tuple[tuple[slice, ...]]:
|
||||
trans_list = []
|
||||
for idx in permutation:
|
||||
trans_list.append(li[idx])
|
||||
if isinstance(li, tuple):
|
||||
return tuple(trans_list)
|
||||
else:
|
||||
return trans_list
|
||||
|
||||
|
||||
def invert_permutation(p: list[int]) -> list[int]:
|
||||
q = [0] * len(p)
|
||||
for i, pi in enumerate(p):
|
||||
q[pi] = i
|
||||
return q
|
||||
@@ -0,0 +1,150 @@
|
||||
# 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 re
|
||||
from enum import Enum, auto
|
||||
|
||||
|
||||
class Token:
|
||||
def __init__(self, type, value):
|
||||
self.type = type
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"Token({self.type}, {self.value!r})"
|
||||
|
||||
|
||||
class TokenType(Enum):
|
||||
IDENTIFIER = auto()
|
||||
NUMBER = auto()
|
||||
COLON = auto()
|
||||
LBRACKET = auto()
|
||||
RBRACKET = auto()
|
||||
COMMA = auto()
|
||||
RARROW = auto()
|
||||
STRING = auto()
|
||||
EQUAL = auto()
|
||||
NEWLINE = auto()
|
||||
EOF = auto()
|
||||
|
||||
|
||||
class Lexer:
|
||||
token_specification = [
|
||||
('RARROW', r'->'),
|
||||
('EQUAL', r'='),
|
||||
('COLON', r':'),
|
||||
('LBRACKET', r'\['),
|
||||
('RBRACKET', r'\]'),
|
||||
('COMMA', r','),
|
||||
('NUMBER', r'\d+'),
|
||||
('STRING', r'"[^"]*"|\'[^\']*\''),
|
||||
('IDENTIFIER', r'[A-Za-z_][A-Za-z\.\$\_\*\d\^T]*'),
|
||||
('SKIP', r'[ \t]+'),
|
||||
('NEWLINE', r'[\r\n]+'),
|
||||
('MISMATCH', r'.'),
|
||||
]
|
||||
|
||||
def __init__(self, context, traceback=None):
|
||||
from .macros import macro_registry
|
||||
|
||||
self.macros = [list(d.values())[1] for d in macro_registry.macros]
|
||||
self.get_token = re.compile(
|
||||
'|'.join(
|
||||
f'(?P<{name}>{regex})'
|
||||
for name, regex in self.token_specification
|
||||
)
|
||||
).match
|
||||
self.context = context
|
||||
self.traceback = traceback
|
||||
|
||||
def tokenize(self, text):
|
||||
pos = 0
|
||||
mo = self.get_token(text, pos)
|
||||
tokens = []
|
||||
if not text.endswith('\n'):
|
||||
text += '\n'
|
||||
while mo is not None:
|
||||
kind = mo.lastgroup
|
||||
value = mo.group()
|
||||
if kind == 'SKIP':
|
||||
pass
|
||||
elif kind == 'MISMATCH':
|
||||
raise RuntimeError(
|
||||
f'Unexpected character {value!r} at position {pos}'
|
||||
)
|
||||
else:
|
||||
tokens.append(Token(TokenType[kind], value))
|
||||
pos = mo.end()
|
||||
mo = self.get_token(text, pos)
|
||||
return tokens
|
||||
|
||||
def apply_macro(self, expression, macro):
|
||||
if isinstance(expression, str):
|
||||
expression = [expression]
|
||||
new_expression = []
|
||||
for expr in expression:
|
||||
results = macro(self.tokenize(expr), expr, self.context)
|
||||
if isinstance(results, str):
|
||||
new_expression.append(results)
|
||||
else:
|
||||
new_expression.extend(results)
|
||||
return new_expression
|
||||
|
||||
def apply_single_macro_to_all(self, expressions, macro):
|
||||
new_expressions = []
|
||||
macro_name = getattr(macro, "__name__", "macro")
|
||||
for expr in expressions:
|
||||
try:
|
||||
results = macro(self.tokenize(expr), expr, self.context)
|
||||
except (AssertionError, ValueError, KeyError, RuntimeError) as e:
|
||||
if self.traceback:
|
||||
chain = self.traceback.build_chain(expr)
|
||||
self.traceback.add_error(
|
||||
error_message=str(e),
|
||||
stage=f"{macro_name}",
|
||||
chain=chain,
|
||||
error_type=type(e).__name__,
|
||||
)
|
||||
self.traceback.print()
|
||||
raise
|
||||
|
||||
if isinstance(results, str):
|
||||
results_list = [results]
|
||||
else:
|
||||
results_list = list(results)
|
||||
|
||||
if self.traceback:
|
||||
if results_list != [expr]:
|
||||
self.traceback.record_children(
|
||||
expr, results_list, macro_name
|
||||
)
|
||||
|
||||
new_expressions.extend(results_list)
|
||||
return new_expressions
|
||||
|
||||
def all_tokens(self, expressions):
|
||||
if self.traceback:
|
||||
self.traceback.register_roots(list(expressions))
|
||||
|
||||
current_expressions = expressions
|
||||
for macro in self.macros:
|
||||
current_expressions = self.apply_single_macro_to_all(
|
||||
current_expressions, macro
|
||||
)
|
||||
|
||||
self.final_expressions = list(current_expressions)
|
||||
tokens = []
|
||||
for expr in current_expressions:
|
||||
tokens.extend(self.tokenize(expr))
|
||||
return tokens
|
||||
@@ -0,0 +1,864 @@
|
||||
# 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 math
|
||||
import re
|
||||
from itertools import product
|
||||
|
||||
from .lexer import Token, TokenType
|
||||
|
||||
|
||||
def macro(name, priority):
|
||||
def decorator(func):
|
||||
macro_registry.register_macro(name, func, priority)
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class MacroRegistry:
|
||||
_instance = None
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not hasattr(self, 'macros'):
|
||||
self.macros = []
|
||||
|
||||
def register_macro(self, name, func, priority):
|
||||
if any(macro['name'] == name for macro in self.macros):
|
||||
raise ValueError(f"Macro '{name}' is already registered.")
|
||||
self.macros.append({'name': name, 'func': func, 'priority': priority})
|
||||
self.macros.sort(key=lambda x: x['priority'], reverse=False)
|
||||
|
||||
|
||||
macro_registry = MacroRegistry()
|
||||
|
||||
GLOBAL_ATTRIBUTE_KEYWORDS = [
|
||||
"axis",
|
||||
'fused_ffn',
|
||||
'fused_qkv_old',
|
||||
'num_heads',
|
||||
'num_key_value_groups',
|
||||
'permute',
|
||||
'dtype',
|
||||
'fused_qkv',
|
||||
'src_dtype',
|
||||
'dst_dtype',
|
||||
]
|
||||
|
||||
EXTRA_SUFFIX = [
|
||||
"^T",
|
||||
]
|
||||
|
||||
|
||||
def extract_axis_and_clean_tokens(tokens):
|
||||
axis = 1
|
||||
for idx, tkn in enumerate(tokens):
|
||||
if tkn.value == "axis" and idx + 2 < len(tokens):
|
||||
axis = int(tokens[idx + 2].value)
|
||||
end_idx = idx + 3
|
||||
if end_idx < len(tokens) - 1:
|
||||
assert tokens[end_idx].value == ",", (
|
||||
f"The different attributes must split by a comma, but now the token is {tokens[end_idx].value}."
|
||||
)
|
||||
end_idx += 1
|
||||
tokens = tokens[:idx] + tokens[end_idx:]
|
||||
break
|
||||
return axis, tokens
|
||||
|
||||
|
||||
# star_macro must be called after layer_id_macro
|
||||
@macro(name='star_macro', priority=3)
|
||||
def star_macro(tokens, expression, context):
|
||||
STAR_TAG = "*"
|
||||
if STAR_TAG not in expression:
|
||||
return expression
|
||||
|
||||
def _sort_keys_by_numeric_part(prefix, suffix, allkeys):
|
||||
pattern = re.compile(rf"{re.escape(prefix)}(\d+){re.escape(suffix)}")
|
||||
filtered_keys = []
|
||||
for key in allkeys:
|
||||
match = pattern.fullmatch(key)
|
||||
if match:
|
||||
num = int(match.group(1))
|
||||
filtered_keys.append((key, num))
|
||||
sorted_keys = sorted(filtered_keys, key=lambda x: x[1])
|
||||
return [key for key, _ in sorted_keys]
|
||||
|
||||
pre_rarrow = True
|
||||
new_tokens = []
|
||||
for token in tokens:
|
||||
if token.type == TokenType.RARROW:
|
||||
pre_rarrow = False
|
||||
if token.type == TokenType.IDENTIFIER and STAR_TAG in token.value:
|
||||
prefix, suffix = token.value.split(STAR_TAG)
|
||||
allkeys = (
|
||||
context.get_all_dst_state_keys()
|
||||
if not pre_rarrow
|
||||
else context.get_all_src_state_keys()
|
||||
)
|
||||
assert len(allkeys) != 0, (
|
||||
f"No keys found with prefix '{prefix}' and suffix '{suffix}' in "
|
||||
f"{'destination_state_shard_info' if not pre_rarrow else 'source_state_shard_info'}, please check!"
|
||||
)
|
||||
keys = list(_sort_keys_by_numeric_part(prefix, suffix, allkeys))
|
||||
for key in keys:
|
||||
new_tokens.append(Token(TokenType.IDENTIFIER, key))
|
||||
if key != keys[-1]:
|
||||
new_tokens.append(Token(TokenType.COMMA, ","))
|
||||
else:
|
||||
new_tokens.append(token)
|
||||
new_expression = "".join([token.value for token in new_tokens])
|
||||
return new_expression
|
||||
|
||||
|
||||
@macro(name='layer_id_offset_macro', priority=1)
|
||||
def layer_id_offset_macro(tokens, expression, context):
|
||||
LAYER_ID_OFFSET_MACRO_TAG = "$LAYER_ID_OFFSET"
|
||||
if LAYER_ID_OFFSET_MACRO_TAG not in expression:
|
||||
return expression
|
||||
|
||||
name_with_layer_id_offset = next(
|
||||
(
|
||||
token.value
|
||||
for token in tokens
|
||||
if token.type == TokenType.IDENTIFIER
|
||||
and LAYER_ID_OFFSET_MACRO_TAG in token.value
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert name_with_layer_id_offset, (
|
||||
"No $LAYER_ID_OFFSET found in NAME tokens.Please check the aoa_config."
|
||||
)
|
||||
assert all(
|
||||
(t.type != TokenType.IDENTIFIER)
|
||||
or (LAYER_ID_OFFSET_MACRO_TAG in t.value)
|
||||
or (t.value in GLOBAL_ATTRIBUTE_KEYWORDS)
|
||||
for t in tokens
|
||||
), (
|
||||
f"All IDENTIFIER tokens must contain {LAYER_ID_OFFSET_MACRO_TAG} when a NAME with it is present, except for GLOBAL_ATTRIBUTE_KEYWORDS."
|
||||
)
|
||||
|
||||
match_layer_id_offset = context.get_num_hidden_layers(
|
||||
name_with_layer_id_offset, LAYER_ID_OFFSET_MACRO_TAG
|
||||
)
|
||||
expanded_expressions = []
|
||||
|
||||
match_layer_id_offset = sorted(match_layer_id_offset)
|
||||
|
||||
for layer_id in match_layer_id_offset:
|
||||
expr = ""
|
||||
before_rarrow = True
|
||||
for token in tokens:
|
||||
if token.type == TokenType.RARROW:
|
||||
before_rarrow = False
|
||||
if before_rarrow:
|
||||
cur_layer_id = layer_id
|
||||
else:
|
||||
cur_layer_id = layer_id - 1
|
||||
if token.type == TokenType.IDENTIFIER:
|
||||
if LAYER_ID_OFFSET_MACRO_TAG in token.value:
|
||||
expr += token.value.replace(
|
||||
LAYER_ID_OFFSET_MACRO_TAG, str(cur_layer_id)
|
||||
)
|
||||
elif token.value not in GLOBAL_ATTRIBUTE_KEYWORDS:
|
||||
expr += f"{token.value}.layer.{cur_layer_id}"
|
||||
else:
|
||||
expr += token.value
|
||||
else:
|
||||
expr += token.value
|
||||
expanded_expressions.append(expr)
|
||||
return expanded_expressions
|
||||
|
||||
|
||||
@macro(name='array_macro', priority=2)
|
||||
def array_macro(tokens, expression, context):
|
||||
if "[" not in expression:
|
||||
return expression
|
||||
new_tokens = []
|
||||
idx = 0
|
||||
while idx < len(tokens):
|
||||
if tokens[idx].type == TokenType.LBRACKET:
|
||||
name = tokens[idx - 1].value
|
||||
assert (
|
||||
tokens[idx + 1].type == TokenType.NUMBER
|
||||
and tokens[idx + 2].type == TokenType.COLON
|
||||
and tokens[idx + 3].type == TokenType.NUMBER
|
||||
and tokens[idx + 4].type == TokenType.RBRACKET
|
||||
), (
|
||||
f"The array macro format is incorrect which is must be like: NAME[START:END], but now the format is {tokens[idx].value}{tokens[idx + 1].value}:{tokens[idx + 3].value}{tokens[idx + 4].value}."
|
||||
)
|
||||
new_tokens.pop()
|
||||
start = int(tokens[idx + 1].value)
|
||||
end = int(tokens[idx + 3].value)
|
||||
for i in range(start, end):
|
||||
new_tokens.append(
|
||||
Token(TokenType.IDENTIFIER, name + "_" + str(i))
|
||||
)
|
||||
if i != end - 1:
|
||||
new_tokens.append(Token(TokenType.COMMA, ","))
|
||||
idx += 5
|
||||
else:
|
||||
new_tokens.append(tokens[idx])
|
||||
idx += 1
|
||||
new_expression = "".join([token.value for token in new_tokens])
|
||||
return new_expression
|
||||
|
||||
|
||||
@macro(name='fused_qkv_old_macro', priority=6)
|
||||
def fused_qkv_old_macro(tokens, expression, context):
|
||||
FUSED_QKV_OLD_TAG = "fused_qkv_old"
|
||||
if not any(tkn.value == FUSED_QKV_OLD_TAG for tkn in tokens):
|
||||
return expression
|
||||
|
||||
axis, tokens = extract_axis_and_clean_tokens(tokens)
|
||||
|
||||
attn_head_num = None
|
||||
num_key_value_groups = None
|
||||
fused_qkv_old_pos = None
|
||||
rarrow_pos = None
|
||||
right_var_end_pos = None
|
||||
|
||||
for idx, token in enumerate(tokens):
|
||||
if token.type == TokenType.IDENTIFIER:
|
||||
if token.value == "num_heads" and idx + 2 < len(tokens):
|
||||
attn_head_num = int(tokens[idx + 2].value)
|
||||
elif token.value == "num_key_value_groups" and idx + 2 < len(
|
||||
tokens
|
||||
):
|
||||
num_key_value_groups = int(tokens[idx + 2].value)
|
||||
elif token.value == FUSED_QKV_OLD_TAG:
|
||||
fused_qkv_old_pos = idx
|
||||
elif token.type == TokenType.RARROW and rarrow_pos is None:
|
||||
rarrow_pos = idx
|
||||
if (
|
||||
right_var_end_pos is None
|
||||
and token.type == TokenType.IDENTIFIER
|
||||
and token.value
|
||||
in {FUSED_QKV_OLD_TAG, "num_heads", "num_key_value_groups"}
|
||||
):
|
||||
right_var_end_pos = idx + 1
|
||||
|
||||
assert attn_head_num and attn_head_num > 0, (
|
||||
f"num_heads must be positive.(got: {attn_head_num})."
|
||||
)
|
||||
assert num_key_value_groups and num_key_value_groups > 0, (
|
||||
f"num_key_value_groups must be positive.(got: {num_key_value_groups})."
|
||||
)
|
||||
assert fused_qkv_old_pos is not None, (
|
||||
f"No fused_qkv_old tag found in expression. The tag must be {FUSED_QKV_OLD_TAG}."
|
||||
)
|
||||
assert rarrow_pos is not None, "No -> found in expression."
|
||||
assert attn_head_num % num_key_value_groups == 0, (
|
||||
f"num_heads ({attn_head_num}) must be divisible by num_key_value_groups ({num_key_value_groups})."
|
||||
)
|
||||
|
||||
results = []
|
||||
num_key_value_heads = num_key_value_groups
|
||||
if rarrow_pos == 1:
|
||||
src_qkv_weight_name = tokens[0].value
|
||||
if fused_qkv_old_pos > 4:
|
||||
dst_qkv_weight_name = None
|
||||
else:
|
||||
dst_qkv_weight_name = tokens[2].value
|
||||
|
||||
if context.aoa_config_reverse:
|
||||
dst_state_shard_num = context.get_src_state_shard_num(
|
||||
dst_qkv_weight_name
|
||||
)
|
||||
src_state_shard_num = (
|
||||
context.get_dst_state_shard_num(src_qkv_weight_name)
|
||||
if src_qkv_weight_name is not None
|
||||
else 1
|
||||
)
|
||||
else:
|
||||
src_state_shard_num = context.get_src_state_shard_num(
|
||||
src_qkv_weight_name
|
||||
)
|
||||
dst_state_shard_num = (
|
||||
context.get_dst_state_shard_num(dst_qkv_weight_name)
|
||||
if dst_qkv_weight_name is not None
|
||||
else 1
|
||||
)
|
||||
|
||||
configs = [
|
||||
(src_state_shard_num, src_qkv_weight_name),
|
||||
(dst_state_shard_num, dst_qkv_weight_name),
|
||||
]
|
||||
|
||||
head_config = [
|
||||
("Q", attn_head_num),
|
||||
("K", num_key_value_heads),
|
||||
("V", num_key_value_heads),
|
||||
]
|
||||
|
||||
def gen_expr(tp_degree, num_heads, tp_rank, comp):
|
||||
start = tp_rank * num_heads // tp_degree
|
||||
count = num_heads // tp_degree
|
||||
return ",".join(
|
||||
f"fused_qkv_old_tmp.{comp}_{i}"
|
||||
for i in range(start, start + count)
|
||||
)
|
||||
|
||||
for idx, (tp_degree, qkv_weight_name) in enumerate(configs):
|
||||
qkv_parts = [
|
||||
gen_expr(tp_degree, n, tp_rank, c)
|
||||
for tp_rank in range(tp_degree)
|
||||
for c, n in head_config
|
||||
]
|
||||
if idx == 0:
|
||||
mapping = (
|
||||
f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis={axis}"
|
||||
)
|
||||
results.append(mapping)
|
||||
elif qkv_weight_name is not None:
|
||||
mapping = (
|
||||
f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis={axis}"
|
||||
)
|
||||
results.append(mapping)
|
||||
|
||||
if fused_qkv_old_pos > 4:
|
||||
|
||||
def _generate_expr(prefix, count, target_name):
|
||||
elements = ",".join(
|
||||
f"fused_qkv_old_tmp.{prefix}_{i}" for i in range(count)
|
||||
)
|
||||
return f"{elements} -> {target_name}, axis={axis}"
|
||||
|
||||
q_name = tokens[2].value
|
||||
k_name = tokens[4].value
|
||||
v_name = tokens[6].value
|
||||
|
||||
results.append(_generate_expr("Q", attn_head_num, q_name))
|
||||
results.append(_generate_expr("K", num_key_value_heads, k_name))
|
||||
results.append(_generate_expr("V", num_key_value_heads, v_name))
|
||||
elif rarrow_pos == 5:
|
||||
q_name = tokens[0].value
|
||||
k_name = tokens[2].value
|
||||
v_name = tokens[4].value
|
||||
dst_qkv_weight_name = tokens[6].value
|
||||
|
||||
fused_qkv_tmp_name = f"{q_name}.{k_name}.{v_name}.tmp"
|
||||
results.append(
|
||||
f"{q_name},{k_name},{v_name} -> {fused_qkv_tmp_name}, axis={axis}"
|
||||
)
|
||||
dst_state_shard_num = context.get_dst_state_shard_num(
|
||||
dst_qkv_weight_name
|
||||
)
|
||||
|
||||
configs = [
|
||||
(1, fused_qkv_tmp_name),
|
||||
(dst_state_shard_num, dst_qkv_weight_name),
|
||||
]
|
||||
|
||||
head_config = [
|
||||
("Q", attn_head_num),
|
||||
("K", num_key_value_heads),
|
||||
("V", num_key_value_heads),
|
||||
]
|
||||
|
||||
def gen_expr(tp_degree, num_heads, tp_rank, comp):
|
||||
start = tp_rank * num_heads // tp_degree
|
||||
count = num_heads // tp_degree
|
||||
return ",".join(
|
||||
f"fused_qkv_old_tmp.{comp}_{i}"
|
||||
for i in range(start, start + count)
|
||||
)
|
||||
|
||||
for idx, (tp_degree, qkv_weight_name) in enumerate(configs):
|
||||
qkv_parts = [
|
||||
gen_expr(tp_degree, n, tp_rank, c)
|
||||
for tp_rank in range(tp_degree)
|
||||
for c, n in head_config
|
||||
]
|
||||
if idx == 0:
|
||||
mapping = (
|
||||
f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis={axis}"
|
||||
)
|
||||
else:
|
||||
mapping = (
|
||||
f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis={axis}"
|
||||
)
|
||||
results.append(mapping)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported fused_qkv_old macro format: {expression}."
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
@macro(name='fused_ffn_macro', priority=6)
|
||||
def fused_ffn_macro(tokens, expression, context):
|
||||
FUSED_FFN_TAG = "fused_ffn"
|
||||
if not any(tkn.value == FUSED_FFN_TAG for tkn in tokens):
|
||||
return expression
|
||||
|
||||
axis, tokens = extract_axis_and_clean_tokens(tokens)
|
||||
|
||||
rarrow_pos = None
|
||||
fused_ffn_pos = None
|
||||
for idx, token in enumerate(tokens):
|
||||
if token.type == TokenType.RARROW and rarrow_pos is None:
|
||||
rarrow_pos = idx
|
||||
elif (
|
||||
token.type == TokenType.IDENTIFIER and token.value == FUSED_FFN_TAG
|
||||
):
|
||||
fused_ffn_pos = idx
|
||||
assert rarrow_pos is not None, "No -> found in expression."
|
||||
assert fused_ffn_pos is not None, (
|
||||
f"No fused_ffn tag found in expression. The tag must be {FUSED_FFN_TAG}."
|
||||
)
|
||||
results = []
|
||||
if rarrow_pos == 1:
|
||||
src_ffn_weight_name = tokens[0].value
|
||||
if fused_ffn_pos == 4:
|
||||
dst_ffn_weight_name = tokens[2].value
|
||||
else:
|
||||
dst_ffn_weight_name = None
|
||||
if context.aoa_config_reverse:
|
||||
dst_state_shard_num = context.get_src_state_shard_num(
|
||||
dst_ffn_weight_name
|
||||
)
|
||||
src_state_shard_num = (
|
||||
context.get_dst_state_shard_num(src_ffn_weight_name)
|
||||
if src_ffn_weight_name is not None
|
||||
else 1
|
||||
)
|
||||
else:
|
||||
src_state_shard_num = context.get_src_state_shard_num(
|
||||
src_ffn_weight_name
|
||||
)
|
||||
dst_state_shard_num = (
|
||||
context.get_dst_state_shard_num(dst_ffn_weight_name)
|
||||
if dst_ffn_weight_name is not None
|
||||
else 1
|
||||
)
|
||||
splited_num = math.lcm(src_state_shard_num, dst_state_shard_num)
|
||||
|
||||
configs = [
|
||||
(src_state_shard_num, src_ffn_weight_name),
|
||||
(dst_state_shard_num, dst_ffn_weight_name),
|
||||
]
|
||||
split_config = [("GATE", splited_num), ("UP", splited_num)]
|
||||
|
||||
def gen_expr(tp_degree, splited_num, tp_rank, comp):
|
||||
return ",".join(
|
||||
f"fused_ffn_tmp.{comp}_{tp_rank * splited_num // tp_degree + idx}"
|
||||
for idx in range(splited_num // tp_degree)
|
||||
)
|
||||
|
||||
for idx, (tp_degree, ffn_weight_name) in enumerate(configs):
|
||||
ffn_parts = [
|
||||
gen_expr(tp_degree, n, tp_rank, c)
|
||||
for tp_rank in range(tp_degree)
|
||||
for c, n in split_config
|
||||
]
|
||||
if idx == 0:
|
||||
results.append(
|
||||
f"{ffn_weight_name} -> {','.join(ffn_parts)}, axis={axis}"
|
||||
)
|
||||
elif ffn_weight_name is not None:
|
||||
results.append(
|
||||
f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis={axis}"
|
||||
)
|
||||
if fused_ffn_pos > 4:
|
||||
|
||||
def _generate_expr(prefix, count, target_name):
|
||||
elements = ",".join(
|
||||
f"fused_ffn_tmp.{prefix}_{i}" for i in range(count)
|
||||
)
|
||||
return f"{elements} -> {target_name}, axis={axis}"
|
||||
|
||||
gate_name = tokens[2].value
|
||||
up_name = tokens[4].value
|
||||
|
||||
results.append(_generate_expr("GATE", splited_num, gate_name))
|
||||
results.append(_generate_expr("UP", splited_num, up_name))
|
||||
|
||||
elif rarrow_pos == 3:
|
||||
gate_name = tokens[0].value
|
||||
up_name = tokens[2].value
|
||||
dst_ffn_weight_name = tokens[4].value
|
||||
|
||||
fused_gate_up_tmp_name = f"{gate_name}.{up_name}.tmp"
|
||||
results.append(
|
||||
f"{gate_name},{up_name} -> {fused_gate_up_tmp_name}, axis={axis}"
|
||||
)
|
||||
dst_state_shard_num = context.get_dst_state_shard_num(
|
||||
dst_ffn_weight_name
|
||||
)
|
||||
|
||||
configs = [
|
||||
(1, fused_gate_up_tmp_name),
|
||||
(dst_state_shard_num, dst_ffn_weight_name),
|
||||
]
|
||||
|
||||
split_config = [
|
||||
("GATE", dst_state_shard_num),
|
||||
("UP", dst_state_shard_num),
|
||||
]
|
||||
|
||||
def gen_expr(tp_degree, splited_num, tp_rank, comp):
|
||||
return ",".join(
|
||||
f"fused_ffn_tmp.{comp}_{tp_rank * splited_num // tp_degree + idx}"
|
||||
for idx in range(splited_num // tp_degree)
|
||||
)
|
||||
|
||||
for idx, (tp_degree, ffn_weight_name) in enumerate(configs):
|
||||
ffn_parts = [
|
||||
gen_expr(tp_degree, n, tp_rank, c)
|
||||
for tp_rank in range(tp_degree)
|
||||
for c, n in split_config
|
||||
]
|
||||
if idx == 0:
|
||||
results.append(
|
||||
f"{ffn_weight_name} -> {','.join(ffn_parts)}, axis={axis}"
|
||||
)
|
||||
else:
|
||||
results.append(
|
||||
f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis={axis}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported fused_ffn macro format: {expression}.")
|
||||
return results
|
||||
|
||||
|
||||
@macro(name='transpose_macro', priority=5)
|
||||
def transpose_macro(tokens, expression, context):
|
||||
TRANSPOSE_TAG = "^T"
|
||||
|
||||
if TRANSPOSE_TAG not in expression:
|
||||
return expression
|
||||
|
||||
transpose_vars = set()
|
||||
new_expression = ""
|
||||
rarrow_pos = None
|
||||
|
||||
for idx, token in enumerate(tokens):
|
||||
if token.type == TokenType.RARROW:
|
||||
rarrow_pos = idx
|
||||
break
|
||||
|
||||
assert rarrow_pos is not None, "No -> found in expression."
|
||||
|
||||
for token in tokens[rarrow_pos + 1 :]:
|
||||
if token.type == TokenType.IDENTIFIER and token.value.endswith(
|
||||
TRANSPOSE_TAG
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot assign to transpose (e.g., 'A -> B^T').\n"
|
||||
"B^T is not a real variable, just a view.\n"
|
||||
"Assign first: A -> B\n"
|
||||
"Then transpose: B^T -> B"
|
||||
)
|
||||
for token in tokens:
|
||||
if token.type == TokenType.IDENTIFIER and token.value.endswith(
|
||||
TRANSPOSE_TAG
|
||||
):
|
||||
var_name = token.value[: -len(TRANSPOSE_TAG)]
|
||||
transpose_vars.add(var_name)
|
||||
new_expression += var_name + "_transpose_tmp"
|
||||
else:
|
||||
new_expression += token.value
|
||||
|
||||
results = [
|
||||
f'{var} -> {var}_transpose_tmp, permute = "[]"'
|
||||
for var in transpose_vars
|
||||
]
|
||||
results.append(new_expression)
|
||||
return results
|
||||
|
||||
|
||||
@macro(name='fused_qkv_macro', priority=6)
|
||||
def fused_qkv_macro(tokens, expression, context):
|
||||
FUSED_QKV_TAG = "fused_qkv"
|
||||
if not any(tkn.value == FUSED_QKV_TAG for tkn in tokens):
|
||||
return expression
|
||||
|
||||
axis, tokens = extract_axis_and_clean_tokens(tokens)
|
||||
|
||||
attn_head_num = num_heads = None
|
||||
num_key_value_groups = None
|
||||
fused_qkv_pos = None
|
||||
rarrow_pos = None
|
||||
|
||||
for idx, token in enumerate(tokens):
|
||||
if token.type == TokenType.IDENTIFIER:
|
||||
if token.value == "num_heads" and idx + 2 < len(tokens):
|
||||
attn_head_num = int(tokens[idx + 2].value)
|
||||
elif token.value == "num_key_value_groups" and idx + 2 < len(
|
||||
tokens
|
||||
):
|
||||
num_key_value_groups = int(tokens[idx + 2].value)
|
||||
elif token.value == FUSED_QKV_TAG:
|
||||
fused_qkv_pos = idx
|
||||
elif token.type == TokenType.RARROW and rarrow_pos is None:
|
||||
rarrow_pos = idx
|
||||
|
||||
assert attn_head_num and attn_head_num > 0, (
|
||||
f"num_heads must be positive (got: {attn_head_num})"
|
||||
)
|
||||
assert num_key_value_groups and num_key_value_groups > 0, (
|
||||
f"num_key_value_groups must be positive (got: {num_key_value_groups})"
|
||||
)
|
||||
assert fused_qkv_pos is not None, (
|
||||
f"No fused_qkv tag found in expression. The tag must be {FUSED_QKV_TAG}."
|
||||
)
|
||||
assert rarrow_pos is not None, "No -> found in expression."
|
||||
assert rarrow_pos == 1 or rarrow_pos == 5, (
|
||||
"Only support q,k,v -> fused_qkv or fused_qkv -> q,k,v patterns"
|
||||
)
|
||||
assert attn_head_num % num_key_value_groups == 0, (
|
||||
f"num_heads ({attn_head_num}) must be divisible by num_key_value_groups ({num_key_value_groups})."
|
||||
)
|
||||
|
||||
num_key_value_heads = attn_head_num // num_key_value_groups
|
||||
|
||||
def make_names(base, n):
|
||||
return [f"{base}{i}" for i in range(n)]
|
||||
|
||||
results = []
|
||||
|
||||
if rarrow_pos == 1:
|
||||
fused_qkv_var = tokens[0].value
|
||||
q_var = tokens[rarrow_pos + 1].value
|
||||
k_var = tokens[rarrow_pos + 3].value
|
||||
v_var = tokens[rarrow_pos + 5].value
|
||||
|
||||
q_names = make_names(q_var, attn_head_num)
|
||||
k_names = make_names(k_var, num_key_value_groups)
|
||||
v_names = make_names(v_var, num_key_value_groups)
|
||||
|
||||
fused_qkv_order = []
|
||||
for g in range(num_key_value_groups):
|
||||
fused_qkv_order.extend(
|
||||
q_names[g * num_key_value_heads : (g + 1) * num_key_value_heads]
|
||||
)
|
||||
fused_qkv_order.append(k_names[g])
|
||||
fused_qkv_order.append(v_names[g])
|
||||
results.append(
|
||||
f"{fused_qkv_var} -> {','.join(fused_qkv_order)}, axis={axis}"
|
||||
)
|
||||
|
||||
results.append(f"{','.join(q_names)} -> {q_var}, axis={axis}")
|
||||
results.append(f"{','.join(k_names)} -> {k_var}, axis={axis}")
|
||||
results.append(f"{','.join(v_names)} -> {v_var}, axis={axis}")
|
||||
|
||||
return results
|
||||
|
||||
elif rarrow_pos == 5:
|
||||
q_var = tokens[0].value
|
||||
k_var = tokens[2].value
|
||||
v_var = tokens[4].value
|
||||
fused_qkv_var = tokens[rarrow_pos + 1].value
|
||||
|
||||
q_names = make_names(q_var, attn_head_num)
|
||||
k_names = make_names(k_var, num_key_value_groups)
|
||||
v_names = make_names(v_var, num_key_value_groups)
|
||||
|
||||
results.append(f"{q_var} -> {','.join(q_names)}, axis={axis}")
|
||||
results.append(f"{k_var} -> {','.join(k_names)}, axis={axis}")
|
||||
results.append(f"{v_var} -> {','.join(v_names)}, axis={axis}")
|
||||
|
||||
fused_qkv_order = []
|
||||
for g in range(num_key_value_groups):
|
||||
fused_qkv_order.extend(
|
||||
q_names[g * num_key_value_heads : (g + 1) * num_key_value_heads]
|
||||
)
|
||||
fused_qkv_order.append(k_names[g])
|
||||
fused_qkv_order.append(v_names[g])
|
||||
results.append(
|
||||
f"{','.join(fused_qkv_order)} -> {fused_qkv_var}, axis={axis}"
|
||||
)
|
||||
return results
|
||||
|
||||
else:
|
||||
return expression
|
||||
|
||||
|
||||
class IDMatcher:
|
||||
def __init__(
|
||||
self,
|
||||
source_keys: list[str],
|
||||
extra_suffixes: list[str],
|
||||
allowed_placeholders: list[str],
|
||||
):
|
||||
self.source_keys = set(source_keys)
|
||||
self.allowed_placeholders = allowed_placeholders
|
||||
# Dynamically build regex pattern from allowed placeholders
|
||||
placeholder_pattern = '|'.join(
|
||||
re.escape(ph) for ph in self.allowed_placeholders
|
||||
)
|
||||
self._placeholder_pattern = re.compile(f'({placeholder_pattern})')
|
||||
self.extra_suffixes = sorted(extra_suffixes, key=lambda x: (-len(x), x))
|
||||
|
||||
def _remove_extra_suffixes(self, key: str) -> str:
|
||||
for sfx in self.extra_suffixes:
|
||||
if key.endswith(sfx):
|
||||
key = key[: -len(sfx)]
|
||||
break
|
||||
return key
|
||||
|
||||
def _pattern_to_regex(self, pattern: str) -> tuple[re.Pattern, list[str]]:
|
||||
placeholders = sorted(set(self._placeholder_pattern.findall(pattern)))
|
||||
regex_str = re.escape(pattern)
|
||||
for ph in placeholders:
|
||||
group_name = ph[1:]
|
||||
regex_str = regex_str.replace(
|
||||
re.escape(ph), f'(?P<{group_name}>\\d+)'
|
||||
)
|
||||
return re.compile(f'^{regex_str}$'), [ph[1:] for ph in placeholders]
|
||||
|
||||
def _substitute_ids(self, pattern: str, id_dict: dict[str, int]) -> str:
|
||||
key = pattern
|
||||
for ph, value in id_dict.items():
|
||||
key = key.replace(f'${ph}', str(value))
|
||||
return key
|
||||
|
||||
def find_matches(self, pattern: str) -> dict[str, list[int]]:
|
||||
pattern = self._remove_extra_suffixes(pattern)
|
||||
regex, ph_names = self._pattern_to_regex(pattern)
|
||||
id_values = {ph: set() for ph in ph_names}
|
||||
for key in self.source_keys:
|
||||
match = regex.match(key)
|
||||
if match:
|
||||
for k, v in match.groupdict().items():
|
||||
id_values[k].add(int(v))
|
||||
return {k: sorted(vs) for k, vs in id_values.items()}
|
||||
|
||||
|
||||
# Global registry for allowed_placeholders
|
||||
_REGISTERED_PLACEHOLDERS = ['$EXPERT_ID', '$LAYER_ID']
|
||||
|
||||
|
||||
# TODO: need to adapt the scene of temp_layers.\$LAYER_ID.weight -> dst_layers.\$LAYER_ID.weight
|
||||
@macro(name='id_macro', priority=1)
|
||||
def id(tokens, expression, context):
|
||||
allowed_placeholders = _REGISTERED_PLACEHOLDERS
|
||||
has_allowed_placeholder = any(
|
||||
ph in expression for ph in allowed_placeholders
|
||||
)
|
||||
if not has_allowed_placeholder:
|
||||
return expression
|
||||
|
||||
if not context.aoa_config_reverse:
|
||||
name_with_id = next(
|
||||
(
|
||||
token.value
|
||||
for token in tokens
|
||||
if token.type == TokenType.IDENTIFIER
|
||||
and any(ph in token.value for ph in allowed_placeholders)
|
||||
),
|
||||
None,
|
||||
)
|
||||
else:
|
||||
flag_right_var = False
|
||||
for token in tokens:
|
||||
if token.type == TokenType.RARROW:
|
||||
flag_right_var = True
|
||||
if token.type == TokenType.IDENTIFIER and any(
|
||||
ph in token.value for ph in allowed_placeholders
|
||||
):
|
||||
if flag_right_var:
|
||||
name_with_id = token.value
|
||||
break
|
||||
|
||||
assert name_with_id is not None, "No $ID found in NAME tokens"
|
||||
all_src_state_keys = context.get_all_src_state_keys()
|
||||
id_matcher = IDMatcher(
|
||||
all_src_state_keys, EXTRA_SUFFIX, allowed_placeholders
|
||||
)
|
||||
valid_id_combos = id_matcher.find_matches(name_with_id)
|
||||
valid_keys = list(valid_id_combos.keys())
|
||||
IDENTIFIER_tokens = []
|
||||
for token in tokens:
|
||||
if token.value in GLOBAL_ATTRIBUTE_KEYWORDS:
|
||||
break
|
||||
if token.type == TokenType.IDENTIFIER:
|
||||
IDENTIFIER_tokens.append(token)
|
||||
|
||||
for token in IDENTIFIER_tokens:
|
||||
assert all(k in token.value for k in valid_keys), (
|
||||
f"The token: {token.value} must contain all of the following keys: {valid_keys}.When use the id macro all IDENTIFIER tokens must contain the same ID placeholders."
|
||||
)
|
||||
|
||||
def dict_cartesian_tuples(d: dict[str, list[int]]):
|
||||
keys = list(d.keys())
|
||||
value_lists = [d[k] for k in keys]
|
||||
for prod in product(*value_lists):
|
||||
yield tuple(zip(keys, prod))
|
||||
|
||||
results = []
|
||||
id_combs = dict_cartesian_tuples(valid_id_combos)
|
||||
id_combs = sorted(id_combs)
|
||||
for id_comb in id_combs:
|
||||
cur_statement = ""
|
||||
for tkn in tokens:
|
||||
tkn_val = tkn.value
|
||||
if tkn.type == TokenType.IDENTIFIER and any(
|
||||
ph in tkn.value for ph in allowed_placeholders
|
||||
):
|
||||
for id_tag, id_val in id_comb:
|
||||
tkn_val = tkn_val.replace("$" + id_tag, str(id_val))
|
||||
cur_statement += tkn_val
|
||||
else:
|
||||
cur_statement += tkn_val
|
||||
results.append(cur_statement)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# This macro processes variable mappings between source and destination states,
|
||||
# but it requires that all expansion macros (layer_id_macro, expert_id_macro,
|
||||
# star_macro, array_macro, etc.) have already been executed to expand template
|
||||
# variables into concrete variable names.
|
||||
@macro(name='get_var_mapping_chain_macro', priority=4)
|
||||
def get_var_mapping_chain_macro(tokens, expression, context):
|
||||
flag_left_var = True
|
||||
left_var_list = []
|
||||
right_var_list = []
|
||||
for tkn in tokens:
|
||||
if tkn.value in GLOBAL_ATTRIBUTE_KEYWORDS:
|
||||
break
|
||||
if tkn.type == TokenType.RARROW:
|
||||
flag_left_var = False
|
||||
if tkn.type == TokenType.IDENTIFIER:
|
||||
extra_suffix_removed_value = tkn.value
|
||||
for sfx in EXTRA_SUFFIX:
|
||||
extra_suffix_removed_value = (
|
||||
extra_suffix_removed_value.removesuffix(sfx)
|
||||
)
|
||||
if flag_left_var:
|
||||
left_var_list.append(extra_suffix_removed_value)
|
||||
else:
|
||||
right_var_list.append(extra_suffix_removed_value)
|
||||
assert len(left_var_list) == 1 or len(right_var_list) == 1, (
|
||||
"Left or right variable must have the only one element,the aoa_statements not support 'multiple var -> multiple var' pattern."
|
||||
)
|
||||
if len(left_var_list) == 1:
|
||||
context.left_var_to_right_var_mapping[left_var_list[0]] = right_var_list
|
||||
for right_var in right_var_list:
|
||||
context.right_var_from_left_var_mapping[right_var] = left_var_list
|
||||
else:
|
||||
context.right_var_from_left_var_mapping[right_var_list[0]] = (
|
||||
left_var_list
|
||||
)
|
||||
for left_var in left_var_list:
|
||||
context.left_var_to_right_var_mapping[left_var] = right_var_list
|
||||
return expression
|
||||
@@ -0,0 +1,142 @@
|
||||
# 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.
|
||||
|
||||
from .lexer import Token, TokenType
|
||||
|
||||
|
||||
class Statement:
|
||||
def __init__(self, left_vars, right_vars, attrs):
|
||||
self.left_vars = left_vars # List[Var]
|
||||
self.right_vars = right_vars # List[Var]
|
||||
self.attrs = attrs # List[Attribute]
|
||||
|
||||
def __repr__(self):
|
||||
return f"Statement({self.left_vars} -> {self.right_vars}, attrs={self.attrs})"
|
||||
|
||||
|
||||
class Var:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
|
||||
def __repr__(self):
|
||||
return self.name
|
||||
|
||||
|
||||
class Attribute:
|
||||
def __init__(self, key, value):
|
||||
self.key = key
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.key}={self.value!r}"
|
||||
|
||||
|
||||
class Parser:
|
||||
"""
|
||||
AOA Grammar
|
||||
PROGRAM ::= { STATEMENT }
|
||||
|
||||
STATEMENT ::= VAR_LIST '->' VAR ',' ATTR_LIST // meige
|
||||
| VAR '->' VAR_LIST ',' ATTR_LIST // split
|
||||
| VAR '->' VAR ',' ATTR_LIST // single variable mapping + attributes
|
||||
| VAR '->' VAR // single variable mapping, rename
|
||||
|
||||
VAR_LIST ::= VAR { ',' VAR }
|
||||
VAR ::= IDENTIFIER
|
||||
ATTR_LIST ::= ATTRIBUTE { ',' ATTRIBUTE }
|
||||
ATTRIBUTE ::= IDENTIFIER '=' VALUE
|
||||
VALUE ::= NUMBER | STRING
|
||||
"""
|
||||
|
||||
def __init__(self, tokens):
|
||||
self.tokens = tokens
|
||||
self.pos = 0
|
||||
|
||||
def at_end(self):
|
||||
return self.peek().type == TokenType.EOF
|
||||
|
||||
def peek(self, offset=0):
|
||||
if self.pos + offset >= len(self.tokens):
|
||||
return Token(TokenType.EOF, '')
|
||||
return self.tokens[self.pos + offset]
|
||||
|
||||
def consume(self, expected_type=None):
|
||||
tok = self.peek()
|
||||
if expected_type and tok.type != expected_type:
|
||||
raise SyntaxError(
|
||||
f'Expected {expected_type}, got {tok.type} at pos {self.pos}'
|
||||
)
|
||||
self.pos += 1
|
||||
return tok
|
||||
|
||||
def expect(self, expected_type):
|
||||
return self.consume(expected_type)
|
||||
|
||||
def skip_newlines(self):
|
||||
while self.peek().type == TokenType.NEWLINE:
|
||||
self.consume()
|
||||
|
||||
def parse_program(self):
|
||||
stmts = []
|
||||
self.skip_newlines()
|
||||
while not self.at_end():
|
||||
stmt = self.parse_statement()
|
||||
stmts.append(stmt)
|
||||
self.skip_newlines()
|
||||
return stmts
|
||||
|
||||
def parse_statement(self):
|
||||
left_vars = [self.parse_var()]
|
||||
while self.peek().type == TokenType.COMMA:
|
||||
self.consume(TokenType.COMMA)
|
||||
left_vars.append(self.parse_var())
|
||||
self.expect(TokenType.RARROW)
|
||||
right_vars = [self.parse_var()]
|
||||
while self.peek().type == TokenType.COMMA:
|
||||
# Lookahead for attribute: IDENT '=' after COMMA means attribute starts
|
||||
if (
|
||||
self.peek(1).type == TokenType.IDENTIFIER
|
||||
and self.peek(2).type == TokenType.EQUAL
|
||||
):
|
||||
break
|
||||
self.consume(TokenType.COMMA)
|
||||
right_vars.append(self.parse_var())
|
||||
attrs = []
|
||||
if self.peek().type == TokenType.COMMA:
|
||||
self.consume(TokenType.COMMA)
|
||||
attrs = self.parse_attr_list()
|
||||
return Statement(left_vars, right_vars, attrs)
|
||||
|
||||
def parse_var(self):
|
||||
name = self.expect(TokenType.IDENTIFIER).value
|
||||
return Var(name)
|
||||
|
||||
def parse_attr_list(self):
|
||||
attrs = [self.parse_attribute()]
|
||||
while self.peek().type == TokenType.COMMA:
|
||||
self.consume(TokenType.COMMA)
|
||||
attrs.append(self.parse_attribute())
|
||||
return attrs
|
||||
|
||||
def parse_attribute(self):
|
||||
key = self.expect(TokenType.IDENTIFIER).value
|
||||
self.expect(TokenType.EQUAL)
|
||||
val_tok = self.consume()
|
||||
if val_tok.type == TokenType.NUMBER:
|
||||
val = int(val_tok.value)
|
||||
elif val_tok.type == TokenType.STRING:
|
||||
val = val_tok.value.strip('"').strip("'")
|
||||
else:
|
||||
raise SyntaxError(f'Unexpected value: {val_tok}')
|
||||
return Attribute(key, val)
|
||||
@@ -0,0 +1,133 @@
|
||||
# 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class AOATraceback:
|
||||
"""
|
||||
When error occurs, print the chain of "original aoa_statement -> ... -> current expression".
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.records: list[dict] = []
|
||||
self.last_error_chain: list[str] = []
|
||||
self.last_error_message: str = ""
|
||||
self.last_error_stage: str = ""
|
||||
self.last_error_type: str = ""
|
||||
self.parent_map: dict[str, str | None] = {}
|
||||
self.child_macro_map: dict[str, str] = {}
|
||||
|
||||
def register_roots(self, expressions: list[str]) -> None:
|
||||
"""Register the original aoa_statements as the root nodes of the chain."""
|
||||
for expr in expressions:
|
||||
self.parent_map.setdefault(expr, None)
|
||||
|
||||
def record_children(
|
||||
self, parent: str, children: list[str], macro_name: str | None = None
|
||||
) -> None:
|
||||
"""Record the children expressions obtained by the parent expression, and mark the macro name used."""
|
||||
macro = macro_name or "Expanded"
|
||||
for child in children:
|
||||
if child == parent:
|
||||
continue
|
||||
self.parent_map[child] = parent
|
||||
self.child_macro_map[child] = macro
|
||||
|
||||
def build_chain(self, expr: str) -> list[str]:
|
||||
"""Build the chain from the root to expr by tracing back from the current expression."""
|
||||
chain: list[str] = []
|
||||
visited = set()
|
||||
cur = expr
|
||||
while cur is not None and cur not in visited:
|
||||
chain.append(cur)
|
||||
visited.add(cur)
|
||||
cur = self.parent_map.get(cur)
|
||||
chain.reverse()
|
||||
return chain
|
||||
|
||||
def add_error(
|
||||
self,
|
||||
error_message: str,
|
||||
stage: str,
|
||||
chain: list[str],
|
||||
error_type: str = "",
|
||||
) -> None:
|
||||
"""Record the error chain and information."""
|
||||
self.last_error_chain = chain
|
||||
self.last_error_message = error_message
|
||||
self.last_error_stage = stage
|
||||
self.last_error_type = error_type or ""
|
||||
self.records.append(
|
||||
{
|
||||
"type": "error",
|
||||
"stage": stage,
|
||||
"message": error_message,
|
||||
"error_type": self.last_error_type,
|
||||
"chain": chain,
|
||||
}
|
||||
)
|
||||
|
||||
def format_traceback(self) -> str:
|
||||
lines: list[str] = []
|
||||
header_text = " AOA Traceback (related chain) "
|
||||
header = f"===={header_text}===="
|
||||
footer = "=" * len(header)
|
||||
|
||||
if self.last_error_chain:
|
||||
lines.append(header)
|
||||
indent_unit = " "
|
||||
|
||||
lines.append("| Origin AOA Statement")
|
||||
origin_expr = self.last_error_chain[0].replace("\n", " ")
|
||||
lines.append(f"|-> {origin_expr}")
|
||||
|
||||
for level, expr in enumerate(self.last_error_chain[1:], start=1):
|
||||
indent = indent_unit * level
|
||||
single_line_expr = expr.replace("\n", " ")
|
||||
macro = self.child_macro_map.get(
|
||||
expr, self.last_error_stage or "Expanded"
|
||||
)
|
||||
lines.append(f"{indent}| {macro}")
|
||||
lines.append(f"{indent}|-> {single_line_expr}")
|
||||
|
||||
if self.last_error_message:
|
||||
err_title = self.last_error_type or "Error"
|
||||
stage_str = (
|
||||
f" [{self.last_error_stage}]"
|
||||
if self.last_error_stage
|
||||
else ""
|
||||
)
|
||||
err_level = len(self.last_error_chain)
|
||||
indent = indent_unit * err_level
|
||||
single_line_msg = self.last_error_message.replace("\n", " ")
|
||||
lines.append(f"{indent}| Error")
|
||||
lines.append(
|
||||
f"{indent}|-> ({err_title}{stage_str}) {single_line_msg}"
|
||||
)
|
||||
|
||||
lines.append(footer)
|
||||
else:
|
||||
lines.append(header)
|
||||
lines.append("(No trace records)")
|
||||
lines.append(footer)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def print(self, logger=None) -> None:
|
||||
text = self.format_traceback()
|
||||
if logger:
|
||||
logger.error(text)
|
||||
else:
|
||||
print(text)
|
||||
Reference in New Issue
Block a user