(dim1, dim2, gtype, device)
| 298 | @pytest.mark.parametrize("device", get_available_devices()) |
| 299 | @pytest.mark.skipif(not get_available_devices(), reason="No device") |
| 300 | def test_global_config(dim1, dim2, gtype, device): |
| 301 | if dim1 == 1 and dim2 == 1: |
| 302 | return |
| 303 | p1 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 |
| 304 | p2 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 |
| 305 | p3 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 |
| 306 | mask = torch.rand_like(p2) < 0.1 |
| 307 | beta1 = 0.9 |
| 308 | beta2 = 0.999 |
| 309 | lr = 0.001 |
| 310 | eps = 1e-8 |
| 311 | |
| 312 | bnb.optim.GlobalOptimManager.get_instance().initialize() |
| 313 | bnb.optim.GlobalOptimManager.get_instance().override_config(p3, "optim_bits", 8) |
| 314 | |
| 315 | bnb.optim.GlobalOptimManager.get_instance().register_parameters([p1, p2, p3]) |
| 316 | p1 = p1.to(device) |
| 317 | p2 = p2.to(device) |
| 318 | p3 = p3.to(device) |
| 319 | |
| 320 | adam2 = bnb.optim.Adam([p1, p2, p3], lr, (beta1, beta2), eps) |
| 321 | |
| 322 | if gtype == torch.float32: |
| 323 | atol, rtol = 1e-6, 1e-5 |
| 324 | else: |
| 325 | atol, rtol = 1e-4, 1e-3 |
| 326 | |
| 327 | for i in range(50): |
| 328 | g1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 + 0.001 |
| 329 | g2 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 + 0.001 |
| 330 | g3 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 + 0.001 |
| 331 | p1.grad = g1 |
| 332 | p2.grad = g2 |
| 333 | p3.grad = g3 |
| 334 | |
| 335 | adam2.step() |
| 336 | |
| 337 | assert adam2.state[p3]["state1"].dtype == torch.uint8 |
| 338 | assert adam2.state[p3]["state2"].dtype == torch.uint8 |
| 339 | |
| 340 | |
| 341 | @pytest.mark.parametrize("device", get_available_devices()) |
nothing calls this directly
no test coverage detected