Add Bagz format support#4496
Open
Marlon666 wants to merge 1 commit into
Open
Conversation
Implements end-to-end support for the Bagz (.bagz) dataset format inside the MaxText training and data processing pipeline. Enables highly efficient, POSIX-compliant file loading optimized specifically for GCS FUSE and local filesystem environments. Scope of Changes: - Input Pipeline: Added `BagzDataSource` to support direct, multiprocess-safe reading of Bagz file shards via PyGrain. - Data Processing: Registered `bagz` as a first-class supported file type across PyGrain data loaders, feature normalizers, and tokenizers. - Tooling: Created `download_hf_dataset_as_bagz.py` with multi-worker support, HuggingFace streaming, and resilient checkpointing for GCSFUSE mounts. Testing: - End-to-End Local Smoke Test: Executed a 3-step offline training loop on a local CPU environment using the newly generated Bagz dataset via `DECOUPLE_GCLOUD=TRUE`. - Parallel Data Generation: Successfully converted and streamed HF datasets into multi-shard Bagz files (38 shards, Salesforce/wikitext). - Resiliency: Verified resilience against rate limits and process crashes with standalone protobuf execution. Test Logs URL: https://paste.googleplex.com/6387838210932736
Marlon666
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
suexu1025,
vipannalla and
xibinliu
as code owners
July 15, 2026 23:16
Marlon666
changed the title
| " controller) enabled (`enable_single_controller=True`)." | ||
| ) | ||
| if self.grain_use_elastic_iterator and self.grain_file_type != "arrayrecord": | ||
| if self.grain_use_elastic_iterator and self.grain_file_type not in ("arrayrecord", "bagz"): |
There was a problem hiding this comment.
Have we tested the elastic iterator with BagZ?
| elastic=elastic, | ||
| ) | ||
| return dataset | ||
| elif data_file_type == "bagz": |
There was a problem hiding this comment.
This logic looks exactly the same as arrayrecord besides create_dataset_from_pattern, should we combine them into one?
| elif data_file_type == "bagz": | ||
| def create_dataset_from_pattern(pattern): | ||
| files = find_data_files(pattern) | ||
| source = input_pipeline_utils.BagzDataSource(files) |
There was a problem hiding this comment.
Any reasons not directly using bagz.Reader here, do we have to add the BagZDataSource? Example code: https://google-grain.readthedocs.io/en/latest/tutorials/data_sources/bagz_data_source_tutorial.html.
| """A fully picklable and fork-safe RandomAccessDataSource for Bagz files.""" | ||
| def __init__(self, paths: list[str] | str): | ||
| if isinstance(paths, (list, tuple)): | ||
| self._path = ",".join(sorted(paths)) |
There was a problem hiding this comment.
Why we have to sort the paths here?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This Pull Request introduces end-to-end support for the Gemini-backed Bagz (
.bagz) dataset format within MaxText's PyGrain input pipeline.Context & Problem Solved:
Previously, MaxText users relying on the ultra-efficient Bagz format for large-scale pre-training faced challenges when bridging cloud storage with PyGrain workers. This PR enables a highly optimized, POSIX-compliant file ingestion path designed explicitly for local filesystem and naturally GCS Fuse mounts. By avoiding complex REST/RPC shims, this approach achieves robust, multiprocess-safe read scalability out-of-the-box.
Scope of Changes:
BagzDataSourcewrapper ininput_pipeline_utils.pyfor parallel and resilient Bagz shard reader instantiation.bagzas a first-class supported file type ingrain_data_processing.py, enabling nativeParseFeaturesandNormalizeFeatureshandling.base.ymlandtypes.pyto seamlessly accept"bagz"as a validgrain_file_type.download_hf_dataset_as_bagz.pywith multi-worker parallelism, HuggingFace streaming, and robust checkpoint recovery.FIXES: (If there is a Github Issue #, paste it here, otherwise leave blank)
Tests
The implementation was successfully validated end-to-end on both Data Generation and Model Training levels:
Salesforce/wikitextdataset (wikitext-103-raw-v1) into 38 contiguous.bagzfile shards using multi-worker execution.DECOUPLE_GCLOUD=TRUE(Log URL attached below).Commands used for local testing: