[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