Add dataset type olmo_grain for AI2 OLMo numpy pretrain mixes

[gagika]

Description

• Adds dataset_type=olmo_grain, a Grain-based input pipeline for AI2's
  pre-tokenized OLMo numpy mixes (e.g. OLMo-mix-0925-official.txt).
  Reads headerless .npy token streams from a gcsfuse mount, applies
  OLMo-core's repeated-n-gram filter, and yields the shapes the MaxText
  pretrain trainer expects.
• Stateless sampler: record at step k is a pure function of
  (seed, shard, k). Resume reads the latest step from
  config.checkpoint_dir and shifts the sampler — no Grain iterator
  state in the checkpoint.
• Ships two data tools (download_olmo_data_to_gcs.py with HTTP-Range
  resume; build_olmo_npy_index.py for header-scan indexing) and two
  launchers (run_olmo3_7b_grain_smoke.sh,
  run_olmo3_7b_grain_resume_test.sh).

Tests

• Unit tests pass (tests/unit/input_pipeline/olmo_*)
• Smoke train: 50 steps, loss 11.99 → 8.93 on v4-8 (4-layer bf16)
• Resume test: Run B picks up at step 50 with loss continuity 8.931 → 8.930

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

authored-by: @aireenmei

[github-actions]

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

[codecov]

Codecov Report

❌ Patch coverage is 79.05759% with 80 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/input_pipeline/olmo_grain_data_processing.py	24.52%	40 Missing ⚠️
src/maxtext/input_pipeline/olmo_data.py	88.63%	14 Missing and 6 partials ⚠️
src/maxtext/input_pipeline/olmo_data_grain.py	87.33%	13 Missing and 6 partials ⚠️
...maxtext/input_pipeline/input_pipeline_interface.py	66.66%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!