From 395289c8ed8863210ff5c34e967f4fc45b2c1fd9 Mon Sep 17 00:00:00 2001 From: Sonica-B Date: Sun, 25 Oct 2026 14:24:23 -0700 Subject: [PATCH] Emit intra_op_parallelism and JSON booleans from ds_nvme_tune ds_nvme_tune prints its recommendation as a ds_config "aio" section, but it spelled the parallelism key num_threads and wrote single_submit and overlap_events as the strings "true"/"false". get_aio_config reads intra_op_parallelism (num_threads was never a key it accepted), so the tuned value was dropped for the default of 1, and the flags reached the aio handle constructors as strings, which declare both parameters as const bool. Build the section from the aio config constants with real booleans, and replace the stale thread_count key in the config reference. Add a test that feeds the printed section back through get_aio_config. Signed-off-by: Sonica-B --- deepspeed/nvme/perf_generate_param.py | 14 ++++--- docs/_pages/config-json.md | 4 +- docs/_tutorials/deepnvme.md | 4 +- tests/unit/v1/nvme/test_ds_nvme_tune.py | 55 +++++++++++++++++++++++++ 4 files changed, 67 insertions(+), 10 deletions(-) create mode 100644 tests/unit/v1/nvme/test_ds_nvme_tune.py diff --git a/deepspeed/nvme/perf_generate_param.py b/deepspeed/nvme/perf_generate_param.py index 99be75c2d919..571d96aeeeaf 100644 --- a/deepspeed/nvme/perf_generate_param.py +++ b/deepspeed/nvme/perf_generate_param.py @@ -10,6 +10,8 @@ import json from .parse_nvme_stats import READ_SPEED, WRITE_SPEED, get_sorted_results from .perf_sweep_utils import BENCH_LOG_DIR, READ_LOG_DIR, WRITE_LOG_DIR +from deepspeed.runtime.swap_tensor.constants import AIO, AIO_BLOCK_SIZE, AIO_QUEUE_DEPTH, AIO_INTRA_OP_PARALLELISM, \ + AIO_SINGLE_SUBMIT, AIO_OVERLAP_EVENTS def parse_arguments(): @@ -41,11 +43,11 @@ def validate_args(args): def convert_to_param(key): assert len(key) == 6 return { - "single_submit": "true" if key[0] == "single" else "false", - "overlap_events": "true" if key[1] == "overlap" else "false", - "num_threads": int(key[5]), - "queue_depth": int(key[3]), - "block_size": int(key[4]) + AIO_SINGLE_SUBMIT: key[0] == "single", + AIO_OVERLAP_EVENTS: key[1] == "overlap", + AIO_INTRA_OP_PARALLELISM: int(key[5]), + AIO_QUEUE_DEPTH: int(key[3]), + AIO_BLOCK_SIZE: int(key[4]) } @@ -72,7 +74,7 @@ def generate_aio_param(read_log_dir, write_log_dir): optimal_perf = value optimal_key = key - aio_param = {"aio": convert_to_param(optimal_key)} + aio_param = {AIO: convert_to_param(optimal_key)} read_perf_keys = {key[1:]: key for key in read_results.keys()} write_perf_keys = {key[1:]: key for key in write_results.keys()} diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 9d310fd90a17..7898d66803a3 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -710,7 +710,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s "aio": { "block_size": 1048576, "queue_depth": 8, - "thread_count": 1, + "intra_op_parallelism": 1, "single_submit": false, "overlap_events": true } @@ -727,7 +727,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s | ---------------- | ------- | | I/O queue depth. | 8 | -***thread_count***: [integer] +***intra_op_parallelism***: [integer] | Description | Default | | ------------------------------------------------------------------------- | ------- | diff --git a/docs/_tutorials/deepnvme.md b/docs/_tutorials/deepnvme.md index 890dc0add791..8a5dabaaf160 100644 --- a/docs/_tutorials/deepnvme.md +++ b/docs/_tutorials/deepnvme.md @@ -232,8 +232,8 @@ Running DeepNVMe performance tuning on ['/local_nvme/'] Best performance (GB/sec): read = 3.69, write = 3.18 { "aio": { - "single_submit": "false", - "overlap_events": "true", + "single_submit": false, + "overlap_events": true, "intra_op_parallelism": 8, "queue_depth": 32, "block_size": 1048576 diff --git a/tests/unit/v1/nvme/test_ds_nvme_tune.py b/tests/unit/v1/nvme/test_ds_nvme_tune.py new file mode 100644 index 000000000000..89215a83c788 --- /dev/null +++ b/tests/unit/v1/nvme/test_ds_nvme_tune.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import json +import os + +import pytest + +from deepspeed.nvme.perf_generate_param import generate_main +from deepspeed.nvme.perf_sweep_utils import READ_LOG_DIR, WRITE_LOG_DIR +from deepspeed.runtime.swap_tensor.aio_config import get_aio_config + +# Log names as written by perf_run_sweep.get_log_file: ___ftd_d_bs_p +FAST_SWEEP_POINT = "single_overlap_ftd1_d32_bs1M_p8" +SLOW_SWEEP_POINT = "block_sequential_ftd1_d128_bs8M_p2" + + +def _write_sweep_logs(log_dir, speeds): + for sub_dir, op, metric in ((READ_LOG_DIR, "read", "E2E Read Speed"), (WRITE_LOG_DIR, "write", "E2E Write Speed")): + os.makedirs(os.path.join(log_dir, sub_dir)) + for sweep_point, speed in speeds.items(): + with open(os.path.join(log_dir, sub_dir, f"{op}_{sweep_point}.txt"), "w") as f: + f.write(f"{metric} = {speed} GB/sec\n") + + +# ds_nvme_tune prints its recommendation as a ds_config "aio" section, so it must survive the config parser +# unchanged: every key spelled as get_aio_config reads it, and the flags as JSON booleans. +@pytest.mark.parametrize("winner, expected", [ + (FAST_SWEEP_POINT, { + "block_size": 1048576, + "queue_depth": 32, + "intra_op_parallelism": 8, + "single_submit": True, + "overlap_events": True + }), + (SLOW_SWEEP_POINT, { + "block_size": 8388608, + "queue_depth": 128, + "intra_op_parallelism": 2, + "single_submit": False, + "overlap_events": False + }), +]) +def test_recommended_aio_section_round_trips_through_config(tmpdir, capsys, winner, expected): + speeds = {FAST_SWEEP_POINT: 1.0, SLOW_SWEEP_POINT: 1.0} + speeds[winner] = 5.0 + _write_sweep_logs(str(tmpdir), speeds) + + generate_main(str(tmpdir)) + printed = capsys.readouterr().out + recommended = json.loads(printed[printed.index("{"):]) + + aio_config = get_aio_config(recommended) + assert {key: aio_config[key] for key in expected} == expected