Загрузка и эффективный по стоимости поиск больших коллекций#

В данном руководстве описан подход к загрузке, индексации и поиску больших объемов данных с минимальными затратами на примере реального набора данных LAION-400M.

Цель данного руководства — показать, какой минимальный объем ресурсов требуется для индексирования и поиска большого набора данных при сохранении разумной задержки поиска и точности.

Набор данных#

Используемый набор данных — LAION-400M, коллекция примерно из 400 миллионов векторов, полученных из изображений, извлеченных из набора данных Common Crawl. Каждый вектор имеет размерность 512 и генерируется с помощью модели CLIP.

Векторы связаны с рядом полей метаданных, таких как url, caption, LICENSE, и так далее.

Общий размер полезной нагрузки составляет около 200 ГБ, а сами векторы занимают 400 ГБ.

Сам набор данных не хранит изображения, а содержит только URL-адреса исходных изображений. На момент написания некоторые из этих URL уже недоступны.

Набор данных доступен в виде 409 фрагментов, каждый из которых содержит примерно 1 миллион векторов. Для загрузки частей набора данных одну за другой будет использован скрипт Python:.

export QDRANT_URL="https://xxxx-xxxx.xxxx.cloud.qdrant.io"
export QDRANT_API_KEY="xxxx-xxxx-xxxx-xxxx"

python upload.py

Аппаратное обеспечение#

Минимальная аппаратная конфигурация для выполнения задачи:

8 ядер CPU;
64 ГБ RAM;
650 ГБ дискового пространства.

Эта конфигурация позволяет индексировать и исследовать набор данных в однопользовательском режиме; задержка достаточно низкая, чтобы строить интерактивные графики и перемещаться по панели инструментов.

Может понадобиться больше ядер CPU и оперативной памяти для производственных конфигураций.

Важно обеспечить высокую пропускную способность сети для данного эксперимента, поэтому клиент и сервер должны находиться в одном регионе.

Загрузка и индексация#

Для загрузки фрагментов набора данных один за другим будет использован скрипт Python:

export QDRANT_URL="https://xxxx-xxxx.xxxx.cloud.qdrant.io"
export QDRANT_API_KEY="xxxx-xxxx-xxxx-xxxx"

python upload.py

Скрипт будет скачивать фрагменты набора данных LAION один за другим и загружать их в Platform V Vector DB (далее - Vector DB). Промежуточные данные не сохраняются на диске, поэтому скрипту требуется немного места на стороне клиента.

Конфигурация коллекции, которую использовали:

client.create_collection(
        QDRANT_COLLECTION_NAME,
        vectors_config=models.VectorParams(
            size=512, # CLIP model output size
            distance=models.Distance.COSINE, # CLIP model uses cosine distance
            datatype=models.Datatype.FLOAT16, # We only need 16 bits for float, otherwise disk usage would be 800Gb instead of 400Gb
            on_disk=True # We don't need original vectors in RAM
        ),
        # Even though CLIP vectors don't work well with binary quantization, out of the box,
        # we can rely on query-time oversampling to get more accurate results
        quantization_config=models.BinaryQuantization(
            binary=models.BinaryQuantizationConfig(
                always_ram=True,
            )
        ),
        optimizers_config=models.OptimizersConfigDiff(
            # Bigger size of segments are desired for faster search
            # However it might be slower for indexing
            max_segment_size=5_000_000, 
        ),
        # Having larger M value is desirable for higher accuracy,
        # but in our case we care more about memory usage
        # We could still achieve reasonable accuracy even with M=6 + oversampling
        hnsw_config=models.HnswConfigDiff(
            m=6, # decrease M for lower memory usage
            on_disk=False
        ),
    )

Есть несколько важных моментов, на которые следует обратить внимание:

Используется тип данных FLOAT16, который позволяет хранить векторы вдвое меньшего размера по сравнению с типом FLOAT32. Для этого набора данных нет значительных потерь точности.
Используется BinaryQuantization совместно с always_ram=True, чтобы включить оверсемплинг во время выполнения запросов. Это позволяет получить точный и ресурсоэффективный поиск, несмотря на то, что вектора CLIP размером 512D плохо работают с двоичной квантизацией из коробки.
Используется HnswConfig совместно с m=6, чтобы уменьшить использование памяти.

