Fix OOM in to_huggingface conversion and stabilize post-training E2E TPU pipelines

[SurbhiJainUSC]

Description

This PR fixes the following issues caught by TPU post-training E2E pipeline:

1. Checkpoint Conversion:

   • Pop transformed parameters from state_dict and trigger gc.collect() per iteration in _transform_weights_to_full_model to prevent host RAM accumulation and OOM-kill on large models (e.g., Llama 3.1 70B).
   • Pop LoRA weights when calculating deltas and delete checkpoint_dict after state dict extraction.
   • Pass target dtype to ocp.ArrayRestoreArgs in load_orbax_checkpoint when weight_dtype is specified, reducing initial checkpoint memory footprint.

2. LLaMA 3.1 70B RL failure:

   • Export required vLLM multiprocessing and IPC flags (VLLM_WORKER_MULTIPROC_METHOD=spawn, VLLM_ENABLE_V1_MULTIPROCESSING=0, GRPC_ENABLE_FORK_SUPPORT=0, PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=upb).
   • Explicitly forward use_pathways=${use_pathways} to train_rl and vllm_decode to prevent falling back to Pathways execution on standard distributed JAX/XPK clusters.
   • Configure chips_per_vm=4 and max_target_length=512 for stable GRPO training on TPU v5p slices.

3. Gemma 3 / Gemma 4 Post-Training E2E Tests:

   • Add export DATASET_PATH=gs://maxtext-dataset to test_gemma3_multimodal_sft.sh.
   • Add missing vLLM runtime flags and use_pathways forwarding to test_gemma3_rl.sh.
   • Update test_gemma4_rl.sh with batch_size=1, chips_per_vm=4, and max_target_length=512.

Tests

E2E tests

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.

[gemini-code-assist]

Code Review

This pull request introduces memory optimizations in the Hugging Face checkpoint conversion process by popping state dictionary keys and invoking garbage collection. It also adds support for specifying a target weight data type when restoring Orbax checkpoints, and updates various TPU end-to-end RL test scripts with new environment variables and execution parameters. Feedback is provided regarding the handling of tuple keys during state dictionary popping to prevent potential downstream crashes if subkeys are missing.

[codecov]

Codecov Report

❌ Patch coverage is 2.32558% with 42 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...c/maxtext/integration/vllm/maxtext_vllm_rollout.py	0.00%	19 Missing ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	5.55%	17 Missing ⚠️
...rc/maxtext/checkpoint_conversion/to_huggingface.py	0.00%	4 Missing ⚠️
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!