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..49ca5edf58 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 @@ -137,6 +137,9 @@ class UVMCacheStatsIndex(enum.IntEnum): class RESParams: res_server_port: int = 0 # the port of the res server res_store_shards: int = 1 # the number of shards to store the raw embeddings + 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 table_names: list[str] = field(default_factory=list) # table names the TBE holds table_offsets: list[int] = field( default_factory=list @@ -1560,6 +1563,9 @@ def __init__( # noqa C901 self.enable_raw_embedding_streaming, self.res_params.res_store_shards, self.res_params.res_server_port, + self.res_params.res_chunk_size, + self.res_params.res_num_consumers, + self.res_params.res_num_copy_threads, self.res_params.table_names, self.res_params.table_offsets, self.res_params.table_sizes, @@ -4529,7 +4535,7 @@ def raw_embedding_stream(self) -> None: f"## uvm_lookup_prefetched_rows {self.timestep} {self.uuid} ##" ): if not self._res_sync_copy and self._res_require_copy: - self._raw_embedding_streamer.join_stream_tensor_copy_thread() + self._raw_embedding_streamer.join_dispatch() prefetched_info = self.prefetched_info_list.pop(0) updated_locations = torch.ops.fbgemm.lxu_cache_lookup( prefetched_info.linear_unique_cache_indices, @@ -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..4bb4dfe3bb 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 @@ -10,11 +10,17 @@ #include #ifdef FBGEMM_FBCODE #include +#include #endif #include +#include #ifdef FBGEMM_FBCODE +namespace folly { +class CPUThreadPoolExecutor; +} // namespace folly + namespace facebook::aiplatform::gmpp::experimental::training_ps { class TrainingPsOdsLogger; } // namespace facebook::aiplatform::gmpp::experimental::training_ps @@ -48,6 +54,9 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { bool enable_raw_embedding_streaming, int64_t res_store_shards, int64_t res_server_port, + int64_t res_chunk_size, + int64_t res_num_consumers, + int64_t res_num_copy_threads, std::vector table_names, std::vector table_offsets, const std::vector& table_sizes); @@ -56,10 +65,12 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { /// 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,10 +97,10 @@ 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(); + void join_dispatch(); #ifdef FBGEMM_FBCODE folly::coro::Task tensor_stream( @@ -97,53 +108,96 @@ 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: drops the ship executor to 0 worker threads so a test can + * read a stable queue size (via get_weights_to_stream_queue_size()). Not + * reversible -- the executor stops shipping after this call. */ void join_weights_stream_thread(); // FOR TESTING: get queue size. uint64_t get_weights_to_stream_queue_size(); #endif private: - std::atomic stop_{false}; std::string unique_id_; bool enable_raw_embedding_streaming_; #ifdef FBGEMM_FBCODE int64_t res_store_shards_; int64_t res_server_port_; + size_t res_chunk_size_; + size_t res_num_consumers_; + size_t res_num_copy_threads_; #endif std::vector table_names_; 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_; + // Named executor that ships enqueued StreamQueueItems to the PS. Push model: + // producers submit one ship task per item and workers wake on submit (no + // polling). Sized to res_num_consumers_. + std::unique_ptr consumer_executor_; + // 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_; + // Persistent size-1 executor that runs the per-iteration dispatch (poll_flag + // + chunked_copy_and_enqueue). Named so its thread is identifiable in traces. + std::unique_ptr dispatch_executor_; + folly::SemiFuture dispatch_future_{folly::makeSemiFuture()}; // 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_chunk_copy_threads(); + // Submit one ship task (blockingWait(tensor_stream(...))) for `item` onto + // consumer_executor_. + void submit_stream_item(StreamQueueItem item); + 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); + + // Waits (spinning) for copy_done_flag to signal the source tensors are safe + // to read, resetting it. Returns false on timeout. Shared by the blocking + // stream() path and dispatch_copy_task. + bool poll_copy_done_flag(const std::optional& copy_done_flag); + + // Coroutine form of the non-blocking dispatch (poll_copy_done_flag + + // chunked_copy_and_enqueue). Args are taken by value so nothing dangles once + // it is scheduled on dispatch_executor_ -- a capturing lambda coroutine would + // risk a use-after-free (clang-tidy cppcoreguidelines-avoid-capturing-lambda- + // coroutines). + folly::coro::Task dispatch_copy_task( + at::Tensor indices, + at::Tensor weights, + std::optional identities, + std::optional runtime_meta, + at::Tensor count, + std::optional copy_done_flag); #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..e09ed887c7 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -7,7 +7,11 @@ */ #ifdef FBGEMM_FBCODE +#include #include +#include +#include +#include #include #include #include "aiplatform/gmpp/experimental/training_ps/TrainingPsOdsLogger.h" @@ -29,7 +33,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,91 +68,120 @@ 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, int64_t res_store_shards [[maybe_unused]], int64_t res_server_port [[maybe_unused]], + int64_t res_chunk_size [[maybe_unused]], + int64_t res_num_consumers [[maybe_unused]], + int64_t res_num_copy_threads [[maybe_unused]], std::vector table_names, std::vector table_offsets, const std::vector& table_sizes) @@ -157,12 +190,29 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( #ifdef FBGEMM_FBCODE res_store_shards_(res_store_shards), res_server_port_(res_server_port), + res_chunk_size_(res_chunk_size), + res_num_consumers_(res_num_consumers), + res_num_copy_threads_(res_num_copy_threads), #endif table_names_(std::move(table_names)), table_offsets_(std::move(table_offsets)), table_sizes_(at::tensor(table_sizes)) { #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,46 +222,53 @@ 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; - } - 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"; - } - }); + // Persistent size-1 executor that runs the non-blocking per-iteration + // dispatch (poll + chunked_copy_and_enqueue) as a coroutine, off the + // trainer thread. Named so its thread is identifiable in traces. + dispatch_executor_ = std::make_unique( + 1, + std::make_unique( + fmt::format("RESDispatch.{}", unique_id_))); + + XLOG(INFO) << "[TBE_ID" << unique_id_ + << "] Starting RES ship executor with " << res_num_consumers_ + << " threads" + << ", chunk_size=" << res_chunk_size_ + << ", copy_threads=" << res_num_copy_threads_; + // Push model: ship tasks are submitted onto this executor (one per enqueued + // StreamQueueItem) and its workers wake on submit -- no polled queue, no + // raw std::thread that could std::terminate on an escaped exception. + // + // Ordering caveat: with res_num_consumers_ > 1, ship tasks run + // 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). res_num_consumers_ == 1 is + // ordered/safe. TODO(T281413204): proper fix is to carry a per-row + // iteration/version and keep-newest in the store. + consumer_executor_ = std::make_unique( + res_num_consumers_, + std::make_unique( + fmt::format("RESShip.{}", unique_id_))); } #endif } RawEmbeddingStreamer::~RawEmbeddingStreamer() { - stop_ = true; #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { - join_stream_tensor_copy_thread(); - join_weights_stream_thread(); + join_dispatch(); + if (dispatch_executor_ != nullptr) { + dispatch_executor_->join(); + } + join_chunk_copy_threads(); + // Producers (dispatch + copy threads) are all joined above, so no further + // ship tasks will be submitted. join() drains the in-flight ones before the + // executor's threads stop, then destroys members in a safe order. + if (consumer_executor_ != nullptr) { + consumer_executor_->join(); + } } #endif } @@ -232,93 +289,216 @@ void RawEmbeddingStreamer::stream( auto rec = torch::autograd::profiler::record_function_enter_new( "## RawEmbeddingStreamer::stream_callback ##"); if (!require_tensor_copy) { - StreamQueueItem stream_item( + submit_stream_item(StreamQueueItem( indices, weights, std::move(identities), std::move(runtime_meta), - count); - weights_to_stream_queue_.enqueue(stream_item); + count)); return; } + auto poll_flag = [this, copy_done_flag]() { + return poll_copy_done_flag(copy_done_flag); + }; + if (blocking_tensor_copy) { - 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 blocking: poll timed out after " - << kCopyDonePollTimeoutUs / 1'000'000 << "s"; - return; - } - } - *ptr = 0; // Reset for next iteration - } else { - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] copy_done_flag not provided, skipping wait (blocking)"; + 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_chunk_copy_threads(); return; } - // Make sure the previous thread is done before starting a new one - 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; - } - } - *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); - }); + // 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_dispatch(); + // Dispatch runs as a coroutine on the persistent size-1 executor; folly + // captures any exception into dispatch_future_, which is logged when the + // future is waited in join_dispatch() (log-and-continue, never terminate). + dispatch_future_ = folly::coro::co_withExecutor( + dispatch_executor_.get(), + dispatch_copy_task( + indices, + weights, + identities, + runtime_meta, + count, + copy_done_flag)) + .start(); rec->record.end(); #endif } -void RawEmbeddingStreamer::join_stream_tensor_copy_thread() { +void RawEmbeddingStreamer::join_dispatch() { #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(); + "## RawEmbeddingStreamer::join_dispatch ##"); + // Wait the previous dispatch. Log-and-continue: an exception the dispatch + // deferred into the future must not escape (would std::terminate the + // trainer). + if (dispatch_future_.valid()) { + try { + std::move(dispatch_future_).get(); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] stream dispatcher caught exception: " << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] stream dispatcher caught unknown exception"; + } + dispatch_future_ = folly::makeSemiFuture(); } + // The real torn-row barrier: iter i's copy must finish reading the source + // cache rows before iter i+1 overwrites them. + join_chunk_copy_threads(); rec->record.end(); #endif } #ifdef FBGEMM_FBCODE +void RawEmbeddingStreamer::join_chunk_copy_threads() { + for (auto& t : chunk_copy_threads_) { + if (t && t->joinable()) { + t->join(); + } + } + chunk_copy_threads_.clear(); +} + +void RawEmbeddingStreamer::submit_stream_item(StreamQueueItem item) { + // Push model: hand the item to a ship worker that wakes on submit. folly + // captures any task exception (so an escaped throw can't std::terminate the + // trainer); we still wrap the body to log a transient tensor_stream failure + // and to keep the depth-gauge / periodic-log behavior of the old consumer. + consumer_executor_->add([this, item = std::move(item)]() mutable { + 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( + "stream_mpmc_depth", + static_cast(consumer_executor_->getTaskQueueSize())); + } + XLOG_EVERY_MS(INFO, 60000) + << "[TBE_ID" << unique_id_ << "] end stream queue size: " + << consumer_executor_->getTaskQueueSize() << " stream takes " + << stop_watch.elapsed().count() << "ms" + << " rows=" << item.indices.size(0); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] ship task caught exception: " << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] ship task caught unknown exception"; + } + }); +} + +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) { + 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()) { + 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); + submit_stream_item(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(); +} + +bool RawEmbeddingStreamer::poll_copy_done_flag( + const std::optional& copy_done_flag) { + 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 poll timed out after " + << kCopyDonePollTimeoutUs / 1'000'000 << "s"; + return false; + } + } + *ptr = 0; + } + return true; +} + +folly::coro::Task RawEmbeddingStreamer::dispatch_copy_task( + at::Tensor indices, + at::Tensor weights, + std::optional identities, + std::optional runtime_meta, + at::Tensor count, + std::optional copy_done_flag) { + if (!poll_copy_done_flag(copy_done_flag)) { + co_return; + } + chunked_copy_and_enqueue( + indices, weights, identities, runtime_meta, count, chunk_copy_threads_); + co_return; +} + folly::coro::Task RawEmbeddingStreamer::tensor_stream( const at::Tensor& indices, const at::Tensor& weights, @@ -439,41 +619,34 @@ 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(); + // TESTING only: drop the ship executor to 0 worker threads so subsequently + // submitted tasks accumulate in its queue (observable via + // get_weights_to_stream_queue_size()) instead of being shipped. Mirrors the + // old "stop the consumer threads" behavior; unlike join() the executor still + // accepts newly submitted tasks. + if (consumer_executor_ != nullptr) { + consumer_executor_->setNumThreads(0); } } uint64_t RawEmbeddingStreamer::get_weights_to_stream_queue_size() { - return weights_to_stream_queue_.size(); + return consumer_executor_ != nullptr ? consumer_executor_->getTaskQueueSize() + : 0; } #endif 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..ab4062aa19 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 @@ -87,6 +87,9 @@ auto raw_embedding_streamer = bool, int64_t, int64_t, + int64_t, + int64_t, + int64_t, std::vector, std::vector, std::vector>(), @@ -96,6 +99,9 @@ auto raw_embedding_streamer = torch::arg("enable_raw_embedding_streaming") = false, torch::arg("res_store_shards") = 0, torch::arg("res_server_port") = 0, + torch::arg("res_chunk_size") = 500000, + torch::arg("res_num_consumers") = 8, + torch::arg("res_num_copy_threads") = 4, torch::arg("table_names") = torch::List(), torch::arg("table_offsets") = torch::List(), torch::arg("table_sizes") = torch::List(), @@ -114,8 +120,6 @@ auto raw_embedding_streamer = torch::arg("blocking_tensor_copy"), torch::arg("copy_done_flag") = std::nullopt, }) - .def( - "join_stream_tensor_copy_thread", - &fbgemm_gpu::RawEmbeddingStreamer::join_stream_tensor_copy_thread); + .def("join_dispatch", &fbgemm_gpu::RawEmbeddingStreamer::join_dispatch); } // namespace 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..f4a0418654 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 @@ -47,6 +47,9 @@ getRawEmbeddingStreamer( enable_raw_embedding_streaming, 3, // res_store_shards 0, // res_server_port + 500000, // res_chunk_size + 8, // res_num_consumers + 4, // res_num_copy_threads table_names, table_offsets, table_sizes); @@ -84,7 +87,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, + /*res_chunk_size=*/0, + /*res_num_consumers=*/8, + /*res_num_copy_threads=*/4, + /*table_names=*/{}, + /*table_offsets=*/{}, + /*table_sizes=*/{})); +} + TEST(RawEmbeddingStreamerTest, TestTensorStream) { std::vector table_names = {"tb1", "tb2", "tb3"}; std::vector table_offsets = {0, 100, 300}; @@ -178,14 +506,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_dispatch() 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(); + streamer->join_dispatch(); 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 +606,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, TestCoSetEmbeddingsThrowPropagates) { +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, 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 +718,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))); } @@ -429,7 +889,10 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopyDoneFlagNonBlockingCopy) { copy_done_flag); // Wait for the async thread to complete - streamer->join_stream_tensor_copy_thread(); + streamer->join_dispatch(); 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..f49360eaa3 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 @@ -92,6 +92,14 @@ EmbeddingKVDB::EmbeddingKVDB( enable_raw_embedding_streaming, res_store_shards, res_server_port, + // 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, std::move(table_names), std::move(table_offsets), table_sizes)) { @@ -332,9 +340,8 @@ void EmbeddingKVDB::stream_sync_cuda() { "## EmbeddingKVDB::stream_sync_cuda ##"); // take reference to self to avoid lifetime issues. auto self = shared_from_this(); - std::function* functor = new std::function([=]() { - self->raw_embedding_streamer_->join_stream_tensor_copy_thread(); - }); + std::function* functor = new std::function( + [=]() { self->raw_embedding_streamer_->join_dispatch(); }); AT_CUDA_CHECK(cudaLaunchHostFunc( at::cuda::getCurrentCUDAStream(), kv_db_utils::cuda_host_func, functor)); rec->record.end();