feat: migrate pipeline to nnx by mesakhcienet · Pull Request #2885 · AI-Hypercomputer/maxtext

mesakhcienet · 2025-12-24T06:35:02Z

Description

implement nnx-based pipeline.

This PR extends PR#2831

Main changes:

nnx_decoders.py: implementing the missing pipeline logic in nnx_decoders.py.
pipeline.py : add a new class NNXPipeline, which is a nnx-based pipeline class.

Tests

we run the pipeline process with command below:

MODEL_NAME=llama2-7b
python -m MaxText.train src/maxtext/configs/base.yml \
    run_name=pipeline_test_${MODEL_NAME}_nnx \
    base_output_directory=/dev/shm/pipeline_test_nnx \
    model_name=${MODEL_NAME}\
    dataset_type=synthetic \
    steps=15 \
    debug_sharding=true \
    per_device_batch_size=2 \
    max_target_length=32 \
    ici_pipeline_parallelism=2 \
    num_pipeline_microbatches=4 \
    num_layers_per_pipeline_stage=2 \
    enable_checkpointing=false \
    enable_nnx=true \
    pure_nnx_decoder=true \
    scan_layers_per_stage=false \
    async_checkpointing=false > nnx-porting-log/pipeline/custom_${MODEL_NAME}.log 2>&1

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

codecov · 2026-01-19T07:41:42Z

Codecov Report

❌ Patch coverage is 39.00000% with 366 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	11.98%	237 Missing and 20 partials ⚠️
src/maxtext/layers/pipeline.py	64.43%	84 Missing and 17 partials ⚠️
src/maxtext/layers/decoders.py	66.66%	5 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

bvandermoon

@gobbleturk what testing do you recommend for migrating pipeline parallelism to NNX? I'll send over an internal doc @hsuan-lun-chiang, @mesakhcienet, and others put together that shows the tests they have already run

bvandermoon · 2026-03-25T21:41:31Z

@gobbleturk what testing do you recommend for migrating pipeline parallelism to NNX? I'll send over an internal doc @hsuan-lun-chiang, @mesakhcienet, and others put together that shows the tests they have already run

@NuojCheng any thoughts here?

NuojCheng

Some additional train compile test for pipeline NNX migration:

Train compile test 1: https://paste.googleplex.com/5960957017849856
Train compile test 2: https://paste.googleplex.com/5749974483730432
Train compile test 3: https://paste.googleplex.com/5201745681711104
If the train compile tests above can pass without getting OOM + current tests in pipeline_parallelism_test.py can all pass, then I think it is good to go! Please ping me if the PR is ready for review.

NuojCheng · 2026-03-25T22:16:56Z

There are also some linen usage in pipeline_utils.py, e.g.

maxtext/src/maxtext/utils/pipeline_utils.py

Line 330 in 77f5334

