Skip to content

测试基建：⚡ 优化 GPU Mode 显存治理 #684

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cangtianhuang merged 9 commits into from
Jul 20, 2026
Merged

测试基建：⚡ 优化 GPU Mode 显存治理 #684

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
7877ab3
⚡️ Fix wrong classify, use report_compare_error
cangtianhuang Jul 16, 2026
ffa1647
🐛 Fix numpy error
cangtianhuang Jul 16, 2026
62a6360
✨ Impl TE moe
cangtianhuang Jul 16, 2026
68d7ceb
⚡️ Update gpu mode
cangtianhuang Jul 17, 2026
821f497
Merge remote-tracking branch 'upstream/main' into opt
cangtianhuang Jul 17, 2026
485b225
⚡️ Update gpu mode
cangtianhuang Jul 17, 2026
1cf7af7
⚡️ Fix accuracy_stable oom
cangtianhuang Jul 20, 2026
b2dcab6
⚡️ Imp gpu mode
cangtianhuang Jul 20, 2026
2a5b5b1
Merge branch 'main' into opt
cangtianhuang Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions engineV2-README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,9 +72,8 @@
| `--accuracy_stable` | bool | 启用稳定性测试（默认 False） |
| `--paddle_custom_device` | bool | 运行Custom Device或者XPU的API与CPU的精度对比（默认 False） |
| `--num_gpus` | int | 使用的 GPU 数量（默认 -1，-1 动态最大） |
| `--num_workers_per_gpu` | int | 每 GPU 的 worker 进程数（默认 1，-1 动态最大） |
| `--num_workers_per_gpu` | int | 每 GPU 的 worker 进程数（默认 1；gpu_mode 下 -1 表示每 GPU 1 个 worker） |
| `--gpu_ids` | str | 使用的 GPU 序号，以逗号分隔或横线范围（默认 ""，"-1" 动态最大） |
| `--required_memory` | float | 每 worker 进程预估使用显存 GB（默认 10.0） |
| `--test_amp` | bool | 启用自动混合精度测试（默认 False） |
| `--test_cpu` | bool | 启用 Paddle CPU 模式测试（默认 False） |
| `--use_cached_numpy` | bool | 启用 Numpy 缓存（默认 False） |
Expand Down
211 changes: 105 additions & 106 deletions engineV2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
)

from tester.api_config.log_writer import *
from tester.runtime_config import TestRuntimeConfig, runtime_config_for_gpu

os.environ["FLAGS_use_system_allocator"] = "1"
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
Expand All @@ -58,6 +59,7 @@
"bitwise_alignment",
"exit_on_error",
"use_gpu_mode",
"runtime_config",
}

DEVICE_TYPE = None
Expand All @@ -71,6 +73,11 @@
FATAL_TORCH_EXIT_CODE = 97
MEMORY_WAIT_SECONDS = 10
MEMORY_WAIT_LOG_INTERVAL = 60
MEMORY_WAIT_MAX_SECONDS = 300


class GpuMemoryDeferred(Exception):
"""Raised when a GPU-mode case should wait for more free memory."""


def cleanup(pool):
Expand Down Expand Up @@ -436,11 +443,6 @@ def validate_gpu_options(options) -> tuple:
f"invalid --num_workers_per_gpu={options.num_workers_per_gpu}: "
"expected -1 or a positive integer"
)
if options.required_memory <= 0:
raise ValueError(
f"invalid --required_memory={options.required_memory:g}: "
"expected a positive number of GiB"
)
return tuple(gpu_ids)


Expand All @@ -457,6 +459,9 @@ def parse_bool(value):

def _prepare_single_config_gpu(options):
if options.test_cpu:
options.gpu_workers_per_gpu_map = {}
options.gpu_total_memory_map = {}
options.runtime_config = TestRuntimeConfig.from_options(options)
return None

gpu_ids = validate_gpu_options(options)
Expand All @@ -465,29 +470,29 @@ def _prepare_single_config_gpu(options):
f"single --api_config run supports exactly one GPU; got {len(gpu_ids)} GPUs: {gpu_ids}"
)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_ids[0])
try:
options.gpu_total_memory_map = {gpu_ids[0]: get_memory_info(gpu_ids[0])[0]}
except Exception:
options.gpu_total_memory_map = {}
options.gpu_workers_per_gpu_map = {gpu_ids[0]: 1}
options.runtime_config = TestRuntimeConfig.from_options(options)
options.runtime_config = runtime_config_for_gpu(options, gpu_ids[0])
return gpu_ids[0]


def check_gpu_memory(gpu_ids, num_workers_per_gpu, required_memory): # required_memory in GB
def check_gpu_memory(gpu_ids, num_workers_per_gpu):
assert isinstance(gpu_ids, tuple) and len(gpu_ids) > 0
available_gpus = []
max_workers_per_gpu = {}

