(dim1, dim2, gtype, optim_name, device)
| 165 | @pytest.mark.parametrize("dim2", [32, 1024, 4097, 1], ids=id_formatter("dim2")) |
| 166 | @pytest.mark.parametrize("device", get_available_devices(), ids=id_formatter("device")) |
| 167 | def test_optimizer32bit(dim1, dim2, gtype, optim_name, device): |
| 168 | if device == "cpu" and optim_name.startswith("paged_"): |
| 169 | pytest.skip("Paged optimizers are not meaningful on CPU") |
| 170 | |
| 171 | if optim_name.startswith("paged_") and sys.platform == "win32": |
| 172 | pytest.skip("Paged optimizers can have issues on Windows.") |
| 173 | |
| 174 | if gtype == torch.bfloat16 and optim_name in ["momentum", "lars", "rmsprop"]: |
| 175 | pytest.skip() |
| 176 | if dim1 == 1 and dim2 == 1: |
| 177 | return |
| 178 | p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 |
| 179 | p2 = p1.clone() |
| 180 | p1 = p1.float() |
| 181 | |
| 182 | torch_optimizer = str2optimizers[optim_name][0]([p1]) |
| 183 | bnb_optimizer = str2optimizers[optim_name][1]([p2]) |
| 184 | |
| 185 | if gtype == torch.float32: |
| 186 | atol, rtol = 1e-6, 1e-5 |
| 187 | elif gtype == torch.bfloat16: |
| 188 | atol, rtol = 1e-3, 1e-2 |
| 189 | else: |
| 190 | atol, rtol = 1e-4, 1e-3 |
| 191 | |
| 192 | for i in range(k): |
| 193 | g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01 |
| 194 | p1.grad = g.clone().float() |
| 195 | p2.grad = g.clone() |
| 196 | |
| 197 | bnb_optimizer.step() |
| 198 | torch_optimizer.step() |
| 199 | |
| 200 | for name1, name2 in str2statenames[optim_name]: |
| 201 | torch.testing.assert_close( |
| 202 | torch_optimizer.state[p1][name1], |
| 203 | bnb_optimizer.state[p2][name2].to(device), |
| 204 | atol=atol, |
| 205 | rtol=rtol, |
| 206 | ) |
| 207 | |
| 208 | # since Lion can have pretty noisy updates where things lie at the boundary |
| 209 | # allow up to 15 errors for Lion |
| 210 | assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=15) |
| 211 | |
| 212 | if i % (k // 5) == 0 and i > 0: |
| 213 | buf = io.BytesIO() |
| 214 | torch.save(bnb_optimizer.state_dict(), buf) |
| 215 | del bnb_optimizer |
| 216 | bnb_optimizer = None |
| 217 | bnb_optimizer = str2optimizers[optim_name][1]([p2]) |
| 218 | buf.seek(0) |
| 219 | bnb_optimizer.load_state_dict(torch.load(buf)) |
| 220 | # since Lion can have pretty noisy updates where things lie at the boundary |
| 221 | # allow up to 10 errors for Lion |
| 222 | assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10) |
| 223 | for name1, name2 in str2statenames[optim_name]: |
| 224 | # since Lion can have pretty noisy updates where things lie at the boundary |
nothing calls this directly
no test coverage detected