Hub documentation
Zephon
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)
breaktokenize 将 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/train | sample-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。