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