Write post-training checkpoints in MaxText's on-disk layout

[ecnal-cienet]

Description

Post-training saved through Tunix's checkpoint manager, which stores nnx.state(model) verbatim under a model_params item. MaxText's on-disk layout is the Linen one: weights in params/params, the optimizer in opt_state and step, and NNX-only state such as rngs in nnx_aux. The two never matched, so a checkpoint from an SFT, DPO, RL or distillation run could not be loaded by pre-training, and loaders had started growing branches to read the post-training shape instead.

This adds MaxTextLayoutCheckpointManager, which converts in both directions and writes the same items tree everything else in MaxText reads. It lives under trainers/post_train rather than common/checkpointing because it subclasses Tunix's manager, and common/checkpointing is imported by pre-training and inference, which run without Tunix installed. Older checkpoints are still in the Tunix layout, so maybe_restore falls back to the base class for those, and load_params_from_path learns to read them by restoring into the NNX state itself.

The adapter level

DPO and RL train through TunixMaxTextAdapter, whose base level would otherwise reach the checkpoint. It is stripped from the weights and from the optimizer accumulators that mirror them, and put back on restore.

Optimizer differences that blocked a full-state resume

• train_dpo passed gradient_accumulation_steps unconditionally, so Tunix wrapped the optimizer in optax.MultiSteps even at 1. It now passes None below 2, as train_sft already did.
• Post-training chained optax.clip_by_global_norm into the optimizer, nesting its state a level deeper than pre-training, which clips raw gradients in its train step. add_gradient_clipping applies the same math inside the update and keeps the optimizer's own state tree.

A DPO checkpoint now resumes into pre-training with its weights, its optimizer state and its step counter.

Checkpoint metadata

A checkpoint also records what the run that wrote it was configured as: scan_layers, and the LoRA settings when there are any. verify_and_sync_scan_layers and lora_utils.sync_lora_metadata read it back to fill in a value the run left at its default, or to reject one that contradicts the checkpoint. Tunix passed none, so post-training checkpoints arrived with an empty dict and got neither behaviour. The manager now takes the run's config and records the same metadata pre-training does, through one shared builder (checkpoint_custom_metadata) so the two cannot drift apart. All four trainers pass their config, distillation the student's.

Reading it back was broken too, and not only for post-training. The metadata belongs to the step, so it sits at <step>/ and not at <step>/items/, while every caller passes load_parameters_path, which by convention points at the item. Both readers have been getting an empty dict from pre-training checkpoints as well and silently doing nothing with it. load_checkpoint_metadata now falls back to the parent directory, so both spellings work.

Demo notebooks

The RL and SFT demo notebooks read the checkpoint back to convert it to HuggingFace format, so the paths they hardcode are updated. Both now point at the items directory rather than model_params. RL also drops the actor level: Tunix appended it only when it owned checkpoint_root_directory, which is now None, so the actor trainer writes straight to checkpoint_dir as SFT and DPO do.

Tests

New unit tests, all runnable on CPU:

• tests/post_training/unit/post_train_checkpointing_test.py — the bulk of the coverage. The saved tree is the MaxText layout and not the Tunix one; the adapter level is stripped from both weights and optimizer accumulators; weights and optimizer restore to what was saved; an absent checkpoint reports step 0; config-derived metadata is stamped the way pre-training does it, with a caller-supplied key winning over it; a missing config and a missing optimizer both still save; scanned and unscanned models each round-trip and are shown to differ on disk; a checkpoint in the old Tunix layout still restores through the base-class fallback; and install replaces the manager, forwards the run config, closes the one it replaced, and restores the step.
• tests/post_training/unit/distillation_checkpointing_test.py — the student saves in MaxText layout and restores, its scan_layers is recorded, and learn_to_init mode leaves the optimizer out.
• tests/unit/checkpointing_test.py — load_checkpoint_metadata falls back to the step directory and stops at the path it was given; checkpoint_custom_metadata records scan_layers either way, records LoRA only once there is a rank, and records nothing without a config.
• tests/unit/optimizers_test.py — add_gradient_clipping produces the same updates as the chained form, clips gradients over the threshold, keeps the unclipped optimizer's state tree, forwards extra args to the inner optimizer, and composes with skip-step-on-spikes.
• tests/post_training/unit/train_{sft,dpo,rl,distill}_test.py — each trainer disables Tunix's own manager and installs this one with its config.

Run the post-training suite:

python3 -m pytest tests/post_training/unit/post_train_checkpointing_test.py tests/post_training/unit/distillation_checkpointing_test.py -v

and the pre-training side:

python3 -m pytest tests/unit/checkpointing_test.py tests/unit/checkpointing_nnx_load_test.py tests/unit/optimizers_test.py -v

The cpu-post-training-unit and cpu-unit CI jobs cover both. The notebook job exercises the changed RL and SFT demo paths end to end, including the checkpoint-to-HuggingFace conversion that reads the new layout back.

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.

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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 58.54093% with 233 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../maxtext/checkpoint_conversion/tunix_to_maxtext.py	19.44%	115 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/tokenizer.py	15.62%	27 Missing ⚠️
src/maxtext/trainers/post_train/checkpointing.py	85.00%	15 Missing and 9 partials ⚠️
...c/maxtext/integration/vllm/maxtext_vllm_rollout.py	36.36%	20 Missing and 1 partial ⚠️
src/maxtext/trainers/post_train/rl/train_rl.py	42.85%	16 Missing ⚠️
.../trainers/post_train/distillation/train_distill.py	65.51%	4 Missing and 6 partials ⚠️
src/maxtext/input_pipeline/input_pipeline_utils.py	25.00%	5 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/hf_data_processing.py	16.66%	4 Missing and 1 partial ⚠️
src/maxtext/utils/lora_utils.py	80.00%	2 Missing and 1 partial ⚠️
...axtext/input_pipeline/synthetic_data_processing.py	60.00%	1 Missing and 1 partial ⚠️
... and 2 more

📢 Thoughts on this report? Let us know!