feat: add NaN capture and replay for fine-tuning

[gabrielfruet]

What has changed and why?

• New NaNCapture debug tool for fine-tuning (train_* task APIs): when a
  NaN/Inf is detected in parameter gradients (scanned after gradient
  accumulation, before clip_gradients/optimizer.step), it captures
  reproducible state and halts training.
• Capture is a single self-contained file
  out_dir/debug/nan_capture/rank{R}/nan_capture.pt holding the model state
  dict, the TrainModel class path + init kwargs (for reconstruction), the
  step's microbatches, and torch/CUDA RNG state. The standard
  checkpoints/last.ckpt is not touched, so resume_interrupted is
  unaffected.
• Adds load_nan_capture(dir) + NaNCaptureState.replay() for zero-setup
  replay: reconstructs the model, restores RNG, and re-runs the triggering
  forward+backward (mirrors the training loop; stops before the optimizer step)
  to reproduce the NaN deterministically in a notebook/REPL.
• Wired into _commands/train_task.py alongside the existing underflow/overflow
  monitor from feat: integrate HF DebugUnderflowOverflow into fine-tuning #814. Enable with debug_args={"nancapture": {"enabled": True}}.

Replay is debug-only — the training loop carries no replay flag.

Closes TRN-2256.

How has it been tested?

• pytest tests/_debug/test_nan_capture.py → 11/11 pass (config, monitor grad
  scan + capture payload, buffer clone/detach/reset, replay roundtrip +
  reproduction).
• make format → clean (no unintended changes).
• mypy on the 4 changed files → clean.

Did you update CHANGELOG.md?

• Yes

Did you update the documentation?

• Not needed (debugging API; docs deferred, matching feat: integrate HF DebugUnderflowOverflow into fine-tuning #814)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1d330952ed

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[gabrielfruet]

/review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 71b3fb4715

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[liopeer]

LGTM! Adressing the comments is optional.