diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py index e76e5d03cd..23c988c32a 100644 --- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py +++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py @@ -144,6 +144,9 @@ class RESParams: table_sizes: list[int] = field( default_factory=list ) # table sizes for the global rows the TBE holds + res_chunk_size: int = 500000 # max rows copied into one enqueued chunk + res_num_consumers: int = 8 # threads draining the stream queue + res_num_copy_threads: int = 4 # parallel chunk-copy threads per stream() call res_enabled_tables: list[str] = field( default_factory=list ) # table names that are enabled for RES (empty means all enabled) @@ -1563,6 +1566,9 @@ def __init__( # noqa C901 self.res_params.table_names, self.res_params.table_offsets, self.res_params.table_sizes, + self.res_params.res_chunk_size, + self.res_params.res_num_consumers, + self.res_params.res_num_copy_threads, ) self._register_res_enabled_feature_mask() logging.info( @@ -4580,7 +4586,7 @@ def raw_embedding_stream(self) -> None: # Lazy resize: runtime_meta shape/dtype is not known until # the first data arrives from the MC module. Must use UVM # (new_unified_tensor) because the C++ RawEmbeddingStreamer - # reads this buffer via raw CPU pointers in tensor_copy(). + # reads this buffer via raw CPU pointers in tensor_copy_chunk(). self.register_buffer( "res_runtime_meta", torch.ops.fbgemm.new_unified_tensor( diff --git a/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h b/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h index 9b0f7be6ee..ffcfea27ca 100644 --- a/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h +++ b/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h @@ -13,6 +13,7 @@ #endif #include +#include #ifdef FBGEMM_FBCODE namespace facebook::aiplatform::gmpp::experimental::training_ps { @@ -50,16 +51,24 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { int64_t res_server_port, std::vector table_names, std::vector table_offsets, - const std::vector& table_sizes); + const std::vector& table_sizes, + int64_t res_chunk_size = 500000, + // TODO(T282801601): 8 was an arbitrary high value picked during + // experimentation; too many consumer threads per TBE may be wasteful -- + // tune via experiments. + int64_t res_num_consumers = 8, + int64_t res_num_copy_threads = 4); ~RawEmbeddingStreamer() override; /// Stream out non-negative elements in and its paired embeddings /// from for the first elements in the tensor. - /// It spins up a thread that will copy all 3 tensors to CPU and inject them - /// into the background queue which will be picked up by another set of thread - /// pools for streaming out to the thrift server (co-located on same host - /// now). + /// It spins up a dispatcher thread that copies the 4 tensors (indices, + /// weights, and the optional identities / runtime_meta) to CPU and injects + /// them into the background queue, which is drained by a pool of consumer + /// threads that stream out to the thrift server (co-located on same host + /// now). The copy is split into <= res_chunk_size-row chunks across up to + /// res_num_copy_threads copy threads. /// /// This is used in cuda stream callback, which doesn't require to be /// serialized with other callbacks, thus a separate thread is used to @@ -86,8 +95,8 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { std::optional copy_done_flag = std::nullopt); /* - * Join the stream tensor copy thread, make sure the thread is properly - * finished before creating new. + * Join the pending dispatch (and the copy threads it spawned), making sure it + * is properly finished before creating new. */ void join_stream_tensor_copy_thread(); @@ -97,20 +106,11 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { const at::Tensor& weights, std::optional identities, std::optional runtime_meta); - /* - * Copy the indices, weights and count tensors and enqueue them for - * asynchronous stream. - */ - void copy_and_enqueue_stream_tensors( - const at::Tensor& indices, - const at::Tensor& weights, - std::optional identities, - std::optional runtime_meta, - const at::Tensor& count); /* - * FOR TESTING: Join the weight stream thread, make sure the thread is - * properly finished for destruction and testing. + * FOR TESTING ONLY: latches stop_ so the consumer threads exit, letting a + * test read a stable queue size. Not reversible -- the streamer stops + * consuming after this call. */ void join_weights_stream_thread(); // FOR TESTING: get queue size. @@ -128,22 +128,53 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { std::vector table_offsets_; at::Tensor table_sizes_; #ifdef FBGEMM_FBCODE - std::unique_ptr weights_stream_thread_; - folly::UMPSCQueue weights_to_stream_queue_; - std::unique_ptr stream_tensor_copy_thread_; + size_t res_chunk_size_; + size_t res_num_consumers_; + size_t res_num_copy_threads_; +#endif +#ifdef FBGEMM_FBCODE + // Multi-threaded consumers for tensor_stream() RPCs. + std::vector> consumer_threads_; + folly::UMPMCQueue weights_to_stream_queue_; + // Copy threads for UVM cache (joined every iteration). Shared by the blocking + // and non-blocking stream() paths; this assumes a given table streams in a + // single mode at a time (blocking OR non-blocking), never concurrently. + std::vector> chunk_copy_threads_; + std::unique_ptr dispatch_thread_; // OBC logger for RES silent-failure counters (res.fail.*). Emits to the // host-level OBC agent, so it reaches ODS from the trainer process without // per-process fb303 scrape config. Only constructed when streaming is on. std::unique_ptr ods_logger_; + + void join_worker_threads(); + void join_consumer_threads(); + void chunked_copy_and_enqueue( + const at::Tensor& indices, + const at::Tensor& weights, + std::optional identities, + std::optional runtime_meta, + const at::Tensor& count, + std::vector>& target_copy_threads); #endif }; -fbgemm_gpu::StreamQueueItem tensor_copy( +fbgemm_gpu::StreamQueueItem tensor_copy_chunk( const at::Tensor& indices, const at::Tensor& weights, std::optional identities, std::optional runtime_meta, - const at::Tensor& count); + int64_t start_row, + int64_t end_row); + +// Tiles [0, num_rows) into per-thread groups of [start, end) chunk ranges, each +// chunk of size <= chunk_size, contiguous and non-overlapping (union is the +// whole range). The outer index is the thread; each inner vector is that +// thread's contiguous band split into chunks. Empty bands produce no group. +// Pure/build-agnostic so it is unit-testable without a GPU or FBGEMM_FBCODE. +// num_threads bounds how the rows are pre-split before chunking, matching +// chunked_copy_and_enqueue's tiling. +std::vector>> +computeChunkRanges(int64_t num_rows, size_t chunk_size, size_t num_threads); } // namespace fbgemm_gpu diff --git a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp index 0a887a4193..3901553567 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -29,7 +29,7 @@ namespace { #ifdef FBGEMM_FBCODE // Timeout for copy_done_flag polling loop (microseconds). -constexpr int64_t kCopyDonePollTimeoutUs = 10'000'000; // 10 seconds +constexpr int64_t kCopyDonePollTimeoutUs = 10'000'000; /* * Get the thrift client to the training parameter server service @@ -64,86 +64,112 @@ inline int64_t get_maybe_uvm_scalar(const at::Tensor& tensor) { } // namespace -fbgemm_gpu::StreamQueueItem tensor_copy( +fbgemm_gpu::StreamQueueItem tensor_copy_chunk( const at::Tensor& indices, const at::Tensor& weights, std::optional identities, std::optional runtime_meta, - const at::Tensor& count) { - auto num_sets = get_maybe_uvm_scalar(count); - auto new_indices = at::empty( - num_sets, at::TensorOptions().device(at::kCPU).dtype(indices.dtype())); + int64_t start_row, + int64_t end_row) { + int64_t n = end_row - start_row; + auto new_indices = + at::empty(n, at::TensorOptions().device(at::kCPU).dtype(indices.dtype())); auto new_weights = at::empty( - {num_sets, weights.size(1)}, + {n, weights.size(1)}, at::TensorOptions().device(at::kCPU).dtype(weights.dtype())); std::optional new_identities = std::nullopt; if (identities.has_value()) { new_identities = at::empty( - {num_sets, identities->size(1)}, + {n, identities->size(1)}, at::TensorOptions().device(at::kCPU).dtype(identities->dtype())); } std::optional new_runtime_meta = std::nullopt; if (runtime_meta.has_value()) { new_runtime_meta = at::empty( - {num_sets, runtime_meta->size(1)}, + {n, runtime_meta->size(1)}, at::TensorOptions().device(at::kCPU).dtype(runtime_meta->dtype())); } auto new_count = at::empty({1}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + // Each tensor is copied under its own dispatch. They are independent copies, + // so there is no need to nest the dispatches (nesting would only reintroduce + // scalar_t shadowing and multiply template instantiations). + FBGEMM_DISPATCH_INTEGRAL_TYPES( + indices.scalar_type(), "tensor_copy_chunk", [&] { + std::copy( + indices.const_data_ptr() + start_row, + indices.const_data_ptr() + end_row, + new_indices.mutable_data_ptr()); + }); FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE( - weights.scalar_type(), "tensor_copy", [&] { - using value_t = scalar_t; - FBGEMM_DISPATCH_INTEGRAL_TYPES( - indices.scalar_type(), "tensor_copy", [&] { - using index_t = scalar_t; - auto indices_addr = indices.const_data_ptr(); - auto new_indices_addr = new_indices.mutable_data_ptr(); - std::copy( - indices_addr, - indices_addr + num_sets, - new_indices_addr); // dst_start - - auto weights_addr = weights.const_data_ptr(); - auto new_weights_addr = new_weights.mutable_data_ptr(); - std::copy( - weights_addr, - weights_addr + num_sets * weights.size(1), - new_weights_addr); // dst_start - if (identities.has_value()) { - FBGEMM_DISPATCH_INTEGRAL_TYPES( - identities->scalar_type(), "tensor_copy", [&] { - using identities_t = scalar_t; - const auto identities_addr = - identities->const_data_ptr(); - auto new_identities_addr = - new_identities->mutable_data_ptr(); - std::copy( - identities_addr, - identities_addr + num_sets * identities->size(1), - new_identities_addr); // dst_start - }); - } - if (runtime_meta.has_value()) { - FBGEMM_DISPATCH_ALL_TYPES( - runtime_meta->scalar_type(), "tensor_copy", [&] { - using runtime_meta_t = scalar_t; - auto runtime_meta_addr = - runtime_meta->const_data_ptr(); - auto new_runtime_meta_addr = - new_runtime_meta->mutable_data_ptr(); - std::copy( - runtime_meta_addr, - runtime_meta_addr + num_sets * runtime_meta->size(1), - new_runtime_meta_addr); // dst_start - }); - } - }); + weights.scalar_type(), "tensor_copy_chunk", [&] { + std::copy( + weights.const_data_ptr() + start_row * weights.size(1), + weights.const_data_ptr() + end_row * weights.size(1), + new_weights.mutable_data_ptr()); }); - *new_count.mutable_data_ptr() = num_sets; + if (identities.has_value()) { + FBGEMM_DISPATCH_INTEGRAL_TYPES( + identities->scalar_type(), "tensor_copy_chunk", [&] { + std::copy( + identities->const_data_ptr() + + start_row * identities->size(1), + identities->const_data_ptr() + + end_row * identities->size(1), + new_identities->mutable_data_ptr()); + }); + } + if (runtime_meta.has_value()) { + FBGEMM_DISPATCH_ALL_TYPES( + runtime_meta->scalar_type(), "tensor_copy_chunk", [&] { + std::copy( + runtime_meta->const_data_ptr() + + start_row * runtime_meta->size(1), + runtime_meta->const_data_ptr() + + end_row * runtime_meta->size(1), + new_runtime_meta->mutable_data_ptr()); + }); + } + *new_count.mutable_data_ptr() = n; return fbgemm_gpu::StreamQueueItem{ new_indices, new_weights, new_identities, new_runtime_meta, new_count}; } +std::vector>> +computeChunkRanges(int64_t num_rows, size_t chunk_size, size_t num_threads) { + // Split [0, num_rows) across up to num_threads contiguous per-thread bands, + // then split each band into <= chunk_size chunks. Returns one inner vector of + // [start, end) chunk ranges per thread (outer index = thread); ranges are + // contiguous, non-overlapping, and their union is the whole range. Empty + // bands produce no group. + std::vector>> thread_chunks; + if (num_rows <= 0 || chunk_size == 0 || num_threads == 0) { + return thread_chunks; + } + // ceil-div (a + b - 1) / b: rounds up so a partial final chunk/band counts. + const size_t n_chunks = + (static_cast(num_rows) + chunk_size - 1) / chunk_size; + const size_t n_threads = std::min(n_chunks, num_threads); + const size_t rows_per_thread = + (static_cast(num_rows) + n_threads - 1) / n_threads; + for (size_t ti = 0; ti < n_threads; ++ti) { + const int64_t thread_start = static_cast(ti * rows_per_thread); + const int64_t thread_end = + std::min(static_cast((ti + 1) * rows_per_thread), num_rows); + std::vector> chunks; + for (int64_t s = thread_start; s < thread_end; + s += static_cast(chunk_size)) { + const int64_t e = + std::min(s + static_cast(chunk_size), thread_end); + chunks.emplace_back(s, e); + } + if (!chunks.empty()) { + thread_chunks.push_back(std::move(chunks)); + } + } + return thread_chunks; +} + RawEmbeddingStreamer::RawEmbeddingStreamer( std::string unique_id, bool enable_raw_embedding_streaming, @@ -151,7 +177,10 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( int64_t res_server_port [[maybe_unused]], std::vector table_names, std::vector table_offsets, - const std::vector& table_sizes) + const std::vector& table_sizes, + int64_t res_chunk_size [[maybe_unused]], + int64_t res_num_consumers [[maybe_unused]], + int64_t res_num_copy_threads [[maybe_unused]]) : unique_id_(std::move(unique_id)), enable_raw_embedding_streaming_(enable_raw_embedding_streaming), #ifdef FBGEMM_FBCODE @@ -160,9 +189,30 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( #endif table_names_(std::move(table_names)), table_offsets_(std::move(table_offsets)), - table_sizes_(at::tensor(table_sizes)) { + table_sizes_(at::tensor(table_sizes)) +#ifdef FBGEMM_FBCODE + , + res_chunk_size_(res_chunk_size), + res_num_consumers_(res_num_consumers), + res_num_copy_threads_(res_num_copy_threads) +#endif +{ #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { + // Fail loud on a misconfigured knob. These are now caller-supplied (were + // compile-time constants), and 0 -- or a negative that wrapped to a huge + // size_t -- would silently break streaming: res_num_consumers=0 spawns no + // drain threads (queue grows unbounded), res_chunk_size=0 / + // res_num_copy_threads=0 make computeChunkRanges return empty (enqueues + // nothing). Reject rather than silently no-op. + TORCH_CHECK( + res_chunk_size > 0 && res_num_consumers > 0 && res_num_copy_threads > 0, + "RES config knobs must be > 0: res_chunk_size=", + res_chunk_size, + ", res_num_consumers=", + res_num_consumers, + ", res_num_copy_threads=", + res_num_copy_threads); XLOG(INFO) << "[TBE_ID" << unique_id_ << "] Raw embedding streaming enabled with res_server_port at" << res_server_port_; @@ -172,36 +222,59 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( ods_logger_ = std::make_unique(); - weights_stream_thread_ = std::make_unique([this] { - while (!stop_) { - auto stream_item_ptr = weights_to_stream_queue_.try_peek(); - if (!stream_item_ptr) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - continue; + XLOG(INFO) << "[TBE_ID" << unique_id_ << "] Starting " << res_num_consumers_ + << " consumer threads" + << ", chunk_size=" << res_chunk_size_ + << ", copy_threads=" << res_num_copy_threads_; + // Ordering caveat: with res_num_consumers_ > 1, these consumers drain the + // queue concurrently, so arrival order at the PS is NOT enqueue (iteration) + // order. The store (TrainingPsHandler) applies same-(fqn,row_id) writes + // arrival-wins with no version compare, so a stale iter-i write can + // transiently clobber a fresh iter-(i+1) one for a hot row (self-heals on + // the row's next in-order update). A single consumer would be ordered/safe. + // TODO(T281413204): proper fix is to carry a per-row iteration/version and + // keep-newest in the store. + for (size_t ci = 0; ci < res_num_consumers_; ++ci) { + consumer_threads_.push_back(std::make_unique([this] { + while (!stop_) { + auto item = weights_to_stream_queue_.try_dequeue(); + if (!item.has_value()) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + continue; + } + if (stop_) { + return; + } + // Guard the per-item body: a transient tensor_stream failure must log + // and move on to the next item, never escape and std::terminate the + // trainer. + try { + folly::stop_watch stop_watch; + folly::coro::blockingWait(tensor_stream( + item->indices, + item->weights, + item->identities, + item->runtime_meta)); + if (ods_logger_) { + ods_logger_->bumpKeyGauge( + "tbe_stream_queue_depth", + static_cast(weights_to_stream_queue_.size())); + } + XLOG_EVERY_MS(INFO, 60000) + << "[TBE_ID" << unique_id_ << "] end stream queue size: " + << weights_to_stream_queue_.size() << " stream takes " + << stop_watch.elapsed().count() << "ms" + << " rows=" << item->indices.size(0); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] consumer thread caught exception: " << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] consumer thread caught unknown exception"; + } } - if (stop_) { - return; - } - auto& indices = stream_item_ptr->indices; - auto& weights = stream_item_ptr->weights; - auto& identities = stream_item_ptr->identities; - auto& runtime_meta = stream_item_ptr->runtime_meta; - folly::stop_watch stop_watch; - folly::coro::blockingWait( - tensor_stream(indices, weights, identities, runtime_meta)); - - weights_to_stream_queue_.dequeue(); - auto post_dequeue_depth = weights_to_stream_queue_.size(); - if (ods_logger_) { - ods_logger_->bumpKeyGauge( - "stream_mpsc_depth", static_cast(post_dequeue_depth)); - } - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] end stream queue size: " << post_dequeue_depth - << " stream takes " << stop_watch.elapsed().count() << "ms"; - } - }); + })); + } } #endif } @@ -211,7 +284,8 @@ RawEmbeddingStreamer::~RawEmbeddingStreamer() { #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { join_stream_tensor_copy_thread(); - join_weights_stream_thread(); + join_worker_threads(); + join_consumer_threads(); } #endif } @@ -241,7 +315,7 @@ void RawEmbeddingStreamer::stream( weights_to_stream_queue_.enqueue(stream_item); return; } - if (blocking_tensor_copy) { + auto poll_flag = [this, copy_done_flag]() { if (copy_done_flag.has_value()) { auto* ptr = static_cast(copy_done_flag->data_ptr()); folly::stop_watch poll_watch; @@ -249,59 +323,58 @@ void RawEmbeddingStreamer::stream( std::this_thread::yield(); if (poll_watch.elapsed().count() > kCopyDonePollTimeoutUs) { LOG(ERROR) << "[TBE_ID" << unique_id_ - << "] copy_done_flag blocking: poll timed out after " + << "] copy_done_flag poll timed out after " << kCopyDonePollTimeoutUs / 1'000'000 << "s"; - return; + return false; } } - *ptr = 0; // Reset for next iteration - } else { - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] copy_done_flag not provided, skipping wait (blocking)"; + *ptr = 0; + } + return true; + }; + + if (blocking_tensor_copy) { + if (!poll_flag()) { + return; } - copy_and_enqueue_stream_tensors( + chunked_copy_and_enqueue( indices, weights, std::move(identities), std::move(runtime_meta), - count); + count, + chunk_copy_threads_); + join_worker_threads(); return; } - // Make sure the previous thread is done before starting a new one + // Non-blocking: join the previous dispatch + copy threads, then spawn new + // ones. The join is the serializer: it guarantees iter i's copy finished + // reading the source cache rows before iter i+1 overwrites them. join_stream_tensor_copy_thread(); - // Cuda dispatches the host callbacks all in the same CPU thread. But the - // callbacks don't need to be serialized. - // So, We need to spin up a new thread to unblock the CUDA stream, so the CUDA - // can continue executing other host callbacks, eg. get/evict. - stream_tensor_copy_thread_ = std::make_unique([this, - copy_done_flag, - indices, - weights, - identities, - runtime_meta, - count]() { - if (copy_done_flag.has_value()) { - auto* ptr = static_cast(copy_done_flag->data_ptr()); - folly::stop_watch poll_watch; - while (*ptr == 0) { - std::this_thread::yield(); - if (poll_watch.elapsed().count() > kCopyDonePollTimeoutUs) { - LOG(ERROR) << "[TBE_ID" << unique_id_ - << "] copy_done_flag non-blocking: poll timed out after " - << kCopyDonePollTimeoutUs / 1'000'000 << "s"; - return; + dispatch_thread_ = std::make_unique( + [this, poll_flag, indices, weights, identities, runtime_meta, count]() { + // Guard the dispatcher body so a copy/enqueue failure logs instead of + // escaping the std::thread and calling std::terminate. + try { + if (!poll_flag()) { + return; + } + chunked_copy_and_enqueue( + indices, + weights, + identities, + runtime_meta, + count, + chunk_copy_threads_); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] stream dispatcher thread caught exception: " + << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] stream dispatcher thread caught unknown exception"; } - } - *ptr = 0; // Reset for next iteration - } else { - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] copy_done_flag not provided, skipping wait (non-blocking)"; - } - copy_and_enqueue_stream_tensors( - indices, weights, identities, runtime_meta, count); - }); + }); rec->record.end(); #endif } @@ -310,15 +383,100 @@ void RawEmbeddingStreamer::join_stream_tensor_copy_thread() { #ifdef FBGEMM_FBCODE auto rec = torch::autograd::profiler::record_function_enter_new( "## RawEmbeddingStreamer::join_stream_tensor_copy_thread ##"); - if (stream_tensor_copy_thread_ != nullptr && - stream_tensor_copy_thread_->joinable()) { - stream_tensor_copy_thread_->join(); + if (dispatch_thread_ != nullptr && dispatch_thread_->joinable()) { + dispatch_thread_->join(); } + join_worker_threads(); rec->record.end(); #endif } #ifdef FBGEMM_FBCODE +void RawEmbeddingStreamer::join_worker_threads() { + for (auto& t : chunk_copy_threads_) { + if (t && t->joinable()) { + t->join(); + } + } + chunk_copy_threads_.clear(); +} + +void RawEmbeddingStreamer::join_consumer_threads() { + for (auto& t : consumer_threads_) { + if (t && t->joinable()) { + t->join(); + } + } + consumer_threads_.clear(); +} + +void RawEmbeddingStreamer::chunked_copy_and_enqueue( + const at::Tensor& indices, + const at::Tensor& weights, + std::optional identities, + std::optional runtime_meta, + const at::Tensor& count, + std::vector>& target_copy_threads) { + auto rec = torch::autograd::profiler::record_function_enter_new( + "## RawEmbeddingStreamer::chunked_copy_and_enqueue ##"); + const auto num_rows = get_maybe_uvm_scalar(count); + const auto thread_chunks = + computeChunkRanges(num_rows, res_chunk_size_, res_num_copy_threads_); + + for (auto& t : target_copy_threads) { // join+clear the previous batch + if (t && t->joinable()) { + t->join(); + } + } + target_copy_threads.clear(); + if (thread_chunks.empty()) { + rec->record.end(); + return; + } + + // One copy thread per pre-computed group. Chunk boundaries and per-thread + // grouping live entirely in computeChunkRanges, so the enqueued row set is + // identical regardless of how threads are laid out here. + target_copy_threads.reserve(thread_chunks.size()); + for (size_t ti = 0; ti < thread_chunks.size(); ++ti) { + target_copy_threads.push_back( + std::make_unique([this, + indices, + weights, + identities, + runtime_meta, + chunks = thread_chunks[ti], + ti]() { + // Guard the copy body so a per-chunk failure logs instead of escaping + // the std::thread and calling std::terminate. + try { + folly::stop_watch thread_watch; + int64_t rows_done = 0; + for (const auto& [s, e] : chunks) { + auto chunk_item = tensor_copy_chunk( + indices, weights, identities, runtime_meta, s, e); + weights_to_stream_queue_.enqueue(std::move(chunk_item)); + rows_done += (e - s); + } + XLOG_EVERY_MS(INFO, 15000) + << "[TBE_ID" << unique_id_ << "] copy_thread tid=" << ti + << " rows=" << rows_done << " chunks=" << chunks.size() + << " copy_ms=" << thread_watch.elapsed().count(); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ << "] copy_thread tid=" << ti + << " caught exception: " << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ << "] copy_thread tid=" << ti + << " caught unknown exception"; + } + })); + } + XLOG_EVERY_MS(INFO, 15000) + << "[RES] chunked_copy tbe=" << unique_id_ << " rows=" << num_rows + << " threads=" << thread_chunks.size(); + rec->record.end(); +} + folly::coro::Task RawEmbeddingStreamer::tensor_stream( const at::Tensor& indices, const at::Tensor& weights, @@ -439,37 +597,23 @@ folly::coro::Task RawEmbeddingStreamer::tensor_stream( try { co_await res_client->co_setEmbeddings(req); } catch (const std::exception& e) { + // A transient per-shard RPC failure must not propagate: it would tear + // down the consumer thread (std::terminate). Log, bump the counter, and + // move on to the next shard. if (ods_logger_) { - ods_logger_->bumpKey("set_embeddings_rpc", 1); + ods_logger_->bumpKey("set_embeddings_rpc_failure", 1); } XLOG(ERR) << "[TBE_ID" << unique_id_ << "] co_setEmbeddings threw on shard " << i << ": " << e.what(); - throw; } } co_return; } -void RawEmbeddingStreamer::copy_and_enqueue_stream_tensors( - const at::Tensor& indices, - const at::Tensor& weights, - std::optional identities, - std::optional runtime_meta, - const at::Tensor& count) { - auto rec = torch::autograd::profiler::record_function_enter_new( - "## RawEmbeddingStreamer::copy_and_enqueue_stream_tensors ##"); - auto stream_item = tensor_copy( - indices, weights, std::move(identities), std::move(runtime_meta), count); - weights_to_stream_queue_.enqueue(stream_item); - rec->record.end(); -} - void RawEmbeddingStreamer::join_weights_stream_thread() { - if (weights_stream_thread_ != nullptr && weights_stream_thread_->joinable()) { - stop_ = true; - weights_stream_thread_->join(); - } + stop_ = true; + join_consumer_threads(); } uint64_t RawEmbeddingStreamer::get_weights_to_stream_queue_size() { diff --git a/fbgemm_gpu/src/split_embeddings_cache/split_embeddings_cache_ops.cpp b/fbgemm_gpu/src/split_embeddings_cache/split_embeddings_cache_ops.cpp index a1294293e9..b83010f0f5 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/split_embeddings_cache_ops.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/split_embeddings_cache_ops.cpp @@ -89,7 +89,10 @@ auto raw_embedding_streamer = int64_t, std::vector, std::vector, - std::vector>(), + std::vector, + int64_t, + int64_t, + int64_t>(), "", { torch::arg("unique_id") = 0, @@ -99,6 +102,9 @@ auto raw_embedding_streamer = torch::arg("table_names") = torch::List(), torch::arg("table_offsets") = torch::List(), torch::arg("table_sizes") = torch::List(), + torch::arg("res_chunk_size") = 500000, + torch::arg("res_num_consumers") = 8, + torch::arg("res_num_copy_threads") = 4, }) .def( "stream", diff --git a/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp b/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp index 50b6f28ada..42d3a351e0 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp @@ -8,7 +8,7 @@ #include #include -#include "deeplearning/fbgemm/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h" // @manual=//deeplearning/fbgemm/fbgemm_gpu/src/split_embeddings_cache:raw_embedding_streamer +#include "fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h" #ifdef FBGEMM_FBCODE #include #include "aiplatform/gmpp/experimental/training_ps/gen-cpp2/TrainingParameterServerService.h" @@ -35,6 +35,11 @@ class MockTrainingParameterServerService }; #endif +// These object-constructing tests run in BOTH targets: the fbcode target links +// the flag-on :raw_embedding_streamer, and the no-fbcode target links the +// flag-off :raw_embedding_streamer_no_fbcode -- so each test's flag view of the +// header layout matches its library, avoiding the sizeof mismatch (an all-off +// test linked against the always-flag-on library would overrun the object). static std::unique_ptr getRawEmbeddingStreamer( const std::string& unique_id, @@ -49,7 +54,10 @@ getRawEmbeddingStreamer( 0, // res_server_port table_names, table_offsets, - table_sizes); + table_sizes, + 500000, // res_chunk_size + 8, // res_num_consumers + 4); // res_num_copy_threads } TEST(RawEmbeddingStreamerTest, TestConstructorAndDestructor) { @@ -84,7 +92,332 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithoutStreaming) { indices, weights, std::nullopt, std::nullopt, count, true, true); } +namespace { +// Row-major tensor with distinct, predictable values so a sliced copy is +// unambiguous: value at [r, c] == r * dim + c. +at::Tensor makeRowMajor(int64_t num_rows, int64_t dim, at::ScalarType dtype) { + return at::arange( + num_rows * dim, at::TensorOptions().device(at::kCPU).dtype(dtype)) + .reshape({num_rows, dim}); +} +} // namespace + +// tensor_copy_chunk is build-agnostic (defined outside FBGEMM_FBCODE). Expected +// tensors are constructed independently via at::slice, never copied from impl +// output. +TEST(RawEmbeddingStreamerTest, TensorCopyChunkFullRange) { + constexpr int64_t kNumRows = 5; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, /*start_row=*/0, kNumRows); + + EXPECT_TRUE(at::equal(item.indices, indices)); + EXPECT_TRUE(at::equal(item.weights, weights)); + EXPECT_FALSE(item.identities.has_value()); + EXPECT_FALSE(item.runtime_meta.has_value()); + const auto expected_count = at::tensor( + {kNumRows}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + EXPECT_TRUE(at::equal(item.count, expected_count)); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkNonZeroStartSlicesCorrectSlice) { + // Guards the start_row*dim / end_row*dim offset arithmetic in the copy. + constexpr int64_t kNumRows = 6; + constexpr int64_t kStart = 2; + constexpr int64_t kEnd = 5; // n == 3 + auto indices = at::tensor( + {10, 20, 30, 40, 50, 60}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_TRUE(at::equal(item.indices, indices.slice(0, kStart, kEnd))); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); + const auto expected_count = at::tensor( + {kEnd - kStart}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + EXPECT_TRUE(at::equal(item.count, expected_count)); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkCopiesIdentitiesAndRuntimeMeta) { + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; // n == 3 + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + auto identities = makeRowMajor(kNumRows, /*dim=*/2, at::kLong); + auto runtime_meta = makeRowMajor(kNumRows, /*dim=*/1, at::kLong); + + auto item = tensor_copy_chunk( + indices, weights, identities, runtime_meta, kStart, kEnd); + + ASSERT_TRUE(item.identities.has_value()); + ASSERT_TRUE(item.runtime_meta.has_value()); + EXPECT_TRUE(at::equal(*item.identities, identities.slice(0, kStart, kEnd))); + EXPECT_TRUE( + at::equal(*item.runtime_meta, runtime_meta.slice(0, kStart, kEnd))); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkAbsentOptionalsStayNullopt) { + constexpr int64_t kNumRows = 4; + auto indices = at::tensor( + {10, 20, 30, 40}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, /*start_row=*/0, kNumRows); + + EXPECT_FALSE(item.identities.has_value()); + EXPECT_FALSE(item.runtime_meta.has_value()); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkInt32IndicesDtype) { + // Coverage for the integral-index dispatch branch with int32 indices. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kInt)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_EQ(item.indices.scalar_type(), at::kInt); + EXPECT_TRUE(at::equal(item.indices, indices.slice(0, kStart, kEnd))); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkFloatRuntimeMeta) { + // tensor_copy_chunk must dispatch runtime_meta over all dtypes + // (FBGEMM_DISPATCH_ALL_TYPES): a float runtime_meta would throw if it were + // dispatched integral-only. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + auto runtime_meta = makeRowMajor(kNumRows, /*dim=*/2, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, runtime_meta, kStart, kEnd); + + ASSERT_TRUE(item.runtime_meta.has_value()); + EXPECT_TRUE( + at::equal(*item.runtime_meta, runtime_meta.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkHalfWeights) { + // Coverage for the half (fp16) weights dispatch branch of + // FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE -- a common quantized serving dtype. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kHalf); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_EQ(item.weights.scalar_type(), at::kHalf); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkByteWeights) { + // Coverage for the byte (int8) weights dispatch branch of + // FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE -- the int8-quantized serving dtype. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kByte); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_EQ(item.weights.scalar_type(), at::kByte); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +namespace { +// computeChunkRanges groups chunks per thread (outer index = thread). Flatten +// to the in-order chunk list so the coverage/contiguity invariants can be +// checked across the whole range regardless of the thread grouping. +std::vector> flatten( + const std::vector>>& + thread_chunks) { + std::vector> ranges; + for (const auto& chunks : thread_chunks) { + ranges.insert(ranges.end(), chunks.begin(), chunks.end()); + } + return ranges; +} + +// Structural invariants computeChunkRanges must always satisfy: ranges are +// contiguous + non-overlapping starting at 0, cover exactly [0, num_rows), and +// every chunk is non-empty and no larger than chunk_size. An off-by-one in the +// tiling arithmetic breaks at least one of these. +void expectValidChunkRanges( + const std::vector>& ranges, + int64_t num_rows, + int64_t chunk_size) { + int64_t cursor = 0; + for (const auto& [start, end] : ranges) { + EXPECT_EQ(start, cursor) << "ranges must be contiguous and non-overlapping"; + EXPECT_GT(end, start) << "no empty ranges"; + EXPECT_LE(end - start, chunk_size) << "each chunk must be <= chunk_size"; + cursor = end; + } + EXPECT_EQ(cursor, num_rows) << "ranges must cover exactly [0, num_rows)"; +} +} // namespace + +// computeChunkRanges (like tensor_copy_chunk) is build-agnostic. Expected +// ranges are constructed independently. +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesExactMultiple) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/8, /*chunk_size=*/4, /*num_threads=*/2)); + const std::vector> expected = {{0, 4}, {4, 8}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/8, /*chunk_size=*/4); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesRemainderChunk) { + // Single thread => pure chunking; last chunk carries the remainder. + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/10, /*chunk_size=*/4, /*num_threads=*/1)); + const std::vector> expected = { + {0, 4}, {4, 8}, {8, 10}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/10, /*chunk_size=*/4); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesCountLessThanChunkSize) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/3, /*chunk_size=*/10, /*num_threads=*/4)); + const std::vector> expected = {{0, 3}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/3, /*chunk_size=*/10); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesSingleChunk) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/5, /*num_threads=*/4)); + const std::vector> expected = {{0, 5}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/5); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesChunkSizeOne) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/4, /*chunk_size=*/1, /*num_threads=*/1)); + const std::vector> expected = { + {0, 1}, {1, 2}, {2, 3}, {3, 4}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/4, /*chunk_size=*/1); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesZeroRowsIsEmpty) { + EXPECT_TRUE( + computeChunkRanges(/*num_rows=*/0, /*chunk_size=*/4, /*num_threads=*/4) + .empty()); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesNumThreadsExceedsNumChunks) { + // n_threads is clamped to n_chunks, so exactly n_chunks groups are emitted, + // one chunk each, with no empty group. + const auto thread_chunks = + computeChunkRanges(/*num_rows=*/6, /*chunk_size=*/3, /*num_threads=*/10); + EXPECT_EQ(thread_chunks.size(), 2u) << "one group per chunk"; + const auto ranges = flatten(thread_chunks); + const std::vector> expected = {{0, 3}, {3, 6}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/6, /*chunk_size=*/3); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesThreadSplitThenChunk) { + // num_threads < num_chunks and rows_per_thread not a multiple of chunk_size: + // rows are pre-split into 2 per-thread bands ([0,50), [50,100)) and each band + // is then chunked by 30, so boundaries land at the thread split (50), not at + // 60. This locks the tiling to the original inline behavior. + const auto thread_chunks = computeChunkRanges( + /*num_rows=*/100, /*chunk_size=*/30, /*num_threads=*/2); + EXPECT_EQ(thread_chunks.size(), 2u) << "one group per thread band"; + const std::vector>> expected_groups = + {{{0, 30}, {30, 50}}, {{50, 80}, {80, 100}}}; + EXPECT_EQ(thread_chunks, expected_groups); + expectValidChunkRanges( + flatten(thread_chunks), /*num_rows=*/100, /*chunk_size=*/30); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesNoEmptyTailRange) { + // rows_per_thread rounds up (ceil(5/4)=2) so the 4th thread's band would be + // [6,5); that empty band must be dropped, leaving 3 non-empty groups. + const auto thread_chunks = + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/1, /*num_threads=*/4); + EXPECT_EQ(thread_chunks.size(), 3u) << "empty trailing band is dropped"; + const auto ranges = flatten(thread_chunks); + const std::vector> expected = { + {0, 1}, {1, 2}, {2, 3}, {3, 4}, {4, 5}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/1); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesOneOverChunk) { + // num_rows == chunk_size + 1: the +1 spills into a second, single-row chunk. + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/4, /*num_threads=*/1)); + const std::vector> expected = {{0, 4}, {4, 5}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/4); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesZeroThreadsOrChunkSizeEmpty) { + // Defensive guards: chunk_size==0 and num_threads==0 would divide by zero in + // the ceil-div tiling, so both must short-circuit to an empty result. + EXPECT_TRUE( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/0, /*num_threads=*/4) + .empty()); + EXPECT_TRUE( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/4, /*num_threads=*/0) + .empty()); +} + #ifdef FBGEMM_FBCODE +TEST(RawEmbeddingStreamerTest, CtorRejectsZeroKnob) { + // A 0-valued RES knob would silently disable streaming (0 consumers never + // drain the queue; res_chunk_size/res_num_copy_threads=0 make chunk ranges + // empty), so the ctor must reject it loudly. The TORCH_CHECK fires before any + // thrift client is created, so no mock server is needed. + EXPECT_ANY_THROW( + fbgemm_gpu::RawEmbeddingStreamer( + "test_zero_knob", + /*enable_raw_embedding_streaming=*/true, + /*res_store_shards=*/3, + /*res_server_port=*/0, + /*table_names=*/{}, + /*table_offsets=*/{}, + /*table_sizes=*/{}, + /*res_chunk_size=*/0, + /*res_num_consumers=*/8, + /*res_num_copy_threads=*/4)); +} + TEST(RawEmbeddingStreamerTest, TestTensorStream) { std::vector table_names = {"tb1", "tb2", "tb3"}; std::vector table_offsets = {0, 100, 300}; @@ -178,14 +511,60 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopy) { indices, weights, std::nullopt, std::nullopt, count, true, true); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); - // Test non-blocking tensor copy + // Test non-blocking tensor copy. The copy runs on the dispatcher, so we must + // join_stream_tensor_copy_thread() before checking the queue -- asserting the + // size before the join would race the background copy. streamer->stream( indices, weights, std::nullopt, std::nullopt, count, true, false); - EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); streamer->join_stream_tensor_copy_thread(); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 2); } +TEST(RawEmbeddingStreamerTest, TestStreamWithCopyZeroCountEnqueuesNothing) { + // count <= 0 drives num_rows == 0, so chunked_copy_and_enqueue early-returns + // and nothing is enqueued. + std::vector table_names = {"tb1", "tb2", "tb3"}; + std::vector table_offsets = {0, 100, 300}; + std::vector table_sizes = {0, 50, 200, 300}; + + auto streamer = getRawEmbeddingStreamer( + "test_zero_count", true, table_names, table_offsets, table_sizes); + + auto mock_service = std::make_shared(); + auto mock_server = + std::make_shared( + mock_service, + "::1", + 0, + facebook::services::TLSConfig::applyDefaultsToThriftServer); + auto& mock_client_factory = + facebook::servicerouter::getMockSRClientFactory(false /* strict */); + mock_client_factory.registerMockService( + "realtime.delta.publish.esr", mock_server); + + auto indices = at::tensor( + {10, 2, 1, 150, 170, 230, 280}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = at::randn( + {indices.size(0), EMBEDDING_DIMENSION}, + at::TensorOptions().device(at::kCPU).dtype(c10::kFloat)); + auto count = + at::tensor({0}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + + // Stop the consumer threads so the queue size is stable to read. + streamer->join_weights_stream_thread(); + + streamer->stream( + indices, + weights, + std::nullopt, + std::nullopt, + count, + /*require_tensor_copy=*/true, + /*blocking_tensor_copy=*/true); + EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 0); +} + TEST(RawEmbeddingStreamerTest, TestStreamE2E) { std::vector table_names = {"tb1", "tb2", "tb3"}; std::vector table_offsets = {0, 100, 300}; @@ -232,12 +611,95 @@ TEST(RawEmbeddingStreamerTest, TestStreamE2E) { streamer->stream( indices, weights, std::nullopt, std::nullopt, count, true, true); - // Make sure dequeue finished - std::this_thread::sleep_for(std::chrono::seconds(1)); + // Bounded wait for the consumer to drain the enqueued item (so + // co_setEmbeddings has run) before stopping the thread -- avoids a + // fixed-sleep flake and the stop_-between-peek-and-process race. + for (int i = 0; i < 1000 && streamer->get_weights_to_stream_queue_size() > 0; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + streamer->join_weights_stream_thread(); +} + +TEST(RawEmbeddingStreamerTest, TestNoCopyMultiItemConsumerDrain) { + // require_tensor_copy=false enqueues one raw item per stream() call (no D2H + // copy, no chunking). Enqueue several and let the N-thread consumer pool + // drain them concurrently through the UMPMC queue: every item must be + // shipped, so co_setEmbeddings runs (kNumItems * shards-per-item) times. + // Exercises the raw-enqueue (require_tensor_copy=false) branch and multi-item + // concurrent drain. Wait on the RPC count BEFORE stopping so no in-flight + // item is dropped. + std::vector table_names = {"tb1", "tb2", "tb3"}; + std::vector table_offsets = {0, 100, 300}; + std::vector table_sizes = {0, 50, 200, 300}; + + // Static storage duration so the co_setEmbeddings coroutine mock below can + // read it WITHOUT capturing -- a capturing coroutine lambda risks + // use-after-free once its closure is destroyed + // (cppcoreguidelines-avoid-capturing-lambda-coroutines). Reset per run. + static std::atomic rpc_count; + rpc_count.store(0); + auto mock_service = std::make_shared(); + auto mock_server = + std::make_shared( + mock_service, + "::1", + 0, + facebook::services::TLSConfig::applyDefaultsToThriftServer); + auto& mock_client_factory = + facebook::servicerouter::getMockSRClientFactory(false /* strict */); + mock_client_factory.registerMockService( + "realtime.delta.publish.esr", mock_server); + + auto counting_response = + [](std::unique_ptr< + aiplatform::gmpp::experimental::training_ps::SetEmbeddingsRequest>) + -> folly::coro::Task> { + rpc_count.fetch_add(1); + co_return std::make_unique< + aiplatform::gmpp::experimental::training_ps::SetEmbeddingsResponse>(); + }; + EXPECT_CALL(*mock_service, co_setEmbeddings(_)) + .WillRepeatedly(folly::coro::gmock_helpers::CoInvoke(counting_response)); + + auto streamer = getRawEmbeddingStreamer( + "test_nocopy_multi_item", true, table_names, table_offsets, table_sizes); + + auto indices = at::tensor( + {10, 2, 1, 150, 170, 230, 280}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = at::randn( + {indices.size(0), EMBEDDING_DIMENSION}, + at::TensorOptions().device(at::kCPU).dtype(c10::kFloat)); + auto count = at::tensor( + {indices.size(0)}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + + constexpr int kNumItems = 3; + // These 7 indices span the 3 tables, so each item ships 3 shards (matches the + // Times(3) in TestStreamE2E for a single item). + constexpr int kShardsPerItem = 3; + for (int i = 0; i < kNumItems; ++i) { + streamer->stream( + indices, + weights, + std::nullopt, + std::nullopt, + count, + /*require_tensor_copy=*/false); + } + // Bounded wait until every item has been shipped, then stop -- waiting on the + // count (not queue size) guarantees no dequeued-but-unprocessed item is + // dropped by stop_. + for (int i = 0; i < 1000 && rpc_count.load() < kNumItems * kShardsPerItem; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_EQ(rpc_count.load(), kNumItems * kShardsPerItem); streamer->join_weights_stream_thread(); } -TEST(RawEmbeddingStreamerTest, TestCoSetEmbeddingsThrowPropagates) { +TEST(RawEmbeddingStreamerTest, TestCoSetEmbeddingsFailureIsSwallowed) { std::vector table_names = {"tb1", "tb2", "tb3"}; std::vector table_offsets = {0, 100, 300}; std::vector table_sizes = {0, 50, 200, 300}; @@ -261,7 +723,9 @@ TEST(RawEmbeddingStreamerTest, TestCoSetEmbeddingsThrowPropagates) { mock_client_factory.registerMockService( "realtime.delta.publish.esr", mock_server); + // Every shard RPC fails. EXPECT_CALL(*mock_service, co_setEmbeddings(_)) + .Times(3) // still attempts all 3 shards despite each failing .WillRepeatedly( folly::coro::gmock_helpers::CoInvoke( [](std::unique_ptrtensor_stream( indices, weights, std::nullopt, std::nullopt))); } @@ -431,5 +896,8 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopyDoneFlagNonBlockingCopy) { // Wait for the async thread to complete streamer->join_stream_tensor_copy_thread(); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); + // poll_flag() must have observed the flag (1) and reset it to 0; without the + // reset the next iteration would stream before the D2H copy finished. + EXPECT_EQ(copy_done_flag.item(), 0); } #endif diff --git a/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.cpp b/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.cpp index 0958049fcc..8eaae2eb97 100644 --- a/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.cpp +++ b/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.cpp @@ -94,7 +94,15 @@ EmbeddingKVDB::EmbeddingKVDB( res_server_port, std::move(table_names), std::move(table_offsets), - table_sizes)) { + table_sizes, + // SSD/kv_db TBEs are intentionally left on the default RES ship/ + // copy knobs: the config layer that makes these tunable is not + // threaded through the kv_db ctor. Plumb res_chunk_size/ + // res_num_consumers/res_num_copy_threads here if SSD-backed + // tables ever need to tune them. + /*res_chunk_size=*/500000, + /*res_num_consumers=*/8, + /*res_num_copy_threads=*/4)) { CHECK(num_shards > 0); if (cache_size_gb > 0) { l2_cache::CacheLibCache::CacheConfig cache_config;