fix(test): use the recipe's attention backend for the HF reload parity check (#3500)
The HF reload phase built its from_pretrained kwargs from scratch and never copied the recipe's attn_implementation, so vanilla HF fell back to its own default. The source-load phase already forwards it, which is why only the reload phase diverged. Attention backends are not bit-identical in bf16. gemma4_2b pins eager while transformers 5.x resolves to sdpa, so the two sides computed different logits from byte-identical weights and the reload reported a checkpoint defect that did not exist. The gap scales with the trained weights, so it moves with world size and passing CI runs are affected too: nemo-ci job 390379862 reports Phase 4 1.168239e-02 against a 2e-2 threshold while Phase 0, which already forwards the backend, is bit-exact. Reproduced locally at 8 GPUs to all 7 digits; at 2 GPUs the same defect crosses the threshold at 3.489082e-02. With the fix both world sizes report 0.0. Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
A
Abhishree Thittenamane committed
8dae92f43e26da2e9723c92ba30f3f069d7aa00b
Parent: 8e53272
Committed by GitHub <noreply@github.com>
on 8/12/2026, 4:54:38 PM