Files
2026-07-13 13:18:33 +08:00

342 lines
13 KiB
C++

// Copyright (c) Microsoft Corporation.
// SPDX-License-Identifier: Apache-2.0
// DeepSpeed Team
/*
Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
*/
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <libaio.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iostream>
#include <limits>
#include <numeric>
#include <string>
#include <vector>
#include "deepspeed_aio_common.h"
using namespace std;
using namespace std::chrono;
#define DEBUG_DS_AIO_PERF 0
#define DEBUG_DS_AIO_SUBMIT_PERF 0
static const std::string c_library_name = "deepspeed_aio";
static void _report_aio_statistics(const char* tag,
const std::vector<std::chrono::duration<double>>& latencies)
__attribute__((unused));
static void _report_aio_statistics(const char* tag,
const std::vector<std::chrono::duration<double>>& latencies)
{
std::vector<double> lat_usec;
for (auto& lat : latencies) { lat_usec.push_back(lat.count() * 1e6); }
const auto min_lat = *(std::min_element(lat_usec.begin(), lat_usec.end()));
const auto max_lat = *(std::max_element(lat_usec.begin(), lat_usec.end()));
const auto avg_lat = std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size();
std::cout << c_library_name << ": latency statistics(usec) " << tag
<< " min/max/avg = " << min_lat << " " << max_lat << " " << avg_lat << std::endl;
}
static void _get_aio_latencies(std::vector<std::chrono::duration<double>>& raw_latencies,
struct deepspeed_aio_latency_t& summary_latencies)
{
std::vector<double> lat_usec;
for (auto& lat : raw_latencies) { lat_usec.push_back(lat.count() * 1e6); }
summary_latencies._min_usec = *(std::min_element(lat_usec.begin(), lat_usec.end()));
summary_latencies._max_usec = *(std::max_element(lat_usec.begin(), lat_usec.end()));
summary_latencies._avg_usec =
std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size();
}
static void _do_io_submit_singles(const int64_t n_iocbs,
const int64_t iocb_index,
std::unique_ptr<aio_context>& aio_ctxt,
std::vector<std::chrono::duration<double>>& submit_times)
{
for (auto i = 0; i < n_iocbs; ++i) {
const auto st = std::chrono::high_resolution_clock::now();
const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, 1, aio_ctxt->_iocbs.data() + i);
submit_times.push_back(std::chrono::high_resolution_clock::now() - st);
#if DEBUG_DS_AIO_SUBMIT_PERF
printf("submit(usec) %f io_index=%lld buf=%p len=%lu off=%llu \n",
submit_times.back().count() * 1e6,
iocb_index,
aio_ctxt->_iocbs[i]->u.c.buf,
aio_ctxt->_iocbs[i]->u.c.nbytes,
aio_ctxt->_iocbs[i]->u.c.offset);
#endif
assert(submit_ret > 0);
}
}
static void _do_io_submit_block(const int64_t n_iocbs,
const int64_t iocb_index,
std::unique_ptr<aio_context>& aio_ctxt,
std::vector<std::chrono::duration<double>>& submit_times)
{
const auto st = std::chrono::high_resolution_clock::now();
const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, n_iocbs, aio_ctxt->_iocbs.data());
submit_times.push_back(std::chrono::high_resolution_clock::now() - st);
#if DEBUG_DS_AIO_SUBMIT_PERF
printf("submit(usec) %f io_index=%lld nr=%lld buf=%p len=%lu off=%llu \n",
submit_times.back().count() * 1e6,
iocb_index,
n_iocbs,
aio_ctxt->_iocbs[0]->u.c.buf,
aio_ctxt->_iocbs[0]->u.c.nbytes,
aio_ctxt->_iocbs[0]->u.c.offset);
#endif
assert(submit_ret > 0);
}
static int _do_io_complete(const int64_t min_completes,
const int64_t max_completes,
std::unique_ptr<aio_context>& aio_ctxt,
std::vector<std::chrono::duration<double>>& reap_times)
{
const auto start_time = std::chrono::high_resolution_clock::now();
int64_t n_completes = io_pgetevents(aio_ctxt->_io_ctxt,
min_completes,
max_completes,
aio_ctxt->_io_events.data(),
nullptr,
nullptr);
reap_times.push_back(std::chrono::high_resolution_clock::now() - start_time);
assert(n_completes >= min_completes);
return n_completes;
}
void do_aio_operation_sequential(const bool read_op,
std::unique_ptr<aio_context>& aio_ctxt,
std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
deepspeed_aio_config_t* config,
deepspeed_aio_perf_t* perf)
{
struct io_prep_context prep_ctxt(read_op, xfer_ctxt, aio_ctxt->_block_size, &aio_ctxt->_iocbs);
const auto num_io_blocks = static_cast<int64_t>(
ceil(static_cast<double>(xfer_ctxt->_num_bytes) / aio_ctxt->_block_size));
#if DEBUG_DS_AIO_PERF
const auto io_op_name = std::string(read_op ? "read" : "write");
std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes
<< " bytes with " << num_io_blocks << " io blocks" << std::endl;
#endif
std::vector<std::chrono::duration<double>> submit_times;
std::vector<std::chrono::duration<double>> reap_times;
const auto max_queue_bytes =
static_cast<int64_t>(aio_ctxt->_queue_depth * aio_ctxt->_block_size);
auto start = std::chrono::high_resolution_clock::now();
for (int64_t iocb_index = 0; iocb_index < num_io_blocks; iocb_index += aio_ctxt->_queue_depth) {
const auto start_offset = iocb_index * aio_ctxt->_block_size;
const auto start_buffer = (char*)xfer_ctxt->_mem_buffer + start_offset;
const auto n_iocbs =
min(static_cast<int64_t>(aio_ctxt->_queue_depth), (num_io_blocks - iocb_index));
const auto num_bytes = min(max_queue_bytes, (xfer_ctxt->_num_bytes - start_offset));
prep_ctxt.prep_iocbs(n_iocbs, num_bytes, start_buffer, start_offset);
if (config->_single_submit) {
_do_io_submit_singles(n_iocbs, iocb_index, aio_ctxt, submit_times);
} else {
_do_io_submit_block(n_iocbs, iocb_index, aio_ctxt, submit_times);
}
_do_io_complete(n_iocbs, n_iocbs, aio_ctxt, reap_times);
}
const std::chrono::duration<double> elapsed = std::chrono::high_resolution_clock::now() - start;
if (perf) {
_get_aio_latencies(submit_times, perf->_submit);
_get_aio_latencies(reap_times, perf->_complete);
perf->_e2e_usec = elapsed.count() * 1e6;
perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9);
}
#if DEBUG_DS_AIO_PERF
_report_aio_statistics("submit", submit_times);
_report_aio_statistics("complete", reap_times);
#endif
#if DEBUG_DS_AIO_PERF
std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6
<< " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl;
#endif
#if DEBUG_DS_AIO_PERF
std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes
<< " bytes " << std::endl;
#endif
}
void do_aio_operation_overlap(const bool read_op,
std::unique_ptr<aio_context>& aio_ctxt,
std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
deepspeed_aio_config_t* config,
deepspeed_aio_perf_t* perf)
{
struct io_prep_generator io_gen(read_op, xfer_ctxt, aio_ctxt->_block_size);
#if DEBUG_DS_AIO_PERF
const auto io_op_name = std::string(read_op ? "read" : "write");
std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes
<< " bytes with " << io_gen._num_io_blocks << " io blocks" << std::endl;
#endif
std::vector<std::chrono::duration<double>> submit_times;
std::vector<std::chrono::duration<double>> reap_times;
auto request_iocbs = aio_ctxt->_queue_depth;
auto n_pending_iocbs = 0;
const auto min_completes = 1;
auto start = std::chrono::high_resolution_clock::now();
while (true) {
const auto n_iocbs = io_gen.prep_iocbs(request_iocbs - n_pending_iocbs, &aio_ctxt->_iocbs);
if (n_iocbs > 0) {
if (config->_single_submit) {
_do_io_submit_singles(
n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times);
} else {
_do_io_submit_block(
n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times);
}
}
n_pending_iocbs += n_iocbs;
assert(n_pending_iocbs <= aio_ctxt->_queue_depth);
if (n_pending_iocbs == 0) { break; }
const auto n_complete =
_do_io_complete(min_completes, n_pending_iocbs, aio_ctxt, reap_times);
n_pending_iocbs -= n_complete;
}
const std::chrono::duration<double> elapsed = std::chrono::high_resolution_clock::now() - start;
if (perf) {
_get_aio_latencies(submit_times, perf->_submit);
_get_aio_latencies(reap_times, perf->_complete);
perf->_e2e_usec = elapsed.count() * 1e6;
perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9);
}
#if DEBUG_DS_AIO_PERF
_report_aio_statistics("submit", submit_times);
_report_aio_statistics("complete", reap_times);
#endif
#if DEBUG_DS_AIO_PERF
std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6
<< " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl;
#endif
#if DEBUG_DS_AIO_PERF
std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes
<< " bytes " << std::endl;
#endif
}
void report_file_error(const char* filename, const std::string file_op, const int error_code)
{
std::string err_msg = file_op + std::string(" failed on ") + std::string(filename) +
" error = " + std::to_string(error_code);
std::cerr << c_library_name << ": " << err_msg << std::endl;
}
int open_file(const char* filename, const bool read_op)
{
const int flags = read_op ? (O_RDONLY | O_DIRECT) : (O_WRONLY | O_CREAT | O_DIRECT);
#if defined(__ENABLE_CANN__)
int* flags_ptr = (int*)&flags;
*flags_ptr = read_op ? (O_RDONLY) : (O_WRONLY | O_CREAT);
#endif
const int mode = 0600;
const auto fd = open(filename, flags, mode);
if (fd == -1) {
const auto error_code = errno;
const auto error_msg = read_op ? " open for read " : " open for write ";
report_file_error(filename, error_msg, error_code);
return -1;
}
return fd;
}
int regular_read(const char* filename, std::vector<char>& buffer)
{
const auto fd = open(filename, O_RDONLY, 0600);
assert(fd != -1);
struct stat fs;
const auto result = fstat(fd, &fs);
assert(result != -1);
int64_t num_bytes = fs.st_size;
buffer.resize(num_bytes);
int64_t read_bytes = 0;
auto r = 0;
do {
const auto buffer_ptr = buffer.data() + read_bytes;
const auto bytes_to_read = num_bytes - read_bytes;
r = read(fd, buffer_ptr, bytes_to_read);
read_bytes += r;
} while (r > 0);
if (read_bytes != num_bytes) {
std::cerr << "read error " << " read_bytes (read) = " << read_bytes
<< " num_bytes (fstat) = " << num_bytes << std::endl;
}
assert(read_bytes == num_bytes);
close(fd);
return 0;
}
static bool _validate_buffer(const char* filename, void* aio_buffer, const int64_t num_bytes)
{
std::vector<char> regular_buffer;
const auto reg_ret = regular_read(filename, regular_buffer);
assert(0 == reg_ret);
std::cout << "regular read of " << filename << " returned " << regular_buffer.size() << " bytes"
<< std::endl;
if (static_cast<int64_t>(regular_buffer.size()) != num_bytes) { return false; }
return (0 == memcmp(aio_buffer, regular_buffer.data(), regular_buffer.size()));
}
bool validate_aio_operation(const bool read_op,
const char* filename,
void* aio_buffer,
const int64_t num_bytes)
{
const auto msg_suffix = std::string("deepspeed_aio_") +
std::string(read_op ? "read()" : "write()") +
std::string("using read()");
if (false == _validate_buffer(filename, aio_buffer, num_bytes)) {
std::cout << "Fail: correctness of " << msg_suffix << std::endl;
return false;
}
std::cout << "Pass: correctness of " << msg_suffix << std::endl;
return true;
}