Uh oh!
There was an error while loading. Please reload this page.
Add Bagz format support - #4496
Conversation
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
| """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)) |
aireenmei
left a comment
There was a problem hiding this comment.
Thanks for adding the feature! Could you also add unit tests in grain_data_processing_test.py? You will need a class setup class similar to _GrainArrayRecordSetup and a test class similar to GrainArrayRecordProcessingTest
Uh oh!
There was an error while loading. Please reload this page.
| # limitations under the License. | ||
| """ | ||
| Download a HuggingFace dataset via streaming and save as Bagz files. |
There was a problem hiding this comment.
Could you add some example commands in the docstring?
There was a problem hiding this comment.
Done, I've added detailed example commands in the docstring covering basic local conversion, GCSFUSE-mounted bucket output with custom prefix/file size, and authenticated access for gated datasets
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: Integrated upstream `grain.BagzDataSource` for 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: - Unit Tests: Added comprehensive unit test suite in `tests/unit/bagz_data_processing_test.py` covering RandomAccessDataSource compatibility, MapDataset, ElasticIterator (single and multi-worker), and end-to-end `get_datasets` integration. - 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).
85206e8 to
7abdb4dCompare…ring Add practical example commands to the docstring of `download_hf_dataset_as_bagz.py`: - Basic local directory conversion (Salesforce/wikitext) - Write-optimized GCSFUSE mounted bucket output with custom shard prefix and file size - Private/gated HuggingFace dataset download using auth token
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: