gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add topology-aware qwen27b 0-8k tuning

Browse Source
This commit is contained in:
Gahow Wang
2026-04-10 17:41:54 +08:00
parent b960607d8f
commit 8d0777e5e2
4 changed files with 144 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/aituner/llm.py
View File

@@ -124,7 +124,15 @@ def _enumerate_parallel_candidates(study: StudySpec) -> list[dict[str, int | boo
for tp in sorted(set(tp_values)):
for dp in sorted(set(dp_values)):
tp_dp_product = tp * dp
if tp_dp_product > study.hardware.gpu_count:
if (
constraints.allowed_tp_dp_products
and tp_dp_product not in constraints.allowed_tp_dp_products
):
continue
if (
not constraints.allowed_tp_dp_products
and tp_dp_product > study.hardware.gpu_count
):
continue
if (
constraints.require_tp_dp_product_equals_gpu_count
@@ -342,7 +350,19 @@ def validate_proposal(proposal: Proposal, study: StudySpec) -> Proposal:
f"expert-parallel-size={effective_ep} not in {constraints.allowed_expert_parallel_sizes}."
)
tp_dp_product = topology["tp_dp_product"]
if tp_dp_product > study.hardware.gpu_count:
if (
constraints.allowed_tp_dp_products
and tp_dp_product not in constraints.allowed_tp_dp_products
):
raise SpecError(
"Proposal violates topology constraints: "
f"tensor-parallel-size * data-parallel-size = {tp_dp_product} not in "
f"{constraints.allowed_tp_dp_products}."
)
if (
not constraints.allowed_tp_dp_products
and tp_dp_product > study.hardware.gpu_count
):
raise SpecError(
"Proposal violates topology constraints: "
f"tensor-parallel-size * data-parallel-size = {tp_dp_product} exceeds "

5
src/aituner/spec.py
View File

@@ -225,6 +225,7 @@ class TopologyConstraintSpec:
require_ep_size_divides_tp_dp_product: bool = False
require_enable_expert_parallel_when_ep_gt_one: bool = True
validate_cuda_graph_sizes_divisible_by_tp_when_tp_ep_reduce_scatter: bool = True
allowed_tp_dp_products: list[int] = field(default_factory=list)
allowed_tensor_parallel_sizes: list[int] = field(default_factory=list)
allowed_data_parallel_sizes: list[int] = field(default_factory=list)
allowed_expert_parallel_sizes: list[int] = field(default_factory=list)
@@ -255,6 +256,10 @@ class TopologyConstraintSpec:
),
context="engine.topology_constraints.validate_cuda_graph_sizes_divisible_by_tp_when_tp_ep_reduce_scatter",
),
allowed_tp_dp_products=_coerce_int_list(
data.get("allowed_tp_dp_products"),
context="engine.topology_constraints.allowed_tp_dp_products",
),
allowed_tensor_parallel_sizes=_coerce_int_list(
data.get("allowed_tensor_parallel_sizes"),
context="engine.topology_constraints.allowed_tensor_parallel_sizes",