From a928b7bef9d2aad2fdc2ac756569d8fae6203992 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 1/7] Chunk the D2H copy across N copy threads (#6068) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2968 Speed up the RES C++ streamer's device->host copy by parallelizing it. Previously stream() copied all updated rows to CPU on a single thread and enqueued one item. This replaces that serial copy with a chunked copy across up to `kNumCopyThreads` threads: tensor_copy_chunk() copies a [start, end) row range, and the per-thread tiling is factored into a pure, unit-testable computeChunkRanges(). Differential Revision: D113594180 --- ...t_table_batched_embeddings_ops_training.py | 4 +- .../raw_embedding_streamer.h | 57 ++- .../raw_embedding_streamer.cpp | 299 +++++++++----- .../split_embeddings_cache_ops.cpp | 4 +- .../tests/raw_embedding_streamer_test.cpp | 372 +++++++++++++++++- .../kv_db_table_batched_embeddings.cpp | 5 +- 6 files changed, 604 insertions(+), 137 deletions(-) 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..42d575e39c 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 @@ -4529,7 +4529,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_thread() prefetched_info = self.prefetched_info_list.pop(0) updated_locations = torch.ops.fbgemm.lxu_cache_lookup( prefetched_info.linear_unique_cache_indices, @@ -4580,7 +4580,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..a9091fcfc6 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 { @@ -56,10 +57,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 the stream thread + /// that streams out to the thrift server (co-located on same host + /// now). The copy is split into <= kChunkSize-row chunks across up to + /// kNumCopyThreads 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 +89,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_thread(); #ifdef FBGEMM_FBCODE folly::coro::Task tensor_stream( @@ -97,16 +100,6 @@ 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 @@ -130,20 +123,44 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { #ifdef FBGEMM_FBCODE std::unique_ptr weights_stream_thread_; folly::UMPSCQueue weights_to_stream_queue_; - std::unique_ptr stream_tensor_copy_thread_; + // 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_chunk_copy_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..b6c1a87b12 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,12 @@ 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; + +// Max rows copied into one enqueued chunk. +constexpr size_t kChunkSize = 500000; +// Parallel chunk-copy threads spawned per non-blocking stream() call. +constexpr size_t kNumCopyThreads = 4; /* * Get the thrift client to the training parameter server service @@ -64,86 +69,114 @@ 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)); FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE( - weights.scalar_type(), "tensor_copy", [&] { + weights.scalar_type(), "tensor_copy_chunk", [&] { using value_t = scalar_t; FBGEMM_DISPATCH_INTEGRAL_TYPES( - indices.scalar_type(), "tensor_copy", [&] { + indices.scalar_type(), "tensor_copy_chunk", [&] { 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(); + indices.const_data_ptr() + start_row, + indices.const_data_ptr() + end_row, + new_indices.mutable_data_ptr()); std::copy( - weights_addr, - weights_addr + num_sets * weights.size(1), - new_weights_addr); // dst_start + weights.const_data_ptr() + + start_row * weights.size(1), + weights.const_data_ptr() + end_row * weights.size(1), + new_weights.mutable_data_ptr()); 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(); + identities->scalar_type(), "tensor_copy_chunk", [&] { + using id_t = scalar_t; std::copy( - identities_addr, - identities_addr + num_sets * identities->size(1), - new_identities_addr); // dst_start + 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", [&] { - 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(); + runtime_meta->scalar_type(), "tensor_copy_chunk", [&] { + using rm_t = scalar_t; std::copy( - runtime_meta_addr, - runtime_meta_addr + num_sets * runtime_meta->size(1), - new_runtime_meta_addr); // dst_start + 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() = num_sets; + *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, @@ -210,7 +243,8 @@ RawEmbeddingStreamer::~RawEmbeddingStreamer() { stop_ = true; #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { - join_stream_tensor_copy_thread(); + join_dispatch_thread(); + join_chunk_copy_threads(); join_weights_stream_thread(); } #endif @@ -241,7 +275,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,76 +283,147 @@ 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_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; + // 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_thread(); + 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 } -void RawEmbeddingStreamer::join_stream_tensor_copy_thread() { +void RawEmbeddingStreamer::join_dispatch_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(); + "## RawEmbeddingStreamer::join_dispatch_thread ##"); + if (dispatch_thread_ != nullptr && dispatch_thread_->joinable()) { + dispatch_thread_->join(); } + 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::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, kChunkSize, kNumCopyThreads); + + 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); + 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(); +} + folly::coro::Task RawEmbeddingStreamer::tensor_stream( const at::Tensor& indices, const at::Tensor& weights, @@ -451,20 +556,6 @@ folly::coro::Task RawEmbeddingStreamer::tensor_stream( 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; 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..ed54e97e28 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 @@ -115,7 +115,7 @@ auto raw_embedding_streamer = torch::arg("copy_done_flag") = std::nullopt, }) .def( - "join_stream_tensor_copy_thread", - &fbgemm_gpu::RawEmbeddingStreamer::join_stream_tensor_copy_thread); + "join_dispatch_thread", + &fbgemm_gpu::RawEmbeddingStreamer::join_dispatch_thread); } // 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..ea6310baa8 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 @@ -84,6 +84,312 @@ 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, TestTensorStream) { std::vector table_names = {"tb1", "tb2", "tb3"}; @@ -178,14 +484,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_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(); + streamer->join_dispatch_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,8 +584,13 @@ 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(); } @@ -429,7 +786,10 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopyDoneFlagNonBlockingCopy) { copy_done_flag); // Wait for the async thread to complete - streamer->join_stream_tensor_copy_thread(); + streamer->join_dispatch_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..bbd9b22c78 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 @@ -332,9 +332,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_thread(); }); AT_CUDA_CHECK(cudaLaunchHostFunc( at::cuda::getCurrentCUDAStream(), kv_db_utils::cuda_host_func, functor)); rec->record.end(); From e7f361ff9f85245ee3c6889c6990d03f9af0f818 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 2/7] Drain the stream queue across N consumer threads (#6069) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2971 Speed up the RES C++ streamer's ship stage by parallelizing it. Previously a single stream thread drained the queue, so the setEmbeddings RPCs to the PS ran one at a time. This diff replaces it with kNumConsumerThreads consumers on a folly::UMPMCQueue that ship concurrently. Differential Revision: D113599253 --- .../raw_embedding_streamer.h | 15 ++- .../raw_embedding_streamer.cpp | 106 ++++++++++++------ .../tests/raw_embedding_streamer_test.cpp | 95 ++++++++++++++-- 3 files changed, 167 insertions(+), 49 deletions(-) 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 a9091fcfc6..7493e9f781 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 @@ -59,8 +59,8 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { /// from for the first elements in the tensor. /// 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 the stream thread - /// that streams out to the thrift server (co-located on same host + /// 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 <= kChunkSize-row chunks across up to /// kNumCopyThreads threads. /// @@ -102,8 +102,9 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { std::optional runtime_meta); /* - * 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. @@ -121,8 +122,9 @@ 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_; + // 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. @@ -136,6 +138,7 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { ods_logger_; void join_chunk_copy_threads(); + void join_consumer_threads(); void chunked_copy_and_enqueue( const at::Tensor& indices, const at::Tensor& weights, 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 b6c1a87b12..1664c264d7 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -33,6 +33,8 @@ constexpr int64_t kCopyDonePollTimeoutUs = 10'000'000; // Max rows copied into one enqueued chunk. constexpr size_t kChunkSize = 500000; +// Threads draining the queue and shipping via co_setEmbeddings. +constexpr size_t kNumConsumerThreads = 8; // Parallel chunk-copy threads spawned per non-blocking stream() call. constexpr size_t kNumCopyThreads = 4; @@ -205,36 +207,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; - } - 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(INFO) << "[TBE_ID" << unique_id_ << "] Starting " + << kNumConsumerThreads << " consumer threads" + << ", chunk_size=" << kChunkSize + << ", copy_threads=" << kNumCopyThreads; + // Ordering caveat: with kNumConsumerThreads > 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 < kNumConsumerThreads; ++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( + "stream_mpmc_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"; + } } - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] end stream queue size: " << post_dequeue_depth - << " stream takes " << stop_watch.elapsed().count() << "ms"; - } - }); + })); + } } #endif } @@ -245,7 +270,7 @@ RawEmbeddingStreamer::~RawEmbeddingStreamer() { if (enable_raw_embedding_streaming_) { join_dispatch_thread(); join_chunk_copy_threads(); - join_weights_stream_thread(); + join_consumer_threads(); } #endif } @@ -361,6 +386,15 @@ void RawEmbeddingStreamer::join_chunk_copy_threads() { 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, @@ -544,23 +578,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::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/tests/raw_embedding_streamer_test.cpp b/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp index ea6310baa8..11785a5b48 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 @@ -594,7 +594,85 @@ TEST(RawEmbeddingStreamerTest, TestStreamE2E) { 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}; @@ -618,7 +696,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))); } From 1e6e42b87b15e510adb8c3a1e33bb1b4b40b1c70 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 3/7] Make RES chunk size + ship/copy thread counts config-tunable (#6075) Summary: X-link: https://github.com/meta-pytorch/torchrec/pull/4464 X-link: https://github.com/facebookresearch/FBGEMM/pull/2976 Make the RES streamer's chunk size and ship/copy thread counts tunable from config instead of compile-time constants. Previously kChunkSize / kNumConsumerThreads / kNumCopyThreads were constexpr in raw_embedding_streamer.cpp, so tuning them meant a base rebuild. This plumbs them as res_chunk_size / res_num_consumers / res_num_copy_threads from TableBatchedEmbeddingConfig through fused_params to the RawEmbeddingStreamer ctor, mirroring the existing res_store_shards. Defaults stay 500000/8/4, so behavior is unchanged until overridden Differential Revision: D113532922 --- ...t_table_batched_embeddings_ops_training.py | 6 +++ .../raw_embedding_streamer.h | 10 ++++- .../raw_embedding_streamer.cpp | 41 ++++++++++++------- .../split_embeddings_cache_ops.cpp | 6 +++ .../tests/raw_embedding_streamer_test.cpp | 22 ++++++++++ .../kv_db_table_batched_embeddings.cpp | 8 ++++ 6 files changed, 77 insertions(+), 16 deletions(-) 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 42d575e39c..12d2fd82e8 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, 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 7493e9f781..e098f25249 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 @@ -49,6 +49,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); @@ -61,8 +64,8 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { /// 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 <= kChunkSize-row chunks across up to - /// kNumCopyThreads threads. + /// 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 @@ -117,6 +120,9 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { #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_; 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 1664c264d7..cfece20124 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -31,13 +31,6 @@ namespace { // Timeout for copy_done_flag polling loop (microseconds). constexpr int64_t kCopyDonePollTimeoutUs = 10'000'000; -// Max rows copied into one enqueued chunk. -constexpr size_t kChunkSize = 500000; -// Threads draining the queue and shipping via co_setEmbeddings. -constexpr size_t kNumConsumerThreads = 8; -// Parallel chunk-copy threads spawned per non-blocking stream() call. -constexpr size_t kNumCopyThreads = 4; - /* * Get the thrift client to the training parameter server service * There is a destruction double free issue when wrapping the member @@ -184,6 +177,9 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( 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) @@ -192,12 +188,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_; @@ -207,11 +220,11 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( ods_logger_ = std::make_unique(); - XLOG(INFO) << "[TBE_ID" << unique_id_ << "] Starting " - << kNumConsumerThreads << " consumer threads" - << ", chunk_size=" << kChunkSize - << ", copy_threads=" << kNumCopyThreads; - // Ordering caveat: with kNumConsumerThreads > 1, these consumers drain the + 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 @@ -219,7 +232,7 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( // 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 < kNumConsumerThreads; ++ci) { + 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(); @@ -404,7 +417,7 @@ void RawEmbeddingStreamer::chunked_copy_and_enqueue( std::vector>& target_copy_threads) { const auto num_rows = get_maybe_uvm_scalar(count); const auto thread_chunks = - computeChunkRanges(num_rows, kChunkSize, kNumCopyThreads); + 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()) { 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 ed54e97e28..ed51afc6af 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(), 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 11785a5b48..3f47afbc7d 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); @@ -391,6 +394,25 @@ TEST(RawEmbeddingStreamerTest, ComputeChunkRangesZeroThreadsOrChunkSizeEmpty) { } #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}; 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 bbd9b22c78..72a8a197b0 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)) { From 86481a41a90b851fd27e20896de462ddb8fd9c57 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 4/7] Un-nest tensor_copy_chunk dispatch lambdas (#6074) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2975 The four tensor copies in tensor_copy_chunk are independent, so dispatch each under its own sibling FBGEMM_DISPATCH_* instead of nesting weights -> indices -> identities -> runtime_meta. Two wins, no behavior change: 1. Readability: the flat structure drops the value_t / index_t / id_t / rm_t aliases that existed only to dodge scalar_t name-shadowing between the nested lambdas -- each copy now just uses scalar_t and reads top-to-bottom instead of 3 levels deep. 2. Fewer template instantiations: nesting stamps each inner copy once per outer type (multiplicative, ~4 x 2); siblings stamp each once per its own type set (additive, 4 + 2) -> smaller binary, faster compile. Differential Revision: D113533821 --- .../raw_embedding_streamer.cpp | 74 +++++++++---------- 1 file changed, 36 insertions(+), 38 deletions(-) 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 cfece20124..909c837f51 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -91,47 +91,45 @@ fbgemm_gpu::StreamQueueItem tensor_copy_chunk( } 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_chunk", [&] { - using value_t = scalar_t; - FBGEMM_DISPATCH_INTEGRAL_TYPES( - indices.scalar_type(), "tensor_copy_chunk", [&] { - using index_t = scalar_t; - std::copy( - indices.const_data_ptr() + start_row, - indices.const_data_ptr() + end_row, - new_indices.mutable_data_ptr()); - 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()); - if (identities.has_value()) { - FBGEMM_DISPATCH_INTEGRAL_TYPES( - identities->scalar_type(), "tensor_copy_chunk", [&] { - using id_t = scalar_t; - 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", [&] { - using rm_t = scalar_t; - 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()); - }); - } - }); + 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()); }); + 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}; From 45897003234b9553938901b56685d9c61f3ff6b5 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 5/7] Replace raw dispatcher thread with a named CPUThreadPoolExecutor Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2978 Replace the per-iteration raw std::thread in the non-blocking stream() dispatch with a persistent named size-1 folly::CPUThreadPoolExecutor + SemiFuture. Kills per-iter thread create/join churn, names the thread in traces, and removes the raw-thread std::terminate risk (folly captures exceptions into the future) Differential Revision: D113669440 --- ...t_table_batched_embeddings_ops_training.py | 2 +- .../raw_embedding_streamer.h | 30 +++- .../raw_embedding_streamer.cpp | 128 ++++++++++++------ .../split_embeddings_cache_ops.cpp | 4 +- .../tests/raw_embedding_streamer_test.cpp | 6 +- .../kv_db_table_batched_embeddings.cpp | 2 +- 6 files changed, 117 insertions(+), 55 deletions(-) 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 12d2fd82e8..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 @@ -4535,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_dispatch_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, 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 e098f25249..ba04e8746d 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,12 +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 @@ -95,7 +100,7 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { * Join the pending dispatch (and the copy threads it spawned), making sure it * is properly finished before creating new. */ - void join_dispatch_thread(); + void join_dispatch(); #ifdef FBGEMM_FBCODE folly::coro::Task tensor_stream( @@ -135,7 +140,10 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { // 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_; + // 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. @@ -152,6 +160,24 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { 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 }; 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 909c837f51..11606f754e 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" @@ -218,6 +222,14 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( ods_logger_ = std::make_unique(); + // 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_num_consumers_ << " consumer threads" << ", chunk_size=" << res_chunk_size_ @@ -279,7 +291,10 @@ RawEmbeddingStreamer::~RawEmbeddingStreamer() { stop_ = true; #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { - join_dispatch_thread(); + join_dispatch(); + if (dispatch_executor_ != nullptr) { + dispatch_executor_->join(); + } join_chunk_copy_threads(); join_consumer_threads(); } @@ -312,21 +327,7 @@ void RawEmbeddingStreamer::stream( return; } 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; - 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; + return poll_copy_done_flag(copy_done_flag); }; if (blocking_tensor_copy) { @@ -346,42 +347,45 @@ void RawEmbeddingStreamer::stream( // 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_thread(); - 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"; - } - }); + 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_dispatch_thread() { +void RawEmbeddingStreamer::join_dispatch() { #ifdef FBGEMM_FBCODE auto rec = torch::autograd::profiler::record_function_enter_new( - "## RawEmbeddingStreamer::join_dispatch_thread ##"); - if (dispatch_thread_ != nullptr && dispatch_thread_->joinable()) { - dispatch_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 @@ -469,6 +473,40 @@ void RawEmbeddingStreamer::chunked_copy_and_enqueue( << " 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, 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 ed51afc6af..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 @@ -120,8 +120,6 @@ auto raw_embedding_streamer = torch::arg("blocking_tensor_copy"), torch::arg("copy_done_flag") = std::nullopt, }) - .def( - "join_dispatch_thread", - &fbgemm_gpu::RawEmbeddingStreamer::join_dispatch_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 3f47afbc7d..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 @@ -507,11 +507,11 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopy) { EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); // Test non-blocking tensor copy. The copy runs on the dispatcher, so we must - // join_dispatch_thread() before checking the queue -- asserting the size + // 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); - streamer->join_dispatch_thread(); + streamer->join_dispatch(); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 2); } @@ -889,7 +889,7 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopyDoneFlagNonBlockingCopy) { copy_done_flag); // Wait for the async thread to complete - streamer->join_dispatch_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. 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 72a8a197b0..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 @@ -341,7 +341,7 @@ void 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_dispatch_thread(); }); + [=]() { self->raw_embedding_streamer_->join_dispatch(); }); AT_CUDA_CHECK(cudaLaunchHostFunc( at::cuda::getCurrentCUDAStream(), kv_db_utils::cuda_host_func, functor)); rec->record.end(); From abae9b863b1553c9da51c676eb7419b3f7417ba1 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 6/7] Replace raw consumer threads + polled queue with a push executor (#6076) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2979 The consumer/ship path drained a hand-rolled folly::UMPMCQueue (weights_to_stream_queue_) with a pool of raw std::thread consumers that ran a 10ms-backoff polling loop try_dequeue(), sleep(10ms) on empty, retry and exited via a stop_ latch. This diff replaces both the queue and the raw thread pool with a persistent folly::CPUThreadPoolExecutor: each ship item is posted with add(...), and the dtor join()s the executor so all pending + in-flight ship tasks drain before teardown. Differential Revision: D113669447 --- .../raw_embedding_streamer.h | 18 +-- .../raw_embedding_streamer.cpp | 125 +++++++++--------- 2 files changed, 72 insertions(+), 71 deletions(-) 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 ba04e8746d..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 @@ -110,16 +110,15 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { std::optional runtime_meta); /* - * 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. + * 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 @@ -133,9 +132,10 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { std::vector table_offsets_; at::Tensor table_sizes_; #ifdef FBGEMM_FBCODE - // Multi-threaded consumers for tensor_stream() RPCs. - std::vector> consumer_threads_; - folly::UMPMCQueue weights_to_stream_queue_; + // 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. @@ -152,7 +152,9 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { ods_logger_; void join_chunk_copy_threads(); - void join_consumer_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, 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 11606f754e..e09ed887c7 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -230,65 +230,32 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( std::make_unique( fmt::format("RESDispatch.{}", unique_id_))); - XLOG(INFO) << "[TBE_ID" << unique_id_ << "] Starting " << res_num_consumers_ - << " consumer threads" + 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_; - // Ordering caveat: with res_num_consumers_ > 1, these consumers drain the - // queue concurrently, so arrival order at the PS is NOT enqueue (iteration) + // 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). 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( - "stream_mpmc_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"; - } - } - })); - } + // 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_dispatch(); @@ -296,7 +263,12 @@ RawEmbeddingStreamer::~RawEmbeddingStreamer() { dispatch_executor_->join(); } join_chunk_copy_threads(); - join_consumer_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 } @@ -317,13 +289,12 @@ 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]() { @@ -401,13 +372,34 @@ void RawEmbeddingStreamer::join_chunk_copy_threads() { chunk_copy_threads_.clear(); } -void RawEmbeddingStreamer::join_consumer_threads() { - for (auto& t : consumer_threads_) { - if (t && t->joinable()) { - t->join(); +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"; } - } - consumer_threads_.clear(); + }); } void RawEmbeddingStreamer::chunked_copy_and_enqueue( @@ -452,7 +444,7 @@ void RawEmbeddingStreamer::chunked_copy_and_enqueue( 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)); + submit_stream_item(std::move(chunk_item)); rows_done += (e - s); } XLOG_EVERY_MS(INFO, 15000) @@ -642,12 +634,19 @@ folly::coro::Task RawEmbeddingStreamer::tensor_stream( } void RawEmbeddingStreamer::join_weights_stream_thread() { - stop_ = true; - join_consumer_threads(); + // 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 From 91f32ee5879c5e68ea5a5159fa1961635a543509 Mon Sep 17 00:00:00 2001 From: Joey Yang Date: Mon, 27 Jul 2026 12:44:49 -0700 Subject: [PATCH 7/7] replace raw copy threads with a persistent copy_executor_ Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2982 Swap the per-stream raw std::thread copy pool for a persistent folly::CPUThreadPoolExecutor; chunk copies fan out + join via collectAllRange; We still join so each copy finishes before the next iteration so this only removes per-stream thread churn. Differential Revision: D113752426 --- .../raw_embedding_streamer.h | 64 +++-- .../raw_embedding_streamer.cpp | 220 ++++++++-------- .../tests/raw_embedding_streamer_test.cpp | 242 ++++++++++-------- 3 files changed, 277 insertions(+), 249 deletions(-) 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 4bb4dfe3bb..aedcc5c3c1 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 @@ -69,8 +69,8 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { /// 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. + /// now). The copy is split into <= res_chunk_size-row chunks that run + /// concurrently across the res_num_copy_threads-worker copy pool. /// /// This is used in cuda stream callback, which doesn't require to be /// serialized with other callbacks, thus a separate thread is used to @@ -97,8 +97,9 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { std::optional copy_done_flag = std::nullopt); /* - * Join the pending dispatch (and the copy threads it spawned), making sure it - * is properly finished before creating new. + * Join the pending dispatch future (which resolves only after all copies for + * the previous iteration have been enqueued), making sure it is properly + * finished before creating new. */ void join_dispatch(); @@ -136,10 +137,12 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { // 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 pool that runs the per-chunk tensor copies (CPU-bound + // std::copy), sized res_num_copy_threads_ so chunks run concurrently. 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::unique_ptr copy_executor_; // 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_; @@ -151,17 +154,33 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder { TrainingPsOdsLogger> 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, + // Copies `count` rows of the source tensors into CPU chunks and enqueues each + // via submit_stream_item, fanning the per-chunk copies out across + // copy_executor_ and awaiting them all (the torn-row barrier). A coroutine so + // the non-blocking dispatch can co_await it; args are by value so nothing + // dangles once a chunk is scheduled on copy_executor_. + folly::coro::Task chunked_copy_and_enqueue( + at::Tensor indices, + at::Tensor weights, std::optional identities, std::optional runtime_meta, - const at::Tensor& count, - std::vector>& target_copy_threads); + at::Tensor count); + // Copies one chunk [start, end) and enqueues it via submit_stream_item. Runs + // to completion on a single copy_executor_ worker; the per-chunk try/catch + // keeps an exception from escaping into collectAllRange (which awaits every + // sibling to completion and then rethrows one exception onto the blocking + // caller / dispatch future). By value for the same reason as + // chunked_copy_and_enqueue. + folly::coro::Task copy_chunk_task( + at::Tensor indices, + at::Tensor weights, + std::optional identities, + std::optional runtime_meta, + int64_t start, + int64_t end); // 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 @@ -191,13 +210,12 @@ fbgemm_gpu::StreamQueueItem tensor_copy_chunk( 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); +// Tiles [0, num_rows) into flat [start, end) chunks of size <= chunk_size, +// contiguous and non-overlapping (their union is the whole range). One task per +// chunk is submitted to copy_executor_, which load-balances them across its +// workers. Pure/build-agnostic so it is unit-testable without a GPU or +// FBGEMM_FBCODE. +std::vector> computeChunks( + int64_t num_rows, + size_t chunk_size); } // 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 e09ed887c7..30d19e0550 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -9,6 +9,8 @@ #ifdef FBGEMM_FBCODE #include #include +#include +#include #include #include #include @@ -139,39 +141,20 @@ fbgemm_gpu::StreamQueueItem tensor_copy_chunk( 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; +std::vector> computeChunks( + int64_t num_rows, + size_t chunk_size) { + // Split [0, num_rows) into flat [start, end) chunks of <= chunk_size rows; + // chunks are contiguous, non-overlapping, and their union is the whole range. + std::vector> chunks; + if (num_rows <= 0 || chunk_size == 0) { + return 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)); - } + for (int64_t s = 0; s < num_rows; s += static_cast(chunk_size)) { + const int64_t e = std::min(s + static_cast(chunk_size), num_rows); + chunks.emplace_back(s, e); } - return thread_chunks; + return chunks; } RawEmbeddingStreamer::RawEmbeddingStreamer( @@ -202,9 +185,10 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( // 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. + // drain threads (queue grows unbounded), res_chunk_size=0 makes + // computeChunks return empty (enqueues nothing), res_num_copy_threads=0 + // sizes copy_executor_ to zero workers (copy tasks never run). 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=", @@ -251,6 +235,15 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( res_num_consumers_, std::make_unique( fmt::format("RESShip.{}", unique_id_))); + + // Persistent pool that runs the per-chunk tensor copies (CPU-bound + // std::copy) off the trainer thread. Sized res_num_copy_threads_ so N + // chunks run concurrently -- the same N-way parallelism the old + // one-std::thread-per-group model gave. + copy_executor_ = std::make_unique( + res_num_copy_threads_, + std::make_unique( + fmt::format("RESCopy.{}", unique_id_))); } #endif } @@ -258,13 +251,19 @@ RawEmbeddingStreamer::RawEmbeddingStreamer( RawEmbeddingStreamer::~RawEmbeddingStreamer() { #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { + // Drain in producer order: dispatch schedules onto copy_executor_, and copy + // groups submit ship tasks onto consumer_executor_, so join dispatch -> + // copy -> consumer. A wrong order risks joining an executor while a + // producer still schedules onto it. 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 + if (copy_executor_ != nullptr) { + copy_executor_->join(); + } + // Producers (dispatch + copy) 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(); @@ -305,19 +304,23 @@ void RawEmbeddingStreamer::stream( if (!poll_flag()) { return; } - chunked_copy_and_enqueue( + // blockingWait blocks the trainer thread until every copy group (bound to + // copy_executor_) has finished and enqueued -- the torn-row barrier. The + // trainer thread is not a copy_executor_ worker, so this cannot self- + // deadlock. The per-group try/catch keeps collectAllRange from rethrowing + // here. + folly::coro::blockingWait(chunked_copy_and_enqueue( indices, weights, std::move(identities), std::move(runtime_meta), - count, - chunk_copy_threads_); - join_chunk_copy_threads(); + count)); return; } - // 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. + // Non-blocking: join the previous dispatch (which awaited its copies on + // copy_executor_) before starting a new one. join_dispatch() is the + // serializer: it guarantees iter i's copies 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 @@ -340,8 +343,11 @@ void RawEmbeddingStreamer::join_dispatch() { #ifdef FBGEMM_FBCODE auto rec = torch::autograd::profiler::record_function_enter_new( "## 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 + // Wait the previous dispatch. The dispatch future resolves only after + // chunked_copy_and_enqueue's collectAllRange completes, so this is also the + // torn-row barrier: iter i's copies finish reading the source cache rows + // before iter i+1 overwrites them. Log-and-continue: an exception the + // dispatch deferred into the future must not escape (would std::terminate the // trainer). if (dispatch_future_.valid()) { try { @@ -355,23 +361,11 @@ void RawEmbeddingStreamer::join_dispatch() { } 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 @@ -402,67 +396,63 @@ void RawEmbeddingStreamer::submit_stream_item(StreamQueueItem item) { }); } -void RawEmbeddingStreamer::chunked_copy_and_enqueue( - const at::Tensor& indices, - const at::Tensor& weights, +folly::coro::Task RawEmbeddingStreamer::copy_chunk_task( + at::Tensor indices, + 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; + int64_t start, + int64_t end) { + // Guard the copy body so a per-chunk failure logs instead of escaping into + // collectAllRange (which would cancel siblings and rethrow onto the blocking + // caller / std::terminate). + try { + auto chunk_item = tensor_copy_chunk( + indices, weights, identities, runtime_meta, start, end); + submit_stream_item(std::move(chunk_item)); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ << "] copy_chunk [" << start << ", " + << end << ") caught exception: " << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ << "] copy_chunk [" << start << ", " + << end << ") caught unknown exception"; } + co_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"; - } - })); - } +folly::coro::Task RawEmbeddingStreamer::chunked_copy_and_enqueue( + at::Tensor indices, + at::Tensor weights, + std::optional identities, + std::optional runtime_meta, + at::Tensor count) { + const auto num_rows = get_maybe_uvm_scalar(count); + const auto chunks = computeChunks(num_rows, res_chunk_size_); XLOG_EVERY_MS(INFO, 15000) << "[RES] chunked_copy tbe=" << unique_id_ << " rows=" << num_rows - << " threads=" << thread_chunks.size(); + << " chunks=" << chunks.size(); + if (chunks.empty()) { + co_return; + } + + // One copy task per chunk. Chunk boundaries live entirely in computeChunks, + // so the enqueued row set is identical regardless of how the tasks run. + // co_withExecutor binds each task to copy_executor_ (a bare Task would + // inherit the size-1 dispatch_executor_ and serialize the copies); + // collectAllRange then runs the tasks across the pool -- the workers pull + // them, so the pool load-balances -- and completes only after every one is + // done (the barrier). + std::vector> tasks; + tasks.reserve(chunks.size()); + for (const auto& [start, end] : chunks) { + tasks.push_back( + folly::coro::co_withExecutor( + copy_executor_.get(), + copy_chunk_task( + indices, weights, identities, runtime_meta, start, end))); + } + co_await folly::coro::collectAllRange(std::move(tasks)); + co_return; } bool RawEmbeddingStreamer::poll_copy_done_flag( @@ -494,8 +484,8 @@ folly::coro::Task RawEmbeddingStreamer::dispatch_copy_task( if (!poll_copy_done_flag(copy_done_flag)) { co_return; } - chunked_copy_and_enqueue( - indices, weights, identities, runtime_meta, count, chunk_copy_threads_); + co_await chunked_copy_and_enqueue( + indices, weights, identities, runtime_meta, count); co_return; } 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 f4a0418654..35e933d83c 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 @@ -250,147 +250,85 @@ TEST(RawEmbeddingStreamerTest, TensorCopyChunkByteWeights) { } 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, +// Structural invariants computeChunks must always satisfy: chunks 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 expectValidChunks( + const std::vector>& chunks, 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"; + for (const auto& [start, end] : chunks) { + EXPECT_EQ(start, cursor) << "chunks must be contiguous and non-overlapping"; + EXPECT_GT(end, start) << "no empty chunks"; 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)"; + EXPECT_EQ(cursor, num_rows) << "chunks 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)); +// computeChunks (like tensor_copy_chunk) is build-agnostic. Expected chunks are +// constructed independently. +TEST(RawEmbeddingStreamerTest, ComputeChunksExactMultiple) { + const auto chunks = computeChunks(/*num_rows=*/8, /*chunk_size=*/4); const std::vector> expected = {{0, 4}, {4, 8}}; - EXPECT_EQ(ranges, expected); - expectValidChunkRanges(ranges, /*num_rows=*/8, /*chunk_size=*/4); + EXPECT_EQ(chunks, expected); + expectValidChunks(chunks, /*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)); +TEST(RawEmbeddingStreamerTest, ComputeChunksRemainderChunk) { + // Last chunk carries the remainder when num_rows isn't a multiple of + // chunk_size. + const auto chunks = computeChunks(/*num_rows=*/10, /*chunk_size=*/4); const std::vector> expected = { {0, 4}, {4, 8}, {8, 10}}; - EXPECT_EQ(ranges, expected); - expectValidChunkRanges(ranges, /*num_rows=*/10, /*chunk_size=*/4); + EXPECT_EQ(chunks, expected); + expectValidChunks(chunks, /*num_rows=*/10, /*chunk_size=*/4); } -TEST(RawEmbeddingStreamerTest, ComputeChunkRangesCountLessThanChunkSize) { - const auto ranges = flatten( - computeChunkRanges(/*num_rows=*/3, /*chunk_size=*/10, /*num_threads=*/4)); +TEST(RawEmbeddingStreamerTest, ComputeChunksCountLessThanChunkSize) { + // num_rows < chunk_size collapses to a single partial chunk. + const auto chunks = computeChunks(/*num_rows=*/3, /*chunk_size=*/10); const std::vector> expected = {{0, 3}}; - EXPECT_EQ(ranges, expected); - expectValidChunkRanges(ranges, /*num_rows=*/3, /*chunk_size=*/10); + EXPECT_EQ(chunks, expected); + expectValidChunks(chunks, /*num_rows=*/3, /*chunk_size=*/10); } -TEST(RawEmbeddingStreamerTest, ComputeChunkRangesSingleChunk) { - const auto ranges = flatten( - computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/5, /*num_threads=*/4)); +TEST(RawEmbeddingStreamerTest, ComputeChunksSingleChunk) { + // num_rows == chunk_size is exactly one full chunk (no empty trailing chunk). + const auto chunks = computeChunks(/*num_rows=*/5, /*chunk_size=*/5); const std::vector> expected = {{0, 5}}; - EXPECT_EQ(ranges, expected); - expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/5); + EXPECT_EQ(chunks, expected); + expectValidChunks(chunks, /*num_rows=*/5, /*chunk_size=*/5); } -TEST(RawEmbeddingStreamerTest, ComputeChunkRangesChunkSizeOne) { - const auto ranges = flatten( - computeChunkRanges(/*num_rows=*/4, /*chunk_size=*/1, /*num_threads=*/1)); +TEST(RawEmbeddingStreamerTest, ComputeChunksChunkSizeOne) { + const auto chunks = computeChunks(/*num_rows=*/4, /*chunk_size=*/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); + EXPECT_EQ(chunks, expected); + expectValidChunks(chunks, /*num_rows=*/4, /*chunk_size=*/1); } -TEST(RawEmbeddingStreamerTest, ComputeChunkRangesOneOverChunk) { +TEST(RawEmbeddingStreamerTest, ComputeChunksOneOverChunk) { // 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 auto chunks = computeChunks(/*num_rows=*/5, /*chunk_size=*/4); const std::vector> expected = {{0, 4}, {4, 5}}; - EXPECT_EQ(ranges, expected); - expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/4); + EXPECT_EQ(chunks, expected); + expectValidChunks(chunks, /*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()); +TEST(RawEmbeddingStreamerTest, ComputeChunksZeroRowsIsEmpty) { + EXPECT_TRUE(computeChunks(/*num_rows=*/0, /*chunk_size=*/4).empty()); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunksZeroChunkSizeIsEmpty) { + // Defensive guard: chunk_size==0 would loop forever (s += 0), so it must + // short-circuit to an empty result. + EXPECT_TRUE(computeChunks(/*num_rows=*/5, /*chunk_size=*/0).empty()); } #ifdef FBGEMM_FBCODE @@ -413,6 +351,88 @@ TEST(RawEmbeddingStreamerTest, CtorRejectsZeroKnob) { /*table_sizes=*/{})); } +TEST(RawEmbeddingStreamerTest, TestMultiChunkFanOutShipsEveryChunk) { + // The stream()/tensor_stream tests all run at the default res_chunk_size (one + // chunk), so the chunked_copy_and_enqueue -> computeChunks -> copy_executor_ + // fan-out -> collectAllRange -> per-chunk submit_stream_item composition is + // otherwise never exercised with >1 chunk. Here res_chunk_size=4 over 10 rows + // tiles into ceil(10/4)=3 chunks; with a single shard each chunk ships + // exactly one co_setEmbeddings, so RPC count == chunk count proves every + // chunk is copied and shipped -- a dropped/duplicated chunk future or an + // off-by-one in the fan-out would change the count. computeChunks' partition + // correctness (no gap/dup/overflow) is covered by the ComputeChunks* unit + // tests above. + std::vector table_names = {"tb1"}; + std::vector table_offsets = {0}; + std::vector table_sizes = {0, 300}; + + // Static storage duration so the co_setEmbeddings coroutine mock can read it + // WITHOUT capturing (avoids the capturing-lambda-coroutine UAF lint). + 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)); + + // res_chunk_size=4 (not the 500000 default) so 10 rows fan out into 3 chunks + // across the 3-worker copy pool; res_store_shards=1 so each chunk ships once. + auto streamer = std::make_unique( + "test_multi_chunk_fanout", + /*enable_raw_embedding_streaming=*/true, + /*res_store_shards=*/1, + /*res_server_port=*/0, + /*res_chunk_size=*/4, + /*res_num_consumers=*/2, + /*res_num_copy_threads=*/3, + table_names, + table_offsets, + table_sizes); + + constexpr int64_t kNumRows = 10; + auto indices = at::arange( + kNumRows, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + auto count = at::tensor( + {kNumRows}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + + streamer->stream( + indices, + weights, + std::nullopt, + std::nullopt, + count, + /*require_tensor_copy=*/true, + /*blocking_tensor_copy=*/true); + + // Wait on the RPC count (not queue size) so no in-flight ship is missed. + constexpr int kExpectedChunks = 3; // ceil(10 / 4) + for (int i = 0; i < 1000 && rpc_count.load() < kExpectedChunks; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_EQ(rpc_count.load(), kExpectedChunks); + streamer->join_weights_stream_thread(); +} + TEST(RawEmbeddingStreamerTest, TestTensorStream) { std::vector table_names = {"tb1", "tb2", "tb3"}; std::vector table_offsets = {0, 100, 300};