From f307942d49105cda88660307498b0ad9f3978568 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Fri, 5 Dec 2025 08:53:20 +0000 Subject: [PATCH 01/14] update usage --- src/MaxText/layers/decoders.py | 12 +++------ src/MaxText/layers/pipeline.py | 47 +++++++++++++++------------------- 2 files changed, 25 insertions(+), 34 deletions(-) diff --git a/src/MaxText/layers/decoders.py b/src/MaxText/layers/decoders.py index 249a65b9da..9fd266c25c 100644 --- a/src/MaxText/layers/decoders.py +++ b/src/MaxText/layers/decoders.py @@ -399,14 +399,10 @@ def setup(self): self.decoder_layer = self.get_decoder_layers() self.norm_layer = self.get_norm_layer(num_features=self.config.emb_dim) if self.config.using_pipeline_parallelism: - nnx_decoder_classes = self.get_nnx_decoder_layers() - if nnx_decoder_classes is not None: - pipeline_stage_module = self.get_pipeline_stage_module(nnx_decoder_classes, use_nnx=True) - else: - pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer, use_nnx=False) + pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer) remat_policy = self.get_remat_policy() self.pipeline_module = pipeline.create_pipeline( - config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy, use_nnx=(nnx_decoder_classes is not None) + config=self.config, mesh=self.mesh, layer=pipeline_stage_module, remat_policy=remat_policy, ) def minimal_policy(self, with_context=False): @@ -685,7 +681,7 @@ def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, me config=cfg, mesh=mesh, name=metadata_axis_name, quant=self.quant, **kwargs # pytype: disable=wrong-keyword-args ) - def get_pipeline_stage_module(self, decoder_blocks, use_nnx=False): + def get_pipeline_stage_module(self, decoder_blocks): """get pipeline stage module Args: @@ -703,7 +699,7 @@ def get_layer_to_pipeline(blocks, cfg): cfg = self.config base_stage = get_layer_to_pipeline(decoder_blocks, cfg) - if use_nnx: + if issubclass(base_stage, nnx.Module): if cfg.num_layers_per_pipeline_stage == 1: return base_stage else: diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 20c9dbe51b..057642bca7 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -32,6 +32,7 @@ from MaxText.sharding import all_gather_over_fsdp from MaxText import max_logging from MaxText.layers import nnx_wrappers +from MaxText.layers.initializers import variable_to_logically_partitioned class Pipeline(nnx.Module): @@ -53,10 +54,10 @@ class Pipeline(nnx.Module): def __init__( self, - layers: Callable | type, + layer: Callable[..., nnx.Module|nn.Module], config: Config, mesh: Mesh, - rngs: nnx.Rngs = None, + rngs: nnx.Rngs, remat_policy: Any = None, ): """Initialize Pipeline with NNX or Linen decoder layers. @@ -88,20 +89,19 @@ def __init__( self.seq_len_axis_name = "activation_length_no_exp" # Detect if layers is a Linen class/instance or NNX class - self._is_linen = (isinstance(layers, type) and issubclass(layers, nn.Module)) or isinstance(layers, nn.Module) - + self._is_linen = (isinstance(layer, type) and issubclass(layer, nn.Module)) or isinstance(layer, nn.Module) if self._is_linen: - if isinstance(layers, nn.Module): - self.layers = layers + if isinstance(layer, nn.Module): + self.layer = layer else: - self.layers = layers(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + self.layer = layer(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) self._linen_variables = None else: # Create num_stages independent NNX instances, stored as attributes for # NNX pytree tracking (not as Python lists). for s in range(self.num_stages): stage_rngs = nnx.Rngs(s) - instance = layers( + instance = layer( config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN, @@ -361,7 +361,7 @@ def _initialize_linen_parameters(self, sample_input, sample_seg_ids, sample_posi linen_rngs = {'params': jax.random.PRNGKey(0), 'dropout': jax.random.PRNGKey(1)} - base_params = self.layers.init( + base_params = self.layer.init( linen_rngs, sample_input, sample_seg_ids, @@ -392,7 +392,7 @@ def _run_stages_linen( ) def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos): - output = self.layers.apply( + output = self.layer.apply( stage_params, stage_input, stage_seg_ids, @@ -606,16 +606,15 @@ class PipelineToLinen(nnx_wrappers.ToLinen): def create_pipeline( config: Config, - layers: Callable | type, + layer: Callable | type, mesh: Mesh, remat_policy: Any = None, - use_nnx: bool = True, -) -> PipelineToLinen: +) -> nnx_wrappers.ToLinen: """Factory function to create a Pipeline wrapped as a Linen module. Args: config: Model configuration - layers: NNX or Linen decoder layer class to use for pipeline stages + layer: NNX or Linen decoder layer class to use for pipeline stages mesh: Device mesh for sharding remat_policy: Remat policy for loop iterations use_nnx: Whether to use NNX pipeline (True) or Linen (False) @@ -623,17 +622,13 @@ def create_pipeline( Returns: PipelineToLinen wrapper around the NNX Pipeline """ - if not use_nnx: - raise ValueError("This implementation only supports NNX pipelines (use_nnx=True)") - - wrapped = PipelineToLinen( + return nnx_wrappers.to_linen( Pipeline, - kwargs={ - 'layers': layers, - 'config': config, - 'mesh': mesh, - 'remat_policy': remat_policy, - } + config=config, + mesh=mesh, + layer=layer, + remat_policy=remat_policy, + name="pipeline_module", + abstract_init=False, + metadata_fn=variable_to_logically_partitioned, ) - - return wrapped From bd540a05d223fc42360c413aa07f58ee2baf623d Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Fri, 5 Dec 2025 10:19:27 +0000 Subject: [PATCH 02/14] first attempt --- src/MaxText/layers/embeddings.py | 1 + src/MaxText/layers/pipeline.py | 42 ++++++++++++++++++++++++++------ 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/src/MaxText/layers/embeddings.py b/src/MaxText/layers/embeddings.py index 76fb1b8a78..5285e530e2 100644 --- a/src/MaxText/layers/embeddings.py +++ b/src/MaxText/layers/embeddings.py @@ -342,6 +342,7 @@ def __call__( "The embedding dims of the rotary position embedding" "must match the hidden dimension of the inputs." ) + position = position[:, :, jnp.newaxis, jnp.newaxis] sinusoid_inp = position / self.timescale sin = jnp.sin(sinusoid_inp).astype(inputs.dtype) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 057642bca7..16ee90d5f2 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -14,7 +14,6 @@ """ Pipeline layer wrapping a decoder layer(s). Supports circular pipelining """ -import functools from typing import Any, Callable import numpy as np @@ -29,8 +28,6 @@ from flax import nnx from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT -from MaxText.sharding import all_gather_over_fsdp -from MaxText import max_logging from MaxText.layers import nnx_wrappers from MaxText.layers.initializers import variable_to_logically_partitioned @@ -52,6 +49,8 @@ class Pipeline(nnx.Module): remat_policy: Remat policy to use for the loop iterations """ + _linen_variables: nnx.Data[dict | None] + def __init__( self, layer: Callable[..., nnx.Module|nn.Module], @@ -89,12 +88,9 @@ def __init__( self.seq_len_axis_name = "activation_length_no_exp" # Detect if layers is a Linen class/instance or NNX class - self._is_linen = (isinstance(layer, type) and issubclass(layer, nn.Module)) or isinstance(layer, nn.Module) + self._is_linen = isinstance(layer, nn.Module) if self._is_linen: - if isinstance(layer, nn.Module): - self.layer = layer - else: - self.layer = layer(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + self.layer = layer self._linen_variables = None else: # Create num_stages independent NNX instances, stored as attributes for @@ -361,6 +357,12 @@ def _initialize_linen_parameters(self, sample_input, sample_seg_ids, sample_posi linen_rngs = {'params': jax.random.PRNGKey(0), 'dropout': jax.random.PRNGKey(1)} + # Ensure inputs to init are 2D + if sample_seg_ids is not None and sample_seg_ids.ndim == 1: + sample_seg_ids = sample_seg_ids[None, :] + if sample_positions is not None and sample_positions.ndim == 1: + sample_positions = sample_positions[None, :] + base_params = self.layer.init( linen_rngs, sample_input, @@ -392,6 +394,12 @@ def _run_stages_linen( ) def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos): + # Ensure per-stage positions and segment_ids are 2D for the layer call + if stage_pos is not None and stage_pos.ndim == 1: + stage_pos = stage_pos[None, :] + if stage_seg_ids is not None and stage_seg_ids.ndim == 1: + stage_seg_ids = stage_seg_ids[None, :] + output = self.layer.apply( stage_params, stage_input, @@ -441,6 +449,12 @@ def _run_stages_vmapped( ) def call_stage(state, stage_input, stage_seg_ids, stage_pos): + + if stage_pos is not None and stage_pos.ndim == 1: + stage_pos = stage_pos[None, :] + if stage_seg_ids is not None and stage_seg_ids.ndim == 1: + stage_seg_ids = stage_seg_ids[None, :] + module = nnx.merge(graphdef, state) output = module(stage_input, stage_seg_ids, stage_pos, deterministic, model_mode) if isinstance(output, tuple): @@ -449,6 +463,10 @@ def call_stage(state, stage_input, stage_seg_ids, stage_pos): if stages_segment_ids is None: def call_stage_no_seg(state, stage_input, stage_pos): + # Ensure per-stage positions is 2D for the layer call + if stage_pos is not None and stage_pos.ndim == 1: + stage_pos = stage_pos[None, :] + module = nnx.merge(graphdef, state) output = module(stage_input, None, stage_pos, deterministic, model_mode) if isinstance(output, tuple): @@ -538,6 +556,14 @@ def __call__( self.config.emb_dim, ) ) + if positions is not None: + positions = positions.reshape( + (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length) + ) + if segment_ids is not None: + segment_ids = segment_ids.reshape( + (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length) + ) if self._is_linen and self._linen_variables is None: example_input = inputs[0] From 0704eb8868c23bb9596d0015a71ada18f62e6a79 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 8 Dec 2025 06:26:25 +0000 Subject: [PATCH 03/14] second attempt --- src/MaxText/layers/decoders.py | 14 --- src/MaxText/layers/pipeline.py | 199 ++++++++++++--------------------- 2 files changed, 72 insertions(+), 141 deletions(-) diff --git a/src/MaxText/layers/decoders.py b/src/MaxText/layers/decoders.py index 9fd266c25c..0306a8b201 100644 --- a/src/MaxText/layers/decoders.py +++ b/src/MaxText/layers/decoders.py @@ -699,20 +699,6 @@ def get_layer_to_pipeline(blocks, cfg): cfg = self.config base_stage = get_layer_to_pipeline(decoder_blocks, cfg) - if issubclass(base_stage, nnx.Module): - if cfg.num_layers_per_pipeline_stage == 1: - return base_stage - else: - return lambda config, mesh, model_mode, rngs, quant=None: SequentialNNXWrapper( - decoder_layer_class=base_stage, - num_decoder_layers=cfg.num_layers_per_pipeline_stage, - config=config, - mesh=mesh, - model_mode=model_mode, - rngs=rngs, - quant=quant - ) - if cfg.set_remat_policy_on_layers_per_stage: policy = self.get_remat_policy() base_stage = self.set_remat_policy([base_stage], policy)[0] diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 16ee90d5f2..5939c46036 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -49,7 +49,7 @@ class Pipeline(nnx.Module): remat_policy: Remat policy to use for the loop iterations """ - _linen_variables: nnx.Data[dict | None] + _linen_variables: nnx.Data[list[dict] | None] def __init__( self, @@ -87,24 +87,11 @@ def __init__( self.batch_axis_name = "activation_batch" self.seq_len_axis_name = "activation_length_no_exp" - # Detect if layers is a Linen class/instance or NNX class - self._is_linen = isinstance(layer, nn.Module) - if self._is_linen: - self.layer = layer - self._linen_variables = None - else: - # Create num_stages independent NNX instances, stored as attributes for - # NNX pytree tracking (not as Python lists). - for s in range(self.num_stages): - stage_rngs = nnx.Rngs(s) - instance = layer( - config=config, - mesh=mesh, - model_mode=MODEL_MODE_TRAIN, - rngs=stage_rngs, - quant=None, - ) - setattr(self, f'stage_{s}', instance) + # TODO: Support NNX layers in addition to Linen layers + self._is_linen = True # Always true given the input + self.layer = layer # Store the module instance + self._linen_variables = None # Will be populated in _initialize_linen_parameters + def need_circ_storage(self): return ( @@ -352,31 +339,27 @@ def permute_output_micro_per_stage_dim(self, output): def _initialize_linen_parameters(self, sample_input, sample_seg_ids, sample_positions, deterministic, model_mode): """Initialize Linen module parameters for all stages.""" - if self._linen_variables is not None: + if not self._is_linen or self._linen_variables is not None: return - linen_rngs = {'params': jax.random.PRNGKey(0), 'dropout': jax.random.PRNGKey(1)} - # Ensure inputs to init are 2D if sample_seg_ids is not None and sample_seg_ids.ndim == 1: sample_seg_ids = sample_seg_ids[None, :] if sample_positions is not None and sample_positions.ndim == 1: sample_positions = sample_positions[None, :] - base_params = self.layer.init( - linen_rngs, - sample_input, - sample_seg_ids, - sample_positions, - deterministic, - model_mode, - ) - stage_params = {} - for stage_idx in range(self.num_stages): - stage_params[f'stage_{stage_idx}'] = jax.tree_util.tree_map(lambda x: x, base_params) + stage_params_list = [] + for _ in range(self.num_stages): + linen_rngs = { + 'params': self.rngs.params(), + 'dropout': self.rngs.dropout() + } + base_params = self.layer.init( + linen_rngs, sample_input, sample_seg_ids, sample_positions, deterministic, model_mode) + stage_params_list.append(base_params) + self._linen_variables = stage_params_list - self._linen_variables = {'params': stage_params} def _run_stages_linen( self, @@ -385,15 +368,16 @@ def _run_stages_linen( stages_positions, deterministic, model_mode, + dropout_key ): """Run stages using Linen module with manual vmap.""" - stage_params_list = [self._linen_variables['params'][f'stage_{i}'] for i in range(self.num_stages)] + stage_params_list = self._linen_variables if self._linen_variables is not None else [] stacked_params = jax.tree_util.tree_map( lambda *xs: jnp.stack(xs, axis=0), *stage_params_list ) - def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos): + def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos, rng_key): # Ensure per-stage positions and segment_ids are 2D for the layer call if stage_pos is not None and stage_pos.ndim == 1: stage_pos = stage_pos[None, :] @@ -407,24 +391,28 @@ def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos): stage_pos, deterministic, model_mode, + rngs={"dropout": rng_key}, ) - if isinstance(output, tuple): - return output[0] - return output + return output[0] if isinstance(output, tuple) else output + + dropout_keys_for_vmap = jax.random.split(dropout_key, self.num_stages) if stages_segment_ids is None: - vmapped_apply = jax.vmap( - lambda p, i, pos: apply_stage(p, i, None, pos), - in_axes=(0, 0, 0), - out_axes=0 + vmap_fn = lambda p, i, pos, rng: apply_stage(p, i, None, pos, rng) + in_axes = (0, 0, 0, 0) + stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( + stacked_params, stages_inputs, stages_positions, dropout_keys_for_vmap ) - stages_outputs = vmapped_apply(stacked_params, stages_inputs, stages_positions) else: - vmapped_apply = jax.vmap(apply_stage, in_axes=(0, 0, 0, 0), out_axes=0) - stages_outputs = vmapped_apply(stacked_params, stages_inputs, stages_segment_ids, stages_positions) + vmap_fn = lambda p, i, s, pos, rng: apply_stage(p, i, s, pos, rng) + in_axes = (0, 0, 0, 0, 0) + stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( + stacked_params, stages_inputs, stages_segment_ids, stages_positions, dropout_keys_for_vmap + ) return stages_outputs + def _run_stages_vmapped( self, stages_inputs, @@ -449,38 +437,35 @@ def _run_stages_vmapped( ) def call_stage(state, stage_input, stage_seg_ids, stage_pos): - if stage_pos is not None and stage_pos.ndim == 1: stage_pos = stage_pos[None, :] if stage_seg_ids is not None and stage_seg_ids.ndim == 1: stage_seg_ids = stage_seg_ids[None, :] - module = nnx.merge(graphdef, state) - output = module(stage_input, stage_seg_ids, stage_pos, deterministic, model_mode) - if isinstance(output, tuple): - return output[0] - return output + output = module(stage_input, stage_seg_ids, stage_pos, deterministic=deterministic, model_mode=model_mode) + return output[0] if isinstance(output, tuple) else output if stages_segment_ids is None: def call_stage_no_seg(state, stage_input, stage_pos): - # Ensure per-stage positions is 2D for the layer call if stage_pos is not None and stage_pos.ndim == 1: stage_pos = stage_pos[None, :] - module = nnx.merge(graphdef, state) - output = module(stage_input, None, stage_pos, deterministic, model_mode) - if isinstance(output, tuple): - return output[0] - return output - - vmapped_call = jax.vmap(call_stage_no_seg, in_axes=(0, 0, 0), out_axes=0) - stages_outputs = vmapped_call(stacked_state, stages_inputs, stages_positions) + output = module(stage_input, None, stage_pos, deterministic=deterministic, model_mode=model_mode) + return output[0] if isinstance(output, tuple) else output + vmap_fn = call_stage_no_seg + in_axes = (0, 0, 0) + stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( + stacked_state, stages_inputs, stages_positions + ) else: - vmapped_call = jax.vmap(call_stage, in_axes=(0, 0, 0, 0), out_axes=0) - stages_outputs = vmapped_call(stacked_state, stages_inputs, stages_segment_ids, stages_positions) - + vmap_fn = call_stage + in_axes = (0, 0, 0, 0) + stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( + stacked_state, stages_inputs, stages_segment_ids, stages_positions + ) return stages_outputs + def run_one_iteration( self, loop_state, @@ -494,33 +479,19 @@ def run_one_iteration( shift = loop_state["shift"] circ_storage = loop_state["circ_storage"] loop_iteration = loop_state["loop_iteration"] - + iter_dropout_key = self.rngs.dropout() + microbatch_ids, _ = self.get_microbatch_and_repeat_ids(loop_iteration) - stages_inputs = self.get_iteration_inputs(loop_iteration, state_io, circ_storage, shift) stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") stages_positions = self.vmap_gather(positions, microbatch_ids, 0) if positions is not None else None stages_segment_ids = self.vmap_gather(segment_ids, microbatch_ids, 0) if segment_ids is not None else None - if self._is_linen: - stages_output = self._run_stages_linen( - stages_inputs, - stages_segment_ids, - stages_positions, - deterministic, - model_mode, - ) - else: - stages_output = self._run_stages_vmapped( - stages_inputs, - stages_segment_ids, - stages_positions, - deterministic, - model_mode, - ) + stages_output = self._run_stages_linen(stages_inputs, stages_segment_ids, stages_positions, deterministic, model_mode, iter_dropout_key) + + new_loop_state = self.get_new_loop_state(stages_output, loop_state) + return new_loop_state - new_state = self.get_new_loop_state(stages_output, loop_state) - return new_state def get_pipeline_remat_policy(self): """Returns the remat policy for pipeline iterations.""" @@ -548,42 +519,18 @@ def __call__( Reshapes inputs into microbatches, runs pipeline iterations with bubble handling, and returns outputs reshaped to original batch size. """ - inputs = inputs.reshape( - ( - self.config.num_pipeline_microbatches, - self.pipeline_microbatch_size, - self.config.max_target_length, - self.config.emb_dim, - ) - ) + inputs = inputs.reshape((self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length, self.config.emb_dim)) if positions is not None: - positions = positions.reshape( - (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length) - ) + positions = positions.reshape((self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length)) if segment_ids is not None: - segment_ids = segment_ids.reshape( - (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length) - ) + segment_ids = segment_ids.reshape((self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length)) - if self._is_linen and self._linen_variables is None: + if self._linen_variables is None: example_input = inputs[0] example_seg_ids = segment_ids[0] if segment_ids is not None else None example_pos = positions[0] if positions is not None else None self._initialize_linen_parameters(example_input, example_seg_ids, example_pos, deterministic, model_mode) - ag_sharding = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec(None, None)) - if positions is not None: - positions = jax.lax.with_sharding_constraint(positions, ag_sharding) - positions = positions.reshape( - (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length) - ) - - if segment_ids is not None: - segment_ids = jax.lax.with_sharding_constraint(segment_ids, ag_sharding) - segment_ids = segment_ids.reshape( - (self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length) - ) - loop_state = self.init_states(inputs) bubble_iterations = self.forwarding_delay * (self.num_stages - 1) @@ -591,37 +538,35 @@ def __call__( total_iterations = real_iterations + bubble_iterations if self.config.scan_pipeline_iterations: - def run_iteration_scannable(loop_state, xs): - return ( - self.run_one_iteration( - loop_state, positions, segment_ids, deterministic, model_mode - ), - None, - ) + def run_iteration_scannable(pipeline_instance, loop_state, xs): + new_loop_state = pipeline_instance.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) + return new_loop_state + remat_fn = run_iteration_scannable if self.config.set_remat_policy_on_pipeline_iterations: - run_iteration_scannable = jax.checkpoint( + remat_fn = nnx.remat( run_iteration_scannable, prevent_cse=False, policy=self.get_pipeline_remat_policy(), ) - loop_state, _ = jax.lax.scan(run_iteration_scannable, loop_state, None, length=total_iterations) + ScannedFn = nnx.scan( + remat_fn, + in_axes=(None, nnx.Carry, None), # Corresponds to (self, loop_state, xs) + out_axes=nnx.Carry, # Corresponds to new_loop_state + length=total_iterations, + ) + loop_state = ScannedFn(self, loop_state, None) else: for _ in range(total_iterations): - loop_state = self.run_one_iteration( - loop_state, positions, segment_ids, deterministic, model_mode - ) + loop_state = self.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) final_output = self.permute_output_micro_per_stage_dim(loop_state["state_io"]) - final_output = jnp.reshape( final_output, (self.config.micro_batch_size_to_train_on, self.config.max_target_length, self.config.emb_dim) ) - return final_output - class PipelineToLinen(nnx_wrappers.ToLinen): """Wrap NNX Pipeline as a Linen module. From a33e97ec12b9aff4d61ced311ad81268ed96c043 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 8 Dec 2025 07:14:26 +0000 Subject: [PATCH 04/14] update pipeline --- src/MaxText/layers/pipeline.py | 44 ++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 5939c46036..b146301b51 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -88,7 +88,6 @@ def __init__( self.seq_len_axis_name = "activation_length_no_exp" # TODO: Support NNX layers in addition to Linen layers - self._is_linen = True # Always true given the input self.layer = layer # Store the module instance self._linen_variables = None # Will be populated in _initialize_linen_parameters @@ -339,7 +338,7 @@ def permute_output_micro_per_stage_dim(self, output): def _initialize_linen_parameters(self, sample_input, sample_seg_ids, sample_positions, deterministic, model_mode): """Initialize Linen module parameters for all stages.""" - if not self._is_linen or self._linen_variables is not None: + if self._linen_variables is not None: return # Ensure inputs to init are 2D @@ -377,37 +376,35 @@ def _run_stages_linen( *stage_params_list ) - def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos, rng_key): - # Ensure per-stage positions and segment_ids are 2D for the layer call - if stage_pos is not None and stage_pos.ndim == 1: - stage_pos = stage_pos[None, :] - if stage_seg_ids is not None and stage_seg_ids.ndim == 1: - stage_seg_ids = stage_seg_ids[None, :] + def apply_stage(stage_idx, stage_params, stage_input, stage_seg_ids, stage_pos): + if stage_pos is not None and stage_pos.ndim == 1: stage_pos = stage_pos[None, :] + if stage_seg_ids is not None and stage_seg_ids.ndim == 1: stage_seg_ids = stage_seg_ids[None, :] + + # Fold in stage index to the iteration key for unique RNG per stage + stage_dropout_key = jax.random.fold_in(dropout_key, stage_idx) output = self.layer.apply( stage_params, stage_input, stage_seg_ids, stage_pos, - deterministic, - model_mode, - rngs={"dropout": rng_key}, + deterministic=deterministic, + model_mode=model_mode, + rngs={"dropout": stage_dropout_key}, ) return output[0] if isinstance(output, tuple) else output - dropout_keys_for_vmap = jax.random.split(dropout_key, self.num_stages) + stage_indices = jnp.arange(self.num_stages) if stages_segment_ids is None: - vmap_fn = lambda p, i, pos, rng: apply_stage(p, i, None, pos, rng) - in_axes = (0, 0, 0, 0) - stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( - stacked_params, stages_inputs, stages_positions, dropout_keys_for_vmap + vmap_fn = lambda idx, p, i, pos: apply_stage(idx, p, i, None, pos) + stages_outputs = jax.vmap(vmap_fn, in_axes=(0, 0, 0, 0), out_axes=0)( + stage_indices, stacked_params, stages_inputs, stages_positions ) else: - vmap_fn = lambda p, i, s, pos, rng: apply_stage(p, i, s, pos, rng) - in_axes = (0, 0, 0, 0, 0) - stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( - stacked_params, stages_inputs, stages_segment_ids, stages_positions, dropout_keys_for_vmap + vmap_fn = lambda idx, p, i, s, pos: apply_stage(idx, p, i, s, pos) + stages_outputs = jax.vmap(vmap_fn, in_axes=(0, 0, 0, 0, 0), out_axes=0)( + stage_indices, stacked_params, stages_inputs, stages_segment_ids, stages_positions ) return stages_outputs @@ -479,7 +476,12 @@ def run_one_iteration( shift = loop_state["shift"] circ_storage = loop_state["circ_storage"] loop_iteration = loop_state["loop_iteration"] - iter_dropout_key = self.rngs.dropout() + + # Get a base key for this iteration + iter_dropout_key_base = self.rngs.dropout() + # Fold in the loop iteration number to make the key unique per iteration + iter_dropout_key = jax.random.fold_in(iter_dropout_key_base, loop_iteration) + microbatch_ids, _ = self.get_microbatch_and_repeat_ids(loop_iteration) stages_inputs = self.get_iteration_inputs(loop_iteration, state_io, circ_storage, shift) From 5c4b2142b1e90e9bd1e300a3a0a00b9a2efcb03c Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 8 Dec 2025 08:20:12 +0000 Subject: [PATCH 05/14] revert nnx.scan --- src/MaxText/layers/pipeline.py | 37 +++++++++++++++------------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index b146301b51..441c8c0a14 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -109,7 +109,7 @@ def iterations_to_complete_first_microbatch(self): + self.iterations_to_complete_first_microbatch_one_repeat() ) - def init_states(self, inputs): + def init_states(self, inputs, initial_dropout_rng): """Initialize pipeline loop state buffers. Assumes inputs are reshaped to [num_microbatches, micro_batch_size, sequence, embed]. @@ -158,6 +158,7 @@ def init_states(self, inputs): "circ_storage_mover": circ_storage_mover, "loop_iteration": 0, "prev_outputs": prev_outputs, + "dropout_rng": initial_dropout_rng # Initialize RNG in state } return init_loop_state @@ -477,10 +478,9 @@ def run_one_iteration( circ_storage = loop_state["circ_storage"] loop_iteration = loop_state["loop_iteration"] - # Get a base key for this iteration - iter_dropout_key_base = self.rngs.dropout() - # Fold in the loop iteration number to make the key unique per iteration - iter_dropout_key = jax.random.fold_in(iter_dropout_key_base, loop_iteration) + dropout_key = loop_state["dropout_rng"] # Get key from state + # Split key for this iteration's use and for the next iteration + iter_dropout_key, next_iter_dropout_key = jax.random.split(dropout_key) microbatch_ids, _ = self.get_microbatch_and_repeat_ids(loop_iteration) @@ -492,6 +492,7 @@ def run_one_iteration( stages_output = self._run_stages_linen(stages_inputs, stages_segment_ids, stages_positions, deterministic, model_mode, iter_dropout_key) new_loop_state = self.get_new_loop_state(stages_output, loop_state) + new_loop_state["dropout_rng"] = next_iter_dropout_key # Store the key for the next loop return new_loop_state @@ -533,32 +534,26 @@ def __call__( example_pos = positions[0] if positions is not None else None self._initialize_linen_parameters(example_input, example_seg_ids, example_pos, deterministic, model_mode) - loop_state = self.init_states(inputs) + loop_state = self.init_states(inputs, self.rngs.dropout()) bubble_iterations = self.forwarding_delay * (self.num_stages - 1) real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats total_iterations = real_iterations + bubble_iterations if self.config.scan_pipeline_iterations: - def run_iteration_scannable(pipeline_instance, loop_state, xs): - new_loop_state = pipeline_instance.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) - return new_loop_state + def run_iteration_scannable(loop_state, xs): + new_loop_state = self.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) + return new_loop_state, None - remat_fn = run_iteration_scannable if self.config.set_remat_policy_on_pipeline_iterations: - remat_fn = nnx.remat( - run_iteration_scannable, - prevent_cse=False, - policy=self.get_pipeline_remat_policy(), + run_iteration_scannable = jax.checkpoint( + run_iteration_scannable, + prevent_cse=False, + policy=self.get_pipeline_remat_policy(), ) - ScannedFn = nnx.scan( - remat_fn, - in_axes=(None, nnx.Carry, None), # Corresponds to (self, loop_state, xs) - out_axes=nnx.Carry, # Corresponds to new_loop_state - length=total_iterations, - ) - loop_state = ScannedFn(self, loop_state, None) + loop_state, _ = jax.lax.scan(run_iteration_scannable, loop_state, None, length=total_iterations) + else: for _ in range(total_iterations): loop_state = self.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) From bf3abf9a3d22f7604f0e1c9c161210b4a3a68605 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 8 Dec 2025 08:30:44 +0000 Subject: [PATCH 06/14] update apply stage --- src/MaxText/layers/pipeline.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 441c8c0a14..c0a8e000f6 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -377,13 +377,9 @@ def _run_stages_linen( *stage_params_list ) - def apply_stage(stage_idx, stage_params, stage_input, stage_seg_ids, stage_pos): + def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos, rng_key): if stage_pos is not None and stage_pos.ndim == 1: stage_pos = stage_pos[None, :] if stage_seg_ids is not None and stage_seg_ids.ndim == 1: stage_seg_ids = stage_seg_ids[None, :] - - # Fold in stage index to the iteration key for unique RNG per stage - stage_dropout_key = jax.random.fold_in(dropout_key, stage_idx) - output = self.layer.apply( stage_params, stage_input, @@ -391,21 +387,21 @@ def apply_stage(stage_idx, stage_params, stage_input, stage_seg_ids, stage_pos): stage_pos, deterministic=deterministic, model_mode=model_mode, - rngs={"dropout": stage_dropout_key}, + rngs={"dropout": rng_key}, ) return output[0] if isinstance(output, tuple) else output - stage_indices = jnp.arange(self.num_stages) + dropout_keys_for_vmap = jax.random.split(dropout_key, self.num_stages) if stages_segment_ids is None: - vmap_fn = lambda idx, p, i, pos: apply_stage(idx, p, i, None, pos) + vmap_fn = lambda p, i, pos, rng: apply_stage(p, i, None, pos, rng) stages_outputs = jax.vmap(vmap_fn, in_axes=(0, 0, 0, 0), out_axes=0)( - stage_indices, stacked_params, stages_inputs, stages_positions + stacked_params, stages_inputs, stages_positions, dropout_keys_for_vmap ) else: - vmap_fn = lambda idx, p, i, s, pos: apply_stage(idx, p, i, s, pos) + vmap_fn = lambda p, i, s, pos, rng: apply_stage(p, i, s, pos, rng) stages_outputs = jax.vmap(vmap_fn, in_axes=(0, 0, 0, 0, 0), out_axes=0)( - stage_indices, stacked_params, stages_inputs, stages_segment_ids, stages_positions + stacked_params, stages_inputs, stages_segment_ids, stages_positions, dropout_keys_for_vmap ) return stages_outputs From 22b26094630c9a53340ad604c6ac77d44b4324a0 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Wed, 10 Dec 2025 08:24:40 +0000 Subject: [PATCH 07/14] update implementation --- src/MaxText/layers/pipeline.py | 1036 +++++++++++++++----------------- 1 file changed, 501 insertions(+), 535 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index c0a8e000f6..d1f442f00e 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -14,560 +14,526 @@ """ Pipeline layer wrapping a decoder layer(s). Supports circular pipelining """ -from typing import Any, Callable +from typing import Any, Optional, Callable -import numpy as np - -from jax import numpy as jnp -from jax.sharding import Mesh import jax -import jax.ad_checkpoint - -from flax.core import meta -from flax import linen as nn +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding, PartitionSpec from flax import nnx - +from flax import linen as nn from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT from MaxText.layers import nnx_wrappers from MaxText.layers.initializers import variable_to_logically_partitioned - - class Pipeline(nnx.Module): - """NNX Module that implements pipelining across stages. - - This module will loop over microbatches and execute the main body with a vmap for both the inputs and weights. - This will produce a pipeline pattern if the stage dimension is sharded. - - Supports circular pipelines, and multiple layers per stage are used when a module that executes multiple layers - is passed as the layers input. - - Attributes: - config: Importantly contains num_pipeline_microbatches, num_pipeline_repeats. - layers: A callable (NNX class or Linen class) that each stage can execute. It can either be a single layer such as a - LlamaDecoderLayer instance or scanned/looped set of decoder layers to execute multiple layers per stage. - mesh: The device mesh of the system. - remat_policy: Remat policy to use for the loop iterations - """ - - _linen_variables: nnx.Data[list[dict] | None] - - def __init__( - self, - layer: Callable[..., nnx.Module|nn.Module], - config: Config, - mesh: Mesh, - rngs: nnx.Rngs, - remat_policy: Any = None, - ): - """Initialize Pipeline with NNX or Linen decoder layers. - - Args: - layers: Either an NNX class (type) or Linen class (type) to instantiate for each stage - config: Model configuration - mesh: Device mesh for sharding - rngs: Optional NNX RNG state (passed by ToLinen wrapper) - remat_policy: Remat policy for loop iterations - """ - self.config = config - self.mesh = mesh - self.rngs = rngs - self.remat_policy = remat_policy - - self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism - self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1 - self.pipeline_microbatch_size = self.config.micro_batch_size_to_train_on // self.config.num_pipeline_microbatches - microbatches_per_stage = self.config.num_pipeline_microbatches // self.num_stages - self.microbatches_per_stage = microbatches_per_stage - self.use_circ_storage = self.need_circ_storage() - - if self.config.expert_shard_attention_option == EP_AS_CONTEXT: - self.batch_axis_name = "activation_batch_no_exp" - self.seq_len_axis_name = "activation_length" - else: - self.batch_axis_name = "activation_batch" - self.seq_len_axis_name = "activation_length_no_exp" - - # TODO: Support NNX layers in addition to Linen layers - self.layer = layer # Store the module instance - self._linen_variables = None # Will be populated in _initialize_linen_parameters - - - def need_circ_storage(self): - return ( - self.config.num_pipeline_repeats > 1 - and self.config.num_pipeline_microbatches > self.num_stages * self.forwarding_delay - ) - - def iterations_to_complete_first_microbatch_one_repeat(self): - """Returns iterations for microbatch 0 to complete one repeat.""" - return self.forwarding_delay * (self.num_stages - 1) - - def iterations_to_complete_first_microbatch(self): - """Returns iterations for microbatch 0 to complete all repeats.""" - return ( - self.config.num_pipeline_microbatches * (self.config.num_pipeline_repeats - 1) - + self.iterations_to_complete_first_microbatch_one_repeat() - ) - - def init_states(self, inputs, initial_dropout_rng): - """Initialize pipeline loop state buffers. - - Assumes inputs are reshaped to [num_microbatches, micro_batch_size, sequence, embed]. - - Returns: - Dictionary containing: - - shift: Buffer for rotating outputs [num_stages, micro_size, sequence, embed] - - prev_outputs: Same shape as shift (only used with pipeline_delay_activation_forwarding) - - state_io: Input/output buffer [num_stages, microbatches/stages, micro_size, sequence, embed] - - circ_storage: Circular storage buffer (only when num_microbatches > num_stages) - - circ_storage_mover: One-iteration delay buffer for circ_storage - - loop_iteration: Iteration counter (starts at 0) - """ - shift = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) - shift = self._with_logical_constraint( - shift, - ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - ) - - if self.config.pipeline_delay_activation_forwarding: - prev_outputs = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) - prev_outputs = self._with_logical_constraint( - prev_outputs, - ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - ) - else: - prev_outputs = None - - state_io = jnp.reshape(inputs, (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:]) - state_io = self._with_logical_constraint( - state_io, - ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - ) - - if self.use_circ_storage: - circ_storage = jnp.zeros((self.num_stages,) + inputs.shape, dtype=inputs.dtype) - circ_storage_mover = shift - else: - circ_storage = None - circ_storage_mover = None - - init_loop_state = { - "state_io": state_io, - "shift": shift, - "circ_storage": circ_storage, - "circ_storage_mover": circ_storage_mover, - "loop_iteration": 0, - "prev_outputs": prev_outputs, - "dropout_rng": initial_dropout_rng # Initialize RNG in state - } - return init_loop_state - - def _with_logical_constraint(self, tensor, logical_axis_names): - """Applies logical sharding constraints to tensor.""" - return nn.with_logical_constraint( - tensor, - logical_axis_names, - rules=self.config.logical_axis_rules, - mesh=self.mesh, - ) - - def get_iteration_inputs(self, loop_iteration, state_io, circ_storage, shift): - """Constructs input array for all stages for this iteration. - - Returns array of shape [stages, micro_size, sequence, embed] with rotated outputs - from previous iteration, except stage 0 which gets new input from state_io or circ_storage. - """ - state_io_batch_idx = loop_iteration % self.microbatches_per_stage - state_io_slice = state_io[:, state_io_batch_idx] - - if self.use_circ_storage: - circ_storage_batch_idx = loop_iteration % self.config.num_pipeline_microbatches - circular_stage_in = circ_storage[:, circ_storage_batch_idx] - else: - circular_stage_in = shift - - first_stage_in = jnp.where(loop_iteration < self.config.num_pipeline_microbatches, state_io_slice, circular_stage_in) - - def select_state_or_input(first_stage_in, shift): - return jnp.where(jax.lax.broadcasted_iota("int32", shift.shape, 0) == 0, first_stage_in, shift) - - stages_in = select_state_or_input(first_stage_in, shift) - stages_in = self._with_logical_constraint( - stages_in, - ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - ) - return stages_in - - def shard_dim_by_stages(self, x, dim: int): - """Shards the specified dimension by stage.""" - dims_mapping = [jax.sharding.PartitionSpec.UNCONSTRAINED] * x.ndim - dims_mapping[dim] = "stage" - dims_mapping = tuple(dims_mapping) - sharding = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec(*dims_mapping)) - return jax.lax.with_sharding_constraint(x, sharding) - - def get_microbatch_and_repeat_ids(self, loop_iteration): - """Gets microbatch and repeat IDs for all stages at this iteration.""" - microbatches_processed = jnp.maximum(loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), 0) - microbatch_ids = microbatches_processed % self.config.num_pipeline_microbatches - repeat_ids = microbatches_processed // self.config.num_pipeline_microbatches - return microbatch_ids, repeat_ids - - def vmap_parallel_gather(self, weights, repeat_ids, repeat_dim_in_weights, stages_dim_in_weights): - """Sharded parallel gather where each stage has its own weights and gets one slice. - - Args: - weights: Per-stage data to gather from. - repeat_ids: Integer tensor of shape [num_stages] with repeat indices per stage. - repeat_dim_in_weights: Dimension where repeat_ids are applied (removed in output). - stages_dim_in_weights: Dimension representing parallel stages. - - Returns: - Per-stage gathered values with repeat_dim_in_weights removed. """ - def _gather_one(x, repeat_id): - return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, repeat_id, 1, repeat_dim_in_weights), repeat_dim_in_weights) - - gathered_weights_stage_dim = 0 - repeat_ids = self.shard_dim_by_stages(repeat_ids, 0) - weights = self.shard_dim_by_stages(weights, stages_dim_in_weights) - stage_weights = jax.vmap(_gather_one, in_axes=(stages_dim_in_weights, 0), out_axes=gathered_weights_stage_dim)( - weights, repeat_ids - ) - stage_weights = self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim) - return stage_weights - - def vmap_gather(self, xs, ids, ids_dim): - """Stage-wise sharded gather with shared input but different offsets per stage. - - Args: - xs: Data shared by all stages. - ids: Integer tensor of shape [num_stages] with offsets per stage. - ids_dim: Dimension where ids are applied (output has [num_stages] size here). - - Returns: - Per-stage gathered values with ids_dim size replaced with [num_stages]. + NNX Implementation of the MaxText Pipeline. + Wraps a Flax Linen Module and executes it using Pipeline Parallelism. """ - def _gather_one(x, i): - return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, i, 1, ids_dim), ids_dim) - ids = self.shard_dim_by_stages(ids, 0) - outs = jax.vmap(_gather_one, in_axes=(None, 0), out_axes=ids_dim)(xs, ids) - return self.shard_dim_by_stages(outs, 0) - - def get_new_loop_state(self, output, loop_state): - """Updates all pipeline buffers after one iteration. + def __init__( + self, + layer: nn.Module, # Instance of the Linen Layer (e.g. DecoderLayer) + config: Any, # MaxText Config object + mesh: Mesh, # JAX Mesh + rngs: nnx.Rngs, # NNX RNGs + remat_policy: Any = None + ): + self.config = config + self.mesh = mesh + self.linen_module = layer + self.remat_policy = remat_policy + + # --- 1. Calculate Pipeline Dimensions --- + self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism + self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1 + + self.total_microbatches = self.config.num_pipeline_microbatches + self.microbatch_size = self.config.micro_batch_size_to_train_on // self.total_microbatches + self.microbatches_per_stage = self.total_microbatches // self.num_stages + + self.use_circ_storage = self._need_circ_storage() + + # Axis Naming + if hasattr(self.config, 'expert_shard_attention_option') and self.config.expert_shard_attention_option == EP_AS_CONTEXT: + self.batch_axis_name = "activation_batch_no_exp" + self.seq_len_axis_name = "activation_length" + else: + self.batch_axis_name = "activation_batch" + self.seq_len_axis_name = "activation_length_no_exp" + + self.input_shape = ( + self.microbatch_size, + self.config.max_target_length, + self.config.emb_dim + ) - Updates shift, state_io, circ_storage, circ_storage_mover, and prev_outputs - to advance the pipeline by one step. - """ - old_state_io = loop_state["state_io"] - old_circ_storage = loop_state["circ_storage"] - old_circ_storage_mover = loop_state["circ_storage_mover"] - loop_iteration = loop_state["loop_iteration"] - old_prev_outputs = loop_state["prev_outputs"] - - def _rotate_right(arr): - last = jax.lax.slice_in_dim(arr, self.num_stages - 1, self.num_stages, axis=0) - except_last = jax.lax.slice_in_dim(arr, 0, self.num_stages - 1, axis=0) - return jnp.concatenate([last, except_last], axis=0) - - def _shift_right(arr): - padding = [[1, 0]] + [[0, 0]] * (arr.ndim - 1) - return jax.lax.slice(jnp.pad(arr, padding), [0] * arr.ndim, arr.shape) - - def _update_shift(output_in): - if self.config.num_pipeline_repeats == 1 or self.use_circ_storage: - return _shift_right(output_in) - else: - return _rotate_right(output_in) - - if self.config.pipeline_delay_activation_forwarding: - new_shift = _update_shift(old_prev_outputs) - new_prev_outputs = output - else: - new_shift = _update_shift(output) - new_prev_outputs = None - - if self.use_circ_storage: - def _rotate_right_and_update(circ_storage_mover_in, circ_storage_in): - rotated = _rotate_right(circ_storage_mover_in) - rotated = jnp.expand_dims(rotated, 1) - offset = ( - loop_iteration - self.iterations_to_complete_first_microbatch_one_repeat() - 1 - ) % self.config.num_pipeline_microbatches - return jax.lax.dynamic_update_slice_in_dim(circ_storage_in, rotated, offset, axis=1) - - new_circ_storage = _rotate_right_and_update(old_circ_storage_mover, old_circ_storage) - new_circ_storage_mover = output - else: - new_circ_storage = None - new_circ_storage_mover = None - - stream_buf_idx = loop_iteration % self.microbatches_per_stage - stream_slice = old_state_io[:, stream_buf_idx] - - def _update_state_io(state_in, stream_slice, output): - padding = [[0, 1]] + [[0, 0]] * (stream_slice.ndim - 1) - stream_slice = jax.lax.slice_in_dim(jnp.pad(stream_slice, padding), 1, stream_slice.shape[0] + 1, axis=0) - stream_slice = jnp.where( - jax.lax.broadcasted_iota("int32", stream_slice.shape, 0) == self.num_stages - 1, output, stream_slice - ) - stream_slice = jnp.expand_dims(stream_slice, 1) - return jax.lax.dynamic_update_slice_in_dim(state_in, stream_slice, stream_buf_idx, axis=1) - - new_state = _update_state_io(old_state_io, stream_slice, output) - - new_loop_state = { - "state_io": new_state, - "shift": new_shift, - "circ_storage": new_circ_storage, - "circ_storage_mover": new_circ_storage_mover, - "loop_iteration": loop_iteration + 1, - "prev_outputs": new_prev_outputs, - } - return new_loop_state - - def permute_output_micro_per_stage_dim(self, output): - """Permutes output to correct microbatch ordering after pipeline completion.""" - microbatch_0_idx = self.iterations_to_complete_first_microbatch() % self.microbatches_per_stage - permutation = ( - np.arange(self.microbatches_per_stage) + microbatch_0_idx - ) % self.microbatches_per_stage - output = output[:, permutation] - return output - - def _initialize_linen_parameters(self, sample_input, sample_seg_ids, sample_positions, deterministic, model_mode): - """Initialize Linen module parameters for all stages.""" - if self._linen_variables is not None: - return - - # Ensure inputs to init are 2D - if sample_seg_ids is not None and sample_seg_ids.ndim == 1: - sample_seg_ids = sample_seg_ids[None, :] - if sample_positions is not None and sample_positions.ndim == 1: - sample_positions = sample_positions[None, :] - - - stage_params_list = [] - for _ in range(self.num_stages): - linen_rngs = { - 'params': self.rngs.params(), - 'dropout': self.rngs.dropout() - } - base_params = self.layer.init( - linen_rngs, sample_input, sample_seg_ids, sample_positions, deterministic, model_mode) - stage_params_list.append(base_params) - self._linen_variables = stage_params_list - - - def _run_stages_linen( - self, - stages_inputs, - stages_segment_ids, - stages_positions, - deterministic, - model_mode, - dropout_key - ): - """Run stages using Linen module with manual vmap.""" - stage_params_list = self._linen_variables if self._linen_variables is not None else [] - stacked_params = jax.tree_util.tree_map( - lambda *xs: jnp.stack(xs, axis=0), - *stage_params_list - ) - - def apply_stage(stage_params, stage_input, stage_seg_ids, stage_pos, rng_key): - if stage_pos is not None and stage_pos.ndim == 1: stage_pos = stage_pos[None, :] - if stage_seg_ids is not None and stage_seg_ids.ndim == 1: stage_seg_ids = stage_seg_ids[None, :] - output = self.layer.apply( - stage_params, - stage_input, - stage_seg_ids, - stage_pos, - deterministic=deterministic, - model_mode=model_mode, - rngs={"dropout": rng_key}, - ) - return output[0] if isinstance(output, tuple) else output - - dropout_keys_for_vmap = jax.random.split(dropout_key, self.num_stages) - - if stages_segment_ids is None: - vmap_fn = lambda p, i, pos, rng: apply_stage(p, i, None, pos, rng) - stages_outputs = jax.vmap(vmap_fn, in_axes=(0, 0, 0, 0), out_axes=0)( - stacked_params, stages_inputs, stages_positions, dropout_keys_for_vmap - ) - else: - vmap_fn = lambda p, i, s, pos, rng: apply_stage(p, i, s, pos, rng) - stages_outputs = jax.vmap(vmap_fn, in_axes=(0, 0, 0, 0, 0), out_axes=0)( - stacked_params, stages_inputs, stages_segment_ids, stages_positions, dropout_keys_for_vmap - ) - - return stages_outputs - - - def _run_stages_vmapped( - self, - stages_inputs, - stages_segment_ids, - stages_positions, - deterministic, - model_mode, - ): - """Run all stages in parallel using JAX vmap over NNX instances.""" - stage_0 = getattr(self, 'stage_0') - graphdef, state_0 = nnx.split(stage_0) - - states = [state_0] - for s in range(1, self.num_stages): - instance = getattr(self, f'stage_{s}') - _, state_s = nnx.split(instance) - states.append(state_s) - - stacked_state = jax.tree_util.tree_map( - lambda *xs: jnp.stack(xs, axis=0), - *states - ) - - def call_stage(state, stage_input, stage_seg_ids, stage_pos): - if stage_pos is not None and stage_pos.ndim == 1: - stage_pos = stage_pos[None, :] - if stage_seg_ids is not None and stage_seg_ids.ndim == 1: - stage_seg_ids = stage_seg_ids[None, :] - module = nnx.merge(graphdef, state) - output = module(stage_input, stage_seg_ids, stage_pos, deterministic=deterministic, model_mode=model_mode) - return output[0] if isinstance(output, tuple) else output - - if stages_segment_ids is None: - def call_stage_no_seg(state, stage_input, stage_pos): - if stage_pos is not None and stage_pos.ndim == 1: - stage_pos = stage_pos[None, :] - module = nnx.merge(graphdef, state) - output = module(stage_input, None, stage_pos, deterministic=deterministic, model_mode=model_mode) - return output[0] if isinstance(output, tuple) else output - vmap_fn = call_stage_no_seg - in_axes = (0, 0, 0) - stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( - stacked_state, stages_inputs, stages_positions - ) - else: - vmap_fn = call_stage - in_axes = (0, 0, 0, 0) - stages_outputs = jax.vmap(vmap_fn, in_axes=in_axes, out_axes=0)( - stacked_state, stages_inputs, stages_segment_ids, stages_positions - ) - return stages_outputs - - - def run_one_iteration( - self, - loop_state, - positions, - segment_ids, - deterministic, - model_mode, - ): - """Run one loop iteration: get inputs, execute stages, update state.""" - state_io = loop_state["state_io"] - shift = loop_state["shift"] - circ_storage = loop_state["circ_storage"] - loop_iteration = loop_state["loop_iteration"] - - dropout_key = loop_state["dropout_rng"] # Get key from state - # Split key for this iteration's use and for the next iteration - iter_dropout_key, next_iter_dropout_key = jax.random.split(dropout_key) - - - microbatch_ids, _ = self.get_microbatch_and_repeat_ids(loop_iteration) - stages_inputs = self.get_iteration_inputs(loop_iteration, state_io, circ_storage, shift) - stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") - stages_positions = self.vmap_gather(positions, microbatch_ids, 0) if positions is not None else None - stages_segment_ids = self.vmap_gather(segment_ids, microbatch_ids, 0) if segment_ids is not None else None - - stages_output = self._run_stages_linen(stages_inputs, stages_segment_ids, stages_positions, deterministic, model_mode, iter_dropout_key) - - new_loop_state = self.get_new_loop_state(stages_output, loop_state) - new_loop_state["dropout_rng"] = next_iter_dropout_key # Store the key for the next loop - return new_loop_state - - - def get_pipeline_remat_policy(self): - """Returns the remat policy for pipeline iterations.""" - if self.config.remat_policy == "custom": - return self.remat_policy - - save_input_policy = jax.checkpoint_policies.save_only_these_names("iteration_input", "decoder_layer_input") - if self.remat_policy is not None: - remat_policy = jax.checkpoint_policies.save_from_both_policies(self.remat_policy, save_input_policy) - else: - remat_policy = save_input_policy - return remat_policy - - def __call__( - self, - inputs: jnp.ndarray, - segment_ids: jnp.ndarray, - positions: jnp.ndarray, - deterministic: bool, - model_mode=MODEL_MODE_TRAIN, - partition_spec=None, - ) -> jnp.ndarray: - """Maps decoder layer inputs to outputs using pipeline parallelism. - - Reshapes inputs into microbatches, runs pipeline iterations with bubble - handling, and returns outputs reshaped to original batch size. - """ - inputs = inputs.reshape((self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length, self.config.emb_dim)) - if positions is not None: - positions = positions.reshape((self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length)) - if segment_ids is not None: - segment_ids = segment_ids.reshape((self.config.num_pipeline_microbatches, self.pipeline_microbatch_size, self.config.max_target_length)) - - if self._linen_variables is None: - example_input = inputs[0] - example_seg_ids = segment_ids[0] if segment_ids is not None else None - example_pos = positions[0] if positions is not None else None - self._initialize_linen_parameters(example_input, example_seg_ids, example_pos, deterministic, model_mode) - - loop_state = self.init_states(inputs, self.rngs.dropout()) - - bubble_iterations = self.forwarding_delay * (self.num_stages - 1) - real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats - total_iterations = real_iterations + bubble_iterations - - if self.config.scan_pipeline_iterations: - def run_iteration_scannable(loop_state, xs): - new_loop_state = self.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) - return new_loop_state, None - - if self.config.set_remat_policy_on_pipeline_iterations: - run_iteration_scannable = jax.checkpoint( - run_iteration_scannable, - prevent_cse=False, - policy=self.get_pipeline_remat_policy(), + # --- 2. Eager Initialization (FIXED for Circular Repeats) --- + init_rng_key = rngs.params() + repeats = self.config.num_pipeline_repeats + + def init_single_stage(key): + # Create dummy inputs + dummy_in = jnp.zeros(self.input_shape, dtype=jnp.float32) + dummy_pos_shape = (self.input_shape[0], self.input_shape[1]) + dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) + dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) + + return self.linen_module.init( + {'params': key}, + dummy_in, + decoder_segment_ids=dummy_segments, + decoder_positions=dummy_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN + ) + + if repeats > 1: + # Case A: Circular Pipeline (Multiple Repeats) + # We need independent weights for [Repeats, Num_Stages] + repeat_keys = jax.random.split(init_rng_key, repeats) + # Split each repeat key into stage keys -> Shape: [Repeats, Stages] + stage_rng_keys = jax.vmap(lambda k: jax.random.split(k, self.num_stages))(repeat_keys) + + # Double vmap: Map over Repeats (Axis 0), then Stages (Axis 0 of inner) + # Output Shape: [Repeats, Stages, ...] + raw_variables = jax.vmap(jax.vmap(init_single_stage))(stage_rng_keys) + else: + # Case B: Standard Pipeline (Single Repeat) + # Shape: [Num_Stages] + stage_rng_keys = jax.random.split(init_rng_key, self.num_stages) + # Output Shape: [Stages, ...] + raw_variables = jax.vmap(init_single_stage)(stage_rng_keys) + + # --- 3. Register Parameters with NNX --- + def _to_nnx_structure(node): + if hasattr(node, 'items'): + return nnx.Dict({k: _to_nnx_structure(v) for k, v in node.items()}) + elif isinstance(node, (list, tuple)): + return nnx.List([_to_nnx_structure(v) for v in node]) + else: + return nnx.Param(node) + + self.stage_params = _to_nnx_structure(raw_variables['params']) + # ========================================================================== + # Helper Methods + # ========================================================================== + + def _need_circ_storage(self): + return (self.config.num_pipeline_repeats > 1 and + self.config.num_pipeline_microbatches > self.num_stages * self.forwarding_delay) + + def iterations_to_complete_first_microbatch_one_repeat(self): + return self.forwarding_delay * (self.num_stages - 1) + + + def init_states(self, inputs): + """Initialize pipeline buffers. + Args: + inputs: Rank 4 array [Total_Microbatches, Micro_Size, Seq, Emb] + """ + # 1. Shift Buffer + # Shape: [Num_Stages, Micro_Size, Seq, Emb] + # (Derived from inputs.shape[1:]) + shift = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) + shift = nn.with_logical_constraint( + shift, + ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), + rules=self.config.logical_axis_rules, + mesh=self.mesh, ) - loop_state, _ = jax.lax.scan(run_iteration_scannable, loop_state, None, length=total_iterations) + # 2. Prev Outputs (for forwarding delay) + if self.config.pipeline_delay_activation_forwarding: + prev_outputs = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) + prev_outputs = nn.with_logical_constraint( + prev_outputs, + ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), + rules=self.config.logical_axis_rules, + mesh=self.mesh, + ) + else: + prev_outputs = None + + # 3. State IO (The Main Buffer) + # Reshape: [Total_Micro, ...] -> [Stages, Micro_Per_Stage, ...] + state_io = jnp.reshape( + inputs, + (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:] + ) + state_io = nn.with_logical_constraint( + state_io, + ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), + rules=self.config.logical_axis_rules, + mesh=self.mesh, + ) - else: - for _ in range(total_iterations): - loop_state = self.run_one_iteration(loop_state, positions, segment_ids, deterministic, model_mode) + # 4. Circular Storage + if self.use_circ_storage: + # Shape: [Num_Stages, Total_Microbatches, Micro_Size, Seq, Emb] + # (Derived from inputs.shape) + circ_storage = jnp.zeros((self.num_stages,) + inputs.shape, dtype=inputs.dtype) + circ_storage_mover = shift + else: + circ_storage = None + circ_storage_mover = None + + return { + "state_io": state_io, + "shift": shift, + "circ_storage": circ_storage, + "circ_storage_mover": circ_storage_mover, + "loop_iteration": jnp.array(0, dtype=jnp.int32), + "prev_outputs": prev_outputs, + "rng_stream": jax.random.PRNGKey(0) + } - final_output = self.permute_output_micro_per_stage_dim(loop_state["state_io"]) - final_output = jnp.reshape( - final_output, (self.config.micro_batch_size_to_train_on, self.config.max_target_length, self.config.emb_dim) - ) - return final_output + def get_iteration_inputs(self, loop_iteration, state_io, circ_storage, shift): + state_io_batch_idx = loop_iteration % self.microbatches_per_stage + state_io_slice = state_io[:, state_io_batch_idx] -class PipelineToLinen(nnx_wrappers.ToLinen): - """Wrap NNX Pipeline as a Linen module. + if self.use_circ_storage: + circ_storage_batch_idx = loop_iteration % self.config.num_pipeline_microbatches + circular_stage_in = circ_storage[:, circ_storage_batch_idx] + else: + circular_stage_in = shift - This allows the NNX Pipeline to be used within the Linen Decoder module. - """ - pass + first_stage_in = jnp.where( + loop_iteration < self.config.num_pipeline_microbatches, + state_io_slice, + circular_stage_in + ) + def select_state_or_input(first_stage_in, shift): + return jnp.where( + jax.lax.broadcasted_iota("int32", shift.shape, 0) == 0, + first_stage_in, + shift + ) + + stages_in = select_state_or_input(first_stage_in, shift) + + stages_in = nn.with_logical_constraint( + stages_in, + ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), + rules=self.config.logical_axis_rules, + mesh=self.mesh, + ) + return stages_in + + def get_new_loop_state(self, output, loop_state): + old_state_io = loop_state["state_io"] + old_circ_storage = loop_state["circ_storage"] + old_circ_storage_mover = loop_state["circ_storage_mover"] + loop_iteration = loop_state["loop_iteration"] + old_prev_outputs = loop_state["prev_outputs"] + + def _rotate_right(arr): + last = jax.lax.slice_in_dim(arr, self.num_stages - 1, self.num_stages, axis=0) + except_last = jax.lax.slice_in_dim(arr, 0, self.num_stages - 1, axis=0) + return jnp.concatenate([last, except_last], axis=0) + + def _shift_right(arr): + padding = [[1, 0]] + [[0, 0]] * (arr.ndim - 1) + return jax.lax.slice(jnp.pad(arr, padding), [0] * arr.ndim, arr.shape) + + def _update_shift(output_in): + if self.config.num_pipeline_repeats == 1 or self.use_circ_storage: + return _shift_right(output_in) + return _rotate_right(output_in) + + if self.config.pipeline_delay_activation_forwarding: + new_shift = _update_shift(old_prev_outputs) + new_prev_outputs = output + else: + new_shift = _update_shift(output) + new_prev_outputs = None + + if self.use_circ_storage: + def _rotate_right_and_update(circ_storage_mover_in, circ_storage_in): + rotated = _rotate_right(circ_storage_mover_in) + rotated = jnp.expand_dims(rotated, 1) + offset = ( + loop_iteration + - self.iterations_to_complete_first_microbatch_one_repeat() + - 1 + ) % self.config.num_pipeline_microbatches + return jax.lax.dynamic_update_slice_in_dim( + circ_storage_in, rotated, offset, axis=1 + ) + + new_circ_storage = _rotate_right_and_update(old_circ_storage_mover, old_circ_storage) + new_circ_storage_mover = output + else: + new_circ_storage = None + new_circ_storage_mover = None + + stream_buf_idx = loop_iteration % self.microbatches_per_stage + stream_slice = old_state_io[:, stream_buf_idx] + + def _update_state_io(state_in, stream_slice, output): + padding = [[0, 1]] + [[0, 0]] * (stream_slice.ndim - 1) + stream_slice = jax.lax.slice_in_dim( + jnp.pad(stream_slice, padding), 1, stream_slice.shape[0] + 1, axis=0 + ) + stream_slice = jnp.where( + jax.lax.broadcasted_iota("int32", stream_slice.shape, 0) == self.num_stages - 1, + output, + stream_slice, + ) + stream_slice = jnp.expand_dims(stream_slice, 1) + return jax.lax.dynamic_update_slice_in_dim( + state_in, stream_slice, stream_buf_idx, axis=1 + ) + + new_state = _update_state_io(old_state_io, stream_slice, output) + + return { + "state_io": new_state, + "shift": new_shift, + "circ_storage": new_circ_storage, + "circ_storage_mover": new_circ_storage_mover, + "loop_iteration": loop_iteration + 1, + "prev_outputs": new_prev_outputs, + } + # ========================================================================== + # FSDP Helpers + # ========================================================================== + + def _all_gather_over_fsdp(self, params, partition_spec): + """Helper to apply FSDP all-gather constraint.""" + def _remove_fsdp_from_spec(spec): + if isinstance(spec, PartitionSpec): + new_spec = [] + for axis in spec: + if isinstance(axis, str) and axis in ("fsdp", "fsdp_transpose"): + new_spec.append(None) + elif isinstance(axis, (list, tuple)): + new_spec.append(tuple(a for a in axis if a not in ("fsdp", "fsdp_transpose"))) + else: + new_spec.append(axis) + return PartitionSpec(*new_spec) + return spec + + def _remove_fsdp_sharding(sharding_tree): + return jax.tree.map( + lambda x: NamedSharding(self.mesh, _remove_fsdp_from_spec(x.spec)) + if isinstance(x, NamedSharding) else x, + sharding_tree + ) + + physical = nn.logical_to_mesh_sharding(partition_spec, mesh=self.mesh, rules=self.config.logical_axis_rules) + physical_no_fsdp = _remove_fsdp_sharding(physical) + return jax.lax.with_sharding_constraint(params, physical_no_fsdp) + + def _to_pure_dict(self, node): + """Recursively converts NNX containers to standard Python dicts/lists.""" + if hasattr(node, 'items'): # Handles nnx.Dict, dict, FrozenDict + return {k: self._to_pure_dict(v) for k, v in node.items()} + elif isinstance(node, (list, tuple)): # Handles nnx.List, list, tuple + return [self._to_pure_dict(v) for v in node] + elif hasattr(node, 'value'): # Handles nnx.Param, nnx.Variable + return node.value + return node # Leaf (e.g. JAX Array/Tracer) + + # ========================================================================== + # Main Execution + # ========================================================================== + def get_microbatch_ids(self, loop_iteration): + """Calculates the microbatch ID for each stage at the current tick.""" + # Calculate how many microbatches each stage has processed effectively + # Stage 0 starts at 0. Stage 1 starts after 'forwarding_delay', etc. + microbatches_processed = jnp.maximum( + loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), + 0 + ) + # Wrap around using modulo to get the ID within the circular buffer + microbatch_ids = microbatches_processed % self.total_microbatches + return microbatch_ids.astype(jnp.int32) + + # Handles logic when num_pipeline_repeats > 1 + def get_microbatch_and_repeat_ids(self, loop_iteration): + """Gets the microbatch_ids and repeat_ids for all stages on this loop_iteration.""" + microbatches_processed = jnp.maximum( + loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), + 0 + ) + microbatch_ids = microbatches_processed % self.total_microbatches + repeat_ids = microbatches_processed // self.total_microbatches + return microbatch_ids.astype(jnp.int32), repeat_ids.astype(jnp.int32) + + # Helper for weight gathering + def shard_dim_by_stages(self, x, dim: int): + dims_mapping = [PartitionSpec.UNCONSTRAINED] * x.ndim + dims_mapping[dim] = "stage" + # We construct the NamedSharding manually to match Linen's logical_to_mesh + # Note: In pure NNX/JAX, we can often just return x if sharding is handled by vmap, + # but we keep the constraint for strict parity. + sharding = NamedSharding(self.mesh, PartitionSpec(*dims_mapping)) + return jax.lax.with_sharding_constraint(x, sharding) + + # Helper for circular weight selection + def vmap_parallel_gather(self, weights, repeat_ids, repeat_dim_in_weights, stages_dim_in_weights): + """Use vmap to implement a sharded parallel gather for weights.""" + def _gather_one(x, repeat_id): + return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, repeat_id, 1, repeat_dim_in_weights), repeat_dim_in_weights) + + gathered_weights_stage_dim = 0 + repeat_ids = self.shard_dim_by_stages(repeat_ids, 0) + weights = self.shard_dim_by_stages(weights, stages_dim_in_weights) + + stage_weights = jax.vmap(_gather_one, in_axes=(stages_dim_in_weights, 0), out_axes=gathered_weights_stage_dim)(weights, repeat_ids) + stage_weights = self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim) + return stage_weights + + # The main function missing from scan_body + def get_current_stage_weights(self, pipeline_weights, loop_iteration): + if self.config.num_pipeline_repeats <= 1: + return pipeline_weights + _, repeat_ids = self.get_microbatch_and_repeat_ids(loop_iteration) + + # Helper to map over the PyTree of weights + def gather_weights_for_stages_in(w): + # Assumes weights are [Repeats, Stages, ...] -> Gather -> [Stages, ...] + return self.vmap_parallel_gather( + w, repeat_ids, repeat_dim_in_weights=0, stages_dim_in_weights=1 + ) + + return jax.tree.map(gather_weights_for_stages_in, pipeline_weights) + + # Output sorting + def permute_output_micro_per_stage_dim(self, output): + microbatch_0_idx = self.iterations_to_complete_first_microbatch_one_repeat() % self.microbatches_per_stage + permutation = (jnp.arange(self.microbatches_per_stage) + microbatch_0_idx) % self.microbatches_per_stage + output = output[:, permutation] + return output + + def __call__( + self, + inputs: jax.Array, + segment_ids: Optional[jax.Array] = None, + positions: Optional[jax.Array] = None, + deterministic: bool = True, + model_mode: str = MODEL_MODE_TRAIN, + partition_spec: Any = None, + ) -> jax.Array: + + # 1. Reshape Inputs + inputs = inputs.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length, + self.config.emb_dim + )) + + if positions is not None: + positions = positions.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length + )) + + if segment_ids is not None: + segment_ids = segment_ids.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length + )) + + # 2. Initialize State + loop_state = self.init_states(inputs) + + # 3. Prepare Weights + # Note: If repeats > 1, param_values MUST have shape [Repeats, Stages, ...] + # Currently __init__ only creates [Stages, ...]. + # For standard training (repeats=1), this is correct. + param_values = self._to_pure_dict(self.stage_params) + + if self.config.pipeline_fsdp_ag_once and partition_spec is not None: + try: + param_values = self._all_gather_over_fsdp(param_values, partition_spec) + except (ValueError, TypeError, KeyError): + pass + + # 4. Scan Loop + def scan_body(carry, _): + iteration = carry['loop_iteration'] + current_rng = carry['rng_stream'] + + step_rng, next_rng = jax.random.split(current_rng) + stage_rngs = jax.random.split(step_rng, self.num_stages) + + stages_inputs = self.get_iteration_inputs( + iteration, + carry['state_io'], + carry['circ_storage'], + carry['shift'] + ) + stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") + + # Position gathering + # [FIX] Use get_microbatch_and_repeat_ids to be consistent with parity logic + mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) + + stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None + stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None + + # [ADDED FOR PARITY] Dynamic Weight Selection (for Circular Pipelines) + # If repeats=1, this simply returns param_values as-is. + current_params = self.get_current_stage_weights(param_values, iteration) + + # VMAP Function + def vmapped_linen_apply(p, x, r, pos, seg): + variables = {'params': p} + rngs_dict = {'dropout': r} if not deterministic else {} + + return self.linen_module.apply( + variables, + x, + decoder_segment_ids=seg, + decoder_positions=pos, + deterministic=deterministic, + rngs=rngs_dict, + model_mode=model_mode + ) + + # Dynamic VMAP Arguments + vmap_args = [current_params, stages_inputs, stage_rngs] + vmap_axes = [0, 0, 0] # Params(0), Inputs(0), Rngs(0) + + vmap_args.append(stages_positions) + vmap_axes.append(0 if stages_positions is not None else None) + + vmap_args.append(stages_segment_ids) + vmap_axes.append(0 if stages_segment_ids is not None else None) + + stages_output = jax.vmap(vmapped_linen_apply, in_axes=tuple(vmap_axes))(*vmap_args) + + if hasattr(self.config, 'scan_layers') and self.config.scan_layers: + if isinstance(stages_output, tuple): + stages_output = stages_output[0] + + new_loop_state = self.get_new_loop_state(stages_output, carry) + new_loop_state['rng_stream'] = next_rng + + return new_loop_state, None + + bubble_iterations = self.forwarding_delay * (self.num_stages - 1) + real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats + total_ticks = real_iterations + bubble_iterations + + final_state, _ = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) + + output = final_state['state_io'] + + # [ADDED FOR PARITY] Permute Output to correct order + output = self.permute_output_micro_per_stage_dim(output) + + output = output.reshape((self.config.micro_batch_size_to_train_on, + self.config.max_target_length, + self.config.emb_dim)) + + return output def create_pipeline( config: Config, layer: Callable | type, From 664c0f1e2f97bb9932edb6684bc5bd0fc103ea3a Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Wed, 10 Dec 2025 10:04:43 +0000 Subject: [PATCH 08/14] update implementation --- src/MaxText/layers/pipeline.py | 329 ++++++++++++++++++++------------- 1 file changed, 199 insertions(+), 130 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index d1f442f00e..eabb139e8f 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -13,7 +13,7 @@ # limitations under the License. """ Pipeline layer wrapping a decoder layer(s). Supports circular pipelining """ - +import functools from typing import Any, Optional, Callable import jax @@ -24,6 +24,8 @@ from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT from MaxText.layers import nnx_wrappers from MaxText.layers.initializers import variable_to_logically_partitioned +from flax.linen import spmd + class Pipeline(nnx.Module): """ NNX Implementation of the MaxText Pipeline. @@ -114,6 +116,9 @@ def _to_nnx_structure(node): return nnx.Param(node) self.stage_params = _to_nnx_structure(raw_variables['params']) + + if 'batch_stats' in raw_variables: + print("Warning: Pipeline detected batch_stats during init. Registering as VariableState.") # ========================================================================== # Helper Methods # ========================================================================== @@ -407,158 +412,222 @@ def permute_output_micro_per_stage_dim(self, output): permutation = (jnp.arange(self.microbatches_per_stage) + microbatch_0_idx) % self.microbatches_per_stage output = output[:, permutation] return output - + def __call__( - self, - inputs: jax.Array, - segment_ids: Optional[jax.Array] = None, - positions: Optional[jax.Array] = None, - deterministic: bool = True, - model_mode: str = MODEL_MODE_TRAIN, - partition_spec: Any = None, - ) -> jax.Array: - - # 1. Reshape Inputs - inputs = inputs.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length, - self.config.emb_dim - )) - - if positions is not None: - positions = positions.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length - )) + self, + inputs: jax.Array, + segment_ids: Optional[jax.Array] = None, + positions: Optional[jax.Array] = None, + deterministic: bool = True, + model_mode: str = MODEL_MODE_TRAIN, + partition_spec: Any = None, + ) -> jax.Array: - if segment_ids is not None: - segment_ids = segment_ids.reshape(( + # 1. Reshape Inputs + inputs = inputs.reshape(( self.total_microbatches, self.microbatch_size, - self.config.max_target_length + self.config.max_target_length, + self.config.emb_dim )) - - # 2. Initialize State - loop_state = self.init_states(inputs) - - # 3. Prepare Weights - # Note: If repeats > 1, param_values MUST have shape [Repeats, Stages, ...] - # Currently __init__ only creates [Stages, ...]. - # For standard training (repeats=1), this is correct. - param_values = self._to_pure_dict(self.stage_params) - - if self.config.pipeline_fsdp_ag_once and partition_spec is not None: - try: - param_values = self._all_gather_over_fsdp(param_values, partition_spec) - except (ValueError, TypeError, KeyError): - pass - - # 4. Scan Loop - def scan_body(carry, _): - iteration = carry['loop_iteration'] - current_rng = carry['rng_stream'] - - step_rng, next_rng = jax.random.split(current_rng) - stage_rngs = jax.random.split(step_rng, self.num_stages) - stages_inputs = self.get_iteration_inputs( - iteration, - carry['state_io'], - carry['circ_storage'], - carry['shift'] - ) - stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") + if positions is not None: + positions = positions.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length + )) + + if segment_ids is not None: + segment_ids = segment_ids.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length + )) + + # 2. Initialize State + loop_state = self.init_states(inputs) - # Position gathering - # [FIX] Use get_microbatch_and_repeat_ids to be consistent with parity logic - mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) + # 3. Prepare Weights + # Convert NNX Params to pure dict for Linen & FSDP + param_values = self._to_pure_dict(self.stage_params) - stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None - stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None - - # [ADDED FOR PARITY] Dynamic Weight Selection (for Circular Pipelines) - # If repeats=1, this simply returns param_values as-is. - current_params = self.get_current_stage_weights(param_values, iteration) - - # VMAP Function - def vmapped_linen_apply(p, x, r, pos, seg): - variables = {'params': p} - rngs_dict = {'dropout': r} if not deterministic else {} + # Apply FSDP All-Gather if requested + if self.config.pipeline_fsdp_ag_once and partition_spec is not None: + try: + param_values = self._all_gather_over_fsdp(param_values, partition_spec) + except (ValueError, TypeError, KeyError): + # Log warning in real code if needed + pass + + # 4. Scan Loop + def scan_body(carry, _): + iteration = carry['loop_iteration'] + current_rng = carry['rng_stream'] + + step_rng, next_rng = jax.random.split(current_rng) + stage_rngs = jax.random.split(step_rng, self.num_stages) - return self.linen_module.apply( - variables, - x, - decoder_segment_ids=seg, - decoder_positions=pos, - deterministic=deterministic, - rngs=rngs_dict, - model_mode=model_mode + stages_inputs = self.get_iteration_inputs( + iteration, + carry['state_io'], + carry['circ_storage'], + carry['shift'] ) + # Checkpoint inputs + stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") + + # Gather Positions & Segments + mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) + + stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None + stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None + + # Dynamic Weight Selection (for Circular Pipelines) + current_params = self.get_current_stage_weights(param_values, iteration) + + # VMAP Function + def vmapped_linen_apply(p, x, r, pos, seg): + variables = {'params': p} + rngs_dict = {'dropout': r} if not deterministic else {} + + return self.linen_module.apply( + variables, + x, + decoder_segment_ids=seg, + decoder_positions=pos, + deterministic=deterministic, + rngs=rngs_dict, + model_mode=model_mode + ) + + # Apply Remat/Checkpointing + if self.remat_policy is not None: + vmapped_linen_apply = jax.checkpoint(vmapped_linen_apply, prevent_cse=False) + + # VMAP Execution + # Build args dynamically handling None + vmap_args = [current_params, stages_inputs, stage_rngs] + vmap_axes = [0, 0, 0] + + vmap_args.append(stages_positions) + vmap_axes.append(0 if stages_positions is not None else None) + + vmap_args.append(stages_segment_ids) + vmap_axes.append(0 if stages_segment_ids is not None else None) + + stages_output = jax.vmap(vmapped_linen_apply, in_axes=tuple(vmap_axes))(*vmap_args) + + if hasattr(self.config, 'scan_layers') and self.config.scan_layers: + if isinstance(stages_output, tuple): + stages_output = stages_output[0] - # Dynamic VMAP Arguments - vmap_args = [current_params, stages_inputs, stage_rngs] - vmap_axes = [0, 0, 0] # Params(0), Inputs(0), Rngs(0) + new_loop_state = self.get_new_loop_state(stages_output, carry) + new_loop_state['rng_stream'] = next_rng + + return new_loop_state, None - vmap_args.append(stages_positions) - vmap_axes.append(0 if stages_positions is not None else None) + bubble_iterations = self.forwarding_delay * (self.num_stages - 1) + real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats + total_ticks = real_iterations + bubble_iterations - vmap_args.append(stages_segment_ids) - vmap_axes.append(0 if stages_segment_ids is not None else None) + final_state, _ = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) - stages_output = jax.vmap(vmapped_linen_apply, in_axes=tuple(vmap_axes))(*vmap_args) - - if hasattr(self.config, 'scan_layers') and self.config.scan_layers: - if isinstance(stages_output, tuple): - stages_output = stages_output[0] + output = final_state['state_io'] - new_loop_state = self.get_new_loop_state(stages_output, carry) - new_loop_state['rng_stream'] = next_rng - - return new_loop_state, None + # Permute output (Sort Microbatches) + output = self.permute_output_micro_per_stage_dim(output) - bubble_iterations = self.forwarding_delay * (self.num_stages - 1) - real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats - total_ticks = real_iterations + bubble_iterations + output = output.reshape((self.config.micro_batch_size_to_train_on, + self.config.max_target_length, + self.config.emb_dim)) + + return output - final_state, _ = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) +def add_stage_axis_to_partitioning(variable, repeats=1): + """ + Metadata function for to_linen. + Applies 'stage' sharding to the correct dimension. + If repeats > 1, the first dimension is 'Repeats' (Replicated/None), + and the second is 'Stage'. + """ + # 1. Try to get existing partitioning + partitioned_obj = variable_to_logically_partitioned(variable) + + # 2. Extract base names + if isinstance(partitioned_obj, nn.LogicallyPartitioned): + base_names = partitioned_obj.names + value = partitioned_obj.value + else: + # Heuristics if metadata lost + value = partitioned_obj + if not hasattr(value, 'ndim'): + return value + + # Base heuristics for inner layers (Embed, MLP) + # Note: These are for the INNER tensor (excluding Stage/Repeat dims) + # We need to deduce what the inner rank is. + # If repeats=1, inner_rank = ndim - 1. + # If repeats>1, inner_rank = ndim - 2. + + # However, checking absolute ndim is safer if we know the structure + # Let's reconstruct based on full ndim. + base_names = None - output = final_state['state_io'] + # 3. Construct new names based on repeats and rank + ndim = value.ndim + + if repeats > 1: + # Expected Structure: [Repeats, Stage, ...] + # We want: (None, 'stage', ...) + + if base_names is not None: + # If we recovered ('fsdp', 'tensor'), just prepend + new_names = (None, 'stage') + base_names + else: + # Heuristics for [Repeats, Stage, In, Out] + if ndim == 4: # [Repeats, Stage, Embed, MLP] + new_names = (None, 'stage', 'fsdp', 'tensor') + elif ndim == 3: # [Repeats, Stage, Bias] + new_names = (None, 'stage', 'fsdp') + else: + new_names = (None, 'stage') + (None,) * (ndim - 2) + else: + # Expected Structure: [Stage, ...] + # We want: ('stage', ...) + + if base_names is not None: + new_names = ('stage',) + base_names + else: + if ndim == 3: # [Stage, Embed, MLP] + new_names = ('stage', 'fsdp', 'tensor') + elif ndim == 2: # [Stage, Bias] + new_names = ('stage', 'fsdp') + else: + new_names = ('stage',) + (None,) * (ndim - 1) - # [ADDED FOR PARITY] Permute Output to correct order - output = self.permute_output_micro_per_stage_dim(output) + return nn.LogicallyPartitioned(value, new_names) - output = output.reshape((self.config.micro_batch_size_to_train_on, - self.config.max_target_length, - self.config.emb_dim)) - - return output def create_pipeline( config: Config, layer: Callable | type, mesh: Mesh, remat_policy: Any = None, ) -> nnx_wrappers.ToLinen: - """Factory function to create a Pipeline wrapped as a Linen module. - - Args: - config: Model configuration - layer: NNX or Linen decoder layer class to use for pipeline stages - mesh: Device mesh for sharding - remat_policy: Remat policy for loop iterations - use_nnx: Whether to use NNX pipeline (True) or Linen (False) - - Returns: - PipelineToLinen wrapper around the NNX Pipeline - """ - return nnx_wrappers.to_linen( - Pipeline, - config=config, - mesh=mesh, - layer=layer, - remat_policy=remat_policy, - name="pipeline_module", - abstract_init=False, - metadata_fn=variable_to_logically_partitioned, - ) + """Factory function to create a Pipeline wrapped as a Linen module.""" + # Determine repeats from config + repeats = getattr(config, 'num_pipeline_repeats', 1) + + # Create specialized metadata function + metadata_fn = functools.partial(add_stage_axis_to_partitioning, repeats=repeats) + + return nnx_wrappers.to_linen( + Pipeline, + config=config, + mesh=mesh, + layer=layer, + remat_policy=remat_policy, + name="pipeline_module", + abstract_init=False, + metadata_fn=metadata_fn, # Pass the partial function + ) From 3aefaf5576842ccb02fb1600e10395c4c36dd32f Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Thu, 11 Dec 2025 01:37:29 +0000 Subject: [PATCH 09/14] initial fix --- src/MaxText/layers/pipeline.py | 75 ++++++++++++++++++++++------------ 1 file changed, 48 insertions(+), 27 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index eabb139e8f..a143cfaabb 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -406,6 +406,30 @@ def gather_weights_for_stages_in(w): return jax.tree.map(gather_weights_for_stages_in, pipeline_weights) + def get_pipeline_remat_policy(self): + """ + Constructs the checkpoint policy. + 1. Always saves 'iteration_input' (required for Pipeline Parallelism to save memory). + 2. Merges with self.remat_policy if provided (e.g. for attention layers). + """ + # If config says "custom", we trust the user provided policy entirely + if self.config.remat_policy == "custom": + return self.remat_policy + + # Base Policy: Save the inputs to the stage. + # This reduces memory from O(N_layers) to O(1) per stage. + save_input_policy = jax.checkpoint_policies.save_only_these_names( + "iteration_input", "decoder_layer_input" + ) + + # Combine with user-provided policy (if any) + if self.remat_policy is not None: + return jax.checkpoint_policies.save_from_both_policies( + self.remat_policy, save_input_policy + ) + return save_input_policy + + # Output sorting def permute_output_micro_per_stage_dim(self, output): microbatch_0_idx = self.iterations_to_complete_first_microbatch_one_repeat() % self.microbatches_per_stage @@ -501,9 +525,25 @@ def vmapped_linen_apply(p, x, r, pos, seg): model_mode=model_mode ) - # Apply Remat/Checkpointing - if self.remat_policy is not None: - vmapped_linen_apply = jax.checkpoint(vmapped_linen_apply, prevent_cse=False) + # 3. APPLY THE SWITCH (set_remat_policy_on_pipeline_iterations) + # This logic mimics the original: + # if self.config.set_remat_policy_on_pipeline_iterations: + # run_iteration_scannable = nn.remat(..., policy=self.get_pipeline_remat_policy()) + print(f"set_remat_policy_on_pipeline_iterations",self.config.set_remat_policy_on_pipeline_iterations) + if self.config.set_remat_policy_on_pipeline_iterations: + # Retrieve the strategy + policy = self.get_pipeline_remat_policy() + + # Logic for CSE (Common Subexpression Elimination) + # prevent_cse defaults to False if scanning, True otherwise to save memory + prevent_cse = not self.config.scan_pipeline_iterations + + # Apply jax.checkpoint (equivalent to nn.remat) + vmapped_linen_apply = jax.checkpoint( + vmapped_linen_apply, + policy=policy, + prevent_cse=prevent_cse + ) # VMAP Execution # Build args dynamically handling None @@ -551,41 +591,25 @@ def add_stage_axis_to_partitioning(variable, repeats=1): If repeats > 1, the first dimension is 'Repeats' (Replicated/None), and the second is 'Stage'. """ - # 1. Try to get existing partitioning partitioned_obj = variable_to_logically_partitioned(variable) - # 2. Extract base names if isinstance(partitioned_obj, nn.LogicallyPartitioned): base_names = partitioned_obj.names value = partitioned_obj.value else: - # Heuristics if metadata lost + # Heuristics value = partitioned_obj if not hasattr(value, 'ndim'): return value - - # Base heuristics for inner layers (Embed, MLP) - # Note: These are for the INNER tensor (excluding Stage/Repeat dims) - # We need to deduce what the inner rank is. - # If repeats=1, inner_rank = ndim - 1. - # If repeats>1, inner_rank = ndim - 2. - - # However, checking absolute ndim is safer if we know the structure - # Let's reconstruct based on full ndim. base_names = None - # 3. Construct new names based on repeats and rank ndim = value.ndim if repeats > 1: - # Expected Structure: [Repeats, Stage, ...] - # We want: (None, 'stage', ...) - + # Structure: [Repeats, Stage, Inner...] if base_names is not None: - # If we recovered ('fsdp', 'tensor'), just prepend new_names = (None, 'stage') + base_names else: - # Heuristics for [Repeats, Stage, In, Out] if ndim == 4: # [Repeats, Stage, Embed, MLP] new_names = (None, 'stage', 'fsdp', 'tensor') elif ndim == 3: # [Repeats, Stage, Bias] @@ -593,21 +617,18 @@ def add_stage_axis_to_partitioning(variable, repeats=1): else: new_names = (None, 'stage') + (None,) * (ndim - 2) else: - # Expected Structure: [Stage, ...] - # We want: ('stage', ...) - + # Structure: [Stage, Inner...] if base_names is not None: new_names = ('stage',) + base_names else: - if ndim == 3: # [Stage, Embed, MLP] + if ndim == 3: new_names = ('stage', 'fsdp', 'tensor') - elif ndim == 2: # [Stage, Bias] + elif ndim == 2: new_names = ('stage', 'fsdp') else: new_names = ('stage',) + (None,) * (ndim - 1) return nn.LogicallyPartitioned(value, new_names) - def create_pipeline( config: Config, layer: Callable | type, From 432ad38794cf6fff02e2eb5abdff7ba24aa3c621 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Thu, 11 Dec 2025 02:15:29 +0000 Subject: [PATCH 10/14] second attempt --- src/MaxText/layers/pipeline.py | 357 ++++++++++++++++----------------- 1 file changed, 174 insertions(+), 183 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index a143cfaabb..647ee554a4 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -34,10 +34,10 @@ class Pipeline(nnx.Module): def __init__( self, - layer: nn.Module, # Instance of the Linen Layer (e.g. DecoderLayer) - config: Any, # MaxText Config object - mesh: Mesh, # JAX Mesh - rngs: nnx.Rngs, # NNX RNGs + layer: nn.Module, + config: Any, + mesh: Mesh, + rngs: nnx.Rngs, remat_policy: Any = None ): self.config = config @@ -45,17 +45,15 @@ def __init__( self.linen_module = layer self.remat_policy = remat_policy - # --- 1. Calculate Pipeline Dimensions --- + # --- Dimensions --- self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1 - self.total_microbatches = self.config.num_pipeline_microbatches self.microbatch_size = self.config.micro_batch_size_to_train_on // self.total_microbatches self.microbatches_per_stage = self.total_microbatches // self.num_stages - self.use_circ_storage = self._need_circ_storage() - # Axis Naming + # --- Axis Naming --- if hasattr(self.config, 'expert_shard_attention_option') and self.config.expert_shard_attention_option == EP_AS_CONTEXT: self.batch_axis_name = "activation_batch_no_exp" self.seq_len_axis_name = "activation_length" @@ -69,17 +67,16 @@ def __init__( self.config.emb_dim ) - # --- 2. Eager Initialization (FIXED for Circular Repeats) --- + # --- Initialization --- init_rng_key = rngs.params() repeats = self.config.num_pipeline_repeats - + + # Helper to init one stage def init_single_stage(key): - # Create dummy inputs dummy_in = jnp.zeros(self.input_shape, dtype=jnp.float32) dummy_pos_shape = (self.input_shape[0], self.input_shape[1]) dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) - return self.linen_module.init( {'params': key}, dummy_in, @@ -89,36 +86,50 @@ def init_single_stage(key): model_mode=MODEL_MODE_TRAIN ) + # 1. Run ONCE without vmap to capture correct Metadata (names) + # This returns LogicallyPartitioned objects with rules like ('embed', 'mlp') + ref_variables = init_single_stage(init_rng_key) + + # 2. Run WITH vmap to get the actual Data (values) if repeats > 1: - # Case A: Circular Pipeline (Multiple Repeats) - # We need independent weights for [Repeats, Num_Stages] repeat_keys = jax.random.split(init_rng_key, repeats) - # Split each repeat key into stage keys -> Shape: [Repeats, Stages] stage_rng_keys = jax.vmap(lambda k: jax.random.split(k, self.num_stages))(repeat_keys) - - # Double vmap: Map over Repeats (Axis 0), then Stages (Axis 0 of inner) - # Output Shape: [Repeats, Stages, ...] raw_variables = jax.vmap(jax.vmap(init_single_stage))(stage_rng_keys) else: - # Case B: Standard Pipeline (Single Repeat) - # Shape: [Num_Stages] stage_rng_keys = jax.random.split(init_rng_key, self.num_stages) - # Output Shape: [Stages, ...] raw_variables = jax.vmap(init_single_stage)(stage_rng_keys) - # --- 3. Register Parameters with NNX --- - def _to_nnx_structure(node): - if hasattr(node, 'items'): - return nnx.Dict({k: _to_nnx_structure(v) for k, v in node.items()}) - elif isinstance(node, (list, tuple)): - return nnx.List([_to_nnx_structure(v) for v in node]) + # 3. Register Parameters with Metadata Injection + # We traverse the 'raw' (data) and 'ref' (metadata) trees together. + def _create_param_with_metadata(raw_node, ref_node): + if hasattr(raw_node, 'items'): + return nnx.Dict({k: _create_param_with_metadata(v, ref_node[k]) for k, v in raw_node.items()}) + elif isinstance(raw_node, (list, tuple)): + return nnx.List([_create_param_with_metadata(v, ref_node[i]) for i, v in enumerate(raw_node)]) else: - return nnx.Param(node) + # Unwrap raw value if it came wrapped + if isinstance(raw_node, nn.LogicallyPartitioned): + actual_value = raw_node.value + else: + actual_value = raw_node + + # Check ref_node for metadata + sharding_kwargs = {} + if hasattr(ref_node, 'names'): + original_names = ref_node.names + if repeats > 1: + new_names = (None, 'stage') + original_names + else: + new_names = ('stage',) + original_names + + # Pass as keyword argument to nnx.Param + sharding_kwargs['sharding_names'] = new_names + + return nnx.Param(actual_value, **sharding_kwargs) + + self.stage_params = _create_param_with_metadata(raw_variables['params'], ref_variables['params']) - self.stage_params = _to_nnx_structure(raw_variables['params']) - if 'batch_stats' in raw_variables: - print("Warning: Pipeline detected batch_stats during init. Registering as VariableState.") # ========================================================================== # Helper Methods # ========================================================================== @@ -407,22 +418,15 @@ def gather_weights_for_stages_in(w): return jax.tree.map(gather_weights_for_stages_in, pipeline_weights) def get_pipeline_remat_policy(self): - """ - Constructs the checkpoint policy. - 1. Always saves 'iteration_input' (required for Pipeline Parallelism to save memory). - 2. Merges with self.remat_policy if provided (e.g. for attention layers). - """ - # If config says "custom", we trust the user provided policy entirely - if self.config.remat_policy == "custom": - return self.remat_policy - - # Base Policy: Save the inputs to the stage. - # This reduces memory from O(N_layers) to O(1) per stage. + # Strict policy: Save ONLY the input tensor named 'iteration_input'. + # Discard everything else (dots, gathers, etc). save_input_policy = jax.checkpoint_policies.save_only_these_names( "iteration_input", "decoder_layer_input" ) - # Combine with user-provided policy (if any) + if self.config.remat_policy == "custom": + return self.remat_policy + if self.remat_policy is not None: return jax.checkpoint_policies.save_from_both_policies( self.remat_policy, save_input_policy @@ -436,168 +440,155 @@ def permute_output_micro_per_stage_dim(self, output): permutation = (jnp.arange(self.microbatches_per_stage) + microbatch_0_idx) % self.microbatches_per_stage output = output[:, permutation] return output - + def __call__( - self, - inputs: jax.Array, - segment_ids: Optional[jax.Array] = None, - positions: Optional[jax.Array] = None, - deterministic: bool = True, - model_mode: str = MODEL_MODE_TRAIN, - partition_spec: Any = None, - ) -> jax.Array: + self, + inputs: jax.Array, + segment_ids: Optional[jax.Array] = None, + positions: Optional[jax.Array] = None, + deterministic: bool = True, + model_mode: str = MODEL_MODE_TRAIN, + partition_spec: Any = None, + ) -> jax.Array: + + # 1. Reshape Inputs + inputs = inputs.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length, + self.config.emb_dim + )) + + if positions is not None: + positions = positions.reshape(( + self.total_microbatches, + self.microbatch_size, + self.config.max_target_length + )) - # 1. Reshape Inputs - inputs = inputs.reshape(( + if segment_ids is not None: + segment_ids = segment_ids.reshape(( self.total_microbatches, self.microbatch_size, - self.config.max_target_length, - self.config.emb_dim + self.config.max_target_length )) + + # 2. Initialize State + loop_state = self.init_states(inputs) + + # 3. Prepare Weights + param_values = self._to_pure_dict(self.stage_params) + + if self.config.pipeline_fsdp_ag_once and partition_spec is not None: + try: + param_values = self._all_gather_over_fsdp(param_values, partition_spec) + except (ValueError, TypeError, KeyError): + pass + + # 4. Scan Loop + def scan_body(carry, _): + iteration = carry['loop_iteration'] + current_rng = carry['rng_stream'] - if positions is not None: - positions = positions.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length - )) - - if segment_ids is not None: - segment_ids = segment_ids.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length - )) - - # 2. Initialize State - loop_state = self.init_states(inputs) + step_rng, next_rng = jax.random.split(current_rng) + stage_rngs = jax.random.split(step_rng, self.num_stages) - # 3. Prepare Weights - # Convert NNX Params to pure dict for Linen & FSDP - param_values = self._to_pure_dict(self.stage_params) + stages_inputs = self.get_iteration_inputs( + iteration, + carry['state_io'], + carry['circ_storage'], + carry['shift'] + ) + # Checkpoint inputs: Crucial for policy detection + stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") - # Apply FSDP All-Gather if requested - if self.config.pipeline_fsdp_ag_once and partition_spec is not None: - try: - param_values = self._all_gather_over_fsdp(param_values, partition_spec) - except (ValueError, TypeError, KeyError): - # Log warning in real code if needed - pass - - # 4. Scan Loop - def scan_body(carry, _): - iteration = carry['loop_iteration'] - current_rng = carry['rng_stream'] - - step_rng, next_rng = jax.random.split(current_rng) - stage_rngs = jax.random.split(step_rng, self.num_stages) - - stages_inputs = self.get_iteration_inputs( - iteration, - carry['state_io'], - carry['circ_storage'], - carry['shift'] - ) - # Checkpoint inputs - stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") - - # Gather Positions & Segments - mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) - - stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None - stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None - - # Dynamic Weight Selection (for Circular Pipelines) - current_params = self.get_current_stage_weights(param_values, iteration) - - # VMAP Function - def vmapped_linen_apply(p, x, r, pos, seg): + # Gather Positions & Segments + mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) + stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None + stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None + + # Dynamic Weight Selection + current_params = self.get_current_stage_weights(param_values, iteration) + + # --- A. Define the VMAP Logic (Pure Execution) --- + def execution_logic(params, inputs, rngs, pos, seg): + # Inner function applied to ONE stage + def stage_fn(p, x, r, po, se): variables = {'params': p} rngs_dict = {'dropout': r} if not deterministic else {} - return self.linen_module.apply( - variables, - x, - decoder_segment_ids=seg, - decoder_positions=pos, - deterministic=deterministic, - rngs=rngs_dict, + variables, x, + decoder_segment_ids=se, decoder_positions=po, + deterministic=deterministic, rngs=rngs_dict, model_mode=model_mode ) - - # 3. APPLY THE SWITCH (set_remat_policy_on_pipeline_iterations) - # This logic mimics the original: - # if self.config.set_remat_policy_on_pipeline_iterations: - # run_iteration_scannable = nn.remat(..., policy=self.get_pipeline_remat_policy()) - print(f"set_remat_policy_on_pipeline_iterations",self.config.set_remat_policy_on_pipeline_iterations) - if self.config.set_remat_policy_on_pipeline_iterations: - # Retrieve the strategy - policy = self.get_pipeline_remat_policy() - - # Logic for CSE (Common Subexpression Elimination) - # prevent_cse defaults to False if scanning, True otherwise to save memory - prevent_cse = not self.config.scan_pipeline_iterations - - # Apply jax.checkpoint (equivalent to nn.remat) - vmapped_linen_apply = jax.checkpoint( - vmapped_linen_apply, - policy=policy, - prevent_cse=prevent_cse - ) - - # VMAP Execution - # Build args dynamically handling None - vmap_args = [current_params, stages_inputs, stage_rngs] - vmap_axes = [0, 0, 0] - - vmap_args.append(stages_positions) - vmap_axes.append(0 if stages_positions is not None else None) - - vmap_args.append(stages_segment_ids) - vmap_axes.append(0 if stages_segment_ids is not None else None) - - stages_output = jax.vmap(vmapped_linen_apply, in_axes=tuple(vmap_axes))(*vmap_args) - if hasattr(self.config, 'scan_layers') and self.config.scan_layers: - if isinstance(stages_output, tuple): - stages_output = stages_output[0] + # Apply VMAP here + # Map axes: Params(0), Inputs(0), Rngs(0), Pos(0/None), Seg(0/None) + vmap_axes = [0, 0, 0] + vmap_args = [params, inputs, rngs] + + vmap_axes.append(0 if pos is not None else None) + vmap_args.append(pos) + + vmap_axes.append(0 if seg is not None else None) + vmap_args.append(seg) + + return jax.vmap(stage_fn, in_axes=tuple(vmap_axes))(*vmap_args) - new_loop_state = self.get_new_loop_state(stages_output, carry) - new_loop_state['rng_stream'] = next_rng + # --- B. Apply Checkpoint to the VMAP (Parity with Linen) --- + # We wrap the entire execution logic (which contains the vmap) + + if self.config.set_remat_policy_on_pipeline_iterations: + policy = self.get_pipeline_remat_policy() - return new_loop_state, None + # Checkpoint the function that DOES the vmap + execution_logic = jax.checkpoint( + execution_logic, + policy=policy, + prevent_cse=False + ) - bubble_iterations = self.forwarding_delay * (self.num_stages - 1) - real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats - total_ticks = real_iterations + bubble_iterations + # --- C. Execute --- + stages_output = execution_logic( + current_params, stages_inputs, stage_rngs, stages_positions, stages_segment_ids + ) + + if hasattr(self.config, 'scan_layers') and self.config.scan_layers: + if isinstance(stages_output, tuple): + stages_output = stages_output[0] - final_state, _ = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) + new_loop_state = self.get_new_loop_state(stages_output, carry) + new_loop_state['rng_stream'] = next_rng + + return new_loop_state, None - output = final_state['state_io'] + bubble_iterations = self.forwarding_delay * (self.num_stages - 1) + real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats + total_ticks = real_iterations + bubble_iterations - # Permute output (Sort Microbatches) - output = self.permute_output_micro_per_stage_dim(output) + final_state, _ = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) - output = output.reshape((self.config.micro_batch_size_to_train_on, - self.config.max_target_length, - self.config.emb_dim)) - - return output + output = final_state['state_io'] + output = self.permute_output_micro_per_stage_dim(output) + output = output.reshape((self.config.micro_batch_size_to_train_on, + self.config.max_target_length, + self.config.emb_dim)) + + return output def add_stage_axis_to_partitioning(variable, repeats=1): """ Metadata function for to_linen. - Applies 'stage' sharding to the correct dimension. - If repeats > 1, the first dimension is 'Repeats' (Replicated/None), - and the second is 'Stage'. """ + # 1. Try to recover metadata partitioned_obj = variable_to_logically_partitioned(variable) if isinstance(partitioned_obj, nn.LogicallyPartitioned): base_names = partitioned_obj.names value = partitioned_obj.value else: - # Heuristics + # 2. Fallback to Heuristics (Crucial for jax.vmap created variables) value = partitioned_obj if not hasattr(value, 'ndim'): return value @@ -606,21 +597,23 @@ def add_stage_axis_to_partitioning(variable, repeats=1): ndim = value.ndim if repeats > 1: - # Structure: [Repeats, Stage, Inner...] + # [Repeats, Stage, Inner...] if base_names is not None: new_names = (None, 'stage') + base_names else: - if ndim == 4: # [Repeats, Stage, Embed, MLP] + # Heuristic for Rank 4 (Repeats, Stage, Embed, MLP) + if ndim == 4: new_names = (None, 'stage', 'fsdp', 'tensor') - elif ndim == 3: # [Repeats, Stage, Bias] + elif ndim == 3: new_names = (None, 'stage', 'fsdp') else: new_names = (None, 'stage') + (None,) * (ndim - 2) else: - # Structure: [Stage, Inner...] + # [Stage, Inner...] if base_names is not None: new_names = ('stage',) + base_names else: + # Heuristic for Rank 3 (Stage, Embed, MLP) if ndim == 3: new_names = ('stage', 'fsdp', 'tensor') elif ndim == 2: @@ -629,6 +622,7 @@ def add_stage_axis_to_partitioning(variable, repeats=1): new_names = ('stage',) + (None,) * (ndim - 1) return nn.LogicallyPartitioned(value, new_names) + def create_pipeline( config: Config, layer: Callable | type, @@ -636,12 +630,9 @@ def create_pipeline( remat_policy: Any = None, ) -> nnx_wrappers.ToLinen: """Factory function to create a Pipeline wrapped as a Linen module.""" - # Determine repeats from config - repeats = getattr(config, 'num_pipeline_repeats', 1) - - # Create specialized metadata function - metadata_fn = functools.partial(add_stage_axis_to_partitioning, repeats=repeats) + repeats = getattr(config, "num_pipeline_repeat", 1) + metadata_fn = functools.partial(add_stage_axis_to_partitioning,repeats=repeats) return nnx_wrappers.to_linen( Pipeline, config=config, @@ -650,5 +641,5 @@ def create_pipeline( remat_policy=remat_policy, name="pipeline_module", abstract_init=False, - metadata_fn=metadata_fn, # Pass the partial function + metadata_fn=metadata_fn, ) From a174af65eaeae5094524175010ffe5e57257283d Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Thu, 11 Dec 2025 07:20:19 +0000 Subject: [PATCH 11/14] third attempt --- src/MaxText/layers/pipeline.py | 243 +++++++++++++++------------------ 1 file changed, 112 insertions(+), 131 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 647ee554a4..84e644e339 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -12,19 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" Pipeline layer wrapping a decoder layer(s). Supports circular pipelining """ +"""Pipeline layer wrapping a decoder layer(s). Supports circular pipelining""" import functools -from typing import Any, Optional, Callable - +from typing import Any, Optional, Callable, Tuple, Union import jax import jax.numpy as jnp from jax.sharding import Mesh, NamedSharding, PartitionSpec from flax import nnx from flax import linen as nn +from flax.linen import spmd + from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT from MaxText.layers import nnx_wrappers from MaxText.layers.initializers import variable_to_logically_partitioned -from flax.linen import spmd + class Pipeline(nnx.Module): """ @@ -45,7 +46,7 @@ def __init__( self.linen_module = layer self.remat_policy = remat_policy - # --- Dimensions --- + # --- 1. Dimensions --- self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1 self.total_microbatches = self.config.num_pipeline_microbatches @@ -67,30 +68,31 @@ def __init__( self.config.emb_dim ) - # --- Initialization --- + # --- 2. Initialization --- init_rng_key = rngs.params() repeats = self.config.num_pipeline_repeats - # Helper to init one stage def init_single_stage(key): - dummy_in = jnp.zeros(self.input_shape, dtype=jnp.float32) + # [FIX] Use bfloat16 for dummy inputs to save memory + dummy_in = jnp.zeros(self.input_shape, dtype=jnp.bfloat16) + dummy_pos_shape = (self.input_shape[0], self.input_shape[1]) dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) + return self.linen_module.init( {'params': key}, dummy_in, - decoder_segment_ids=dummy_segments, - decoder_positions=dummy_positions, - deterministic=True, - model_mode=MODEL_MODE_TRAIN + dummy_segments, + dummy_positions, + True, + MODEL_MODE_TRAIN ) - # 1. Run ONCE without vmap to capture correct Metadata (names) - # This returns LogicallyPartitioned objects with rules like ('embed', 'mlp') + # Step A: Init ONCE to get Metadata ref_variables = init_single_stage(init_rng_key) - # 2. Run WITH vmap to get the actual Data (values) + # Step B: Init WITH VMAP to get Data if repeats > 1: repeat_keys = jax.random.split(init_rng_key, repeats) stage_rng_keys = jax.vmap(lambda k: jax.random.split(k, self.num_stages))(repeat_keys) @@ -99,21 +101,24 @@ def init_single_stage(key): stage_rng_keys = jax.random.split(init_rng_key, self.num_stages) raw_variables = jax.vmap(init_single_stage)(stage_rng_keys) - # 3. Register Parameters with Metadata Injection - # We traverse the 'raw' (data) and 'ref' (metadata) trees together. + print(f"raw_variables : {raw_variables}") + # --- 3. Register Parameters --- def _create_param_with_metadata(raw_node, ref_node): if hasattr(raw_node, 'items'): return nnx.Dict({k: _create_param_with_metadata(v, ref_node[k]) for k, v in raw_node.items()}) elif isinstance(raw_node, (list, tuple)): return nnx.List([_create_param_with_metadata(v, ref_node[i]) for i, v in enumerate(raw_node)]) else: - # Unwrap raw value if it came wrapped + # Unwrap LogicallyPartitioned if isinstance(raw_node, nn.LogicallyPartitioned): actual_value = raw_node.value else: actual_value = raw_node - # Check ref_node for metadata + # Cast float32 -> bfloat16 to save memory + if actual_value.dtype == jnp.float32: + actual_value = actual_value.astype(jnp.bfloat16) + sharding_kwargs = {} if hasattr(ref_node, 'names'): original_names = ref_node.names @@ -121,15 +126,20 @@ def _create_param_with_metadata(raw_node, ref_node): new_names = (None, 'stage') + original_names else: new_names = ('stage',) + original_names - - # Pass as keyword argument to nnx.Param sharding_kwargs['sharding_names'] = new_names return nnx.Param(actual_value, **sharding_kwargs) + + # Depending on how the inner module is wrapped (Scan vs Sequential), 'params' might be the root or nested. + if 'params' in raw_variables: + params_dict = raw_variables['params'] + ref_dict = ref_variables['params'] + else: + # Fallback: Assume the whole dict is params (common in some scan wrappings) + params_dict = raw_variables + ref_dict = ref_variables - self.stage_params = _create_param_with_metadata(raw_variables['params'], ref_variables['params']) - - + self.stage_params = _create_param_with_metadata(params_dict, ref_dict) # ========================================================================== # Helper Methods # ========================================================================== @@ -141,15 +151,9 @@ def _need_circ_storage(self): def iterations_to_complete_first_microbatch_one_repeat(self): return self.forwarding_delay * (self.num_stages - 1) - def init_states(self, inputs): - """Initialize pipeline buffers. - Args: - inputs: Rank 4 array [Total_Microbatches, Micro_Size, Seq, Emb] - """ - # 1. Shift Buffer - # Shape: [Num_Stages, Micro_Size, Seq, Emb] - # (Derived from inputs.shape[1:]) + """Initialize pipeline buffers.""" + # Shift Buffer shift = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) shift = nn.with_logical_constraint( shift, @@ -158,7 +162,7 @@ def init_states(self, inputs): mesh=self.mesh, ) - # 2. Prev Outputs (for forwarding delay) + # Prev Outputs if self.config.pipeline_delay_activation_forwarding: prev_outputs = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) prev_outputs = nn.with_logical_constraint( @@ -170,8 +174,7 @@ def init_states(self, inputs): else: prev_outputs = None - # 3. State IO (The Main Buffer) - # Reshape: [Total_Micro, ...] -> [Stages, Micro_Per_Stage, ...] + # State IO state_io = jnp.reshape( inputs, (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:] @@ -183,10 +186,8 @@ def init_states(self, inputs): mesh=self.mesh, ) - # 4. Circular Storage + # Circular Storage if self.use_circ_storage: - # Shape: [Num_Stages, Total_Microbatches, Micro_Size, Seq, Emb] - # (Derived from inputs.shape) circ_storage = jnp.zeros((self.num_stages,) + inputs.shape, dtype=inputs.dtype) circ_storage_mover = shift else: @@ -227,7 +228,6 @@ def select_state_or_input(first_stage_in, shift): ) stages_in = select_state_or_input(first_stage_in, shift) - stages_in = nn.with_logical_constraint( stages_in, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), @@ -313,11 +313,10 @@ def _update_state_io(state_in, stream_slice, output): } # ========================================================================== - # FSDP Helpers + # Parity Helpers # ========================================================================== def _all_gather_over_fsdp(self, params, partition_spec): - """Helper to apply FSDP all-gather constraint.""" def _remove_fsdp_from_spec(spec): if isinstance(spec, PartitionSpec): new_spec = [] @@ -341,35 +340,17 @@ def _remove_fsdp_sharding(sharding_tree): physical = nn.logical_to_mesh_sharding(partition_spec, mesh=self.mesh, rules=self.config.logical_axis_rules) physical_no_fsdp = _remove_fsdp_sharding(physical) return jax.lax.with_sharding_constraint(params, physical_no_fsdp) - + def _to_pure_dict(self, node): - """Recursively converts NNX containers to standard Python dicts/lists.""" - if hasattr(node, 'items'): # Handles nnx.Dict, dict, FrozenDict + if hasattr(node, 'items'): return {k: self._to_pure_dict(v) for k, v in node.items()} - elif isinstance(node, (list, tuple)): # Handles nnx.List, list, tuple + elif isinstance(node, (list, tuple)): return [self._to_pure_dict(v) for v in node] - elif hasattr(node, 'value'): # Handles nnx.Param, nnx.Variable + elif hasattr(node, 'value'): return node.value - return node # Leaf (e.g. JAX Array/Tracer) + return node - # ========================================================================== - # Main Execution - # ========================================================================== - def get_microbatch_ids(self, loop_iteration): - """Calculates the microbatch ID for each stage at the current tick.""" - # Calculate how many microbatches each stage has processed effectively - # Stage 0 starts at 0. Stage 1 starts after 'forwarding_delay', etc. - microbatches_processed = jnp.maximum( - loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), - 0 - ) - # Wrap around using modulo to get the ID within the circular buffer - microbatch_ids = microbatches_processed % self.total_microbatches - return microbatch_ids.astype(jnp.int32) - - # Handles logic when num_pipeline_repeats > 1 def get_microbatch_and_repeat_ids(self, loop_iteration): - """Gets the microbatch_ids and repeat_ids for all stages on this loop_iteration.""" microbatches_processed = jnp.maximum( loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), 0 @@ -378,19 +359,13 @@ def get_microbatch_and_repeat_ids(self, loop_iteration): repeat_ids = microbatches_processed // self.total_microbatches return microbatch_ids.astype(jnp.int32), repeat_ids.astype(jnp.int32) - # Helper for weight gathering def shard_dim_by_stages(self, x, dim: int): dims_mapping = [PartitionSpec.UNCONSTRAINED] * x.ndim dims_mapping[dim] = "stage" - # We construct the NamedSharding manually to match Linen's logical_to_mesh - # Note: In pure NNX/JAX, we can often just return x if sharding is handled by vmap, - # but we keep the constraint for strict parity. sharding = NamedSharding(self.mesh, PartitionSpec(*dims_mapping)) return jax.lax.with_sharding_constraint(x, sharding) - # Helper for circular weight selection def vmap_parallel_gather(self, weights, repeat_ids, repeat_dim_in_weights, stages_dim_in_weights): - """Use vmap to implement a sharded parallel gather for weights.""" def _gather_one(x, repeat_id): return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, repeat_id, 1, repeat_dim_in_weights), repeat_dim_in_weights) @@ -402,45 +377,40 @@ def _gather_one(x, repeat_id): stage_weights = self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim) return stage_weights - # The main function missing from scan_body def get_current_stage_weights(self, pipeline_weights, loop_iteration): if self.config.num_pipeline_repeats <= 1: return pipeline_weights - _, repeat_ids = self.get_microbatch_and_repeat_ids(loop_iteration) - # Helper to map over the PyTree of weights + _, repeat_ids = self.get_microbatch_and_repeat_ids(loop_iteration) def gather_weights_for_stages_in(w): - # Assumes weights are [Repeats, Stages, ...] -> Gather -> [Stages, ...] - return self.vmap_parallel_gather( - w, repeat_ids, repeat_dim_in_weights=0, stages_dim_in_weights=1 - ) + return self.vmap_parallel_gather(w, repeat_ids, repeat_dim_in_weights=0, stages_dim_in_weights=1) return jax.tree.map(gather_weights_for_stages_in, pipeline_weights) def get_pipeline_remat_policy(self): - # Strict policy: Save ONLY the input tensor named 'iteration_input'. - # Discard everything else (dots, gathers, etc). + if self.config.remat_policy == "custom": + return self.remat_policy + save_input_policy = jax.checkpoint_policies.save_only_these_names( "iteration_input", "decoder_layer_input" ) - if self.config.remat_policy == "custom": - return self.remat_policy - if self.remat_policy is not None: return jax.checkpoint_policies.save_from_both_policies( self.remat_policy, save_input_policy ) return save_input_policy - - # Output sorting def permute_output_micro_per_stage_dim(self, output): microbatch_0_idx = self.iterations_to_complete_first_microbatch_one_repeat() % self.microbatches_per_stage permutation = (jnp.arange(self.microbatches_per_stage) + microbatch_0_idx) % self.microbatches_per_stage output = output[:, permutation] return output - + + # ========================================================================== + # Main Execution + # ========================================================================== + def __call__( self, inputs: jax.Array, @@ -451,7 +421,6 @@ def __call__( partition_spec: Any = None, ) -> jax.Array: - # 1. Reshape Inputs inputs = inputs.reshape(( self.total_microbatches, self.microbatch_size, @@ -473,11 +442,20 @@ def __call__( self.config.max_target_length )) - # 2. Initialize State loop_state = self.init_states(inputs) - # 3. Prepare Weights param_values = self._to_pure_dict(self.stage_params) + + # Cast weights to compute_dtype (bfloat16) to save memory + # MaxText defaults to 'bfloat16'. Using float32 for weights fills HBM instantly. + compute_dtype = getattr(self.config, 'compute_dtype', jnp.bfloat16) + if isinstance(compute_dtype, str): + # Map string config (e.g., 'bfloat16') to JAX dtype + dtype_map = {'bfloat16': jnp.bfloat16, 'float32': jnp.float32, 'float16': jnp.float16} + compute_dtype = dtype_map.get(compute_dtype, jnp.bfloat16) + + # Helper to cast leaves + param_values = jax.tree.map(lambda x: x.astype(compute_dtype), param_values) if self.config.pipeline_fsdp_ag_once and partition_spec is not None: try: @@ -485,7 +463,6 @@ def __call__( except (ValueError, TypeError, KeyError): pass - # 4. Scan Loop def scan_body(carry, _): iteration = carry['loop_iteration'] current_rng = carry['rng_stream'] @@ -499,57 +476,50 @@ def scan_body(carry, _): carry['circ_storage'], carry['shift'] ) - # Checkpoint inputs: Crucial for policy detection + # MARKER for Remat Policy stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") - # Gather Positions & Segments mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None - # Dynamic Weight Selection current_params = self.get_current_stage_weights(param_values, iteration) - # --- A. Define the VMAP Logic (Pure Execution) --- + # 1. DEFINE EXECUTION LOGIC (Pure) def execution_logic(params, inputs, rngs, pos, seg): - # Inner function applied to ONE stage - def stage_fn(p, x, r, po, se): + def stage_fn(p, x, r, pos, seg): variables = {'params': p} rngs_dict = {'dropout': r} if not deterministic else {} return self.linen_module.apply( - variables, x, - decoder_segment_ids=se, decoder_positions=po, - deterministic=deterministic, rngs=rngs_dict, - model_mode=model_mode + variables, + x, + seg, pos, + deterministic, + model_mode, + rngs=rngs_dict, ) - # Apply VMAP here - # Map axes: Params(0), Inputs(0), Rngs(0), Pos(0/None), Seg(0/None) vmap_axes = [0, 0, 0] vmap_args = [params, inputs, rngs] - vmap_axes.append(0 if pos is not None else None) vmap_args.append(pos) - vmap_axes.append(0 if seg is not None else None) vmap_args.append(seg) return jax.vmap(stage_fn, in_axes=tuple(vmap_axes))(*vmap_args) - # --- B. Apply Checkpoint to the VMAP (Parity with Linen) --- - # We wrap the entire execution logic (which contains the vmap) - + # 2. APPLY CHECKPOINT TO VMAP (Logic matches Linen run_iteration_scannable) if self.config.set_remat_policy_on_pipeline_iterations: policy = self.get_pipeline_remat_policy() + prevent_cse = not self.config.scan_pipeline_iterations - # Checkpoint the function that DOES the vmap execution_logic = jax.checkpoint( execution_logic, - policy=policy, + # policy=policy, prevent_cse=False ) - # --- C. Execute --- + # 3. EXECUTE stages_output = execution_logic( current_params, stages_inputs, stage_rngs, stages_positions, stages_segment_ids ) @@ -577,62 +547,73 @@ def stage_fn(p, x, r, po, se): return output +# ========================================================================== +# Factory Functions & Metadata +# ========================================================================== + def add_stage_axis_to_partitioning(variable, repeats=1): """ - Metadata function for to_linen. + Metadata function. Idempotent check for 'stage' axis. """ - # 1. Try to recover metadata partitioned_obj = variable_to_logically_partitioned(variable) if isinstance(partitioned_obj, nn.LogicallyPartitioned): - base_names = partitioned_obj.names + names = partitioned_obj.names value = partitioned_obj.value else: - # 2. Fallback to Heuristics (Crucial for jax.vmap created variables) value = partitioned_obj if not hasattr(value, 'ndim'): return value - base_names = None + names = None + # Check if 'stage' is already present (Idempotency Check) + is_stage_present = False + if names is not None and len(names) > 0: + if repeats > 1: + if len(names) > 1 and names[1] == 'stage': + is_stage_present = True + else: + if names[0] == 'stage': + is_stage_present = True + + if is_stage_present: + return nn.LogicallyPartitioned(value, names) + + # Fallback/Construction ndim = value.ndim - if repeats > 1: - # [Repeats, Stage, Inner...] - if base_names is not None: - new_names = (None, 'stage') + base_names + if names is not None: + new_names = (None, 'stage') + names else: - # Heuristic for Rank 4 (Repeats, Stage, Embed, MLP) - if ndim == 4: + if ndim == 4: new_names = (None, 'stage', 'fsdp', 'tensor') - elif ndim == 3: + elif ndim == 3: new_names = (None, 'stage', 'fsdp') else: - new_names = (None, 'stage') + (None,) * (ndim - 2) + new_names = (None, 'stage') + (None,) * (max(0, ndim - 2)) else: - # [Stage, Inner...] - if base_names is not None: - new_names = ('stage',) + base_names + if names is not None: + new_names = ('stage',) + names else: - # Heuristic for Rank 3 (Stage, Embed, MLP) if ndim == 3: new_names = ('stage', 'fsdp', 'tensor') elif ndim == 2: new_names = ('stage', 'fsdp') else: - new_names = ('stage',) + (None,) * (ndim - 1) + new_names = ('stage',) + (None,) * (max(0, ndim - 1)) return nn.LogicallyPartitioned(value, new_names) + def create_pipeline( config: Config, layer: Callable | type, mesh: Mesh, remat_policy: Any = None, ) -> nnx_wrappers.ToLinen: - """Factory function to create a Pipeline wrapped as a Linen module.""" + repeats = getattr(config, 'num_pipeline_repeats', 1) + metadata_fn = functools.partial(add_stage_axis_to_partitioning, repeats=repeats) - repeats = getattr(config, "num_pipeline_repeat", 1) - metadata_fn = functools.partial(add_stage_axis_to_partitioning,repeats=repeats) return nnx_wrappers.to_linen( Pipeline, config=config, @@ -642,4 +623,4 @@ def create_pipeline( name="pipeline_module", abstract_init=False, metadata_fn=metadata_fn, - ) + ) \ No newline at end of file From a21eb3c5a55c03aadd971a1d600c9ada3cb0f2f0 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Thu, 11 Dec 2025 10:42:14 +0000 Subject: [PATCH 12/14] attempt five --- src/MaxText/layers/decoders.py | 1 + src/MaxText/layers/pipeline.py | 669 ++++++++++++--------------------- 2 files changed, 245 insertions(+), 425 deletions(-) diff --git a/src/MaxText/layers/decoders.py b/src/MaxText/layers/decoders.py index 0306a8b201..b2daf65d30 100644 --- a/src/MaxText/layers/decoders.py +++ b/src/MaxText/layers/decoders.py @@ -712,6 +712,7 @@ def get_layer_to_pipeline(blocks, cfg): "layers_per_stage", self.mesh, in_axes_tuple=(nn.broadcast,) * 4, + model_mode=self.model_mode, ) else: stage_module = SequentialBlockDecoderLayers( diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index 84e644e339..d7826f2a40 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -12,31 +12,69 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pipeline layer wrapping a decoder layer(s). Supports circular pipelining""" +"""Pipeline layer wrapping a decoder layer(s). Supports circular pipelining. +NNX Implementation. +""" import functools -from typing import Any, Optional, Callable, Tuple, Union +from typing import Any, Optional, Callable + import jax import jax.numpy as jnp from jax.sharding import Mesh, NamedSharding, PartitionSpec + from flax import nnx from flax import linen as nn -from flax.linen import spmd from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT from MaxText.layers import nnx_wrappers from MaxText.layers.initializers import variable_to_logically_partitioned +# ============================================================================== +# Shared Sharding Logic (Heuristics) +# ============================================================================== + +def _infer_partition_names(value, repeats=1, base_names=None): + """Helper to infer logical axis names (PartitionSpec) for pipeline weights.""" + ndim = value.ndim + + if repeats > 1: + # Structure: [Repeats, Stage, ...] + if base_names is not None: + new_names = (None, 'stage') + base_names + else: + # Heuristics for [Repeats, Stage, ...] + if ndim == 4: # [Repeats, Stage, Embed, MLP] + new_names = (None, 'stage', 'fsdp', 'tensor') + elif ndim == 3: # [Repeats, Stage, Bias] + new_names = (None, 'stage', 'fsdp') + else: + new_names = (None, 'stage') + (None,) * (ndim - 2) + else: + # Structure: [Stage, ...] + if base_names is not None: + new_names = ('stage',) + base_names + else: + # Heuristics for [Stage, ...] + if ndim == 3: # [Stage, Embed, MLP] + new_names = ('stage', 'fsdp', 'tensor') + elif ndim == 2: # [Stage, Bias] + new_names = ('stage', 'fsdp') + else: + new_names = ('stage',) + (None,) * (ndim - 1) + + return new_names + +# ============================================================================== +# Pipeline Class +# ============================================================================== class Pipeline(nnx.Module): - """ - NNX Implementation of the MaxText Pipeline. - Wraps a Flax Linen Module and executes it using Pipeline Parallelism. - """ + """NNX Implementation of the MaxText Pipeline.""" def __init__( self, layer: nn.Module, - config: Any, + config: Config, mesh: Mesh, rngs: nnx.Rngs, remat_policy: Any = None @@ -46,7 +84,7 @@ def __init__( self.linen_module = layer self.remat_policy = remat_policy - # --- 1. Dimensions --- + # --- Dimensions --- self.num_stages = self.config.ici_pipeline_parallelism * self.config.dcn_pipeline_parallelism self.forwarding_delay = 2 if self.config.pipeline_delay_activation_forwarding else 1 self.total_microbatches = self.config.num_pipeline_microbatches @@ -54,7 +92,7 @@ def __init__( self.microbatches_per_stage = self.total_microbatches // self.num_stages self.use_circ_storage = self._need_circ_storage() - # --- Axis Naming --- + # Axis Naming if hasattr(self.config, 'expert_shard_attention_option') and self.config.expert_shard_attention_option == EP_AS_CONTEXT: self.batch_axis_name = "activation_batch_no_exp" self.seq_len_axis_name = "activation_length" @@ -62,88 +100,93 @@ def __init__( self.batch_axis_name = "activation_batch" self.seq_len_axis_name = "activation_length_no_exp" - self.input_shape = ( - self.microbatch_size, - self.config.max_target_length, - self.config.emb_dim - ) + self.input_shape = (self.microbatch_size, self.config.max_target_length, self.config.emb_dim) - # --- 2. Initialization --- + # --- Eager Initialization with JIT + Out Shardings --- init_rng_key = rngs.params() repeats = self.config.num_pipeline_repeats - - def init_single_stage(key): - # [FIX] Use bfloat16 for dummy inputs to save memory - dummy_in = jnp.zeros(self.input_shape, dtype=jnp.bfloat16) - - dummy_pos_shape = (self.input_shape[0], self.input_shape[1]) - dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) - dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) - - return self.linen_module.init( - {'params': key}, - dummy_in, - dummy_segments, - dummy_positions, - True, - MODEL_MODE_TRAIN - ) - - # Step A: Init ONCE to get Metadata - ref_variables = init_single_stage(init_rng_key) - - # Step B: Init WITH VMAP to get Data - if repeats > 1: - repeat_keys = jax.random.split(init_rng_key, repeats) - stage_rng_keys = jax.vmap(lambda k: jax.random.split(k, self.num_stages))(repeat_keys) - raw_variables = jax.vmap(jax.vmap(init_single_stage))(stage_rng_keys) - else: - stage_rng_keys = jax.random.split(init_rng_key, self.num_stages) - raw_variables = jax.vmap(init_single_stage)(stage_rng_keys) - - print(f"raw_variables : {raw_variables}") - # --- 3. Register Parameters --- - def _create_param_with_metadata(raw_node, ref_node): - if hasattr(raw_node, 'items'): - return nnx.Dict({k: _create_param_with_metadata(v, ref_node[k]) for k, v in raw_node.items()}) - elif isinstance(raw_node, (list, tuple)): - return nnx.List([_create_param_with_metadata(v, ref_node[i]) for i, v in enumerate(raw_node)]) + + # 1. Define the Generator (Pure JAX, no side effects) + def generate_vars_fn(key): + def init_single_stage(k): + # Create dummy inputs locally to ensure shape inference works inside JIT + dummy_in = jnp.zeros(self.input_shape, dtype=jnp.float32) + dummy_pos_shape = (self.input_shape[0], self.input_shape[1]) + dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) + dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) + + return self.linen_module.init( + {'params': k}, + dummy_in, + dummy_segments, + dummy_positions, + True, + MODEL_MODE_TRAIN + ) + + if repeats > 1: + repeat_keys = jax.random.split(key, repeats) + # Outer vmap over repeats, Inner vmap over stages + stage_rng_keys = jax.vmap(lambda k: jax.random.split(k, self.num_stages))(repeat_keys) + return jax.vmap(jax.vmap(init_single_stage))(stage_rng_keys) else: - # Unwrap LogicallyPartitioned - if isinstance(raw_node, nn.LogicallyPartitioned): - actual_value = raw_node.value - else: - actual_value = raw_node - - # Cast float32 -> bfloat16 to save memory - if actual_value.dtype == jnp.float32: - actual_value = actual_value.astype(jnp.bfloat16) - - sharding_kwargs = {} - if hasattr(ref_node, 'names'): - original_names = ref_node.names - if repeats > 1: - new_names = (None, 'stage') + original_names - else: - new_names = ('stage',) + original_names - sharding_kwargs['sharding_names'] = new_names - - return nnx.Param(actual_value, **sharding_kwargs) + stage_rng_keys = jax.random.split(key, self.num_stages) + return jax.vmap(init_single_stage)(stage_rng_keys) + + # 2. Abstract Evaluation (Zero Memory) + # Determine the shape and structure of the variables without allocating data + abstract_variables = jax.eval_shape(generate_vars_fn, init_rng_key) + + # 3. Compute Target Sharding (Heuristics) + # We assume the output structure matches what we expect for partitioning + def get_target_sharding(abstract_leaf): + # abstract_leaf is a ShapeDtypeStruct. We check its ndim. + logical_names = _infer_partition_names(abstract_leaf, repeats=repeats) - # Depending on how the inner module is wrapped (Scan vs Sequential), 'params' might be the root or nested. - if 'params' in raw_variables: - params_dict = raw_variables['params'] - ref_dict = ref_variables['params'] - else: - # Fallback: Assume the whole dict is params (common in some scan wrappings) - params_dict = raw_variables - ref_dict = ref_variables + if logical_names: + mesh_axes = nn.logical_to_mesh_axes(logical_names, self.config.logical_axis_rules) + return NamedSharding(self.mesh, PartitionSpec(*mesh_axes)) + + # Default fallback: Replicated + return NamedSharding(self.mesh, PartitionSpec()) + + sharding_tree = jax.tree.map(get_target_sharding, abstract_variables) + + # 4. JIT Compile with Out Shardings (Direct Allocation) + # This tells XLA to compile a kernel that writes outputs directly to their + # final sharded destination, bypassing the single-device OOM bottleneck. + sharded_variables = jax.jit(generate_vars_fn, out_shardings=sharding_tree)(init_rng_key) - self.stage_params = _create_param_with_metadata(params_dict, ref_dict) + # --- Register with NNX --- + self.layers = self._to_nnx_structure(sharded_variables) + + # Handle batch_stats if present + if 'batch_stats' in sharded_variables: + pass # ========================================================================== # Helper Methods # ========================================================================== + def _to_nnx_structure(self, node): + if isinstance(node, (dict, nn.FrozenDict)): + return nnx.Dict({k: self._to_nnx_structure(v) for k, v in node.items()}) + elif isinstance(node, (list, tuple)): + return nnx.List([self._to_nnx_structure(v) for v in node]) + else: + return nnx.Param(node) + + def _to_pure_dict(self, node): + if hasattr(node, 'items'): return {k: self._to_pure_dict(v) for k, v in node.items()} + elif isinstance(node, (list, tuple)): return [self._to_pure_dict(v) for v in node] + elif hasattr(node, 'value'): return node.value + return node + + def _with_logical_constraint(self, x, axis_names): + if axis_names is None: return x + mesh_axes = nn.logical_to_mesh_axes(axis_names, self.config.logical_axis_rules) + sharding = NamedSharding(self.mesh, PartitionSpec(*mesh_axes)) + return jax.lax.with_sharding_constraint(x, sharding) + def _need_circ_storage(self): return (self.config.num_pipeline_repeats > 1 and self.config.num_pipeline_microbatches > self.num_stages * self.forwarding_delay) @@ -151,89 +194,46 @@ def _need_circ_storage(self): def iterations_to_complete_first_microbatch_one_repeat(self): return self.forwarding_delay * (self.num_stages - 1) + # ========================================================================== + # Buffer Management + # ========================================================================== + def init_states(self, inputs): - """Initialize pipeline buffers.""" - # Shift Buffer shift = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) - shift = nn.with_logical_constraint( - shift, - ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - rules=self.config.logical_axis_rules, - mesh=self.mesh, - ) - - # Prev Outputs + shift = self._with_logical_constraint(shift, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed")) + if self.config.pipeline_delay_activation_forwarding: prev_outputs = jnp.zeros((self.num_stages,) + inputs.shape[1:], dtype=inputs.dtype) - prev_outputs = nn.with_logical_constraint( - prev_outputs, - ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - rules=self.config.logical_axis_rules, - mesh=self.mesh, - ) + prev_outputs = self._with_logical_constraint(prev_outputs, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed")) else: prev_outputs = None - - # State IO - state_io = jnp.reshape( - inputs, - (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:] - ) - state_io = nn.with_logical_constraint( - state_io, - ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - rules=self.config.logical_axis_rules, - mesh=self.mesh, - ) - - # Circular Storage + + state_io = jnp.reshape(inputs, (self.num_stages, self.microbatches_per_stage) + inputs.shape[1:]) + state_io = self._with_logical_constraint(state_io, ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed")) + if self.use_circ_storage: circ_storage = jnp.zeros((self.num_stages,) + inputs.shape, dtype=inputs.dtype) circ_storage_mover = shift else: circ_storage = None circ_storage_mover = None - - return { - "state_io": state_io, - "shift": shift, - "circ_storage": circ_storage, - "circ_storage_mover": circ_storage_mover, - "loop_iteration": jnp.array(0, dtype=jnp.int32), - "prev_outputs": prev_outputs, - "rng_stream": jax.random.PRNGKey(0) - } + + return {"state_io": state_io, "shift": shift, "circ_storage": circ_storage, "circ_storage_mover": circ_storage_mover, "loop_iteration": jnp.array(0, dtype=jnp.int32), "prev_outputs": prev_outputs, "rng_stream": jax.random.PRNGKey(0)} def get_iteration_inputs(self, loop_iteration, state_io, circ_storage, shift): state_io_batch_idx = loop_iteration % self.microbatches_per_stage state_io_slice = state_io[:, state_io_batch_idx] - + if self.use_circ_storage: circ_storage_batch_idx = loop_iteration % self.config.num_pipeline_microbatches circular_stage_in = circ_storage[:, circ_storage_batch_idx] else: circular_stage_in = shift - - first_stage_in = jnp.where( - loop_iteration < self.config.num_pipeline_microbatches, - state_io_slice, - circular_stage_in - ) - - def select_state_or_input(first_stage_in, shift): - return jnp.where( - jax.lax.broadcasted_iota("int32", shift.shape, 0) == 0, - first_stage_in, - shift - ) - - stages_in = select_state_or_input(first_stage_in, shift) - stages_in = nn.with_logical_constraint( - stages_in, - ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed"), - rules=self.config.logical_axis_rules, - mesh=self.mesh, - ) + + first_stage_in = jnp.where(loop_iteration < self.config.num_pipeline_microbatches, state_io_slice, circular_stage_in) + + stages_in = jnp.where(jax.lax.broadcasted_iota("int32", shift.shape, 0) == 0, first_stage_in, shift) + stages_in = self._with_logical_constraint(stages_in, ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed")) return stages_in def get_new_loop_state(self, output, loop_state): @@ -242,122 +242,53 @@ def get_new_loop_state(self, output, loop_state): old_circ_storage_mover = loop_state["circ_storage_mover"] loop_iteration = loop_state["loop_iteration"] old_prev_outputs = loop_state["prev_outputs"] - + def _rotate_right(arr): last = jax.lax.slice_in_dim(arr, self.num_stages - 1, self.num_stages, axis=0) except_last = jax.lax.slice_in_dim(arr, 0, self.num_stages - 1, axis=0) return jnp.concatenate([last, except_last], axis=0) - def _shift_right(arr): padding = [[1, 0]] + [[0, 0]] * (arr.ndim - 1) return jax.lax.slice(jnp.pad(arr, padding), [0] * arr.ndim, arr.shape) - def _update_shift(output_in): if self.config.num_pipeline_repeats == 1 or self.use_circ_storage: return _shift_right(output_in) return _rotate_right(output_in) - + if self.config.pipeline_delay_activation_forwarding: new_shift = _update_shift(old_prev_outputs) new_prev_outputs = output else: new_shift = _update_shift(output) new_prev_outputs = None - + if self.use_circ_storage: - def _rotate_right_and_update(circ_storage_mover_in, circ_storage_in): - rotated = _rotate_right(circ_storage_mover_in) - rotated = jnp.expand_dims(rotated, 1) - offset = ( - loop_iteration - - self.iterations_to_complete_first_microbatch_one_repeat() - - 1 - ) % self.config.num_pipeline_microbatches - return jax.lax.dynamic_update_slice_in_dim( - circ_storage_in, rotated, offset, axis=1 - ) - - new_circ_storage = _rotate_right_and_update(old_circ_storage_mover, old_circ_storage) + rotated = _rotate_right(old_circ_storage_mover) + rotated = jnp.expand_dims(rotated, 1) + offset = (loop_iteration - self.iterations_to_complete_first_microbatch_one_repeat() - 1) % self.config.num_pipeline_microbatches + new_circ_storage = jax.lax.dynamic_update_slice_in_dim(old_circ_storage, rotated, offset, axis=1) new_circ_storage_mover = output else: new_circ_storage = None new_circ_storage_mover = None - + stream_buf_idx = loop_iteration % self.microbatches_per_stage stream_slice = old_state_io[:, stream_buf_idx] - - def _update_state_io(state_in, stream_slice, output): - padding = [[0, 1]] + [[0, 0]] * (stream_slice.ndim - 1) - stream_slice = jax.lax.slice_in_dim( - jnp.pad(stream_slice, padding), 1, stream_slice.shape[0] + 1, axis=0 - ) - stream_slice = jnp.where( - jax.lax.broadcasted_iota("int32", stream_slice.shape, 0) == self.num_stages - 1, - output, - stream_slice, - ) - stream_slice = jnp.expand_dims(stream_slice, 1) - return jax.lax.dynamic_update_slice_in_dim( - state_in, stream_slice, stream_buf_idx, axis=1 - ) - - new_state = _update_state_io(old_state_io, stream_slice, output) - - return { - "state_io": new_state, - "shift": new_shift, - "circ_storage": new_circ_storage, - "circ_storage_mover": new_circ_storage_mover, - "loop_iteration": loop_iteration + 1, - "prev_outputs": new_prev_outputs, - } + padding = [[0, 1]] + [[0, 0]] * (stream_slice.ndim - 1) + stream_slice = jax.lax.slice_in_dim(jnp.pad(stream_slice, padding), 1, stream_slice.shape[0] + 1, axis=0) + stream_slice = jnp.where(jax.lax.broadcasted_iota("int32", stream_slice.shape, 0) == self.num_stages - 1, output, stream_slice) + stream_slice = jnp.expand_dims(stream_slice, 1) + new_state = jax.lax.dynamic_update_slice_in_dim(old_state_io, stream_slice, stream_buf_idx, axis=1) + + return {"state_io": new_state, "shift": new_shift, "circ_storage": new_circ_storage, "circ_storage_mover": new_circ_storage_mover, "loop_iteration": loop_iteration + 1, "prev_outputs": new_prev_outputs} # ========================================================================== - # Parity Helpers + # Logic for VMAP and Weight Gathering # ========================================================================== - def _all_gather_over_fsdp(self, params, partition_spec): - def _remove_fsdp_from_spec(spec): - if isinstance(spec, PartitionSpec): - new_spec = [] - for axis in spec: - if isinstance(axis, str) and axis in ("fsdp", "fsdp_transpose"): - new_spec.append(None) - elif isinstance(axis, (list, tuple)): - new_spec.append(tuple(a for a in axis if a not in ("fsdp", "fsdp_transpose"))) - else: - new_spec.append(axis) - return PartitionSpec(*new_spec) - return spec - - def _remove_fsdp_sharding(sharding_tree): - return jax.tree.map( - lambda x: NamedSharding(self.mesh, _remove_fsdp_from_spec(x.spec)) - if isinstance(x, NamedSharding) else x, - sharding_tree - ) - - physical = nn.logical_to_mesh_sharding(partition_spec, mesh=self.mesh, rules=self.config.logical_axis_rules) - physical_no_fsdp = _remove_fsdp_sharding(physical) - return jax.lax.with_sharding_constraint(params, physical_no_fsdp) - - def _to_pure_dict(self, node): - if hasattr(node, 'items'): - return {k: self._to_pure_dict(v) for k, v in node.items()} - elif isinstance(node, (list, tuple)): - return [self._to_pure_dict(v) for v in node] - elif hasattr(node, 'value'): - return node.value - return node - def get_microbatch_and_repeat_ids(self, loop_iteration): - microbatches_processed = jnp.maximum( - loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), - 0 - ) - microbatch_ids = microbatches_processed % self.total_microbatches - repeat_ids = microbatches_processed // self.total_microbatches - return microbatch_ids.astype(jnp.int32), repeat_ids.astype(jnp.int32) + microbatches_processed = jnp.maximum(loop_iteration - self.forwarding_delay * jnp.arange(self.num_stages), 0) + return microbatches_processed % self.total_microbatches, (microbatches_processed // self.total_microbatches).astype(jnp.int32) def shard_dim_by_stages(self, x, dim: int): dims_mapping = [PartitionSpec.UNCONSTRAINED] * x.ndim @@ -368,259 +299,147 @@ def shard_dim_by_stages(self, x, dim: int): def vmap_parallel_gather(self, weights, repeat_ids, repeat_dim_in_weights, stages_dim_in_weights): def _gather_one(x, repeat_id): return jnp.squeeze(jax.lax.dynamic_slice_in_dim(x, repeat_id, 1, repeat_dim_in_weights), repeat_dim_in_weights) - gathered_weights_stage_dim = 0 repeat_ids = self.shard_dim_by_stages(repeat_ids, 0) weights = self.shard_dim_by_stages(weights, stages_dim_in_weights) - stage_weights = jax.vmap(_gather_one, in_axes=(stages_dim_in_weights, 0), out_axes=gathered_weights_stage_dim)(weights, repeat_ids) - stage_weights = self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim) - return stage_weights + return self.shard_dim_by_stages(stage_weights, gathered_weights_stage_dim) def get_current_stage_weights(self, pipeline_weights, loop_iteration): - if self.config.num_pipeline_repeats <= 1: - return pipeline_weights - + if self.config.num_pipeline_repeats <= 1: return pipeline_weights _, repeat_ids = self.get_microbatch_and_repeat_ids(loop_iteration) - def gather_weights_for_stages_in(w): - return self.vmap_parallel_gather(w, repeat_ids, repeat_dim_in_weights=0, stages_dim_in_weights=1) - - return jax.tree.map(gather_weights_for_stages_in, pipeline_weights) - - def get_pipeline_remat_policy(self): - if self.config.remat_policy == "custom": - return self.remat_policy - - save_input_policy = jax.checkpoint_policies.save_only_these_names( - "iteration_input", "decoder_layer_input" - ) - - if self.remat_policy is not None: - return jax.checkpoint_policies.save_from_both_policies( - self.remat_policy, save_input_policy - ) - return save_input_policy + return jax.tree.map(lambda w: self.vmap_parallel_gather(w, repeat_ids, 0, 1), pipeline_weights) def permute_output_micro_per_stage_dim(self, output): microbatch_0_idx = self.iterations_to_complete_first_microbatch_one_repeat() % self.microbatches_per_stage permutation = (jnp.arange(self.microbatches_per_stage) + microbatch_0_idx) % self.microbatches_per_stage - output = output[:, permutation] - return output + return output[:, permutation] # ========================================================================== - # Main Execution + # Integration Methods # ========================================================================== - def __call__( - self, - inputs: jax.Array, - segment_ids: Optional[jax.Array] = None, - positions: Optional[jax.Array] = None, - deterministic: bool = True, - model_mode: str = MODEL_MODE_TRAIN, - partition_spec: Any = None, - ) -> jax.Array: - - inputs = inputs.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length, - self.config.emb_dim - )) - - if positions is not None: - positions = positions.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length - )) - - if segment_ids is not None: - segment_ids = segment_ids.reshape(( - self.total_microbatches, - self.microbatch_size, - self.config.max_target_length - )) + def get_pipeline_remat_policy(self): + if self.config.remat_policy == "custom": return self.remat_policy + save_input_policy = jax.checkpoint_policies.save_only_these_names("iteration_input", "decoder_layer_input") + if self.remat_policy is not None: return jax.checkpoint_policies.save_from_both_policies(self.remat_policy, save_input_policy) + return save_input_policy + + def get_weight_sharding(self, *init_args): + variables = self.layers + repeats = self.config.num_pipeline_repeats + def _infer_partition_spec(node): + if hasattr(node, 'value'): + names = _infer_partition_names(node.value, repeats=repeats, base_names=None) + return nn.LogicallyPartitioned(node.value, names).get_partition_spec() + return None + specs = jax.tree.map(_infer_partition_spec, variables, is_leaf=lambda x: hasattr(x, 'value')) + return {'params': {'layers': specs['params']}} + + def _all_gather_over_fsdp(self, params, partition_spec): + def _remove_fsdp_from_spec(spec): + if isinstance(spec, PartitionSpec): + new_spec = [] + for axis in spec: + if isinstance(axis, str) and axis in ("fsdp", "fsdp_transpose"): new_spec.append(None) + elif isinstance(axis, (list, tuple)): new_spec.append(tuple(a for a in axis if a not in ("fsdp", "fsdp_transpose"))) + else: new_spec.append(axis) + return PartitionSpec(*new_spec) + return spec + def _remove_fsdp_sharding(sharding_tree): + return jax.tree.map(lambda x: NamedSharding(self.mesh, _remove_fsdp_from_spec(x.spec)) if isinstance(x, NamedSharding) else x, sharding_tree) + physical = nn.logical_to_mesh_sharding(partition_spec, mesh=self.mesh, rules=self.config.logical_axis_rules) + physical_no_fsdp = _remove_fsdp_sharding(physical) + return jax.lax.with_sharding_constraint(params, physical_no_fsdp) + + # ========================================================================== + # Main Call (__call__) + # ========================================================================== + def __call__(self, inputs: jax.Array, segment_ids: Optional[jax.Array] = None, positions: Optional[jax.Array] = None, deterministic: bool = True, model_mode: str = MODEL_MODE_TRAIN, partition_spec: Any = None) -> jax.Array: + inputs = inputs.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length, self.config.emb_dim)) + if positions is not None: positions = positions.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length)) + if segment_ids is not None: segment_ids = segment_ids.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length)) + loop_state = self.init_states(inputs) + layer_variables = self._to_pure_dict(self.layers) - param_values = self._to_pure_dict(self.stage_params) - - # Cast weights to compute_dtype (bfloat16) to save memory - # MaxText defaults to 'bfloat16'. Using float32 for weights fills HBM instantly. compute_dtype = getattr(self.config, 'compute_dtype', jnp.bfloat16) if isinstance(compute_dtype, str): - # Map string config (e.g., 'bfloat16') to JAX dtype - dtype_map = {'bfloat16': jnp.bfloat16, 'float32': jnp.float32, 'float16': jnp.float16} - compute_dtype = dtype_map.get(compute_dtype, jnp.bfloat16) + compute_dtype = {'bfloat16': jnp.bfloat16, 'float32': jnp.float32, 'float16': jnp.float16}.get(compute_dtype, jnp.bfloat16) + layer_variables = jax.tree.map(lambda x: x.astype(compute_dtype), layer_variables) - # Helper to cast leaves - param_values = jax.tree.map(lambda x: x.astype(compute_dtype), param_values) - if self.config.pipeline_fsdp_ag_once and partition_spec is not None: - try: - param_values = self._all_gather_over_fsdp(param_values, partition_spec) - except (ValueError, TypeError, KeyError): - pass + try: + if "params" in partition_spec and "layers" in partition_spec["params"]: + params_only = layer_variables['params'] + params_spec = partition_spec['params']['layers']['params'] + layer_variables['params'] = self._all_gather_over_fsdp(params_only, params_spec) + except (KeyError, TypeError): pass def scan_body(carry, _): iteration = carry['loop_iteration'] current_rng = carry['rng_stream'] - step_rng, next_rng = jax.random.split(current_rng) stage_rngs = jax.random.split(step_rng, self.num_stages) - - stages_inputs = self.get_iteration_inputs( - iteration, - carry['state_io'], - carry['circ_storage'], - carry['shift'] - ) - # MARKER for Remat Policy + stages_inputs = self.get_iteration_inputs(iteration, carry['state_io'], carry['circ_storage'], carry['shift']) stages_inputs = jax.ad_checkpoint.checkpoint_name(stages_inputs, "iteration_input") - mb_ids, _ = self.get_microbatch_and_repeat_ids(iteration) stages_positions = jnp.take(positions, mb_ids, axis=0) if positions is not None else None stages_segment_ids = jnp.take(segment_ids, mb_ids, axis=0) if segment_ids is not None else None + current_vars = self.get_current_stage_weights(layer_variables, iteration) - current_params = self.get_current_stage_weights(param_values, iteration) - - # 1. DEFINE EXECUTION LOGIC (Pure) - def execution_logic(params, inputs, rngs, pos, seg): - def stage_fn(p, x, r, pos, seg): - variables = {'params': p} + def execution_logic(vars_in, inputs, rngs, pos, seg): + def stage_fn(v, x, r, po, se): rngs_dict = {'dropout': r} if not deterministic else {} - return self.linen_module.apply( - variables, - x, - seg, pos, - deterministic, - model_mode, - rngs=rngs_dict, - ) - + mutables = ['aux_loss', 'intermediates'] + return self.linen_module.apply(v, x, se, po, deterministic, model_mode, rngs=rngs_dict, mutable=mutables) vmap_axes = [0, 0, 0] - vmap_args = [params, inputs, rngs] - vmap_axes.append(0 if pos is not None else None) - vmap_args.append(pos) - vmap_axes.append(0 if seg is not None else None) - vmap_args.append(seg) - + vmap_args = [vars_in, inputs, rngs] + vmap_axes.append(0 if pos is not None else None); vmap_args.append(pos) + vmap_axes.append(0 if seg is not None else None); vmap_args.append(seg) return jax.vmap(stage_fn, in_axes=tuple(vmap_axes))(*vmap_args) - # 2. APPLY CHECKPOINT TO VMAP (Logic matches Linen run_iteration_scannable) if self.config.set_remat_policy_on_pipeline_iterations: policy = self.get_pipeline_remat_policy() - prevent_cse = not self.config.scan_pipeline_iterations - - execution_logic = jax.checkpoint( - execution_logic, - # policy=policy, - prevent_cse=False - ) + execution_logic = jax.checkpoint(execution_logic, policy=policy, prevent_cse=not self.config.scan_pipeline_iterations) - # 3. EXECUTE - stages_output = execution_logic( - current_params, stages_inputs, stage_rngs, stages_positions, stages_segment_ids - ) - + stages_output, stages_mutables = execution_logic(current_vars, stages_inputs, stage_rngs, stages_positions, stages_segment_ids) if hasattr(self.config, 'scan_layers') and self.config.scan_layers: - if isinstance(stages_output, tuple): - stages_output = stages_output[0] - + if isinstance(stages_output, tuple): stages_output = stages_output[0] + new_loop_state = self.get_new_loop_state(stages_output, carry) new_loop_state['rng_stream'] = next_rng - - return new_loop_state, None + return new_loop_state, stages_mutables bubble_iterations = self.forwarding_delay * (self.num_stages - 1) real_iterations = self.config.num_pipeline_microbatches * self.config.num_pipeline_repeats total_ticks = real_iterations + bubble_iterations - - final_state, _ = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) - + final_state, stacked_mutables = jax.lax.scan(scan_body, loop_state, None, length=total_ticks) output = final_state['state_io'] output = self.permute_output_micro_per_stage_dim(output) - output = output.reshape((self.config.micro_batch_size_to_train_on, - self.config.max_target_length, - self.config.emb_dim)) - + output = output.reshape((self.config.micro_batch_size_to_train_on, self.config.max_target_length, self.config.emb_dim)) return output -# ========================================================================== -# Factory Functions & Metadata -# ========================================================================== +# ============================================================================== +# Factory +# ============================================================================== def add_stage_axis_to_partitioning(variable, repeats=1): - """ - Metadata function. Idempotent check for 'stage' axis. - """ + """Metadata helper for ToLinen.""" partitioned_obj = variable_to_logically_partitioned(variable) - if isinstance(partitioned_obj, nn.LogicallyPartitioned): - names = partitioned_obj.names + base_names = partitioned_obj.names value = partitioned_obj.value else: value = partitioned_obj - if not hasattr(value, 'ndim'): - return value - names = None - - # Check if 'stage' is already present (Idempotency Check) - is_stage_present = False - if names is not None and len(names) > 0: - if repeats > 1: - if len(names) > 1 and names[1] == 'stage': - is_stage_present = True - else: - if names[0] == 'stage': - is_stage_present = True - - if is_stage_present: - return nn.LogicallyPartitioned(value, names) - - # Fallback/Construction - ndim = value.ndim - if repeats > 1: - if names is not None: - new_names = (None, 'stage') + names - else: - if ndim == 4: - new_names = (None, 'stage', 'fsdp', 'tensor') - elif ndim == 3: - new_names = (None, 'stage', 'fsdp') - else: - new_names = (None, 'stage') + (None,) * (max(0, ndim - 2)) - else: - if names is not None: - new_names = ('stage',) + names - else: - if ndim == 3: - new_names = ('stage', 'fsdp', 'tensor') - elif ndim == 2: - new_names = ('stage', 'fsdp') - else: - new_names = ('stage',) + (None,) * (max(0, ndim - 1)) - + if not hasattr(value, 'ndim'): return value + base_names = None + + new_names = _infer_partition_names(value, repeats=repeats, base_names=base_names) return nn.LogicallyPartitioned(value, new_names) - -def create_pipeline( - config: Config, - layer: Callable | type, - mesh: Mesh, - remat_policy: Any = None, -) -> nnx_wrappers.ToLinen: +def create_pipeline(config: Config, layer: Callable | type, mesh: Mesh, remat_policy: Any = None) -> nnx_wrappers.ToLinen: repeats = getattr(config, 'num_pipeline_repeats', 1) metadata_fn = functools.partial(add_stage_axis_to_partitioning, repeats=repeats) - - return nnx_wrappers.to_linen( - Pipeline, - config=config, - mesh=mesh, - layer=layer, - remat_policy=remat_policy, - name="pipeline_module", - abstract_init=False, - metadata_fn=metadata_fn, - ) \ No newline at end of file + return nnx_wrappers.to_linen(Pipeline, config=config, mesh=mesh, layer=layer, remat_policy=remat_policy, name="pipeline_module", abstract_init=False, metadata_fn=metadata_fn) \ No newline at end of file From 0cf2334a86d3394251273c16c131b737a5bbfe81 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Fri, 12 Dec 2025 03:21:02 +0000 Subject: [PATCH 13/14] change name layer --- src/MaxText/layers/pipeline.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/src/MaxText/layers/pipeline.py b/src/MaxText/layers/pipeline.py index d7826f2a40..2720fa53fd 100644 --- a/src/MaxText/layers/pipeline.py +++ b/src/MaxText/layers/pipeline.py @@ -73,7 +73,7 @@ class Pipeline(nnx.Module): def __init__( self, - layer: nn.Module, + layers: nn.Module, config: Config, mesh: Mesh, rngs: nnx.Rngs, @@ -81,7 +81,7 @@ def __init__( ): self.config = config self.mesh = mesh - self.linen_module = layer + self.layers = layers self.remat_policy = remat_policy # --- Dimensions --- @@ -115,7 +115,7 @@ def init_single_stage(k): dummy_positions = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) dummy_segments = jnp.zeros(dummy_pos_shape, dtype=jnp.int32) - return self.linen_module.init( + return self.layers.init( {'params': k}, dummy_in, dummy_segments, @@ -158,7 +158,7 @@ def get_target_sharding(abstract_leaf): sharded_variables = jax.jit(generate_vars_fn, out_shardings=sharding_tree)(init_rng_key) # --- Register with NNX --- - self.layers = self._to_nnx_structure(sharded_variables) + self.stage_params = self._to_nnx_structure(sharded_variables) # Handle batch_stats if present if 'batch_stats' in sharded_variables: @@ -326,7 +326,7 @@ def get_pipeline_remat_policy(self): return save_input_policy def get_weight_sharding(self, *init_args): - variables = self.layers + variables = self.stage_params repeats = self.config.num_pipeline_repeats def _infer_partition_spec(node): if hasattr(node, 'value'): @@ -334,7 +334,7 @@ def _infer_partition_spec(node): return nn.LogicallyPartitioned(node.value, names).get_partition_spec() return None specs = jax.tree.map(_infer_partition_spec, variables, is_leaf=lambda x: hasattr(x, 'value')) - return {'params': {'layers': specs['params']}} + return {'params': {'stage_params': specs['params']}} def _all_gather_over_fsdp(self, params, partition_spec): def _remove_fsdp_from_spec(spec): @@ -362,7 +362,7 @@ def __call__(self, inputs: jax.Array, segment_ids: Optional[jax.Array] = None, p if segment_ids is not None: segment_ids = segment_ids.reshape((self.total_microbatches, self.microbatch_size, self.config.max_target_length)) loop_state = self.init_states(inputs) - layer_variables = self._to_pure_dict(self.layers) + layer_variables = self._to_pure_dict(self.stage_params) compute_dtype = getattr(self.config, 'compute_dtype', jnp.bfloat16) if isinstance(compute_dtype, str): @@ -371,9 +371,9 @@ def __call__(self, inputs: jax.Array, segment_ids: Optional[jax.Array] = None, p if self.config.pipeline_fsdp_ag_once and partition_spec is not None: try: - if "params" in partition_spec and "layers" in partition_spec["params"]: + if "params" in partition_spec and "stage_params" in partition_spec["params"]: params_only = layer_variables['params'] - params_spec = partition_spec['params']['layers']['params'] + params_spec = partition_spec['params']['stage_params']['params'] layer_variables['params'] = self._all_gather_over_fsdp(params_only, params_spec) except (KeyError, TypeError): pass @@ -393,7 +393,7 @@ def execution_logic(vars_in, inputs, rngs, pos, seg): def stage_fn(v, x, r, po, se): rngs_dict = {'dropout': r} if not deterministic else {} mutables = ['aux_loss', 'intermediates'] - return self.linen_module.apply(v, x, se, po, deterministic, model_mode, rngs=rngs_dict, mutable=mutables) + return self.layers.apply(v, x, se, po, deterministic, model_mode, rngs=rngs_dict, mutable=mutables) vmap_axes = [0, 0, 0] vmap_args = [vars_in, inputs, rngs] vmap_axes.append(0 if pos is not None else None); vmap_args.append(pos) @@ -421,6 +421,8 @@ def stage_fn(v, x, r, po, se): output = output.reshape((self.config.micro_batch_size_to_train_on, self.config.max_target_length, self.config.emb_dim)) return output + + # ============================================================================== # Factory # ============================================================================== @@ -439,7 +441,7 @@ def add_stage_axis_to_partitioning(variable, repeats=1): new_names = _infer_partition_names(value, repeats=repeats, base_names=base_names) return nn.LogicallyPartitioned(value, new_names) -def create_pipeline(config: Config, layer: Callable | type, mesh: Mesh, remat_policy: Any = None) -> nnx_wrappers.ToLinen: +def create_pipeline(config: Config, layers: Callable | type, mesh: Mesh, remat_policy: Any = None) -> nnx_wrappers.ToLinen: repeats = getattr(config, 'num_pipeline_repeats', 1) metadata_fn = functools.partial(add_stage_axis_to_partitioning, repeats=repeats) - return nnx_wrappers.to_linen(Pipeline, config=config, mesh=mesh, layer=layer, remat_policy=remat_policy, name="pipeline_module", abstract_init=False, metadata_fn=metadata_fn) \ No newline at end of file + return nnx_wrappers.to_linen(Pipeline, config=config, mesh=mesh, layers=layers, remat_policy=remat_policy, name="pipeline_module", abstract_init=False, metadata_fn=metadata_fn) \ No newline at end of file From 4ca0e628caab516d6567498ca24a80f420847fd8 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Fri, 12 Dec 2025 03:25:02 +0000 Subject: [PATCH 14/14] fix decoders --- src/MaxText/layers/decoders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/layers/decoders.py b/src/MaxText/layers/decoders.py index b2daf65d30..94579a6d34 100644 --- a/src/MaxText/layers/decoders.py +++ b/src/MaxText/layers/decoders.py @@ -402,7 +402,7 @@ def setup(self): pipeline_stage_module = self.get_pipeline_stage_module(self.decoder_layer) remat_policy = self.get_remat_policy() self.pipeline_module = pipeline.create_pipeline( - config=self.config, mesh=self.mesh, layer=pipeline_stage_module, remat_policy=remat_policy, + config=self.config, mesh=self.mesh, layers=pipeline_stage_module, remat_policy=remat_policy, ) def minimal_policy(self, with_context=False):