Return True only if CUDA is present and has free memory with headroom.
(num_bytes: int)
| 301 | gen = torch.Generator().manual_seed(8) |
| 302 | dy = torch.randn(out.shape, generator=gen, dtype=out.dtype) |
| 303 | out.backward(dy) |
| 304 | |
| 305 | assert torch.isfinite(hidden.grad).all() and torch.isfinite(weight.grad).all() |
| 306 | assert hidden.grad.shape == hidden.shape and weight.grad.shape == weight.shape |
| 307 | exp_h = dy @ weight.detach() # [.., V] @ [V, K] -> [.., K] |
| 308 | exp_w = dy.reshape(-1, _VOCAB).t() @ hidden.detach().reshape(-1, _HIDDEN) # [V, K] |
| 309 | torch.testing.assert_close(hidden.grad, exp_h, rtol=1e-5, atol=1e-5) |
| 310 | torch.testing.assert_close(weight.grad, exp_w, rtol=1e-5, atol=1e-5) |
| 311 | |
| 312 | |
| 313 | # Registry dispatch -- "lm_head" resolves to NativeLMHeadOp. |
| 314 | def test_registry_dispatches_native_lm_head_op(): |
nothing calls this directly
no outgoing calls
no test coverage detected