AI-Hypercomputer / AI-Hypercomputer/maxtext

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

Offen
#4,611 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug
Vorherrschende Sprache
Python
Sterne
2.4k
Forks
607
Ø Merge
2 T. 19 Std.
Gemergte PRs (30 T.)
158

Beschreibung

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

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.