Цель данной конфигурации — обеспечить, чтобы компонент поиска Prefetch никогда не нуждался в загрузке данных с диска, а минимальная версия векторов и векторного индекса всегда находилась в оперативной памяти. Второй этап поиска может явно определить, сколько раз можно загружать данные с диска.

В примере процесс загрузки происходил со скоростью 5000 точек в секунду. Процесс индексирования проходил параллельно с загрузкой и происходил со скоростью примерно 4000 точек в секунду.

Использование памяти#

На высоком уровне использование памяти состоит из трех компонентов:

Системная память - 8,34 ГБ — память, зарезервированная для внутренних систем и ОС, она не зависит от размера набора данных.
Память данных - 39,27 ГБ — резидентная память процесса Vector DB, ее нельзя выгрузить, и процесс Vector DB завершится аварийно, если превысит лимит.
Кеш-память - 14,54 ГБ — дисковый кеш, который использует Vector DB. Он необходим для быстрого поиска, но его можно освободить при необходимости.

Больше всего интересна память данных и кеша. Посмотрите, что именно хранится в этих компонентах.

В текущей ситуации Vector DB использует память для хранения следующих компонентов:

Хранение векторов;
Хранение векторного индекса;
Хранение информации об идентификаторах (ID) и версиях точек.

Обратите внимание, что индексы полезной нагрузки выходят за рамки данного руководства. Если используете индексы полезной нагрузки в коллекции, может потребоваться скорректировать оценки соответствующим образом.

Размер векторов#

В текущей ситуации в оперативной памяти хранятся только квантифицированные векторы , поэтому довольно легко рассчитать требуемый размер:

400_000_000 * 512d / 8 bits / 1024 (Kb) / 1024 (Mb) / 1024 (Gb) = 23.84GB

Размер векторного индекса#

Векторный индекс немного сложнее, так как он не является простой матрицей.

Внутренне он хранится как список соединений в графе, причем каждое соединение представляет собой целое число длиной 4 байта.

Количество соединений определяется параметром M индекса HNSW, и в текущем случае оно составляет 6 на верхнем уровне и 2 x M на нулевом уровне.

Это дает следующую оценку:

400_000_000 * (6 * 2) * 4 bytes / 1024 (Kb) / 1024 (Mb) / 1024 (Gb) = 17.881Gb

Индекс HNSW в Vector DB хранится как маппинг и может быть выгружен из оперативной памяти по мере необходимости. Таким образом, потребление памяти HNSW относится к категории Cache memory.

Размер ID и версий#

Vector DB должен хранить дополнительную информацию о каждой точке, такую как идентификатор и версия. Эта информация необходима при каждом запросе, поэтому очень важно держать ее в оперативной памяти для быстрого доступа.

Рассмотрим внутреннюю структуру Vector DB, чтобы понять, сколько памяти требуется для этой информации.

// This is s simplified version of the IdTracker struct
// It omits all optimizations and small details,
// but gives a good estimation of memory usage
IdTracker {
    // Mapping of internal id to version (u64), compressed to 4 bytes
    // Required for versioning and conflict resolution between segments
    internal_to_version, // 400M x 4 = 1.5Gb

    // Mapping of external id to internal id, 4 bytes per point.
    // Required to determine original point ID after search inside the segment
    internal_to_external: Vec<u128>, // 400M x 16 = 6.4Gb

    // Mapping of external id to internal id. For numeric ids it uses 8 bytes,
    //  UUIDs are stored as 16 bytes.
    // Required to determine sequential point ID inside the segment
    external_to_internal: Vec<u64, u32>, // 400M x (8 + 4) = 4.5Gb
}

Общее потребление памяти IdTracker в текущем случае составляет приблизительно 12.4Gb.

Таким образом, ожидаемое общее использование оперативной памяти сервером Vector DB в текущем случае составляет около 23.84Gb + 17.881Gb + 12.4Gb = 54.121Gb, что очень близко к фактическому использованию памяти, которое ранее наблюдалось: 39.27Gb + 14.54Gb = 53.81Gb.

Пришлось применить некоторые упрощения к оценкам, но они достаточно хороши, чтобы понять использование памяти сервером Vector DB.

Поиск#

Данные эталонной истины#

Для более точной оценки производительности поиска необходимо сравнить результаты поиска с данными эталонной истины. К сожалению, набор данных LAION не содержит пригодных данных эталонной истины, поэтому необходимо создать их самостоятельно.

