fix(rerank) Fix rerank finetune optimizer precision

[shan-nvidia]

Summary

Fix the reranker fine-tune recipe so local torch AdamW training preserves fp32 optimizer-state behavior and can resume without regressing optimizer metadata to bf16.

The investigation found that the recipe path could create Adam moment state from bf16 model parameters, while the reference POC behavior kept optimizer state in fp32. This PR makes that behavior explicit for the local recipe path and adds regression coverage for resume.

What Changed

• Keep local fine-tune parameters and torch AdamW optimizer state fp32 when force_fp32_parameters=true.
• Use native torch.optim.AdamW(fused=True) for local mode instead of requiring Transformer Engine.
• Load the reranker with torch_dtype: float32 and configure an fp32 FSDP MixedPrecisionPolicy.
• Add explicit resume handling so restore loads model state, casts model parts to fp32, then loads optimizer/scheduler state.
• Assert after explicit resume that optimizer checkpoint metadata still records fp32 Adam state.
• Pin Stage 3 eval tokenizer EOS behavior to match training (add_eos_token=false).
• Document the fp32 optimizer-state behavior and train/eval tokenizer invariant in the rerank README.
• Keep checkpoint frequency at the recipe default of 100 steps.

Root Cause

AdamW moment tensors (exp_avg, exp_avg_sq, and step) followed the dtype of the parameters used to construct/load optimizer state. With bf16 parameters, the optimizer checkpoint metadata recorded bf16 moment state, which changed training behavior relative to the fp32-state POC.

Resume had a second failure mode: casting parameters after the normal Automodel restore path was too late, because optimizer state had already been loaded. The new restore hook makes the load order explicit.

Validation

Manual fine-tune/eval runs confirmed that the fixed train curve and eval results match the reference POC behavior.

[rnyak]

@shan-nvidia @oliverholworthy we should set "master_weights: true" here. that will give you proper loss curve. If it is False by default, you might get lower accuracy.

"master_weights: true" tells Transformer Engine’s FusedAdam optimizer to keep a separate higher-precision copy of the model weights, usually FP32, inside the optimizer.

I see you also changed the   target: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam  to target: torch.optim.AdamW . In that case we would not need master_weights: true, but was that really necessary?

Please see the example config: https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/retrieval/cross_encoder/llama3_2_1b.yaml#L115C3-L115C17

[shan-nvidia]

Thanks @rnyak. Automodel uses TE Adam whereas recipe uses torch Adam. I think torch Adam doesn't have a master_weights flag. So the fix is sort of mimicking the master_weights with FP32 behavior.

I tried to switch to TE Adam for recipe but got dep version issue (on rapids lab server) on local mode. Here are some details:
transformer-engine[pytorch]<=2.11.0 (added on this branch) cannot install in your local uv environment because no prebuilt TE wheel exists for torch 2.10. I verified the actual assets on the TE GitHub release pages:

TE version	Available CUDA‑12 wheels (cp312, x86_64)
v2.11	+cu12torch2.8.0+cu129 only
v2.12	+cu12torch2.8.0+cu129 only
v2.13	+cu12torch2.8.0+cu129 only
v2.14 / v2.14.1	+cu12torch2.8.0+cu129 only
v2.15	+cu12torch2.8.0+cu129 only
What's happening at install time:

_execute_with_uv_torch in src/nemotron/kit/uv_local.py runs uv run --with torch ... with UV_TORCH_BACKEND=auto. Your driver advertises CUDA 12.9, so uv resolves torch to 2.10.0+cu129.
TE 2.11's setup.py then guesses the wheel URL transformer_engine_torch-2.11.0+cu12torch2.10.0+cu128cxx11abiTRUE-...whl — that file doesn't exist (correct asset is +cu12torch2.8.0+cu129).
It falls back to source build, which immediately fails on cudnn.h: No such file or directory because cuDNN dev headers aren't on your bare host (they live in the nvcr.io/nvidia/pytorch:25.12-py3 container the recipe declares).
The mismatch isn't really cudnn — it's that TE 2.11 only ships a wheel for torch 2.8.0, while uv's auto backend pulled torch 2.10.0. The cudnn error is just the second symptom after the wheel‑lookup misses.

nemo-automodel @ 897ebedf declares torch>=2.6.0,<=2.10.0 and transformer-engine[pytorch]<=2.11.0, so it has the same shape; nemo-automodel relies on the NGC container to either ship TE prebuilt or to provide cudnn for source builds.

@oliverholworthy let me know if there is a better fix.