Hub documentation

Zephon

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Zephon

Zephon 是一个用于基础模型训练的数据加载器,由 DatologyAI 构建。 它在训练运行期间对数据进行分词、打包和混合,并且在更改 GPU 数量时提供相同的批次(弹性确定性)。 Zephon 使用 hf:// 路径直接从 Hugging Face Hub 读取数据集。

安装

hf 额外依赖项添加了对 Hub 数据集和 Parquet 文件的支持。以下示例还需要 transformers 用于分词:

pip install "zephon[hf]" transformers

从 Hub 加载数据集

将 hf:// 路径传递给 Dataset.from_path。该路径指定了数据集、配置和分片,例如 HuggingFaceFW/fineweb-edu 的 sample-10BT 配置的 train 分片。 Zephon 通过其 shard cache 读取 Hub 数据集,因此必须启用缓存。请将缓存放置在快速的本地磁盘上,并设置大小限制:

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 列转换为 token ID。pack_flat 将分词后的文档拼接为长度为 max_length 个 token 的序列。请将 max_length 设置为上下文长度加一:当 return_labels=True 时,to_training 会将标签向后移动一个位置,因此 input_ids 和 labels 的长度均为上下文长度。

Pipeline 是一个 Python 可迭代对象,因此您可以在任何训练循环中使用它。Zephon 会在读取之前下载完整的分片,因此第一个批次可能需要一些时间。

要读取私有或受门禁的数据集,请先使用您的 Hugging Face 账户登录,例如使用 hf auth login。Zephon 还会读取 HF_TOKEN 环境变量。

路径格式

Zephon 的 Hub 数据集 hf:// 路径格式如下:

hf://{org}/{name}[@[{revision}][~{source}]]/[{config}/]{split}
  • revision:分支、标签或提交。默认值为 main。Zephon 在加载数据集时会将该路径固定到某个提交,以确保所有 rank 读取相同的文件。
  • source:当值为 original 时,读取上传到仓库的文件;当值为 parquet 时,读取 Hub 的自动 Parquet 转换。请将 source 放在 @ 之后,例如 @~parquet。默认情况下,如果 Zephon 支持上传文件的格式,则读取这些文件;否则使用 Parquet 转换。
  • config:数据集配置。当数据集包含多个配置时,请包含此项。
路径读取
hf://HuggingFaceFW/fineweb-edu/sample-10BT/trainsample-10BT 配置的 train 拆分,来自 main 分支
hf://HuggingFaceFW/fineweb-edu@v1.4.0/sample-10BT/train来自 v1.4.0 分支的相同拆分
hf://yahma/alpaca-cleaned@~parquet/train自动 Parquet 转换中的 train 拆分

Zephon 路径用于指定数据集和拆分,而非文件。它们不使用 pandas、Polars 和 DuckDB 所用的 hf://datasets/... 文件路径。请使用 hf://HuggingFaceFW/fineweb-edu/sample-10BT/train,而不是 hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train。

混合数据集

要在多个 Hub 数据集上进行训练,请将它们添加到工作源中,并指定每个数据集的比例。此示例将 FineWeb-Edu 与英文维基百科混合,目标按 token 数达到 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。

了解更多

Update on GitHub