for gpu_id in gpu_ids:
try:
total_memory, used_memory = get_memory_info(gpu_id)
free_memory = total_memory - used_memory
max_workers = int(free_memory // required_memory)
if max_workers >= 1:
available_gpus.append(gpu_id)
max_workers_per_gpu[gpu_id] = (
max_workers
if num_workers_per_gpu == -1
else min(max_workers, num_workers_per_gpu)
)
get_memory_info(gpu_id)
except pynvml.NVMLError as e:
print(f"[warn] Failed to check GPU {gpu_id}: {e!s}", flush=True)
continue
available_gpus.append(gpu_id)
max_workers_per_gpu[gpu_id] = 1 if num_workers_per_gpu == -1 else num_workers_per_gpu

return available_gpus, max_workers_per_gpu

Expand Down Expand Up @@ -570,24 +575,7 @@ def run_test_case(api_config_str, options):
flush=True,
)

last_memory_log_time = 0
while True:
total_memory, used_memory = get_memory_info(gpu_id)
free_memory = total_memory - used_memory

if free_memory >= options.required_memory:
break

now = time.time()
if now - last_memory_log_time >= MEMORY_WAIT_LOG_INTERVAL:
print(
f"{datetime.now()} device {gpu_id} Free: {free_memory:.1f} GB, "
f"Required: {options.required_memory:.1f} GB. ",
"Waiting for available memory...",
flush=True,
)
last_memory_log_time = now
time.sleep(MEMORY_WAIT_SECONDS)
runtime_config = runtime_config_for_gpu(options, gpu_id)

if options.show_runtime_status:
total_memory, used_memory_before = get_memory_info(gpu_id)
Expand All @@ -608,6 +596,7 @@ def run_test_case(api_config_str, options):

test_class = _select_test_class(options)
kwargs = {k: v for k, v in vars(options).items() if k in VALID_TEST_ARGS}
kwargs["runtime_config"] = runtime_config
case = test_class(api_config, **kwargs)
try:
case.test()
Expand Down Expand Up @@ -783,20 +772,14 @@ def main():
"--num_workers_per_gpu",
type=int,
default=1,
help="Workers per GPU. Use -1 to maximize workers based on free memory.",
help="Workers per GPU. In gpu_mode, -1 uses one worker per GPU.",
)
parser.add_argument(
"--gpu_ids",
type=str,
default="",
help="GPU IDs to use, e.g. '0', '0,2', '0-3'. Use '-1' for all GPUs.",
)
parser.add_argument(
"--required_memory",
type=float,
default=10.0,
help="Minimum free memory required per worker, in GiB.",
)
parser.add_argument(
"--test_cpu",
type=parse_bool,
Expand Down Expand Up @@ -974,12 +957,10 @@ def main():
options.use_cached_numpy = False
os.environ["USE_CACHED_NUMPY"] = str(options.use_cached_numpy)
os.environ["USE_GPU_MODE"] = str(options.use_gpu_mode)
os.environ["SKIP_GPU_CLEANUP"] = str(options.use_gpu_mode)
if options.use_gpu_mode:
print(
"[gpu_mode] enabled: GPU tensor generation, GPU compare, and allocator reuse are active.",
flush=True,
)
print("[gpu_mode] enabled: use GPU tensors and comparison.", flush=True)
elif options.use_cached_numpy:
print("[use_cached_numpy] enabled: reuse cached NumPy inputs.", flush=True)
if options.bitwise_alignment:
options.atol = 0.0
options.rtol = 0.0
Expand Down Expand Up @@ -1046,6 +1027,7 @@ def main():
bitwise_alignment=options.bitwise_alignment,
exit_on_error=options.exit_on_error,
use_gpu_mode=options.use_gpu_mode,
runtime_config=options.runtime_config,
)
else:
case = test_class(api_config, test_amp=options.test_amp)
Expand Down Expand Up @@ -1141,18 +1123,22 @@ def main():
print(all_case, "cases will be tested.", flush=True)
del api_config_count, dup_case, finish_case

# validate GPU memory
available_gpus, max_workers_per_gpu = check_gpu_memory(
gpu_ids, options.num_workers_per_gpu, options.required_memory
)
# validate GPU visibility and derive per-GPU worker counts
available_gpus, max_workers_per_gpu = check_gpu_memory(gpu_ids, options.num_workers_per_gpu)
if not available_gpus:
print(
f"No GPUs with sufficient memory available. Current memory constraint is {options.required_memory} GB.",
flush=True,
)
print("No usable GPUs available.", flush=True)
return

