AI-Hypercomputer / AI-Hypercomputer/maxtext

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

Offen
#4,888 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
2.4k
Forks
607
Ø Merge
2 T. 19 Std.
Gemergte PRs (30 T.)
158

Beschreibung

### 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 #4886).
- [ ] Update Cloud Composer / Airflow DAG definitions to target the CPU node selector.

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.