Datasets streaming=True: как обучать модели на терабайтных данных без скачивания — разгон в 100 раз

Ваша модель простаивает, пока диск скрипит. Пора это остановить

Вы запускаете тренировку. Видеокарты готовы, код отлажен, оптимизатор жаждет градиентов. А потом… тишина. Десять минут. Двадцать. Система молча пытается впихнуть 800 ГБ The Pile в вашу оперативку. Дисковый кэш раскаляется, swap умирает, а вы смотрите на пустую консоль и думаете о смысле жизни.

Классический load_dataset без флага streaming — это архаичный обряд инициации, через который проходят все. Сначала скачать. Потом распаковать. Потом конвертировать в Arrow. Потом загрузить в RAM. И только потом, может быть, начать обучение. Для терабайтных датасетов это путь в никуда.

streaming=True переворачивает этот процесс с ног на голову. Никакого скачивания целиком. Никакой гигантской RAM. Данные текут к модели тонким ручейком, по мере необходимости. Как водопровод, а не как затопление подвала.

Под капотом: почему это работает в 100 раз быстрее

Забудьте про магию. Здесь работает простая механика. Когда вы вызываете load_dataset("HuggingFaceM4/FineVisionMax", streaming=True), происходит следующее:

1. Hugging Face Hub получает запрос на метаданные датасета (несколько килобайт).
2. Создается ленивый итератор, который знает, как и откуда брать данные.
3. Когда тренировочный цикл запрашивает следующий батч, итератор скачивает ровно один чанк (например, 1000 примеров) с удаленного хранилища или читает его с локального кэша.
4. Данные декодируются, преобразуются, отдаются модели и тут же забываются. Следующий чанк вытесняет предыдущий из памяти.

Ключевое отличие — состояние системы. Обычная загрузка: «все данные здесь и сейчас». Streaming: «данные там, я принесу, что нужно, и выброшу упаковку».

1 Сначала — как сломать всё. Типичные ошибки

Прежде чем делать правильно, посмотрите, как можно убить всю производительность одним движением. Я собрал это из обгоревших остатков чужих пайплайнов.

Ошибка №1: Превращение итератора в список. Это фатально. Весь смысл streaming — в ленивой загрузке. Если вы делаете list(dataset) или dataset = list(dataset), вы приказываете системе: «Все-таки загрузи всё в память, пожалуйста». Она послушает.

# КАТЕГОРИЧЕСКИ НЕ ДЕЛАТЬ
dataset = load_dataset("big_dataset", streaming=True)
dataset_list = list(dataset)  # Прощай, оперативная память!

Ошибка №2: Shuffle без buffer_size. В streaming-режиме нельзя перемешать датасет, которого нет в памяти. .shuffle() по умолчанию попытается это сделать и либо упадет, либо тихо проигнорируется. Нужен buffer_size — размер окна, в пределах которого происходит перемешивание.

# ПЛОХО
dataset = dataset.shuffle()  # Не сработает или сломается

# ХОРОШО
dataset = dataset.shuffle(buffer_size=10000)  # Перемешивает в окне 10к примеров

Ошибка №3: Игнорирование кэша. Каждый чанк по умолчанию скачивается заново при каждом новом проходе (epoch). Для датасета в 100 ГБ это тысячи гигабайт трафика. Нужно настроить кэширование, о чем ниже.

2 Правильный пайплайн: от загрузки до DataLoader

Вот рабочий код, который мы использовали при дообучении SmolLM3 на датасете HuggingFaceM4/FineVisionMax. Это не абстрактный пример — это скрипт, который работал на машине с 16 ГБ ОЗУ и датасетом в ~400 ГБ.

from datasets import load_dataset, load_from_disk
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer

# 1. Загружаем в streaming режиме
# cache_dir — критически важен для кэширования чанков на диск
dataset = load_dataset(
    "HuggingFaceM4/FineVisionMax",
    streaming=True,
    cache_dir="/path/to/large_cache"  # Должен быть на быстром SSD!
)

# Получаем splits как итераторы
train_stream = dataset["train"]
val_stream = dataset["validation"]

# 2. Применяем трансформации ЛЕНИВО
# Каждая операция возвращает новый итератор, а не обрабатывает данные
tokenizer = AutoTokenizer.from_pretrained("microsoft/SmolLM3-135M-Instruct")

def tokenize_function(examples):
    # Обрабатываем один чанк данных
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=2048
    )

# .map применяется к чанкам на лету
tokenized_train = train_stream.map(
    tokenize_function,
    batched=True,
    batch_size=1000  # Размер чанка для обработки
)

# 3. Перемешиваем в скользящем окне
tokenized_train_shuffled = tokenized_train.shuffle(buffer_size=50000)

