From 59e7fc92daf7a1993b86597859158ddc125394ca Mon Sep 17 00:00:00 2001 From: Matt Davidow Date: Fri, 7 Nov 2025 16:37:58 -0800 Subject: [PATCH] Fix rl and integration imports PiperOrigin-RevId: 829620709 --- src/MaxText/integration/tunix/tunix_adapter.py | 2 +- src/MaxText/integration/tunix/weight_mapping/__init__.py | 4 ++-- src/MaxText/rl/train_rl.py | 5 ++++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/MaxText/integration/tunix/tunix_adapter.py b/src/MaxText/integration/tunix/tunix_adapter.py index 499adda94..4d54f3aba 100644 --- a/src/MaxText/integration/tunix/tunix_adapter.py +++ b/src/MaxText/integration/tunix/tunix_adapter.py @@ -26,7 +26,7 @@ from jax import Array from flax import nnx from MaxText.layers.models import Transformer -from maxtext.src.maxtext.integration.tunix.utils import VllmWeightMapping +from MaxText.integration.tunix.utils import VllmWeightMapping from MaxText.utils.ckpt_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS # pylint: disable=ungrouped-imports diff --git a/src/MaxText/integration/tunix/weight_mapping/__init__.py b/src/MaxText/integration/tunix/weight_mapping/__init__.py index c1dc34a88..d250ee2fe 100644 --- a/src/MaxText/integration/tunix/weight_mapping/__init__.py +++ b/src/MaxText/integration/tunix/weight_mapping/__init__.py @@ -19,8 +19,8 @@ model name. This allows for easy extension to support new models. """ -from maxtext.src.maxtext.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING -from maxtext.src.maxtext.integration.tunix.weight_mapping.qwen3 import QWEN3_VLLM_MAPPING +from MaxText.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING +from MaxText.integration.tunix.weight_mapping.qwen3 import QWEN3_VLLM_MAPPING class StandaloneVllmWeightMapping: diff --git a/src/MaxText/rl/train_rl.py b/src/MaxText/rl/train_rl.py index 2a2bbeeae..fadda5f3d 100644 --- a/src/MaxText/rl/train_rl.py +++ b/src/MaxText/rl/train_rl.py @@ -75,7 +75,10 @@ os.environ["SKIP_JAX_PRECOMPILE"] = "1" -from MaxText import max_logging, max_utils, maxtext_utils, pyconfig +from MaxText import max_logging +from MaxText import max_utils +from MaxText import maxtext_utils +from MaxText import pyconfig from MaxText import model_creation_utils from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter from MaxText.rl.evaluate_rl import evaluate