[NNX] NNX migration prep (6/N): NNX-native DPO

[ecnal-cienet]

Summary

Implements NNX-native DPO. The pure_nnx=True training path no longer raises NotImplementedError on use_dpo runs.

The Linen DPO overlay pattern (model.apply(params=..., reference_params=...)) does not translate to NNX modules, which carry their parameters internally. Instead, the policy and reference models are held as separate nnx.Module instances on TrainStateNNX, and a new dpo_loss_fn_nnx runs both forwards with stop_gradient on the reference logits.

This is the next step in the NNX migration after feat/nnx-correctness-fixes — it closes the remaining hard NotImplementedError on the NNX path and unblocks default-flipping pure_nnx=True for users who run DPO.

What's in this PR

TrainStateNNX (src/maxtext/layers/train_state_nnx.py)

• Optional reference_model: nnx.Module field. apply_gradients continues to update only self.model, leaving self.reference_model bit-identical across steps.

dpo_utils.py (src/maxtext/trainers/post_train/dpo/dpo_utils.py)

• New dpo_loss_fn_nnx(policy_model, config, data, dropout_rng, params, reference_model, is_train=True) mirroring the Linen dpo_loss_fn signature so it slots into gradient_accumulation_loss_and_grad's dispatcher (dropout_rng / params slots are unused for NNX; reference_model is the single extra_dpo_args entry).
• Reference forward wrapped in jax.lax.stop_gradient. With nnx.value_and_grad(..., argnums=0) over the policy, no gradient flows to the reference model's nnx.Param leaves.
• Both dpo_loss_fn (Linen) and dpo_loss_fn_nnx (NNX) now include indexer_loss=0.0 / mtp_loss=0.0 in their aux dicts so the gradient-accumulation aux pytree shape matches the non-DPO loss_fn.
• The Linen _split_dpo_state / _merge_dpo_state overlay pattern has no NNX counterpart by design — NNX holds the reference as a sibling field on TrainStateNNX, and apply_gradients already only touches self.model, so no split/merge is needed.

train.py (src/maxtext/trainers/pre_train/train.py)

• NotImplementedError removed from train_step and eval_step NNX branches; both dispatch to dpo_loss_fn_nnx when use_dpo, with state.reference_model passed as extra_dpo_args[0].
• diff_wrapper picks _loss_fn / extra_dpo_args from the per-path init block so GA and non-GA NNX paths route DPO identically.
• _split_dpo_state checkpoint-save stripping is now Linen-only (NNX saves the whole TrainStateNNX including reference_model; the step-0 reload later overwrites reference_model).

train_utils.py (src/maxtext/utils/train_utils.py)

• NNX init_state_fn materializes a frozen reference model alongside the policy when config.use_dpo. Both are constructed via _create_model_partial() with config.init_weights_seed, so they start identical (standard DPO practice) until the step-0 reload.
• Step-0 checkpoint reload writes step0_state["model"] into state["reference_model"]. Linen path unchanged.

Tests

• New tests/unit/dpo_nnx_test.py (7 tests):
  • TrainStateNNX(reference_model=...) init / hasattr semantics
  • apply_gradients leaves the reference model bit-identical
  • dpo_loss_fn_nnx aux key set
  • Identical policy/reference yields loss = log(2) and reward_accuracy = 0.0 (strict > on equal logratios)
  • dropout_rng / params slots are signature-compat only — passing arbitrary values does not change the result
  • nnx.value_and_grad(..., argnums=0) over the policy yields finite grads on policy params only
• tests/unit/train_nnx_test.py: dropped two stale negative tests:
  • vocab_tiling_raises_not_implemented — vocab tiling on NNX was implemented in feat/nnx-correctness-fixes
  • train_step_dpo_raises_for_nnx — DPO on NNX is implemented here

Linen invariant

Linen DPO is behaviorally unchanged: the only Linen-side change is that dpo_loss_fn now includes indexer_loss=0.0 / mtp_loss=0.0 in aux. Existing callers either ignore those keys or read them with .get(..., 0.0), so this is a no-op for downstream code. The NNX non-DPO path is unchanged: every new code path is gated on config.use_dpo.

Stats

• 4 source files + 2 test files modified, +412 / −41 lines
• All 27 tests pass across train_nnx_test, dpo_nnx_test (new), train_state_nnx_test, train_utils_nnx_test, gradient_accumulation_nnx_test

Followups (not in this PR)

• DPO + NNX with gradient_accumulation_steps > 1 is wired up via the same dispatcher (extra_dpo_args carries reference_model and gradient_accumulation_loss_and_grad's NNX branch already calls grad_func(local_model, config, data, None, None, *extra_dpo_args, is_train=True) against dpo_loss_fn_nnx's signature) but is not exercised by an explicit GA-DPO test. Tracked as a followup.
• Step-0 abstract-state shape mismatch (DPO-shape abstract state used to restore a non-DPO step-0 checkpoint) is mirrored from the Linen path's existing behavior; revisit only if integration testing surfaces issues.

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.