Zephon

Zephon est un chargeur de données pour l’entraînement de modèles de fondation, développé par DatologyAI. Il tokenise, agrège et mélange les données pendant l’exécution de l’entraînement, et il fournit les mêmes lots lorsque le nombre de GPU est modifié (déterminisme élastique). Zephon lit les jeux de données directement depuis le Hugging Face Hub via des chemins hf://.

Installation

L’extra hf ajoute la prise en charge des jeux de données du Hub et des fichiers Parquet. Les exemples ci-dessous nécessitent également transformers pour la tokenisation :

pip install "zephon[hf]" transformers

Charger un jeu de données depuis le Hub

Passez un chemin hf:// à Dataset.from_path. Le chemin désigne un jeu de données, une configuration et un split, par exemple le split train de la configuration sample-10BT de HuggingFaceFW/fineweb-edu. Zephon lit les jeux de données du Hub via son cache de shards, il faut donc activer le cache. Placez le cache sur un disque local rapide et définissez une limite de taille :

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 convertit la colonne text en identifiants de jetons. pack_flat assemble les documents tokenisés en séquences de max_length jetons. Définissez max_length sur la longueur du contexte plus un : avec return_labels=True, to_training décale les étiquettes d’une position, de sorte que input_ids et labels aient tous deux la longueur du contexte.

Un Pipeline est un itérable Python, vous pouvez donc l'utiliser dans n'importe quelle boucle d'entraînement. Zephon télécharge un shard complet avant d'y lire, ce qui peut prendre un certain temps pour le premier lot.

Pour lire des jeux de données privés ou protégés, connectez-vous d’abord avec votre compte Hugging Face, par exemple avec hf auth login. Zephon lit également la variable d’environnement HF_TOKEN.

Format du chemin

Les chemins hf:// de Zephon pour les jeux de données du Hub ont ce format :

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

• revision : une branche, un tag ou un commit. La valeur par défaut est main. Zephon fixe le chemin vers un commit lors du chargement du jeu de données, afin que tous les rangs lisent les mêmes fichiers.
• source : original lit les fichiers téléversés dans le dépôt, et parquet lit la conversion automatique en Parquet du Hub. Placez la source après le @, par exemple @~parquet. Par défaut, Zephon lit les fichiers téléversés s’il prend en charge leur format, et utilise la conversion Parquet dans le cas contraire.
• config : la configuration du jeu de données. À inclure lorsque le jeu de données comporte plus d’une configuration.

Les chemins Zephon désignent un jeu de données et un split, et non un fichier. Ils n’utilisent pas les chemins de fichiers hf://datasets/... utilisés par pandas, Polars et DuckDB. Utilisez hf://HuggingFaceFW/fineweb-edu/sample-10BT/train, et non hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train.

Mélanger des jeux de données

Pour entraîner sur plusieurs jeux de données du Hub, ajoutez-les à la source de travail avec la proportion de chacun. Cet exemple mélange FineWeb-Edu avec Wikipédia en anglais, en visant un mélange de 70/30 par jetons :

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

Pour mélanger par échantillons plutôt que par jetons, omettez token_estimation.

En savoir plus

• Documentation Zephon
• Article de blog de lancement
• Article Zephon
• DatologyAI sur le Hub