# 4. Превращаем в бесконечный итератор (для многопроходного обучения)
tokenized_train_shuffled = tokenized_train_shuffled.repeat()

# 5. Создаем DataLoader
# Здесь итератор уже готов, DataLoader просто запрашивает у него батчи
train_dataloader = DataLoader(
    tokenized_train_shuffled,
    batch_size=16,
    collate_fn=lambda batch: {
        "input_ids": torch.stack([torch.tensor(item["input_ids"]) for item in batch]),
        "attention_mask": torch.stack([torch.tensor(item["attention_mask"]) for item in batch])
    }
)

# 6. Тренировочный цикл
for epoch in range(num_epochs):
    for batch_idx, batch in enumerate(train_dataloader):
        if batch_idx % 1000 == 0:
            print(f"Epoch {epoch}, batch {batch_idx}: Memory used {torch.cuda.memory_allocated() / 1e9:.2f} GB")
        # ... ваша логика обучения ...
        # Когда батч обработан, Python может его собрать как мусор, память освобождается

Нюансы, о которых молчит документация

После десятка таких пайплайнов понимаешь, что devil is in the details.

Кэширование — ваш друг и враг. По умолчанию datasets кэширует скачанные чанки в ~/.cache/huggingface/datasets. Для терабайтного датасета это может убить ваш системный диск. Всегда указывайте cache_dir на том разделе, где есть место. И следите за ним — старые кэши не удаляются автоматически.

Скорость зависит от сети и диска. Первый проход по датасету будет медленным, если чанки не закэшированы. Данные качаются из интернета. Второй и последующие проходы — в разы быстрее, если кэш на SSD. Если вы в облаке (AWS, GCP), размещайте кэш на временном SSD-диске (ephemeral storage), он быстрее и часто бесплатен.

Не все операции поддерживаются в streaming. .sort() — нет. .unique() — нет. Любая операция, требующая глобального взгляда на данные, недоступна. Это плата за работу с бесконечным потоком.

Реальный кейс: SmolLM3 и терабайты мультимодальных данных

Когда мы дообучали SmolLM3 на HuggingFaceM4/FineVisionMax, альтернативы streaming не было. Датасет — сотни гигабайт изображений и текста. Железо — несколько машин с 32 ГБ ОЗУ каждая.

Без streaming пришлось бы:

• Скачать 400+ ГБ на каждый узел (полдня).
• Выделить под данные >500 ГБ RAM (невозможно).
• Смириться с тем, что 90% времени тренировки — это ожидание загрузки.

Со streaming:

• Метаданные скачались за 2 секунды.
• Первый батч пошел в модель через 30 секунд (скачался и закэшировался первый чанк).
• Пиковое использование RAM не превысило 8 ГБ (батч + буферы).
• Дисковый кэш постепенно заполнялся в фоне, к третьей эпохе данные уже были локально, скорость упиралась только в GPU.

Разница в эффективности использования ресурсов — на порядки. Это как сравнивать доставку пиццы курьером (streaming) и покупку целой пиццерии для одного ужина (обычная загрузка).

Что делать, если streaming тормозит

Иногда даже с правильным кодом скорость не радует. Вот checklist для поиска узких мест:

1. Сеть. Используйте dataset = dataset.with_format("torch") ДО цикла. Это предотвратит конвертацию типов на лету для каждого батча, которая может создавать лишнюю нагрузку.
2. Диск. Кэш на HDD? Переместите на NVMe. Используйте /tmp если данные временные.
3. CPU. Операции в .map() слишком тяжелые? Увеличьте batch_size в map, чтобы амортизировать накладные расходы. Или используйте num_proc для распараллеливания (осторожно, с streaming это может быть tricky).
4. Шаблон доступа. DataLoader с num_workers > 0 может конфликтовать с streaming-итератором. Начните с num_workers=0, если есть проблемы.

Если вы работаете в распределенной среде (например, как в случае с DGX Spark), каждая нода будет кэшировать данные независимо. Убедитесь, что у всех нод достаточно дискового пространства.

Streaming — это не серебряная пуля (но почти)

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

Но для всего, что больше размера RAM — это единственный цивилизованный способ. Особенно если вы собираетесь использовать техники вроде Entropy-Adaptive Finetuning, где данные динамически фильтруются, и их полная загрузка избыточна.

Мой прогноз: через год флаг streaming=True станет де-факто стандартом для любой неигрушечной тренировки. Зачем хранить то, что можно прочитать? Зачем качать то, что нужно только кусочками? Логика streaming настолько же очевидна, как и логика пагинации в базе данных вместо выгрузки всех записей разом.

Попробуйте на своем следующем проекте. Замените одну строчку в загрузке данных. И посмотрите, как исчезает мучительное ожидание, а ваши GPU наконец-то получают работу, для которой их покупали.