AI-Hypercomputer / AI-Hypercomputer/maxtext

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

Aberta
#4,888 0 comentários 0 reações 0 responsáveis Ver no GitHub
Linguagem predominante
Python
Estrelas
2.4k
Forks
607
Merge médio
2d 19h
PRs com merge (30d)
158

Descrição

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

Guia de contribuição

Abrir o guia de contribuição

Avaliação

Esta issue ainda não foi avaliada.

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.