Fix on-the-fly model conversion in Pathways by forwarding storage flags

[darisoy]

Description

This PR ensures that on-the-fly model checkpoint conversion (Hugging Face to MaxText) during a Pathways run correctly respects the storage compatibility flags configured by the parent process. It explicitly forwards the checkpoint_storage_use_ocdbt and checkpoint_storage_use_zarr3 configurations to the spawned to_maxtext subprocess.

Context & Problem Solved

According to the official Hugging Face to MaxText Conversion Guide, when manually converting checkpoints for a Pathways run, users are instructed to set the Pathways compatibility flags to False (abstracted via USE_PATHWAYS=1 which disables ocdbt and zarr3 storage formats).

However, MaxText also supports automatic, on-the-fly checkpoint conversion during training startup (enabled by default in RL configs via convert_checkpoint_if_possible: True in post_train/rl.yml).

How it is broken today:
When running a Pathways RL job with on-the-fly conversion enabled, from_pretrained in model_creation_utils.py spawns the to_maxtext.py subprocess but fails to forward the parent process's checkpoint_storage_use_ocdbt and checkpoint_storage_use_zarr3 configurations.

As a result, the subprocess defaults to saving the converted checkpoint with ocdbt=True and zarr3=True (the McJAX defaults from base.yml). Because Pathways does not support these formats, the parent training process immediately fails with a restore error when attempting to load the newly converted parameters.

Solution

This PR programmatically forwards checkpoint_storage_use_ocdbt and checkpoint_storage_use_zarr3 from the parent configuration to the to_maxtext subprocess command line, ensuring that automatically converted checkpoints are instantly compatible with the running Pathways environment.

FIXES: b/497915531

Tests

Verified the changes by running the MaxText unit tests inside the TPU VM environment.
Ran unit tests forcing CPU execution and simulating 16 devices (to satisfy test sharding divisibility requirements):

JAX_PLATFORMS=cpu XLA_FLAGS=--xla_force_host_platform_device_count=16 pytest tests/unit/maxtext_utils_test.py

Results: All relevant checkpointing and model setup tests passed successfully.

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!