SIGN IN SIGN UP

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