MB-71397: Fix CPU OOM during GPU Train

faiss/gpu/GpuIndexIVF.cu

@@ -207,6 +207,8 @@ void GpuIndexIVF::copyFrom(const faiss::IndexIVF* index) {
                 "that already contains a GPU coarse (level 1) quantizer "
                 "is not currently supported");
     }
+    // copy the num vectors to use to train the encoder from the source index
+    encoderTrainingVecs_ = index->train_encoder_num_vectors();
 
     // Validate equality
     FAISS_ASSERT(is_trained == index->is_trained);
@@ -254,6 +256,10 @@ void GpuIndexIVF::copyTo(faiss::IndexIVF* index) const {
     index->make_direct_map(false);
 }
 
+idx_t GpuIndexIVF::train_encoder_num_vectors() const {
+    return encoderTrainingVecs_;
+}
+
 idx_t GpuIndexIVF::getNumLists() const {
     return nlist;
 }

faiss/gpu/GpuIndexIVF.h

@@ -81,6 +81,10 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface {
     /// Returns the number of inverted lists we're managing
     virtual idx_t getNumLists() const;
 
+    /// can be redefined by subclasses to indicate how many training vectors
+    /// they need
+    virtual idx_t train_encoder_num_vectors() const;
+
     /// Returns the number of vectors present in a particular inverted list
     virtual idx_t getListLength(idx_t listId) const;
 
@@ -148,6 +152,9 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface {
 
     /// For a trained/initialized index, this is a reference to the base class
     std::shared_ptr<IVFBase> baseIndex_;
+
+    /// Number of vectors to use to train the encoder (if applicable)
+    idx_t encoderTrainingVecs_ = 0;
 };
 
 } // namespace gpu

faiss/gpu/GpuIndexIVFScalarQuantizer.cu

@@ -12,6 +12,7 @@
 #include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
 #include <faiss/gpu/impl/IVFFlat.cuh>
 #include <faiss/gpu/utils/CopyUtils.cuh>
+#include <faiss/utils/utils.h>
 #include <limits>
 
 namespace faiss {
@@ -219,14 +220,19 @@ void GpuIndexIVFScalarQuantizer::reset() {
 
 void GpuIndexIVFScalarQuantizer::trainResiduals_(idx_t n, const float* x) {
     // The input is already guaranteed to be on the CPU
+    idx_t max_nt = train_encoder_num_vectors();
+    if (max_nt <= 0) {
+        max_nt = (size_t)1 << 35;
+    }
+    TransformedVectors tv(x, fvecs_maybe_subsample(d, (size_t*)&n, max_nt, x, verbose));
     if (!by_residual) {
-        sq.train(n, x);
+        sq.train(n, tv.x);
     } else {
         std::vector<idx_t> assign(n);
-        quantizer->assign(n, x, assign.data());
+        quantizer->assign(n, tv.x, assign.data());
 
         std::vector<float> residuals(n * d);
-        quantizer->compute_residual_n(n, x, residuals.data(), assign.data());
+        quantizer->compute_residual_n(n, tv.x, residuals.data(), assign.data());
 
         sq.train(n, residuals.data());
     }