Fix CI notebook execution and add SFT training step verification

[SurbhiJainUSC]

Description

This PR hardens our Jupyter Notebook CI pipeline against silent failures, reorganizes our post-training integration tests, and includes several bug fixes for sharding specifications and gradient accumulation in the SFT trainer. It also optimizes the Jupyter notebook testing workflow (run_jupyter_notebooks.yml) by transitioning from sequential execution of all notebooks on a single runner to dynamic matrix parallel execution. It also adds intelligent PR diff detection, ensuring that Pull Requests only run the specific notebooks that were modified.

Why is this change required?

1. Previously, if a code change caused the training loop to crash internally or exit prematurely, the notebook execution could still finish successfully without propagating a failing exit code to the CI runner. A major vulnerability in our CI was recently exposed by PR fix(post_train): Fix Tunix SFT signatures, Qwix LoRA mesh sharding, and scale Qwen/LLaMA parallelism #4866 (the Tunix SFT signatures failure). Because the training loop failed silently, the CI passed, and the regression was merged into main. If the verification steps in this PR had been present in main, we would have immediately detected failures in PR fix(post_train): Fix Tunix SFT signatures, Qwix LoRA mesh sharding, and scale Qwen/LLaMA parallelism #4866. By explicitly asserting that the trainer completes the expected number of steps, we ensure that CI only passes if the training loop is actually functional, eliminating the risk of "false positive" CI passes on broken notebooks.
2. Long CI Duration: Previously, all 7 active example notebooks (src/maxtext/examples/*.ipynb) executed sequentially in a single job on a single TPU v6e-8 runner. Total execution time took 60–90+ minutes.
3. Resource Inefficiency on PRs: A PR that only touched a single notebook (e.g. sft_qwen3_demo.ipynb) was forced to execute the entire suite of 7 notebooks sequentially.
4. Cumulative Flakiness & Contamination: Sequential execution inside a single container accumulated checkpoint artifacts and cache files, leading to potential disk space issues or memory fragmentation.

Key Changes:

Notebook CI Hardening

• Explicit Step Verification: Added a strict post-training check (trainer.train_steps < config.steps) to sft_llama3_demo_tpu.ipynb, sft_qwen3_demo.ipynb, and lora_llama3_demo.ipynb. It throws a RuntimeError if training aborts early.
• Faster CI Execution: Reduced the number of target steps to 5 in the Llama 3, Qwen 3, Gemma 3 multimodal, and Llama 3 LoRA SFT demos. This accelerates the CI pipeline while still proving that the forward/backward passes and optimizer steps execute correctly.
• Notebook Flag Handling: Added absl.flags.mark_as_parsed() to the notebooks to prevent UnparsedFlagAccessErrors during execution.
• Workflow Cleanup: Cleaned up run_jupyter_notebooks.yml by removing the maxtext_installed conditional blocks.
• Lightweight Discovery: Added a discover_notebooks job that runs git diff against origin/${GITHUB_BASE_REF} to identify only the .ipynb files that were added or modified in the PR. On Scheduled Runs / Workflow Edits: Falls back to discovering all active notebooks.
• Parallel Matrix Execution: Converted the run job into a matrix job parameterized by matrix: { notebook: ${{ fromJson(...) }} } with fail-fast: false. Each runner executes exactly one notebook in isolation.
• Clean Container Isolation: Each notebook executes in a fresh, isolated container with dedicated disk and memory space.

Test Reorganization & CI Pipelines

• Extracted SFT LoRA integration tests from tests/integration/lora_e2e_nnx_test.py into a dedicated tests/post_training/integration/lora_e2e_nnx_test.py file.
• Updated ci_pipeline.yml, run_ci_tests.yml, and run_tests_coordinator.yml to explicitly execute cpu-post-training-integration and tpu-post-training-integration flavors, ensuring these tests are not skipped.

SFT & Core Fixes

• SFT Trainer: Updated train_sft.py to properly handle gradient accumulation and removed the legacy inspection of Tunix's grad_norm signature.
• Sharding Specifications: Fixed an issue where partition specs containing local_layers exceeded array dimensions. Added stripping logic in sharding.py, maxtext_utils.py, lora_utils.py, and nnx_scan.py.
• Quantization Cleanup: Added logic to quantizations.py and lora_utils.py to delete qwix_rngs metadata from the graph.
• vLLM Integration: Removed the unused swap_space argument from maxtext_vllm_rollout.py.

Tests

CI tests and verified E2E Airflow tests

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.

[gemini-code-assist]

Code Review

This pull request updates the SFT LLaMA 3 demo notebook for TPU by reducing the training steps from 100 to 5 and adding a verification check to ensure the expected number of training steps successfully executed. The review feedback identifies an issue with significant trailing whitespace in the newly added notebook lines, which violates PEP 8, and provides a suggestion to clean up the code.

[codecov]

Codecov Report

❌ Patch coverage is 52.08333% with 46 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/lora_utils.py	40.00%	15 Missing ⚠️
src/maxtext/trainers/post_train/sft/train_sft.py	33.33%	5 Missing and 3 partials ⚠️
src/maxtext/layers/nnx_scan.py	77.41%	3 Missing and 4 partials ⚠️
src/maxtext/utils/sharding.py	12.50%	6 Missing and 1 partial ⚠️
src/maxtext/utils/maxtext_utils.py	0.00%	4 Missing and 1 partial ⚠️
src/maxtext/layers/quantizations.py	73.33%	3 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!