| 550 | @pytest.mark.parametrize("optim_name", optimizer_names_benchmark, ids=id_formatter("opt")) |
| 551 | @pytest.mark.benchmark |
| 552 | def test_benchmark_blockwise(dim1, dim2, gtype, optim_name, device): |
| 553 | if dim1 == 1 and dim2 == 1: |
| 554 | return |
| 555 | p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 |
| 556 | |
| 557 | bnb_optimizer = str2optimizers[optim_name][1]([p1]) |
| 558 | |
| 559 | g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01 |
| 560 | p1.grad = g |
| 561 | total_steps = 500 |
| 562 | for i in range(total_steps): |
| 563 | if i == total_steps // 5: |
| 564 | # 100 iterations for burn-in |
| 565 | sync_gpu(p1) |
| 566 | t0 = time.time() |
| 567 | |
| 568 | bnb_optimizer.step() |
| 569 | |
| 570 | sync_gpu(p1) |
| 571 | s = time.time() - t0 |
| 572 | print("") |
| 573 | params = (total_steps - total_steps // 5) * dim1 * dim2 |
| 574 | print(optim_name, gtype, s, params, s / params) |
| 575 | # assert s < 3.9 |
| 576 | |
| 577 | |
| 578 | ademamix_state_dict_opts = [ |