AI-Hypercomputer / AI-Hypercomputer/maxtext

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

未关闭
#4,611 1 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看
bug
主要语言
Python
星标
2.4k
派生
607
平均合并
2 天 19 小时
30 天内合并 PR
158

描述

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

贡献指南

打开贡献指南

评估

这个 Issue 还没有评估数据。

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。