Migrate Pipeline to use NNX module.

[ecnal-cienet]

Description

Migrate Pipelinegit to use NNX module.

Environment

Machine type : v6e-8 TPU VM
Software machine : v2-alpha-tpuv6e
Python Version : 3.12.11

Test Command

Run train command to train llama2-7b for 10 steps:

python3 -m MaxText.train src/MaxText/configs/base.yml \
run_name=pipeline_test_llama2_7b_nnx \
base_output_directory=gs://wanglance-maxtext/pipeline_test \
model_name=llama2-7b \
dataset_type=synthetic \
steps=10 \
per_device_batch_size=2 \
max_target_length=1024 \
ici_pipeline_parallelism=4 \
num_pipeline_microbatches=8 \
num_layers_per_pipeline_stage=2 \
enable_checkpointing=false \
scan_layers_per_stage=false \
async_checkpointing=false

Logs:
Linen, before migration
NNX, after migration
WebDiff

---

Test Suite

Run test command to execute Pipeline test suite:

pytest --cache-clear -s tests/pipeline_parallelism_test.py

Logs:
Linen, before migration
NNX, after migration
WebDiff

---

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.