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

Constrain harness topology by visible GPUs

Browse Source
This commit is contained in:
Gahow Wang
2026-05-13 01:25:31 +08:00
parent fb6d74a18c
commit 5c2958e6c1
3 changed files with 143 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/aituner/harness.py
View File

@@ -1210,17 +1210,18 @@ def _legal_topology_points(
) -> list[dict[str, Any]]:
constraints = study.engine.topology_constraints
tunable = set(study.engine.tunable_flags)
effective_gpu_count = _effective_gpu_count(study)
if constraints is not None and constraints.allowed_tensor_parallel_sizes:
tp_values = sorted(set(constraints.allowed_tensor_parallel_sizes))
elif "tensor-parallel-size" in tunable:
tp_values = [value for value in [1, 2, 4, 8] if value <= study.hardware.gpu_count]
tp_values = [value for value in [1, 2, 4, 8] if value <= effective_gpu_count]
else:
tp_values = [current_tp]
if constraints is not None and constraints.allowed_data_parallel_sizes:
dp_values = sorted(set(constraints.allowed_data_parallel_sizes))
elif "data-parallel-size" in tunable:
dp_values = [value for value in [1, 2, 4, 8] if value <= study.hardware.gpu_count]
dp_values = [value for value in [1, 2, 4, 8] if value <= effective_gpu_count]
else:
dp_values = [current_dp]
@@ -1243,14 +1244,16 @@ def _legal_topology_points(
continue
if (
constraints.require_tp_dp_product_equals_gpu_count
and tp_dp_product != study.hardware.gpu_count
and tp_dp_product != effective_gpu_count
):
continue
elif tp_dp_product > study.hardware.gpu_count:
elif tp_dp_product > effective_gpu_count:
continue
if constraints is not None and not constraints.allowed_tp_dp_products:
if tp_dp_product > study.hardware.gpu_count:
if tp_dp_product > effective_gpu_count:
continue
if tp_dp_product > effective_gpu_count:
continue
for ep in ep_values:
enable_ep = current_enable_ep or ep > 1
if constraints is not None:
@@ -1330,6 +1333,18 @@ def _normalized_topology_flags(flags: dict[str, Any]) -> dict[str, Any]:
}
def _effective_gpu_count(study: StudySpec) -> int:
visible = str(study.engine.base_envs.get("CUDA_VISIBLE_DEVICES") or "").strip()
if not visible:
return study.hardware.gpu_count
if visible.lower() in {"none", "void", "-1"}:
return 0
devices = [item.strip() for item in visible.split(",") if item.strip()]
if not devices:
return study.hardware.gpu_count
return min(study.hardware.gpu_count, len(devices))
def _score_topology_candidate(
top_bottleneck: str,
bottleneck_hypotheses: list[dict[str, Any]],
@@ -1621,6 +1636,7 @@ def _round_up_to_multiple(value: int, multiple: int) -> int:
def _next_allowed_tp(study: StudySpec, *, current_tp: int, current_dp: int) -> int | None:
constraints = study.engine.topology_constraints
effective_gpu_count = _effective_gpu_count(study)
if constraints is not None and constraints.allowed_tensor_parallel_sizes:
candidates = sorted({int(item) for item in constraints.allowed_tensor_parallel_sizes})
else:
@@ -1629,7 +1645,7 @@ def _next_allowed_tp(study: StudySpec, *, current_tp: int, current_dp: int) -> i
if candidate <= current_tp:
continue
tp_dp_product = candidate * current_dp
if tp_dp_product > study.hardware.gpu_count:
if tp_dp_product > effective_gpu_count:
continue
if constraints is not None:
if (
@@ -1639,7 +1655,7 @@ def _next_allowed_tp(study: StudySpec, *, current_tp: int, current_dp: int) -> i
continue
if (
constraints.require_tp_dp_product_equals_gpu_count
and tp_dp_product != study.hardware.gpu_count
and tp_dp_product != effective_gpu_count
):
continue
return candidate