(dim1, dim2, gtype, optim_name, device)
| 385 | @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1")) |
| 386 | @pytest.mark.parametrize("device", get_available_devices()) |
| 387 | def test_optimizer8bit(dim1, dim2, gtype, optim_name, device): |
| 388 | torch.set_printoptions(precision=6) |
| 389 | |
| 390 | if dim1 == 1 and dim2 == 1: |
| 391 | return |
| 392 | |
| 393 | p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 |
| 394 | p2 = p1.clone() |
| 395 | p1 = p1.float() |
| 396 | blocksize = 256 |
| 397 | |
| 398 | torch_optimizer = str2optimizers[optim_name][0]([p1]) |
| 399 | bnb_optimizer = str2optimizers[optim_name][1]([p2]) |
| 400 | |
| 401 | if gtype == torch.float32: |
| 402 | atol, rtol = 3e-3, 1e-3 |
| 403 | patol, prtol = 1e-5, 1e-3 |
| 404 | elif gtype == torch.bfloat16: |
| 405 | atol, rtol = 3e-3, 1e-3 |
| 406 | patol, prtol = 1e-4, 1e-2 |
| 407 | else: |
| 408 | atol, rtol = 3e-3, 1e-3 |
| 409 | patol, prtol = 1e-5, 1e-3 |
| 410 | |
| 411 | errors = [] |
| 412 | relerrors = [] |
| 413 | |
| 414 | for i in range(50): |
| 415 | g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01 |
| 416 | p1.grad = g.clone().float() |
| 417 | p2.grad = g.clone() |
| 418 | |
| 419 | torch_optimizer.step() |
| 420 | bnb_optimizer.step() |
| 421 | |
| 422 | # since Lion can have pretty noisy updates where things lie at the boundary |
| 423 | # assert_most_approx_close(p1, p2.float(), patol, prtol, max_error_count=0) |
| 424 | |
| 425 | dequant_states = [] |
| 426 | for name1, name2, qmap, max_val in str2statenames[optim_name]: |
| 427 | ## For AdEMAMix, we need to dequantize [p2][name2][0] and [p2][name2][1] |
| 428 | ## separately and then stack them. The qmap is shared, but absmax is also stacked. |
| 429 | if optim_name == "ademamix8bit_blockwise" and name1 == "m1_m2": |
| 430 | m1 = F.dequantize_blockwise( |
| 431 | code=bnb_optimizer.state[p2][qmap], |
| 432 | absmax=bnb_optimizer.state[p2][max_val][0], |
| 433 | A=bnb_optimizer.state[p2][name2][0], |
| 434 | blocksize=blocksize, |
| 435 | ) |
| 436 | m2 = F.dequantize_blockwise( |
| 437 | code=bnb_optimizer.state[p2][qmap], |
| 438 | absmax=bnb_optimizer.state[p2][max_val][1], |
| 439 | A=bnb_optimizer.state[p2][name2][1], |
| 440 | blocksize=blocksize, |
| 441 | ) |
| 442 | |
| 443 | s1 = torch.stack((m1, m2)) |
| 444 | else: |
nothing calls this directly
no test coverage detected