AI-Hypercomputer / AI-Hypercomputer/maxtext
[bug] Dynamic TPU Metric Protobuf Import Causes Linkage Panic ('Check failed: GeneratedDatabase()->Add(...)')
- Langage dominant
- Python
- Étoiles
- 2.4k
- Forks
- 607
- Merge moyen
- 2 j 19 h
- PR mergées (30 j)
- 158
Description
### 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`.
Guide de contribution
Ouvrir le guide de contribution
Évaluation
Cette issue n'a pas encore été évaluée.