342 lines
13 KiB
C++
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;
|
|
}
|