Zephon

Zephon es un cargador de datos para el entrenamiento de modelos de base, desarrollado por DatologyAI. Tokeniza, empaqueta y mezcla datos mientras se ejecuta el entrenamiento, y proporciona los mismos lotes al cambiar el número de GPUs (determinismo elástico). Zephon lee conjuntos de datos directamente desde Hugging Face Hub mediante rutas hf://.

Instalación

El extra hf añade soporte para conjuntos de datos del Hub y archivos Parquet. Los ejemplos a continuación también necesitan transformers para la tokenización:

pip install "zephon[hf]" transformers

Cargar un conjunto de datos desde el Hub

Pase una ruta hf:// a Dataset.from_path. La ruta identifica un conjunto de datos, una configuración y una división, por ejemplo, la división train de la configuración sample-10BT de HuggingFaceFW/fineweb-edu. Zephon lee los conjuntos de datos del Hub a través de su shard cache, por lo que debe habilitar la caché. Coloque la caché en un disco local rápido y establezca un límite de tamaño:

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 convierte la columna text en IDs de tokens. pack_flat une los documentos tokenizados en secuencias de max_length tokens. Establece max_length a la longitud de contexto más uno: con return_labels=True, to_training desplaza las etiquetas una posición, de modo que input_ids y labels tengan ambos la longitud de contexto.

Un Pipeline es un iterable de Python, por lo que puedes usarlo en cualquier bucle de entrenamiento. Zephon descarga un shard completo antes de leerlo, por lo que el primer lote puede tardar un poco.

Para leer conjuntos de datos privados o con acceso restringido, inicie sesión primero con su cuenta de Hugging Face, por ejemplo, con hf auth login. Zephon también lee la variable de entorno HF_TOKEN.

Formato de ruta

Las rutas hf:// de Zephon para conjuntos de datos del Hub tienen este formato:

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

• revision: una rama, etiqueta o commit. El valor predeterminado es main. Zephon fija la ruta a un commit cuando carga el conjunto de datos, de modo que todos los rangos lean los mismos archivos.
• source: original lee los archivos subidos al repositorio, y parquet lee la conversión automática a Parquet del Hub. Coloca la fuente después de @, por ejemplo @~parquet. De forma predeterminada, Zephon lee los archivos subidos cuando admite su formato, y utiliza la conversión a Parquet si no lo hace.
• config: la configuración del conjunto de datos. Inclúyala cuando el conjunto de datos tenga más de una configuración.

Las rutas de Zephon nombran un conjunto de datos y una división, no un archivo. No utilizan las rutas de archivo hf://datasets/... que usan pandas, Polars y DuckDB. Use hf://HuggingFaceFW/fineweb-edu/sample-10BT/train, no hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train.

Mezclar conjuntos de datos

Para entrenar con varios conjuntos de datos del Hub, añádalos a la fuente de trabajo con la proporción de cada uno. Este ejemplo combina FineWeb-Edu con Wikipedia en inglés, con un objetivo de mezcla de 70/30 por tokens:

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)
)

Para mezclar por muestras en su lugar, omita token_estimation.

Más información

• Documentación de Zephon
• Publicación del blog de lanzamiento
• Artículo de Zephon
• DatologyAI en el Hub