From 495a7b2ccdb6ab8ecfeb2eafbf53f19f5b20bac8 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 15 Jun 2026 06:41:16 +0000 Subject: [PATCH] grpo_trainer: route goodput/vertex imports through gcloud_stub for decoupled mode (cherry picked from commit 83448872f2a8ccd812703cd5284c3543deffa23a) --- src/maxtext/experimental/rl/grpo_trainer.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/maxtext/experimental/rl/grpo_trainer.py b/src/maxtext/experimental/rl/grpo_trainer.py index f6038ae018..41b20e4e04 100644 --- a/src/maxtext/experimental/rl/grpo_trainer.py +++ b/src/maxtext/experimental/rl/grpo_trainer.py @@ -67,7 +67,9 @@ import transformers -from ml_goodput_measurement.src.goodput import GoodputRecorder +from maxtext.common.goodput import goodput as _goodput_module + +GoodputRecorder = _goodput_module.GoodputRecorder import maxtext as mt from maxtext.configs import pyconfig @@ -87,7 +89,7 @@ from maxtext.experimental.rl import grpo_input_pipeline from maxtext.experimental.rl import grpo_utils from maxtext.common.metric_logger import MetricLogger -from maxtext.common.vertex_tensorboard import VertexTensorboardManager +from maxtext.common.gcloud_stub import vertex_tensorboard_modules from maxtext.utils import exceptions from maxtext.utils import gcs_utils from maxtext.utils import max_logging @@ -98,6 +100,8 @@ from maxtext.utils import sharding from maxtext.utils import train_utils +VertexTensorboardManager, _vertex_tb_is_stub = vertex_tensorboard_modules() + # pylint: disable=too-many-positional-arguments # -----------------------------------------------------------------------------