total_workers = sum(max_workers_per_gpu.values())
gpu_total_memory_map = {}
for gpu_id in available_gpus:
try:
gpu_total_memory_map[gpu_id] = get_memory_info(gpu_id)[0]
except Exception:
pass
options.gpu_workers_per_gpu_map = dict(max_workers_per_gpu)
options.gpu_total_memory_map = gpu_total_memory_map
options.runtime_config = TestRuntimeConfig.from_options(options)
print(
f"Using {len(available_gpus)} GPU(s) with max workers per GPU: {max_workers_per_gpu}. Total workers: {total_workers}.",
flush=True,
Expand Down Expand Up @@ -1192,7 +1178,8 @@ def cleanup_handler(*args):
for batch_start in range(0, len(api_configs), BATCH_SIZE):
batch = api_configs[batch_start : batch_start + BATCH_SIZE]
futures = {}
for config in batch:

def schedule_config(config):
# timeout = estimate_timeout(config)
timeout = options.timeout
future = pool.schedule(
Expand All @@ -1202,67 +1189,79 @@ def cleanup_handler(*args):
)
futures[future] = config

for future in as_completed(futures):
config = futures[future]
checkpoint_ready = True
try:
future.result()
if options.show_runtime_status or tested_case % 10000 == 0:
print(f"[info] Test case succeeded for {config}", flush=True)
except TimeoutError as err:
write_terminal_log("timeout", config)
print(
f"[timeout] {config}: {err}",
flush=True,
)
except ProcessExpired as err:
# CUDA, OOM, and Torch fatal errors may also expire the subprocess;
# classify them by dedicated terminal log types instead of falling through to crash.
if err.exitcode == FATAL_CUDA_EXIT_CODE:
write_terminal_log("paddle_cuda", config)
print(
f"[paddle_cuda] {config}: {err}",
flush=True,
)
elif err.exitcode == FATAL_OOM_EXIT_CODE:
write_terminal_log("oom", config)
print(
f"[oom] {config}: {err}",
flush=True,
)
elif err.exitcode == FATAL_TORCH_EXIT_CODE:
write_terminal_log("torch_error", config)
for config in batch:
schedule_config(config)

while futures:
done_futures = list(as_completed(futures))
for future in done_futures:
config = futures.pop(future)
checkpoint_ready = True
try:
future.result()
if options.show_runtime_status or tested_case % 10000 == 0:
print(f"[info] Test case succeeded for {config}", flush=True)
except TimeoutError as err:
write_terminal_log("timeout", config)
print(
f"[torch_error] {config}: {err}",
f"[timeout] {config}: {err}",
flush=True,
)
elif err.exitcode in (-signal.SIGKILL, -signal.SIGTERM):
except ProcessExpired as err:
# CUDA, OOM, and Torch fatal errors may also expire the subprocess;
# classify them by dedicated terminal log types instead of falling through to crash.
if err.exitcode == FATAL_CUDA_EXIT_CODE:
write_terminal_log("paddle_cuda", config)
print(
f"[paddle_cuda] {config}: {err}",
flush=True,
)
elif err.exitcode == FATAL_OOM_EXIT_CODE:
write_terminal_log("oom", config)
print(
f"[oom] {config}: {err}",
flush=True,
)
elif err.exitcode == FATAL_TORCH_EXIT_CODE:
write_terminal_log("torch_error", config)
print(
f"[torch_error] {config}: {err}",
flush=True,
)
elif err.exitcode in (-signal.SIGKILL, -signal.SIGTERM):
checkpoint_ready = False
print(
f"[warn] Worker was externally killed for {config} "
f"(exit={err.exitcode}); case will be retried on next run.",
flush=True,
)
else:
write_terminal_log("paddle_crash", config)
print(
f"[paddle_crash] {config}: {err}",
flush=True,
)
except GpuMemoryDeferred as err:
checkpoint_ready = False
print(
f"[warn] Worker was externally killed for {config} "
f"(exit={err.exitcode}); case will be retried on next run.",
f"[gpu_mode] Deferred {config}: {err}",
flush=True,
)
else:
write_terminal_log("paddle_crash", config)
schedule_config(config)
except Exception as err:
write_terminal_log("config_parse", config)
print(
f"[paddle_crash] {config}: {err}",
flush=True,
)
except Exception as err:
write_terminal_log("config_parse", config)
print(
f"[config_parse] {config}: {err}",
flush=True,
)
checkpoint_ready = False # checkpoint already written by write_terminal_log
if checkpoint_ready:
tested_case += 1
if options.show_runtime_status or tested_case % 10000 == 0:
print(
f"[{tested_case}/{all_case}] Testing {config}",
f"[config_parse] {config}: {err}",
flush=True,
)
checkpoint_ready = False
if checkpoint_ready:
tested_case += 1
if options.show_runtime_status or tested_case % 10000 == 0:
print(
f"[{tested_case}/{all_case}] Testing {config}",
flush=True,
)
aggregate_logs()
pool.close()
pool.join()
Expand Down
Loading
Loading