AI-Hypercomputer / AI-Hypercomputer/maxtext

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

Đang mở
#4,611 1 bình luận 0 reaction 0 người được giao Xem trên GitHub
bug
Ngôn ngữ chính
Python
Star
2.4k
Fork
607
Merge trung bình
2 ngày 19 giờ
Pull request đã merge (30 ngày)
158

Mô tả

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

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.