Function _run

tests/test_determinism.py:232–243

()

    x_data = torch.randn(batch, seqlen, model.d_model, device=device, dtype=dtype)

    def _run() -> dict[str, torch.Tensor]:
        _set_deterministic(False)
        model.zero_grad(set_to_none=True)
        x = x_data.clone().requires_grad_(True)
        y = model(x)
        (y.float().square().mean()).backward()
        torch.cuda.synchronize()
        grads = {"input": x.grad.detach().float().clone()}
        for name, p in model.named_parameters():
            if p.grad is not None:
                grads[name] = p.grad.detach().float().clone()
        return grads

    _run()  # warmup
    ref = _run()