return nn.scan(
maxtext/src/maxtext/utils/pipeline_utils.py

Line 307 in 77f5334

return nn.remat(

I don't see them get updated in this PR but I think they probably should be updated?

Another thing is the usage of function in

maxtext/src/maxtext/utils/pipeline_utils.py

Lines 151 to 162 in 77f5334

    
           # TODO(chengnuojin) Remove this function and its usage after pipeline nnx migration 
        
           def remove_logically_partition(weights): 
        
             """Removes LogicallyPartitioned wrapper from weights.""" 
        
             def _remove_logically_partition_leaf(v): 
        
               return getattr(v, "value") if isinstance(v, LogicallyPartitioned) else v 
        
             return jax.tree.map( 
        
                 _remove_logically_partition_leaf, 
        
                 weights, 
        
                 is_leaf=lambda v: isinstance(v, LogicallyPartitioned), 
        
             )

. I suspect NNX migration can help us get rid of using this function since it is mostly dealing with linen wrapper troubles. Take a look on this if you can. Thank you for the hard work!

mesakhcienet · 2026-03-30T07:12:35Z

There are also some linen usage in pipeline_utils.py, e.g.

maxtext/src/maxtext/utils/pipeline_utils.py

Line 330 in 77f5334

return nn.scan(

maxtext/src/maxtext/utils/pipeline_utils.py

Line 307 in 77f5334

return nn.remat(

I don't see them get updated in this PR but I think they probably should be updated?

Another thing is the usage of function in

maxtext/src/maxtext/utils/pipeline_utils.py

Lines 151 to 162 in 77f5334

# TODO(chengnuojin) Remove this function and its usage after pipeline nnx migration

def remove_logically_partition(weights):

"""Removes LogicallyPartitioned wrapper from weights."""

def _remove_logically_partition_leaf(v):

return getattr(v, "value") if isinstance(v, LogicallyPartitioned) else v

return jax.tree.map(

_remove_logically_partition_leaf,

weights,

is_leaf=lambda v: isinstance(v, LogicallyPartitioned),

)

. I suspect NNX migration can help us get rid of using this function since it is mostly dealing with linen wrapper troubles. Take a look on this if you can. Thank you for the hard work!

As far as I know, the current objective is to migrate the Linen pipeline to NNX while preserving the current Linen version. Please advise if any additional progress is required at this time. Thanks!

NuojCheng · 2026-03-30T15:46:01Z

There are also some linen usage in pipeline_utils.py, e.g.

maxtext/src/maxtext/utils/pipeline_utils.py

Line 330 in 77f5334

return nn.scan(

maxtext/src/maxtext/utils/pipeline_utils.py

Line 307 in 77f5334

return nn.remat(

I don't see them get updated in this PR but I think they probably should be updated?
Another thing is the usage of function in

maxtext/src/maxtext/utils/pipeline_utils.py

Lines 151 to 162 in 77f5334

# TODO(chengnuojin) Remove this function and its usage after pipeline nnx migration

def remove_logically_partition(weights):

"""Removes LogicallyPartitioned wrapper from weights."""

def _remove_logically_partition_leaf(v):

return getattr(v, "value") if isinstance(v, LogicallyPartitioned) else v

return jax.tree.map(

_remove_logically_partition_leaf,

weights,

is_leaf=lambda v: isinstance(v, LogicallyPartitioned),

)

. I suspect NNX migration can help us get rid of using this function since it is mostly dealing with linen wrapper troubles. Take a look on this if you can. Thank you for the hard work!

As far as I know, the current objective is to migrate the Linen pipeline to NNX while preserving the current Linen version. Please advise if any additional progress is required at this time. Thanks!

Shouldn't we have a nnx version of functions in pipeline_utils.py as well?

bvandermoon · 2026-03-30T22:47:11Z

There are also some linen usage in pipeline_utils.py, e.g.

maxtext/src/maxtext/utils/pipeline_utils.py

Line 330 in 77f5334

return nn.scan(

maxtext/src/maxtext/utils/pipeline_utils.py

Line 307 in 77f5334

return nn.remat(

I don't see them get updated in this PR but I think they probably should be updated?
Another thing is the usage of function in

maxtext/src/maxtext/utils/pipeline_utils.py

Lines 151 to 162 in 77f5334

# TODO(chengnuojin) Remove this function and its usage after pipeline nnx migration

def remove_logically_partition(weights):

"""Removes LogicallyPartitioned wrapper from weights."""

def _remove_logically_partition_leaf(v):

return getattr(v, "value") if isinstance(v, LogicallyPartitioned) else v

return jax.tree.map(

_remove_logically_partition_leaf,

weights,

is_leaf=lambda v: isinstance(v, LogicallyPartitioned),

)

. I suspect NNX migration can help us get rid of using this function since it is mostly dealing with linen wrapper troubles. Take a look on this if you can. Thank you for the hard work!

As far as I know, the current objective is to migrate the Linen pipeline to NNX while preserving the current Linen version. Please advise if any additional progress is required at this time. Thanks!

Are we able to bridge the NNX version back to Linen at a higher layer? If so, then I think we could get rid of the old Linen code that is no longer used

cgarciae · 2026-03-30T23:42:39Z

src/maxtext/layers/nnx_decoders.py

+      new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+      return new_carry, nnx.state(layer)
+
+    final_carry, scanned_state = jax.lax.scan(layer_fn, inputs, (params, state))


you can alternatively use nnx.scan here which already does the moveaxis for you.

it also avoids the use of split and merge

thank you for your review,
as far as I know, we should avoid the nnx.scan and nnx.remat as mentioned in the internal docs chat here

mesakhcienet · 2026-03-31T06:41:29Z

There are also some linen usage in pipeline_utils.py, e.g.

maxtext/src/maxtext/utils/pipeline_utils.py

Line 330 in 77f5334

return nn.scan(

maxtext/src/maxtext/utils/pipeline_utils.py

Line 307 in 77f5334

return nn.remat(

I don't see them get updated in this PR but I think they probably should be updated?
Another thing is the usage of function in

maxtext/src/maxtext/utils/pipeline_utils.py

Lines 151 to 162 in 77f5334

# TODO(chengnuojin) Remove this function and its usage after pipeline nnx migration

def remove_logically_partition(weights):

"""Removes LogicallyPartitioned wrapper from weights."""

def _remove_logically_partition_leaf(v):

return getattr(v, "value") if isinstance(v, LogicallyPartitioned) else v

return jax.tree.map(

_remove_logically_partition_leaf,

weights,

is_leaf=lambda v: isinstance(v, LogicallyPartitioned),

)

. I suspect NNX migration can help us get rid of using this function since it is mostly dealing with linen wrapper troubles. Take a look on this if you can. Thank you for the hard work!

As far as I know, the current objective is to migrate the Linen pipeline to NNX while preserving the current Linen version. Please advise if any additional progress is required at this time. Thanks!

Are we able to bridge the NNX version back to Linen at a higher layer? If so, then I think we could get rid of the old Linen code that is no longer used

@bvandermoon
Yes we are able to brigde the NNX version. Meanwhile, we are considering two solutions and would like your input on which direction to take:

Option 1: If we use nnx_wrappers.ToLinen to wrap NNX model back to Linen at a higher layer, we can completely remove the old Linen code, but we will need to conduct more tests.

Option 2: Delay the full migration until nnx_decoders.py is stable. In the meantime, we retain the current Linen pipeline and use the enable_nnx and pure_nnx_decoder flags to trigger the NNX version for testing. After the migration is stable, we can remove the flags and linen version of pipeline.

Please let me know which of these two solutions you prefer, thank you.

@NuojCheng
Apologize i missed the first questions you asked.

The NNX pipeline classes (NNXPipeline, NNXCircularPipeline) already handle these internally with JAX-native equivalents:

nn.remat→ jax.checkpoint (pipeline.py L1551, L1870)
nn.scan → jax.lax.scan + nnx.split/nnx.merge (pipeline.py L1556, L1876)
remove_logically_partition()→ inline jax.tree.map unboxing (pipeline.py L1519-1522)

So no NNX versions of those functions are needed — the NNX path bypasses them entirely. maybe you have some suggestions or any part that i am wrong? Please let me know. Thank you.

cgarciae · 2026-03-31T22:05:21Z

src/maxtext/layers/nnx_decoders.py

+      scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
+      scanned_state = nnx.State.merge(scanned_params, scanned_other)
+
+    self.scanned_layers = nnx.merge(graphdef, scanned_state)


Suggested change

self.scanned_layers = nnx.merge(graphdef, scanned_state)

nnx.update(self.scanned_layers, scanned_state)

bvandermoon · 2026-04-02T00:40:29Z

There are also some linen usage in pipeline_utils.py, e.g.

maxtext/src/maxtext/utils/pipeline_utils.py

Line 330 in 77f5334

return nn.scan(

maxtext/src/maxtext/utils/pipeline_utils.py

Line 307 in 77f5334

return nn.remat(

I don't see them get updated in this PR but I think they probably should be updated?
Another thing is the usage of function in

maxtext/src/maxtext/utils/pipeline_utils.py

Lines 151 to 162 in 77f5334

# TODO(chengnuojin) Remove this function and its usage after pipeline nnx migration

def remove_logically_partition(weights):

"""Removes LogicallyPartitioned wrapper from weights."""

def _remove_logically_partition_leaf(v):

return getattr(v, "value") if isinstance(v, LogicallyPartitioned) else v

return jax.tree.map(

_remove_logically_partition_leaf,

weights,

is_leaf=lambda v: isinstance(v, LogicallyPartitioned),

)

. I suspect NNX migration can help us get rid of using this function since it is mostly dealing with linen wrapper troubles. Take a look on this if you can. Thank you for the hard work!

As far as I know, the current objective is to migrate the Linen pipeline to NNX while preserving the current Linen version. Please advise if any additional progress is required at this time. Thanks!

Are we able to bridge the NNX version back to Linen at a higher layer? If so, then I think we could get rid of the old Linen code that is no longer used

@bvandermoon Yes we are able to brigde the NNX version. Meanwhile, we are considering two solutions and would like your input on which direction to take:

Option 1: If we use nnx_wrappers.ToLinen to wrap NNX model back to Linen at a higher layer, we can completely remove the old Linen code, but we will need to conduct more tests.

Option 2: Delay the full migration until nnx_decoders.py is stable. In the meantime, we retain the current Linen pipeline and use the enable_nnx and pure_nnx_decoder flags to trigger the NNX version for testing. After the migration is stable, we can remove the flags and linen version of pipeline.

Please let me know which of these two solutions you prefer, thank you.

@NuojCheng Apologize i missed the first questions you asked.

The NNX pipeline classes (NNXPipeline, NNXCircularPipeline) already handle these internally with JAX-native equivalents:

nn.remat→ jax.checkpoint (pipeline.py L1551, L1870)

nn.scan → jax.lax.scan + nnx.split/nnx.merge (pipeline.py L1556, L1876)

remove_logically_partition()→ inline jax.tree.map unboxing (pipeline.py L1519-1522)

So no NNX versions of those functions are needed — the NNX path bypasses them entirely. maybe you have some suggestions or any part that i am wrong? Please let me know. Thank you.

Thank you @mesakhcienet. Let's go with option 1 please. That way we can continue running unit tests along the way, and we don't need to worry about the Linen/NNX versions diverging before the migration is fully done

NuojCheng · 2026-04-02T18:31:37Z

src/maxtext/layers/pipeline.py

    if bsw_pps is not None:

-      @jax.shard_map(mesh=self.mesh, in_specs=((bsw_pps, bsw_pps), P("stage")), out_specs=bsw_pps, check_vma=True)
+      @jax.shard_map(mesh=self.mesh, in_specs=((bsw_pps, bsw_pps), P("stage")), out_specs=bsw_pps, check_vma=False)


please set check_vma=True in pipeline code

This reverts commit 6d6bea2.

refactor: replace Linen pipeline classes with to_linen_class wrappers Remove PipelineBaseLinen, Pipeline, CircularPipeline (~740 lines). Add Pipeline = to_linen_class(NNXPipeline) and CircularPipeline = to_linen_class(NNXCircularPipeline). Update create_pipeline to accept stage_factory instead of layers. test: update pipeline test to use NNX stage factory refactor: switch Linen decoder pipeline path to NNX stage factory Replace get_pipeline_stage_module with _get_nnx_decoder_block_classes and _build_nnx_pipeline_stage. Delete SequentialBlockDecoderLayers. Pipeline setup now uses stage_factory callable (rngs -> NNX module).

…teration run_one_iteration now: 1. Fetches params from BSW (params-only, matching shard_map specs) 2. Gathers metrics/mutables directly for current repeat 3. Merges into full state for forward pass 4. Scatter-updates only non-params back (params static in scan) Fixes ValueError: pytree structure error in shard_map where out_specs had None (leaf) at RNG paths but BSW had RngCount (pytree node). fix(pipeline): pass params-only to weight_prefetching in scan_body BSW prefetching only needs parameters. Non-param state (metrics, mutables) is now passed separately to run_one_iteration for direct gathering. fix(pipeline): clean up run_one_iteration consistency - Filter to non-params in num_pipeline_repeats == 1 path for consistency - Remove redundant get_microbatch_and_repeat_ids call

…opout state

…gnment

mesakhcienet changed the title ~~core: migrate pipeline to nnx~~ feat: migrate pipeline to nnx Dec 24, 2025

mesakhcienet force-pushed the test/pipeline-scan-nnx branch 8 times, most recently from 6875da8 to f34b1a3 Compare January 15, 2026 23:43

mesakhcienet force-pushed the test/pipeline-scan-nnx branch 4 times, most recently from 12a3907 to 2c16599 Compare January 28, 2026 08:04

mesakhcienet force-pushed the test/pipeline-scan-nnx branch 2 times, most recently from 64dc147 to 9e4518e Compare February 2, 2026 01:58

mesakhcienet force-pushed the test/pipeline-scan-nnx branch from 631a73e to ac97a1d Compare March 2, 2026 08:48

mesakhcienet changed the base branch from main to xibin/nnx_all March 2, 2026 08:48

ecnal-cienet force-pushed the xibin/nnx_all branch 12 times, most recently from 1849f0b to 669dc01 Compare March 3, 2026 19:59

bvandermoon reviewed Mar 25, 2026

View reviewed changes

NuojCheng reviewed Mar 25, 2026

View reviewed changes

cgarciae reviewed Mar 30, 2026

View reviewed changes

cgarciae reviewed Mar 31, 2026

View reviewed changes

NuojCheng reviewed Apr 2, 2026

View reviewed changes

mesakhcienet added 18 commits April 7, 2026 09:42

feat: implement nnx-based pipeline

a673751

fix: update

6d6bea2

Revert "fix: update"

db1ccca

This reverts commit 6d6bea2.

fix: update

29781b6

fix: update

c205f3f

update

9624929

fix: update implementation

c66213e

Update pipeline.py

34b8d52

Update pipeline.py

934cdbb

Update pipeline.py

e96ff0b

fix(pipeline): adopt two-spec pattern in NNXCircularPipeline.__call__

6e724b9

revert(pipeline): go back to vmap-only get_current_weights_from_bsw

aee8747

fix(pipeline): re-enable shard_map in get_current_weights_from_bsw

7334a1b

fix(pipeline): NNXPipeline all_gather_over_fsdp tree mismatch with dr…

6badb5d

…opout state

fix(pipeline): NNXPipeline get_current_repeat_from_stages spec misali…

307edcd

…gnment

fix(pipeline): use spmd_axis_name

b2bf791

	self.scanned_layers = nnx.merge(graphdef, scanned_state)
	nnx.update(self.scanned_layers, scanned_state)

Conversation

mesakhcienet commented Dec 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

codecov bot commented Jan 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

bvandermoon left a comment

Choose a reason for hiding this comment

Uh oh!

bvandermoon commented Mar 25, 2026

Uh oh!

NuojCheng left a comment

Choose a reason for hiding this comment

Uh oh!

NuojCheng commented Mar 25, 2026

Uh oh!

mesakhcienet commented Mar 30, 2026

Uh oh!

NuojCheng commented Mar 30, 2026

Uh oh!

bvandermoon commented Mar 30, 2026

Uh oh!

cgarciae Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

cgarciae Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

mesakhcienet Mar 31, 2026

Choose a reason for hiding this comment

Uh oh!

mesakhcienet commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cgarciae Mar 31, 2026

Choose a reason for hiding this comment

Uh oh!

bvandermoon commented Apr 2, 2026

Uh oh!

NuojCheng Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

mesakhcienet commented Dec 24, 2025 •

edited

Loading

codecov bot commented Jan 19, 2026 •

edited

Loading

mesakhcienet commented Mar 31, 2026 •

edited

Loading