Add Bagz format support

[Marlon666]

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:

• Input Pipeline: Added the BagzDataSource wrapper in input_pipeline_utils.py for parallel and resilient Bagz shard reader instantiation.
• Grain Integration: Registered bagz as a first-class supported file type in grain_data_processing.py, enabling native ParseFeatures and NormalizeFeatures handling.
• Configurations: Expanded base.yml and types.py to seamlessly accept "bagz" as a valid grain_file_type.
• Tooling: Created download_hf_dataset_as_bagz.py with 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:

1. Parallel Shard Generation: Successfully streamed and converted the Salesforce/wikitext dataset (wikitext-103-raw-v1) into 38 contiguous .bagz file shards using multi-worker execution.
2. End-to-End Model Training: Executed a 3-step offline training loop on a CPU-only environment using the generated local Bagz dataset via DECOUPLE_GCLOUD=TRUE (Log URL attached below).

Commands used for local testing:

# Data Generation
python tools/data_generation/download_hf_dataset_as_bagz.py \
  --dataset Salesforce/wikitext \
  --config '{"name": "wikitext-103-raw-v1"}' \
  --output ~/datasets/wikitext_bagz \
  --file-size-mb 10 \
  --workers 2

# MaxText CPU Smoke Test
export DECOUPLE_GCLOUD=TRUE
python -m maxtext.trainers.pre_train.train \
  src/maxtext/configs/base.yml \
  run_name="bagz_cpu_smoke_test" \
  steps=3 \
  dataset_type="grain" \
  grain_file_type="bagz" \
  grain_train_files="$HOME/datasets/wikitext_bagz/*.bagz" \
  grain_eval_files="$HOME/datasets/wikitext_bagz/*.bagz" \
  tokenize_eval_data=True \
  per_device_batch_size=1 \
  max_target_length=64 \
  enable_checkpointing=false \
  override_model_config=true \
  base_num_decoder_layers=1 \
  base_emb_dim=32 \
  base_mlp_dim=32 \
  head_dim=8 \
  base_num_kv_heads=1 \
  base_num_query_heads=1

Test Logs URL: https://paste.googleplex.com/6387838210932736

[lepan-google]

Why we have to sort the paths here?

[aireenmei]

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

[aireenmei]

Could you add some example commands in the docstring?

[Marlon666]

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