Zephon

Zephon là trình tải dữ liệu cho việc huấn luyện mô hình nền tảng, được phát triển bởi DatologyAI. Nó thực hiện token hóa, đóng gói và trộn dữ liệu trong khi quá trình huấn luyện đang chạy, và nó tạo ra các lô dữ liệu giống nhau khi bạn thay đổi số lượng GPU (tính tất định co giãn). Zephon đọc bộ dữ liệu trực tiếp từ Hugging Face Hub bằng các đường dẫn hf://.

Cài đặt

Extra hf bổ sung hỗ trợ cho các bộ dữ liệu Hub và tệp Parquet. Các ví dụ bên dưới cũng cần transformers để thực hiện tokenization:

pip install "zephon[hf]" transformers

Tải một bộ dữ liệu từ Hub

Truyền đường dẫn hf:// vào Dataset.from_path. Đường dẫn này chỉ định một dataset, một cấu hình và một split, ví dụ như split train của cấu hình sample-10BT của HuggingFaceFW/fineweb-edu. Zephon đọc các dataset trên Hub thông qua shard cache của nó, vì vậy bạn phải bật cache. Đặt cache trên ổ đĩa cục bộ tốc độ cao và đặt giới hạn kích thước:

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 chuyển cột text thành các token ID. pack_flat ghép các tài liệu đã tokenize thành các chuỗi có max_length token. Đặt max_length bằng độ dài ngữ cảnh cộng thêm một: với return_labels=True, to_training dịch nhãn sang một vị trí, nên cả input_ids và labels đều có độ dài bằng độ dài ngữ cảnh.

Một Pipeline là một đối tượng lặp được trong Python, vì vậy bạn có thể sử dụng nó trong bất kỳ vòng lặp huấn luyện nào. Zephon tải xuống một shard hoàn chỉnh trước khi đọc từ nó, do đó batch đầu tiên có thể mất một chút thời gian.

Để đọc các bộ dữ liệu riêng tư hoặc có kiểm soát, hãy đăng nhập bằng tài khoản Hugging Face của bạn trước, ví dụ bằng hf auth login. Zephon cũng đọc biến môi trường HF_TOKEN.

Định dạng đường dẫn

Đường dẫn hf:// của Zephon cho các bộ dữ liệu Hub có định dạng này:

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

• revision: một nhánh, thẻ hoặc commit. Mặc định là main. Zephon cố định đường dẫn đến một commit khi tải bộ dữ liệu, do đó tất cả các rank đọc cùng các tệp.
• source: original đọc các tệp đã tải lên kho, còn parquet đọc chuyển đổi Parquet tự động của Hub. Đặt nguồn sau @, ví dụ @~parquet. Mặc định, Zephon đọc các tệp đã tải lên nếu hỗ trợ định dạng của chúng, và dùng bản chuyển đổi Parquet nếu không.
• config: cấu hình của bộ dữ liệu. Bao gồm tham số này khi bộ dữ liệu có nhiều hơn một cấu hình.

Đường dẫn Zephon chỉ định một dataset và một split, không phải một file. Chúng không sử dụng đường dẫn file hf://datasets/... mà pandas, Polars và DuckDB dùng. Hãy dùng hf://HuggingFaceFW/fineweb-edu/sample-10BT/train, không dùng hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train.

Trộn các bộ dữ liệu

Để huấn luyện trên nhiều bộ dữ liệu Hub, hãy thêm chúng vào nguồn công việc với tỷ lệ của từng bộ. Ví dụ này trộn FineWeb-Edu với Wikipedia tiếng Anh, nhắm đến hỗn hợp 70/30 theo token:

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

Để trộn theo mẫu, hãy bỏ qua token_estimation.

Tìm hiểu thêm

• Tài liệu Zephon
• Bài đăng blog ra mắt
• Bài báo Zephon
• DatologyAI trên Hub