diff --git a/docs/models.md b/docs/models.md index ff3a3a8ff..2ede50ed7 100644 --- a/docs/models.md +++ b/docs/models.md @@ -153,6 +153,45 @@ model, model_path = AutoModel.from_pretrained( ) ``` +##### Pipeline parallel MaxText models + +Tunix can use MaxText's pipeline runtime for trainer forward and backward +passes. Use `MaxTextPipelineConfig` to keep the MaxText config and JAX mesh in +sync. The helper builds a mesh with MaxText's `stage` and `tensor` axis names; +a native Tunix mesh using an axis such as `tp` does not enable MaxText pipeline +parallelism. + +For example, this layout assigns the 36 Qwen3-8B decoder layers to two pipeline +stages, with tensor parallelism degree four inside every stage: + +```python +from tunix.models import automodel +from tunix.models import maxtext_parallelism + +parallelism = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=2, + tensor_parallelism=4, + num_layers_per_pipeline_stage=18, + num_pipeline_microbatches=4, + pipeline_parallel_layers=36, +) +mesh = parallelism.create_mesh() +parallelism.validate_batch_size(global_batch_size=16) + +model, model_path = automodel.AutoModel.from_pretrained( + model_id="Qwen/Qwen3-8B", + mesh=mesh, + model_source=automodel.ModelSource.MAXTEXT, + model_path="gs://my-bucket/qwen3-8b-maxtext-checkpoint", + maxtext_pipeline_config=parallelism, + per_device_batch_size=2, +) +``` + +On eight devices, another valid Qwen3-8B layout is PP4 x TP2 with nine layers +per stage and eight pipeline microbatches. Pipeline and tensor degrees multiply: +PP8 x TP8 therefore requires 64 devices rather than eight. + To load a Maxtext model when launching training via a shell script, append the corresponding override arguments directly to the script execution: ```bash diff --git a/tests/models/automodel_test.py b/tests/models/automodel_test.py index 90d3d00ed..ad9d7488f 100644 --- a/tests/models/automodel_test.py +++ b/tests/models/automodel_test.py @@ -8,6 +8,7 @@ from absl.testing import parameterized import jax from tunix.models import automodel +from tunix.models import maxtext_parallelism from tunix.models import naming @@ -212,6 +213,91 @@ class MockMaxTextConfig: mock_config, mesh=mock_mesh, wrap_with_tunix_adapter=True ) + @mock.patch( + "tunix.models.automodel.download_model", + return_value="qwen3-8b", + ) + def test_from_pretrained_maxtext_pipeline_config(self, mock_download): + del mock_download + m_maxtext = types.ModuleType("maxtext") + m_configs = types.ModuleType("maxtext.configs") + m_pyconfig = types.ModuleType("maxtext.configs.pyconfig") + m_types = types.ModuleType("maxtext.configs.types") + m_utils = types.ModuleType("maxtext.utils") + m_creation = types.ModuleType("maxtext.utils.model_creation_utils") + + config = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=2, + tensor_parallelism=4, + num_layers_per_pipeline_stage=18, + num_pipeline_microbatches=4, + pipeline_parallel_layers=36, + ) + mesh = types.SimpleNamespace( + shape=dict( + zip( + maxtext_parallelism.MAXTEXT_MESH_AXIS_NAMES, + config.mesh_axis_shapes, + ) + ) + ) + + class MockMaxTextConfig: + model_fields = { + "skip_jax_distributed_system": True, + **{key: True for key in config.as_maxtext_kwargs()}, + } + + mock_config = mock.MagicMock() + m_pyconfig.initialize = mock.MagicMock(return_value=mock_config) + m_creation.from_pretrained = mock.MagicMock() + m_types.MaxTextConfig = MockMaxTextConfig + m_pyconfig.__file__ = "/opt/maxtext/configs/pyconfig.py" + m_configs.pyconfig = m_pyconfig + m_configs.types = m_types + m_utils.model_creation_utils = m_creation + m_maxtext.configs = m_configs + m_maxtext.utils = m_utils + + with mock.patch.dict( + "sys.modules", + { + "maxtext": m_maxtext, + "maxtext.configs": m_configs, + "maxtext.configs.pyconfig": m_pyconfig, + "maxtext.configs.types": m_types, + "maxtext.utils": m_utils, + "maxtext.utils.model_creation_utils": m_creation, + }, + ): + automodel.AutoModel.from_pretrained( + "Qwen/Qwen3-8B", + mesh=mesh, + model_source=automodel.ModelSource.MAXTEXT, + maxtext_pipeline_config=config, + ) + + called_argv = m_pyconfig.initialize.call_args[0][0] + self.assertEqual(called_argv[1], "/opt/maxtext/configs/base.yml") + self.assertIn("ici_pipeline_parallelism=2", called_argv) + self.assertIn("ici_tensor_parallelism=4", called_argv) + self.assertIn("num_layers_per_pipeline_stage=18", called_argv) + self.assertIn("num_pipeline_microbatches=4", called_argv) + + def test_rejects_maxtext_pipeline_config_for_native_model(self): + config = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=2, + tensor_parallelism=4, + ) + + with self.assertRaisesRegex(ValueError, "only supported"): + automodel.AutoModel.from_pretrained( + "Qwen/Qwen3-8B", + mesh=mock.MagicMock(), + model_source=automodel.ModelSource.HUGGINGFACE, + maxtext_pipeline_config=config, + ) + @parameterized.named_parameters(*_get_all_models_test_parameters()) def test_obtain_model_params_valid(self, model_name: str): automodel.call_model_config(model_name) @@ -433,5 +519,6 @@ class FakeConfig: self.assertEqual(called_config.flash_attention_block_size, 512) self.assertFalse(hasattr(called_config, "invalid_param")) + if __name__ == "__main__": absltest.main() diff --git a/tests/models/maxtext_parallelism_test.py b/tests/models/maxtext_parallelism_test.py new file mode 100644 index 000000000..e8728c1f0 --- /dev/null +++ b/tests/models/maxtext_parallelism_test.py @@ -0,0 +1,168 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import collections +import types + +from absl.testing import absltest +from absl.testing import parameterized +from tunix.models import maxtext_parallelism + + +class MaxTextPipelineConfigTest(parameterized.TestCase): + + @parameterized.named_parameters( + dict( + testcase_name="pp2_tp4_qwen3_8b", + pipeline_parallelism=2, + tensor_parallelism=4, + layers_per_stage=18, + microbatches=4, + expected_shapes=(1, 1, 2, 1, 1, 1, 1, 4, 1, 1, 1, 1), + ), + dict( + testcase_name="pp4_tp2_qwen3_8b", + pipeline_parallelism=4, + tensor_parallelism=2, + layers_per_stage=9, + microbatches=8, + expected_shapes=(1, 1, 4, 1, 1, 1, 1, 2, 1, 1, 1, 1), + ), + ) + def test_hybrid_layout( + self, + pipeline_parallelism, + tensor_parallelism, + layers_per_stage, + microbatches, + expected_shapes, + ): + config = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=pipeline_parallelism, + tensor_parallelism=tensor_parallelism, + num_layers_per_pipeline_stage=layers_per_stage, + num_pipeline_microbatches=microbatches, + pipeline_parallel_layers=36, + ) + + self.assertEqual(config.required_device_count, 8) + self.assertEqual(config.mesh_axis_shapes, expected_shapes) + self.assertEqual( + config.as_maxtext_kwargs()["ici_pipeline_parallelism"], + pipeline_parallelism, + ) + self.assertEqual( + config.as_maxtext_kwargs()["ici_tensor_parallelism"], + tensor_parallelism, + ) + config.validate_batch_size(16) + + def test_mesh_axes_match_maxtext_default_order(self): + self.assertEqual( + maxtext_parallelism.MAXTEXT_MESH_AXIS_NAMES, + ( + "diloco", + "data", + "stage", + "fsdp", + "fsdp_transpose", + "context", + "context_autoregressive", + "tensor", + "tensor_transpose", + "tensor_sequence", + "expert", + "autoregressive", + ), + ) + + def test_validate_exact_mesh(self): + config = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=2, + tensor_parallelism=4, + num_layers_per_pipeline_stage=18, + num_pipeline_microbatches=4, + ) + mesh = types.SimpleNamespace( + shape=collections.OrderedDict( + zip( + maxtext_parallelism.MAXTEXT_MESH_AXIS_NAMES, + config.mesh_axis_shapes, + ) + ) + ) + + config.validate_mesh(mesh) + + def test_rejects_tunix_tp_axis_names(self): + config = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=2, + tensor_parallelism=4, + ) + mesh = types.SimpleNamespace( + shape=collections.OrderedDict((("fsdp", 1), ("tp", 8))) + ) + + with self.assertRaisesRegex(ValueError, "Missing axes"): + config.validate_mesh(mesh) + + @parameterized.named_parameters( + dict( + testcase_name="one_pipeline_stage", + kwargs={"pipeline_parallelism": 1}, + error="at least 2", + ), + dict( + testcase_name="microbatches_not_divisible", + kwargs={ + "pipeline_parallelism": 4, + "num_pipeline_microbatches": 6, + }, + error="must be divisible", + ), + dict( + testcase_name="layers_not_divisible", + kwargs={ + "pipeline_parallelism": 4, + "num_layers_per_pipeline_stage": 3, + "pipeline_parallel_layers": 28, + }, + error="must be divisible", + ), + dict( + testcase_name="conflicting_fsdp_modes", + kwargs={ + "pipeline_parallelism": 2, + "pipeline_fsdp_ag_once": True, + "pipeline_fsdp_ag_per_repeat": True, + }, + error="mutually exclusive", + ), + ) + def test_invalid_config(self, kwargs, error): + with self.assertRaisesRegex(ValueError, error): + maxtext_parallelism.MaxTextPipelineConfig(**kwargs) + + def test_rejects_incompatible_batch_size(self): + config = maxtext_parallelism.MaxTextPipelineConfig( + pipeline_parallelism=4, + num_pipeline_microbatches=8, + ) + + with self.assertRaisesRegex(ValueError, "global_batch_size"): + config.validate_batch_size(12) + + +if __name__ == "__main__": + absltest.main() diff --git a/tunix/models/automodel.py b/tunix/models/automodel.py index 5c41a6318..63ef0204d 100644 --- a/tunix/models/automodel.py +++ b/tunix/models/automodel.py @@ -18,19 +18,30 @@ import gc import importlib import os +from pathlib import Path import shutil from typing import Any + from absl import logging from flax import nnx import jax import jax.numpy as jnp from orbax import checkpoint as ocp +from tunix.models import maxtext_parallelism from tunix.models import naming - _BASE_MODULE_PATH = 'tunix.models' # pylint: disable=invalid-name +def _maxtext_base_config_path(pyconfig_module: Any) -> str: + """Finds base.yml next to the imported MaxText pyconfig module.""" + module_file = getattr(pyconfig_module, '__file__', None) + if module_file is None: + # Some internal loaders and unit-test module shims do not expose __file__. + return 'src/maxtext/configs/base.yml' + return str(Path(module_file).resolve().with_name('base.yml')) + + class ModelModule(enum.Enum): """Specifies the type of model module to import.""" @@ -101,7 +112,9 @@ def call_model_config(model_name: str) -> Any: f"for model '{model_name}'. Target object type: {type(target_obj)}" ) - method_to_call = getattr(target_obj, config_id) # pyrefly: ignore[bad-argument-type] + method_to_call = getattr( + target_obj, config_id + ) # pyrefly: ignore[bad-argument-type] if not callable(method_to_call): raise TypeError( @@ -184,7 +197,8 @@ def _nnx_convert_and_reload() -> tuple[nnx.Module, Any]: logging.warning( 'model_path is not provided. Inferring from model_name. This may lead' ' to incorrect results if the model_name (%s) is not a standard Gemma' - ' model name.', model_name + ' model name.', + model_name, ) naming_info = naming.ModelNaming(model_name=model_name) version_dashed = None @@ -198,7 +212,9 @@ def _nnx_convert_and_reload() -> tuple[nnx.Module, Any]: else: # gemma dir_name = version_dashed - params_path = os.path.join(ckpt_path, dir_name) # pyrefly: ignore[no-matching-overload] + params_path = os.path.join( + ckpt_path, dir_name + ) # pyrefly: ignore[no-matching-overload] model, params = create_gemma_model_from_params(params_path, model_name) @@ -310,11 +326,15 @@ def download_model( if model_source == ModelSource.KAGGLE: from tunix.oss import utils as oss_utils # pylint: disable=g-import-not-at-top - return oss_utils.kaggle_pipeline(model_id_or_path, model_download_path) # pyrefly: ignore[bad-argument-type] + return oss_utils.kaggle_pipeline( + model_id_or_path, model_download_path + ) # pyrefly: ignore[bad-argument-type] elif model_source == ModelSource.HUGGINGFACE: from tunix.oss import utils as oss_utils # pylint: disable=g-import-not-at-top - return oss_utils.hf_pipeline(model_id_or_path, model_download_path) # pyrefly: ignore[bad-argument-type] + return oss_utils.hf_pipeline( + model_id_or_path, model_download_path + ) # pyrefly: ignore[bad-argument-type] elif model_source in (ModelSource.GCS, ModelSource.MAXTEXT): return model_id_or_path elif model_source == ModelSource.INTERNAL: @@ -332,7 +352,7 @@ def create_model_from_safe_tensors( model_config: Any, mesh: jax.sharding.Mesh, dtype: jnp.dtype | None = None, - mode: str = "auto", + mode: str = 'auto', ) -> Any: """Dynamically imports the correct module and calls `create_model_from_safe_tensors` based on the model_name. @@ -355,7 +375,11 @@ def create_model_from_safe_tensors( """ naming_info = naming.ModelNaming(model_name=model_name) if naming_info.model_family in ( - 'gemma', 'gemma1p1', 'gemma2', 'gemma3', 'gemma4' + 'gemma', + 'gemma1p1', + 'gemma2', + 'gemma3', + 'gemma4', ): params_module = get_model_module(model_name, ModelModule.PARAMS_SAFETENSORS) else: @@ -396,6 +420,9 @@ def from_pretrained( model_source: ModelSource = ModelSource.HUGGINGFACE, model_path: str | None = None, model_download_path: str | None = None, + maxtext_pipeline_config: ( + maxtext_parallelism.MaxTextPipelineConfig | None + ) = None, **kwargs, ) -> tuple[nnx.Module, str | None]: """Loads a pretrained model from a given identifier. @@ -419,6 +446,9 @@ def from_pretrained( model_download_path: The local directory where the model should be downloaded. The corresponding model_source will handle `None` cases differently. + maxtext_pipeline_config: Optional validated single-slice pipeline and + tensor parallelism layout for ``ModelSource.MAXTEXT``. The supplied + mesh must match this configuration. **kwargs: Additional keyword arguments passed to the underlying model creation functions. - For ModelSource.KAGGLE, Gemma models: `intermediate_ckpt_dir` , `rng_seed` @@ -432,6 +462,15 @@ def from_pretrained( model_params: Any = None naming_info = naming.ModelNaming(model_id=model_id) + if ( + maxtext_pipeline_config is not None + and model_source != ModelSource.MAXTEXT + ): + raise ValueError( + 'maxtext_pipeline_config is only supported with ' + 'model_source=ModelSource.MAXTEXT.' + ) + # Download the model if model_path: model_id_or_path = model_path @@ -452,6 +491,21 @@ def from_pretrained( # Case 1: MaxText models if model_source == ModelSource.MAXTEXT: + if maxtext_pipeline_config is not None: + maxtext_pipeline_config.validate_mesh(mesh) + pipeline_overrides = maxtext_pipeline_config.as_maxtext_kwargs() + conflicting_overrides = { + key: {'config': value, 'kwargs': kwargs[key]} + for key, value in pipeline_overrides.items() + if key in kwargs and kwargs[key] != value + } + if conflicting_overrides: + raise ValueError( + 'MaxText pipeline overrides conflict with explicit kwargs: ' + f'{conflicting_overrides}.' + ) + kwargs.update(pipeline_overrides) + try: import maxtext.configs.pyconfig as pyconfig # pylint: disable=g-import-not-at-top # pytype: disable=import-error from maxtext.configs.types import MaxTextConfig # pylint: disable=g-import-not-at-top # pytype: disable=import-error @@ -464,7 +518,7 @@ def from_pretrained( # We provide load_parameters_path instead of model_path since that's what maxtext expects. argv = [ '', - 'src/maxtext/configs/base.yml', + _maxtext_base_config_path(pyconfig), f'model_name={naming_info.model_name}', ] @@ -548,7 +602,8 @@ def from_pretrained( ) elif model_source == ModelSource.INTERNAL: model, model_params = create_gemma_model_from_params( - params_path=resolved_model_path, model_name=naming_info.model_name # pyrefly: ignore[bad-argument-type] + params_path=resolved_model_path, + model_name=naming_info.model_name, # pyrefly: ignore[bad-argument-type] ) else: raise NotImplementedError( @@ -566,7 +621,9 @@ def from_pretrained( # Common path for all other native Tunix models -- create model from safe tensors if not model_params: # pick corresponding config based on model version - model_params = call_model_config(naming_info.model_name) # pyrefly: ignore[bad-argument-type] + model_params = call_model_config( + naming_info.model_name + ) # pyrefly: ignore[bad-argument-type] # Get load_dtype explicitly from kwargs load_dtype_str = kwargs.get('load_dtype') @@ -574,8 +631,8 @@ def from_pretrained( load_dtype = getattr(jnp, load_dtype_str) except AttributeError: raise ValueError( - f"Invalid load_dtype: {load_dtype_str}. Must be a valid" - " jax.numpy type." + f'Invalid load_dtype: {load_dtype_str}. Must be a valid' + ' jax.numpy type.' ) except TypeError: load_dtype = load_dtype_str @@ -584,17 +641,27 @@ def from_pretrained( # use_flash_attention, flash_attention_block_size). if dataclasses.is_dataclass(model_params): valid_fields = {f.name for f in dataclasses.fields(model_params)} - overrides = {k: v for k, v in kwargs.items() if k in valid_fields and v is not None} - if 'remat_config' in overrides and isinstance(overrides['remat_config'], str): - model_module = get_model_module(naming_info.model_name, ModelModule.MODEL) + overrides = { + k: v + for k, v in kwargs.items() + if k in valid_fields and v is not None + } + if 'remat_config' in overrides and isinstance( + overrides['remat_config'], str + ): + model_module = get_model_module( + naming_info.model_name, ModelModule.MODEL + ) if hasattr(model_module, 'RematConfig'): remat_cfg_str = overrides['remat_config'] try: - overrides['remat_config'] = getattr(model_module.RematConfig, remat_cfg_str) + overrides['remat_config'] = getattr( + model_module.RematConfig, remat_cfg_str + ) except AttributeError: raise ValueError( - f"Invalid remat_config: {remat_cfg_str}. Must be a valid" - " RematConfig type." + f'Invalid remat_config: {remat_cfg_str}. Must be a valid' + ' RematConfig type.' ) if 'dtype' in overrides: dtype_str = overrides['dtype'] @@ -602,7 +669,7 @@ def from_pretrained( overrides['dtype'] = getattr(jnp, dtype_str) except AttributeError: raise ValueError( - f"Invalid dtype: {dtype_str}. Must be a valid jax.numpy type." + f'Invalid dtype: {dtype_str}. Must be a valid jax.numpy type.' ) except TypeError: pass diff --git a/tunix/models/maxtext_parallelism.py b/tunix/models/maxtext_parallelism.py new file mode 100644 index 000000000..c38390456 --- /dev/null +++ b/tunix/models/maxtext_parallelism.py @@ -0,0 +1,248 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pipeline-parallel configuration helpers for MaxText-backed models.""" + +from collections.abc import Sequence +import dataclasses +import math +from typing import Any + +from tunix.utils import mesh as mesh_lib + +# Keep this in the same order as MaxText's default ``mesh_axes``. Including +# singleton axes matters because MaxText logical partition rules may refer to +# them even when their parallelism degree is one. +MAXTEXT_MESH_AXIS_NAMES = ( + "diloco", + "data", + "stage", + "fsdp", + "fsdp_transpose", + "context", + "context_autoregressive", + "tensor", + "tensor_transpose", + "tensor_sequence", + "expert", + "autoregressive", +) + + +@dataclasses.dataclass(frozen=True, slots=True) +class MaxTextPipelineConfig: + """Single-slice MaxText pipeline and tensor parallelism configuration. + + This helper deliberately configures only ICI parallelism. Multi-slice DCN + layouts and custom MaxText mesh rules should continue to be configured + directly in MaxText. + + Attributes: + pipeline_parallelism: Number of pipeline stages (the ``stage`` mesh axis). + tensor_parallelism: Tensor-parallel degree within each pipeline stage. + data_parallelism: Data-parallel degree. + fsdp_parallelism: FSDP degree. + num_layers_per_pipeline_stage: Decoder layers executed by a stage during + one pipeline repeat. + num_pipeline_microbatches: Number of microbatches in a forward pass. When + set, it must be a multiple of ``pipeline_parallelism``. + pipeline_parallel_layers: Optional number of decoder layers assigned to + the pipeline. ``None`` lets MaxText use all decoder layers. + pipeline_delay_activation_forwarding: Whether MaxText should delay + activation forwarding to expose communication/computation overlap. + pipeline_fsdp_ag_once: Whether MaxText should all-gather FSDP weights once + before pipeline execution. + pipeline_fsdp_ag_per_repeat: Whether MaxText should prefetch FSDP weights + before every circular-pipeline repeat. + """ + + pipeline_parallelism: int + tensor_parallelism: int = 1 + data_parallelism: int = 1 + fsdp_parallelism: int = 1 + num_layers_per_pipeline_stage: int = 1 + num_pipeline_microbatches: int | None = None + pipeline_parallel_layers: int | None = None + pipeline_delay_activation_forwarding: bool = False + pipeline_fsdp_ag_once: bool = False + pipeline_fsdp_ag_per_repeat: bool = False + + def __post_init__(self): + degrees = { + "pipeline_parallelism": self.pipeline_parallelism, + "tensor_parallelism": self.tensor_parallelism, + "data_parallelism": self.data_parallelism, + "fsdp_parallelism": self.fsdp_parallelism, + "num_layers_per_pipeline_stage": self.num_layers_per_pipeline_stage, + } + for name, value in degrees.items(): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}.") + + if self.pipeline_parallelism < 2: + raise ValueError( + "pipeline_parallelism must be at least 2; use ordinary MaxText " + "parallelism when no pipeline stages are required." + ) + + if self.num_pipeline_microbatches is not None: + if ( + not isinstance(self.num_pipeline_microbatches, int) + or isinstance(self.num_pipeline_microbatches, bool) + or self.num_pipeline_microbatches <= 0 + ): + raise ValueError( + "num_pipeline_microbatches must be a positive integer or None, " + f"got {self.num_pipeline_microbatches!r}." + ) + if self.num_pipeline_microbatches % self.pipeline_parallelism: + raise ValueError( + "num_pipeline_microbatches must be divisible by " + f"pipeline_parallelism ({self.pipeline_parallelism}), got " + f"{self.num_pipeline_microbatches}." + ) + if ( + self.pipeline_delay_activation_forwarding + and self.num_pipeline_microbatches < 2 * self.pipeline_parallelism + ): + raise ValueError( + "pipeline_delay_activation_forwarding requires at least twice " + "as many microbatches as pipeline stages." + ) + + if self.pipeline_parallel_layers is not None: + if ( + not isinstance(self.pipeline_parallel_layers, int) + or isinstance(self.pipeline_parallel_layers, bool) + or self.pipeline_parallel_layers <= 0 + ): + raise ValueError( + "pipeline_parallel_layers must be a positive integer or None, " + f"got {self.pipeline_parallel_layers!r}." + ) + layers_per_repeat = ( + self.pipeline_parallelism * self.num_layers_per_pipeline_stage + ) + if self.pipeline_parallel_layers % layers_per_repeat: + raise ValueError( + "pipeline_parallel_layers must be divisible by pipeline stages " + "times layers per stage; got " + f"{self.pipeline_parallel_layers} % {layers_per_repeat}." + ) + + if self.pipeline_fsdp_ag_once and self.pipeline_fsdp_ag_per_repeat: + raise ValueError( + "pipeline_fsdp_ag_once and pipeline_fsdp_ag_per_repeat are " + "mutually exclusive." + ) + + @property + def required_device_count(self) -> int: + """Number of accelerator devices required by this ICI layout.""" + return math.prod(( + self.pipeline_parallelism, + self.tensor_parallelism, + self.data_parallelism, + self.fsdp_parallelism, + )) + + @property + def mesh_axis_shapes(self) -> tuple[int, ...]: + """MaxText-compatible mesh shape, including singleton logical axes.""" + axis_sizes = { + "data": self.data_parallelism, + "stage": self.pipeline_parallelism, + "fsdp": self.fsdp_parallelism, + "tensor": self.tensor_parallelism, + } + return tuple(axis_sizes.get(axis, 1) for axis in MAXTEXT_MESH_AXIS_NAMES) + + def create_mesh(self, devices: Sequence[Any] | None = None): + """Creates a MaxText-compatible JAX mesh for this configuration.""" + return mesh_lib.create_mesh( + self.mesh_axis_shapes, + MAXTEXT_MESH_AXIS_NAMES, + devices=devices, + ) + + def validate_mesh(self, mesh: Any) -> None: + """Validates that an existing mesh implements this exact ICI layout.""" + shape = getattr(mesh, "shape", None) + if shape is None or not hasattr(shape, "get"): + raise ValueError("mesh must expose a mapping-like shape attribute.") + + missing_axes = [ + axis for axis in MAXTEXT_MESH_AXIS_NAMES if shape.get(axis) is None + ] + if missing_axes: + raise ValueError( + "MaxText pipeline meshes must include all logical mesh axes, even " + f"when their size is one. Missing axes: {missing_axes}." + ) + + mismatches = { + axis: {"expected": expected, "actual": int(shape.get(axis))} + for axis, expected in zip( + MAXTEXT_MESH_AXIS_NAMES, self.mesh_axis_shapes + ) + if int(shape.get(axis)) != expected + } + if mismatches: + raise ValueError( + f"mesh shape does not match MaxTextPipelineConfig: {mismatches}." + ) + + mesh_size = math.prod(int(shape.get(axis)) for axis in shape) + if mesh_size != self.required_device_count: + raise ValueError( + f"mesh uses {mesh_size} devices, but this configuration requires " + f"{self.required_device_count}." + ) + + def validate_batch_size(self, global_batch_size: int) -> None: + """Validates pipeline microbatch divisibility for a global batch.""" + if not isinstance(global_batch_size, int) or global_batch_size <= 0: + raise ValueError( + "global_batch_size must be a positive integer, got " + f"{global_batch_size!r}." + ) + if ( + self.num_pipeline_microbatches is not None + and global_batch_size % self.num_pipeline_microbatches + ): + raise ValueError( + f"global_batch_size ({global_batch_size}) must be divisible by " + "num_pipeline_microbatches " + f"({self.num_pipeline_microbatches})." + ) + + def as_maxtext_kwargs(self) -> dict[str, Any]: + """Returns validated ``AutoModel.from_pretrained`` MaxText overrides.""" + overrides: dict[str, Any] = { + "ici_pipeline_parallelism": self.pipeline_parallelism, + "ici_tensor_parallelism": self.tensor_parallelism, + "ici_data_parallelism": self.data_parallelism, + "ici_fsdp_parallelism": self.fsdp_parallelism, + "num_layers_per_pipeline_stage": self.num_layers_per_pipeline_stage, + "pipeline_delay_activation_forwarding": ( + self.pipeline_delay_activation_forwarding + ), + "pipeline_fsdp_ag_once": self.pipeline_fsdp_ag_once, + "pipeline_fsdp_ag_per_repeat": self.pipeline_fsdp_ag_per_repeat, + } + if self.num_pipeline_microbatches is not None: + overrides["num_pipeline_microbatches"] = self.num_pipeline_microbatches + if self.pipeline_parallel_layers is not None: + overrides["pipeline_parallel_layers"] = self.pipeline_parallel_layers + return overrides