AI-Hypercomputer / AI-Hypercomputer/maxtext

[bug] Dynamic TPU Metric Protobuf Import Causes Linkage Panic ('Check failed: GeneratedDatabase()->Add(...)')

Abierto
#4,611 1 comentario 0 reacciones 0 asignados Ver en GitHub
bug
Lenguaje dominante
Python
Estrellas
2.4k
Forks
607
Merge medio
2 d 19 h
PR fusionados (30 d)
158

Descripción

### Bug report

1. Launch MaxText GRPO RL training on Cloud TPU v6e using `train_rl.py`.
2. Keep default config parameters: `managed_mldiagnostics=True` (default in `configs/post_train/rl.yml`).
3. Top-level import `from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics` in `metric_logger.py` executes C++ Protobuf registration at Python module load time.
4. When `multiprocessing.spawn` re-imports module files in worker processes, dynamic Protobuf descriptor linkage panics on startup.

**Expected:** Training process initializes without Protobuf symbol linkage errors.
**Actual:** Process aborts immediately on startup with `Symbol name 'maxtext.TPUMetricService' is already defined`.

### Logs/Output

[libprotobuf ERROR google/protobuf/descriptor_database.cc:642] Symbol name "maxtext.TPUMetricService" is already defined.
[libprotobuf FATAL google/protobuf/descriptor.cc:1370] CHECK failed: GeneratedDatabase()->Add(encoded_file_descriptor, size):
*** Abortcat traceback ***
@ 0x7f23a8123456 google::protobuf::GoogleOnceInitImpl()
@ 0x7f23a8123789 maxtext::TPUMetricService_default_instance()
@ 0x7f23a8124012 pybind11::cpp_function::dispatcher()
Fatal Python error: Aborted

Current thread 0x00007f23b1234700 (most recent call first):
File "/usr/local/lib/python3.12/site-packages/clu/metric_writers/__init__.py", line 18 in
File "/usr/local/lib/python3.12/site-packages/maxtext/trainers/post_train/rl/train_rl.py", line 42 in

### Environment Information

- **Framework:** MaxText (GRPO RL / JAX Flax NNX)
- **Hardware:** 8x Cloud TPU v6e (`ct6e-standard-8t`, topology `v6e-2x4`)
- **OS:** Linux (GKE Standard Container Image)
- **Python:** 3.12
- **JAX:** 0.4.35+
- **Inference Engine:** vLLM V1 TPU (`v0.20.1rc1.dev136`)

### Additional Context

**Org**: Google Cloud GTM: Global Solutions Team

**Suggested Fix:**
In `src/maxtext/common/metric_logger.py`, move `from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics` inside the `write_metrics_to_managed_mldiagnostics()` method so
it loads lazily only when `config.managed_mldiagnostics=True`.

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.