Expose enable_continuous_profiling CLI flag in MaxText RL

[youchunni]

Description

Expose the enable_continuous_profiling CLI configuration flag in MaxText RL workflow to support saving trace profiles larger than 2GB as .xplane.riegeli.

Context & Solution:

When profiling large model workloads or multi-step training runs, monolithic .xplane.pb trace files can exceed Protocol Buffer's 2GB wire format limit, leading to failed file upload.

Setting enable_continuous_profiling=True configures the profiler stream smaller trace chunks into Riegeli record containers (.xplane.riegeli), bypassing the 2GB single-protobuf size limitation.

Note: Direct streaming of .xplane.riegeli files to GCS paths (gs://) is currently not supported. Profiles should be saved locally on the host path and uploaded to GCS.

Specific Changes:

• Added enable_continuous_profiling (boolean, default False) to src/maxtext/configs/base.yml and src/maxtext/configs/types.py.

• Updated src/maxtext/trainers/post_train/rl/train_rl.py to forward enable_continuous_profiling to the profiler configuration. Related Tunix PR: Support enable_continuous_profiling option in the Tunix JAX profiler wrapper google/tunix#2036

• Added documentation notes in base.yml clarifying that continuous profiling requires JAX >= 0.11.1.

Tests

• Verified on Cloud TPU v6e-1 that setting enable_continuous_profiling=True enables continuous trace collection and produces .xplane.riegeli container files on local path when the profile is >2GB. See screenshot:

• Verified that default enable_continuous_profiling=False retains existing JAX profiling behavior.

Checklist

• 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.

[gemini-code-assist]

Code Review

This pull request introduces a new configuration option, enable_continuous_profiling, to support saving traces larger than 2GB in the RL/Tunix profiler workflow. The reviewer noted a typo in the JAX version specified in the configuration comment (suggesting 0.4.11 instead of 0.11.1) and recommended clarifying that this option is currently only supported in the RL/Tunix profiler workflow.