Zephon

Zephon — это загрузчик данных для обучения базовых моделей, разработанный командой DatologyAI. Он выполняет токенизацию, упаковку и смешивание данных в процессе обучения и обеспечивает получение одинаковых батчей при изменении количества GPU (эластичная детерминированность). Zephon считывает наборы данных напрямую из Hugging Face Hub по путям hf://.

Установка

Дополнение hf добавляет поддержку наборов данных Hub и файлов Parquet. Для токенизации в примерах ниже также требуется transformers:

pip install "zephon[hf]" transformers

Загрузка набора данных из Hub

Передайте путь hf:// в Dataset.from_path. Путь указывает набор данных, конфигурацию и разбиение, например, разбиение train конфигурации sample-10BT набора HuggingFaceFW/fineweb-edu. Zephon читает наборы данных Hub через свой shard cache, поэтому необходимо включить кэш. Разместите кэш на быстром локальном диске и установите ограничение по размеру:

from zephon import Pipeline
from zephon.io import CacheOptions, Dataset, StoreOptions
from zephon.work import MixtureSpec, StaticMixtureWorkSource

dataset = Dataset.from_path("fineweb_edu", "hf://HuggingFaceFW/fineweb-edu/sample-10BT/train")
work_source = StaticMixtureWorkSource(
    datasets=[dataset],
    mixture=MixtureSpec({"fineweb_edu": 1.0}),
    seed=42,
)
cache = CacheOptions(enabled=True, root="/tmp/zephon-cache", limit_bytes=50 << 30)  # 50 GiB

pipeline = (
    Pipeline(work_source)
    .options(io_options=StoreOptions(cache=cache))
    .tokenize(tokenizer_id="gpt2", field="text")
    .pack_flat(max_length=2049, algorithm="wrap")
    .batch(microbatch_size=8)
)

for batch in pipeline:
    tensors = batch.to_training(return_labels=True)
    print(tensors["input_ids"].shape, tensors["labels"].shape)  # (8, 2048) (8, 2048)
    break

tokenize преобразует столбец text в идентификаторы токенов. pack_flat объединяет токенизированные документы в последовательности из max_length токенов. Установите max_length равным длине контекста плюс один: при return_labels=True to_training сдвигает метки на одну позицию, поэтому input_ids и labels имеют одинаковую длину контекста.

Pipeline — это итерируемый объект Python, поэтому его можно использовать в любом цикле обучения. Zephon загружает полный шард перед тем, как читать из него, поэтому первый батч может потребовать некоторого времени.

Чтобы читать частные или закрытые наборы данных, сначала войдите в свой аккаунт Hugging Face, например, с помощью hf auth login. Zephon также читает переменную окружения HF_TOKEN.

Формат пути

Пути hf:// для наборов данных Hub в Zephon имеют следующий формат:

hf://{org}/{name}[@[{revision}][~{source}]]/[{config}/]{split}

• revision: ветка, тег или коммит. По умолчанию — main. Zephon фиксирует путь к коммиту при загрузке набора данных, поэтому все ранги читают одни и те же файлы.
• source: original читает файлы, загруженные в репозиторий, а parquet — автоматическое преобразование в Parquet на Hub. Укажите источник после @, например @~parquet. По умолчанию Zephon читает загруженные файлы, если поддерживает их формат, и использует преобразование в Parquet в противном случае.
• config: конфигурация набора данных. Указывайте её, если у набора данных есть несколько конфигураций.

Пути Zephon указывают на набор данных и его разбиение, а не на файл. Они не используют файловые пути вида hf://datasets/..., которые применяются в pandas, Polars и DuckDB. Используйте hf://HuggingFaceFW/fineweb-edu/sample-10BT/train, а не hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train.

Смешивание наборов данных

Чтобы обучать на нескольких наборах данных из Hub, добавьте их в источник данных с указанием доли каждого. В этом примере FineWeb-Edu смешивается с английской Википедией, целевая пропорция по токенам — 70/30:

from zephon import Pipeline
from zephon.io import CacheOptions, Dataset, StoreOptions
from zephon.work import MixtureSpec, StaticMixtureWorkSource, TokenEstimation

work_source = StaticMixtureWorkSource(
    datasets=[
        Dataset.from_path("fineweb_edu", "hf://HuggingFaceFW/fineweb-edu/sample-10BT/train"),
        Dataset.from_path("wikipedia", "hf://wikimedia/wikipedia/20231101.en/train"),
    ],
    mixture=MixtureSpec({"fineweb_edu": 0.7, "wikipedia": 0.3}),
    seed=42,
    token_estimation=TokenEstimation(primer="measure", calibration_samples=2048, calibration_shards_max=4),
)
cache = CacheOptions(enabled=True, root="/tmp/zephon-cache", limit_bytes=50 << 30)  # 50 GiB

pipeline = (
    Pipeline(work_source)
    .options(io_options=StoreOptions(cache=cache))
    .tokenize(tokenizer_id="gpt2", field="text")
    .pack_flat(max_length=2049, algorithm="wrap")
    .batch(microbatch_size=8)
)

Чтобы смешивать по образцам, опустите token_estimation.

Подробнее

• Документация Zephon
• Блог-пост о запуске
• Статья Zephon
• DatologyAI на Hub