Zephon

Zephon é um carregador de dados para o treinamento de modelos de fundação, desenvolvido pela DatologyAI. Ele realiza a tokenização, o empacotamento e a mistura de dados durante a execução do treinamento e fornece os mesmos lotes quando você altera o número de GPUs (determinismo elástico). O Zephon lê conjuntos de dados diretamente do Hugging Face Hub com caminhos hf://.

Instalação

O extra hf adiciona suporte para conjuntos de dados do Hub e arquivos Parquet. Os exemplos abaixo também precisam de transformers para tokenização:

pip install "zephon[hf]" transformers

Carregue um conjunto de dados do Hub

Passe um caminho hf:// para Dataset.from_path. O caminho identifica um conjunto de dados, uma configuração e uma divisão, por exemplo, a divisão train da configuração sample-10BT de HuggingFaceFW/fineweb-edu. O Zephon lê conjuntos de dados do Hub por meio de seu shard cache, portanto, você deve habilitar o cache. Coloque o cache em um disco local rápido e defina um limite de tamanho:

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 converte a coluna text em IDs de tokens. pack_flat junta os documentos tokenizados em sequências de max_length tokens. Defina max_length como o comprimento do contexto mais um: com return_labels=True, to_training desloca os rótulos em uma posição, de modo que input_ids e labels tenham ambos o comprimento do contexto.

Uma Pipeline é um iterável em Python, portanto, você pode usá-la em qualquer loop de treinamento. O Zephon baixa um shard completo antes de ler dele, então o primeiro lote pode levar algum tempo.

Para ler conjuntos de dados privados ou restritos, faça login com sua conta do Hugging Face primeiro, por exemplo, com hf auth login. O Zephon também lê a variável de ambiente HF_TOKEN.

Formato do caminho

Os caminhos hf:// do Zephon para conjuntos de dados do Hub têm o seguinte formato:

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

• revision: um branch, tag ou commit. O padrão é main. O Zephon fixa o caminho em um commit ao carregar o conjunto de dados, para que todos os ranks leiam os mesmos arquivos.
• source: original lê os arquivos enviados para o repositório, e parquet lê a conversão automática para Parquet do Hub. Coloque a origem após o @, por exemplo @~parquet. Por padrão, o Zephon lê os arquivos enviados quando suporta o formato deles, e usa a conversão para Parquet quando não suporta.
• config: a configuração do conjunto de dados. Inclua-o quando o conjunto de dados tiver mais de uma configuração.

Os caminhos Zephon nomeiam um conjunto de dados e uma divisão, não um arquivo. Eles não usam os caminhos de arquivo hf://datasets/... usados por pandas, Polars e DuckDB. Use hf://HuggingFaceFW/fineweb-edu/sample-10BT/train, e não hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train.

Misturar conjuntos de dados

Para treinar em vários conjuntos de dados do Hub, adicione-os à fonte de trabalho com a proporção de cada um. Este exemplo mistura FineWeb-Edu com a Wikipédia em inglês, visando uma mistura 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 misturar por amostras, omita token_estimation.

Saiba mais

• Documentação do Zephon
• Post de lançamento no blog
• Artigo Zephon
• DatologyAI no Hub