feat: migrate pipeline to nnx

[mesakhcienet]

Description

implement nnx-based pipeline.

This PR extends PR#2831

Main changes:

1. nnx_decoders.py: implementing the missing pipeline logic in nnx_decoders.py.
2. 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]

Codecov Report

❌ Patch coverage is 14.51187% with 324 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	11.24%	270 Missing and 22 partials ⚠️
src/maxtext/layers/decoders.py	37.77%	25 Missing and 3 partials ⚠️
src/maxtext/layers/nnx_wrappers.py	20.00%	4 Missing ⚠️

📢 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]

@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]

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]

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]

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]

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

[mesakhcienet]

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.

[bvandermoon]

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

[mesakhcienet]

@bvandermoon

Update on the pipeline migration approach:

After further investigation, we've adjusted from the original Option 1 (remove Linen entirely) to a hybrid approach. Here's why:

The current branch keeps both Linen (PipelineLinen/CircularPipelineLinen) and NNX (NNXPipeline/NNXCircularPipeline) pipeline classes in pipeline.py, gated by a use_nnx_pipeline config flag (default: False).

Reasoning:

• The NNX pipeline wrapped via ToLinen still needs validation against the full pipeline test suite (circular pipeline, FSDP+AG, multi-repeat, etc.) to confirm numerical equivalence with the original Linen version.
• Keeping the Linen version as a reference allows us to A/B test and verify correctness before removing it.
• The Linen classes are verified line-by-line identical to main branch -- no divergence risk.
• nnx_decoders.py uses create_nnx_pipeline() directly for the pure NNX path.
• decoders.py uses create_pipeline() which dispatches on the flag -- False gives native Linen, True gives NNX-wrapped-in-ToLinen.

Plan to converge to Option 1:
Once the NNX pipeline passes all integration tests, we'll remove the Linen classes and flip the default, effectively completing Option 1.

Let me know if you'd prefer we prioritize removing the Linen code sooner. Thank you.

[bvandermoon]

@bvandermoon

Update on the pipeline migration approach:

After further investigation, we've adjusted from the original Option 1 (remove Linen entirely) to a hybrid approach. Here's why:

The current branch keeps both Linen (PipelineLinen/CircularPipelineLinen) and NNX (NNXPipeline/NNXCircularPipeline) pipeline classes in pipeline.py, gated by a use_nnx_pipeline config flag (default: False).

Reasoning:

• The NNX pipeline wrapped via ToLinen still needs validation against the full pipeline test suite (circular pipeline, FSDP+AG, multi-repeat, etc.) to confirm numerical equivalence with the original Linen version.
• Keeping the Linen version as a reference allows us to A/B test and verify correctness before removing it.
• The Linen classes are verified line-by-line identical to main branch -- no divergence risk.
• nnx_decoders.py uses create_nnx_pipeline() directly for the pure NNX path.
• decoders.py uses create_pipeline() which dispatches on the flag -- False gives native Linen, True gives NNX-wrapped-in-ToLinen.

Plan to converge to Option 1: Once the NNX pipeline passes all integration tests, we'll remove the Linen classes and flip the default, effectively completing Option 1.

Let me know if you'd prefer we prioritize removing the Linen code sooner. Thank you.

Thanks @mesakhcienet. I feel we should push strongly to remove the Linen logic if possible. Even though the NNX and Linen versions are equivalent now, that could change if someone updates the Linen version and misses that the NNX one needs to be updated also. Is testing differences between Linen/NNX the main reason to fork these two? Maybe we could test before/after this commit if needed?