Как реализовать обучение модели XGBoost при данных, не помещающихся в память

Что проверяют

Знаете ли вы, что у XGBoost есть собственный режим внешней памяти, и понимаете ли, чем он отличается от распределённого обучения. Следом спрашивают, почему нельзя просто обучать модель по кускам, как нейросеть мини-батчами.

Как отвечать

Сначала честная диагностика: часто данные не влезают только в сыром виде.

• QuantileDMatrix: хранит признаки уже квантованными в бины гистограммы (по умолчанию max_bin=256), а не сырые значения, и заметно экономит память по сравнению с DMatrix. Работает с методом hist, который в актуальных версиях используется по умолчанию;
• типы и отбор: float32 вместо float64 ещё в pandas, категориальные признаки через category и enable_categorical=True, отказ от бесполезных колонок, даунсэмплинг мажоритарного класса с поправочными весами.

Если этого мало, путей три:

• внешняя память: наследуетесь от xgboost.DataIter, реализуете next(), который отдаёт очередной кусок через input_data, и reset(), а из итератора строите ExtMemQuantileDMatrix (в свежих версиях) или QuantileDMatrix. Куски читаются с диска, например из Parquet, скорость упирается в диск, и NVMe здесь заметно лучше;
• Dask: модуль xgboost.dask и DaskXGBClassifier при запущенном dask.distributed.Client. Данные раскладываются по воркерам, но суммарно должны поместиться в память кластера;
• Spark: xgboost.spark.SparkXGBClassifier как обычный estimator PySpark, если данные и так живут в Spark.

Для обучения на GPU задают device="cuda", внешняя память там тоже поддерживается.

Где спотыкаются

Предлагают мини-батчи с дообучением через xgb_model. Каждый новый кусок только добавляет деревья поверх старых, а старые построены на части данных и не пересматриваются. Это не эквивалент обучения на всём датасете, качество обычно хуже.

Считают, что Dask решает проблему на одной машине с малой памятью. Он масштабирует на кластер, а на одном узле без внешней памяти упадёт так же.

Забывают про копирование. При сборке матрицы из DataFrame на пике в памяти лежат и фрейм, и его копия внутри XGBoost, поэтому данные либо читают сразу итератором, либо освобождают фрейм сразу после построения матрицы.