SIGN IN SIGN UP

[OLMo] Fix OOM in logits tests by adding torch.no_grad() (#47986)

* [OLMo] Fix OOM in logits tests by adding torch.no_grad()

test_model_7b_logits and test_model_7b_twin_2t_logits call model()
directly without torch.no_grad(), causing PyTorch to retain intermediate
activations for potential backward passes. On A10G (22 GB), this extra
memory pressure causes OOM when accelerate tries to move CPU-offloaded
layers to GPU during the forward pass. Adding torch.no_grad() resolves
the issue, consistent with how generate() handles inference internally.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Also fix test_model_1b_logits device mismatch and missing no_grad

OLMo-1B fits entirely on GPU with device_map="auto", so the input
tensor must be moved to model.device. Also wrap with torch.no_grad()
and move output to CPU for comparison, consistent with the 7B tests.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
Y
Yih-Dar committed
c7ad72879438c7d83955394312cea9f437e9f43b
Parent: ce5c8f5
Committed by GitHub <noreply@github.com> on 8/15/2026, 1:02:49 PM