Для этого нужно выполнить полный сканирующий поиск по каждому вектору в наборе данных и сохранить результаты в отдельном файле. К сожалению, этот процесс занимает много времени и требует больших ресурсов, поэтому ограничьте количество запросов до 100. Скачайте готовый файл эталонной истины и скрипт для его генерации (требуется машина с памятью 512 ГБ и около 20 часов времени выполнения).

Файл эталонной истины содержит 100 запросов, каждый из которых имеет 50 результатов. Для создания запросов использовались первые 100 векторов самого набора данных.

Обратите внимание, что этот набор данных содержит значительное количество точных дубликатов, поэтому порядок результатов может различаться при разных запусках.

Запрос поиска#

Чтобы точно контролировать степень избыточности выборки, будет использоваться следующий запрос поиска:

limit = 50
rescore_limit = 1000 # oversampling factor is 20

query = vectors[query_id] # One of existing vectors

response = client.query_points(
        collection_name=QDRANT_COLLECTION_NAME,
        query=query,
        limit=limit,
        # Go to disk 
        search_params=models.SearchParams(
            quantization=models.QuantizationSearchParams(
                rescore=True,
            ),
        ),
        # Prefetch is performed using only in-RAM data,
        # so querying even large amount of data is fast
        prefetch=models.Prefetch(
            query=query,
            limit=rescore_limit,
            params=models.SearchParams(
                quantization=models.QuantizationSearchParams(
                    # Avoid rescoring in prefetch
                    # We should do it explicitly on the second stage
                    rescore=False,
                ),
            )
        )
    )

Этот запрос состоит из двух этапов:

Первый этап — предварительная выборка, выполняемая только с использованием данных в оперативной памяти. Он выполняется быстро и позволяет получить большое количество тенантов.
Второй этап — пересчет, который выполняется с полными векторами, хранящимися на дисках.

Используя двухэтапный поиск можно точно контролировать объем данных, загружаемых с диска, и обеспечить баланс между скоростью поиска и точностью.

Полностью код процесса поиска можно найти в файле eval.py.

Настройка производительности#

Одна важная настройка производительности, которая полезнА для данного набора данных, заключается во включении асинхронного ввода-вывода в Vector DB.

По умолчанию Vector DB использует синхронный ввод-вывод, что хорошо подходит для наборов данных в памяти, но может стать узким местом, когда нужно прочитать большой объем данных с диска.

Асинхронный ввод-вывод (io_uring) позволяет отправлять параллельные запросы на диск и полностью задействовать пропускную способность диска.

Это именно то, что нужно при выполнении крупномасштабного повторного расчета с исходными векторами.

Вместо того чтобы читать векторы один за другим и ждать ответа от диска тысячу раз, можно отправить тысяча запросов на диск и дождаться завершения всех них. Это позволяет максимально использовать пропускную способность диска и получать результаты быстрее.

Чтобы включить асинхронный ввод-вывод в Vector DB необходимо установить следующую переменную окружения:

QDRANT__STORAGE__PERFORMANCE__ASYNC_SCORER=true

Или задать параметр в конфигурационном файле:

storage:
  performance:
    async_scorer: true

Выполнение поисковых запросов#

Когда все приготовления завершены можно запустить поисковые запросы и оценить полученные результаты.

Полный код процесса поиска можно найти в файле eval.py.

Этот скрипт выполнит 100 поисковых запросов с настроенным коэффициентом избыточной выборки и сравнит результаты с эталонной истиной.

python eval.py --rescore_limit 1000

В тестовых запросах были получены следующие результаты:

Ограничение пересчета	Точность@50	Время на запрос
1000	75.2%	0.7с
5000	81.0%	2.2с

Дополнительные эксперименты с m=16 показали, что можно достичь точности 85% при использовании rescore_limit=1000, хотя это потребует немного больше памяти.

Заключение#

В данном руководстве продемонстрировано как эффективно загрузить, проиндексировать и искать большие наборы данных в Vector DB. Бинарная квантизация может применяться даже к векторам размером 512D, если она сочетается с избыточной выборкой во время запроса.

Vector DB позволяет точно управлять тем, где находится каждая часть хранилища, что позволяет добиться хорошего баланса между скоростью поиска и потреблением памяти.