[Feature Proposal]: Transition MaxText Checkpoint Conversion to Autoscaling CPU Node Pool (n2-highmem-64)

[RexBearIU]

Summary / Motivation

Currently, MaxText checkpoint conversion (to_maxtext.py) and forward-pass logit parity verification tests are often dispatched to TPU slices (e.g., v5e / v5p).

Checkpoint conversion (weight reshaping, tensor transpositions, format transformation) is fundamentally CPU and I/O-bound rather than accelerator-compute-bound. Running these tests on TPU VMs creates several bottlenecks:

1. Ties up scarce TPU accelerator capacity that should be reserved for distributed model training and hardware benchmarking.
2. Introduces long queue wait times and preemption in Kueue/XPK.
3. High compute cost for workloads that do not utilize matrix multiplication units.

---

Proposed Solution: Standardize on Unified Autoscaling CPU Node Pool (n2-highmem-64)

We propose standardizing all MaxText checkpoint conversion and logit verification pipelines onto an Autoscaling n2-highmem-64 CPU Node Pool:

1. Single Standardized VM Type:

   • n2-highmem-64 (64 vCPUs, 512 GB RAM per node).
   • Provides universal coverage for all models from 2B to 70B without special-casing or infrastructure fragmentation (including dual-model on-the-fly PyTorch + MaxText logit checks for 70B models).

2. Scale-to-Zero (min-nodes=0, max-nodes=10):

   • Spins up 1 dedicated node per model on-demand when CI jobs run.
   • Runs all models in parallel with zero CPU thread or memory contention.
   • Automatically scales back down to 0 nodes when idle (zero idle cost).

---

Expected Benefits

1. Pipeline Execution Time: Identical runtime to TPU VMs, with full concurrency across all models, instant provisioning, and zero TPU queue delays.
2. Infrastructure Cost Reduction: >85% compute cost reduction compared to running conversion on TPU slices.
3. Zero Maintenance Overhead & Future-Proof: A single, consistent node pool configuration across all MaxText conversion DAGs that automatically scales as new models and PRs are added.

---

Action Items

• Configure autoscaling n2-highmem-64 node pool (min-nodes=0, max-nodes=10) on GKE.
• Ensure all conversion test scripts pass hardware=cpu and skip_jax_distributed_system=True (see PR fix(testing): add CPU execution flags and standardize E2E checkpoint conversion tests #4886).
• Update Cloud Composer / Airflow DAG definitions to target the CPU node selector.