| """Block-wise 8-bit AdamW in pure torch: codebook, quantization, the optimizer, |
| resume, the factory, and the 2-rank gloo/FSDP2 behaviour. |
| |
| Run: .venv/bin/python -m pytest tests/test_adamw8bit.py -q -m "not gpu and not oracle" -p no:cacheprovider |
| |
| Why this optimizer exists (measured, do not re-derive): every shipped full-FT |
| config names `adamw8bit`; fp32 AdamW does not fit a 23GB A10G and |
| `OffloadAdamW` overshoots its pinned-staging cap 4x at 7.25B params / 4 ranks |
| and projects ~105GB host RSS against ~85GB available. See adamw8bit.py. |
| |
| House conventions: each regression is NAMED AFTER THE INCIDENT/FAILURE CLASS it |
| prevents, and each assertion is one a mutant that reintroduces the bug fails |
| (the mutation each test kills is stated in its docstring). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import math |
| import os |
| import socket |
| import sys |
| import traceback |
|
|
| import pytest |
| import torch |
| import torch.multiprocessing as mp |
|
|
| from align.train import quant8 |
| from align.train.adamw8bit import ( |
| BlockwiseAdamW8bit, |
| adamw_update_, |
| stochastic_round_to_bf16, |
| ) |
|
|
| HP = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1) |
| B = 2048 |
|
|
|
|
| def _randn(*shape, seed=0, dtype=torch.float32): |
| g = torch.Generator().manual_seed(seed) |
| return torch.randn(*shape, generator=g, dtype=torch.float32).to(dtype) |
|
|
|
|
| def _param(*shape, seed=0, dtype=torch.float32): |
| return torch.nn.Parameter(_randn(*shape, seed=seed, dtype=dtype)) |
|
|
|
|
| def _set_grads(params, step, seed=7, scale=1.0, dtype=None): |
| g = torch.Generator().manual_seed(seed + step) |
| for p in params: |
| gr = torch.randn(p.shape, generator=g, dtype=torch.float32) * scale |
| p.grad = gr.to(dtype or p.dtype) |
|
|
|
|
| |
| |
| |
|
|
|
|
| @pytest.mark.parametrize("signed", [True, False]) |
| def test_codebook_is_a_sorted_256_entry_map_covering_the_unit_range(signed): |
| """Kills: a miscounted decade (255/257 entries -> uint8 overflow or a dead |
| code), a duplicate/unsorted entry (searchsorted requires strictly |
| increasing input and returns garbage otherwise), a map that does not reach |
| +-1 (the block's absmax element would clip), a missing exact zero (the |
| all-zero state could not be represented).""" |
| cb = quant8.create_dynamic_map(signed=signed) |
| assert cb.shape == (256,) and cb.dtype == torch.float32 |
| assert bool(torch.isfinite(cb).all()) |
| assert bool((cb[1:] > cb[:-1]).all()), "codebook must be STRICTLY increasing" |
| assert 0.0 in cb.tolist() |
| assert cb[quant8.zero_code(signed)].item() == 0.0 |
| assert cb.max().item() == 1.0 |
| if signed: |
| assert cb.min().item() <= -0.99, "signed map must reach (nearly) -1" |
| assert (cb < 0).sum().item() == (cb > 0).sum().item() - 1, ( |
| "signed map is symmetric except for the extra exact +1.0" |
| ) |
| else: |
| assert cb.min().item() == 0.0 and bool((cb >= 0).all()) |
|
|
|
|
| def test_codebook_is_dynamic_not_linear(): |
| """The point of the map: relative resolution across decades. A linear grid |
| has 1/255 absolute spacing, so values below ~4e-3 of absmax collapse to |
| zero; the dynamic map must still resolve 1e-5 of absmax to within a few |
| percent. Kills: replacing create_dynamic_map with linspace(-1, 1, 256).""" |
| for signed in (True, False): |
| cb = quant8.create_dynamic_map(signed) |
| codes = quant8.nearest_code(torch.tensor([1e-5, 1e-3, 0.5]), cb).long() |
| got = cb[codes] |
| rel = ((got - torch.tensor([1e-5, 1e-3, 0.5])).abs() / torch.tensor([1e-5, 1e-3, 0.5])) |
| assert bool((rel < 0.45).all()), (signed, got, rel) |
| assert rel[1] < 0.1 and rel[2] < 0.02 |
| |
| assert quant8.create_dynamic_map(False)[1].item() < 1e-6 |
|
|
|
|
| @pytest.mark.parametrize("signed", [True, False]) |
| def test_nearest_code_is_truly_the_nearest_including_edges_and_midpoints(signed): |
| """Kills: an off-by-one in the searchsorted/midpoint arithmetic (picking |
| the code ABOVE every x, or the wrong side of a midpoint), edge handling |
| (x below cb[0] / above cb[-1] must resolve to the end codes, not index -1 |
| or 256), and NaN-through-searchsorted. Brute-force argmin is the oracle.""" |
| cb = quant8.create_dynamic_map(signed) |
| lo = -1.5 if signed else -0.2 |
| g = torch.Generator().manual_seed(0) |
| x = torch.cat( |
| [ |
| torch.rand(20000, generator=g) * (1.5 - lo) + lo, |
| cb, |
| (cb[1:] + cb[:-1]) / 2 + 1e-9, |
| (cb[1:] + cb[:-1]) / 2 - 1e-9, |
| torch.tensor([-1.0, 0.0, 1.0, -1e-30, 1e-30]), |
| ] |
| ) |
| got = quant8.nearest_code(x, cb) |
| assert got.dtype == torch.uint8 |
| want = (x.unsqueeze(1) - cb.unsqueeze(0)).abs().argmin(dim=1) |
| |
| |
| |
| d_got = (x - cb[got.long()]).abs() |
| d_want = (x - cb[want]).abs() |
| assert bool((d_got <= d_want * 1.001 + 1e-9).all()), "nearest_code returned a farther code" |
|
|
|
|
| |
| |
| |
|
|
|
|
| @pytest.mark.parametrize("signed", [True, False]) |
| def test_roundtrip_error_is_bounded_on_gaussian_blocks(signed): |
| """Quantization error, measured as the optimizer sees it: relative to the |
| block's absmax (the scale the codes are normalised by). Median < 1e-2 and |
| max < 0.1 are the specified bounds; the measured values are ~2e-3 and |
| ~7e-3 (the top bin is 0.014 wide), so the bound has 10x headroom and a |
| broken codebook (linear, shifted, unsorted) blows through it. |
| |
| ELEMENT-wise relative error on the large elements (|x| > 0.1 absmax, where |
| the optimizer's update is decided) is also pinned: median < 2e-2 (measured |
| ~1.3e-2; a 64-bin decade cannot do better than ~half a bin of 0.014/0.2). |
| """ |
| x = _randn(8 * B, seed=1) |
| if not signed: |
| x = x.pow(2) |
| codes, absmax = quant8.quantize_blockwise(x, B, signed) |
| assert codes.dtype == torch.uint8 and codes.numel() == 8 * B and absmax.shape == (8,) |
| y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, signed) |
| scale = absmax.repeat_interleave(B) |
| abs_err = (x - y).abs() / scale |
| assert abs_err.median().item() < 1e-2 |
| assert abs_err.max().item() < 0.1 |
| big = x.abs() > 0.1 * scale |
| rel = ((x - y).abs() / x.abs())[big] |
| assert rel.median().item() < 2e-2 and rel.max().item() < 0.1 |
|
|
|
|
| def test_blockwise_isolates_an_outlier_to_its_own_block(): |
| """The reason for BLOCK-wise: one huge element must not coarsen the other |
| blocks. Kills: a single global absmax (every other block's median error |
| would be ~1e4x worse).""" |
| x = _randn(4 * B, seed=2) |
| x[5] = 1e4 |
| codes, absmax = quant8.quantize_blockwise(x, B, True) |
| y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, True) |
| clean = slice(B, 4 * B) |
| err = ((x[clean] - y[clean]).abs() / x[clean].abs().clamp_min(1e-3)).median() |
| assert err.item() < 2e-2, err |
| assert absmax[0].item() == 1e4 and absmax[1].item() < 6 |
|
|
|
|
| def test_zero_blocks_and_partial_last_block_roundtrip_exactly_without_nan(): |
| """Kills: dividing by a zero absmax (NaN codes), and a partial final block |
| that leaks its zero padding into the dequantized values or the numel.""" |
| x = torch.zeros(3 * B + 77) |
| x[B : B + 5] = torch.tensor([1.0, -2.0, 3.0, 0.0, 0.5]) |
| codes, absmax = quant8.quantize_blockwise(x, B, True) |
| assert codes.numel() == 4 * B |
| y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, True) |
| assert y.shape == x.shape and bool(torch.isfinite(y).all()) |
| assert torch.equal(y[:B], torch.zeros(B)) and torch.equal(y[2 * B :], torch.zeros(B + 77)) |
| assert absmax[0].item() == 0.0 and absmax[1].item() == 3.0 |
|
|
|
|
| def test_dequantize_indexes_with_a_gather_not_a_uint8_mask(): |
| """THE UINT8 INDEX TRAP: `cb[codes]` with a uint8 `codes` is a boolean-mask |
| index in torch, not a gather. dequantize must return cb[code] * absmax for |
| codes 0 and 255 (the two ends), not a mask-selected junk tensor.""" |
| cb = quant8.codebook(True, "cpu") |
| codes = torch.tensor([[0, 255, 127, 1]], dtype=torch.uint8) |
| out = quant8.dequantize_rows(codes, torch.tensor([2.0]), cb) |
| assert torch.equal(out, torch.stack([cb[0], cb[255], cb[127], cb[1]]).unsqueeze(0) * 2.0) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_fp32_fallback_for_small_tensors_is_bit_exact_vs_torch_adamw(): |
| """Tensors under `min_8bit_size` keep fp32 moments, and must match |
| `torch.optim.AdamW` BITWISE over many steps (decoupled decay, betas, |
| bias correction, eps all included). Kills: any reordering of the op |
| sequence (lerp_ vs hand EMA, folding bias-correction differently), a |
| dropped weight_decay, a mis-forwarded beta2 (0.95 vs torch's 0.999).""" |
| a = [_param(4, 5, seed=s) for s in range(3)] + [_param(4095, seed=9)] |
| b = [torch.nn.Parameter(p.detach().clone()) for p in a] |
| mine = BlockwiseAdamW8bit(a, **HP) |
| ref = torch.optim.AdamW(b, **HP) |
| assert all("m" in s for s in mine._state), "all four are below min_8bit_size=4096" |
| for step in range(40): |
| _set_grads(a, step) |
| for pa, pb in zip(a, b): |
| pb.grad = pa.grad.clone() |
| assert mine.step() is True |
| ref.step() |
| for pa, pb in zip(a, b): |
| assert torch.equal(pa, pb) |
|
|
|
|
| def test_threshold_is_below_min_8bit_size_inclusive_at_the_boundary(): |
| """bnb's rule: numel < min_8bit_size stays fp32; numel == min is 8-bit. |
| Kills: `<=` vs `<` drift.""" |
| opt = BlockwiseAdamW8bit([_param(4095), _param(4096)], lr=1e-3) |
| assert "m" in opt._state[0] and "m_codes" in opt._state[1] |
| assert opt._state[1]["m_codes"].dtype == torch.uint8 |
|
|
|
|
| def _quadratic(opt_fn, n=8192, steps=300, seed=1): |
| g = torch.Generator().manual_seed(seed) |
| t = torch.randn(n, generator=g) |
| curv = torch.rand(n, generator=g) * 9 + 1 |
| x = torch.nn.Parameter(torch.zeros(n)) |
| opt = opt_fn([x]) |
| loss0 = 0.5 * (curv * t**2).sum().item() |
| for _ in range(steps): |
| x.grad = None |
| loss = 0.5 * (curv * (x - t) ** 2).sum() |
| loss.backward() |
| opt.step() |
| return 0.5 * (curv * (x.detach() - t) ** 2).sum().item(), loss0 |
|
|
|
|
| def test_300_steps_on_a_convex_quadratic_converge_like_torch_adamw(): |
| """Both reach <0.5% of the initial loss, and the 8-bit endpoint is within |
| 1e-3 * L0 of fp32 AdamW's. Kills: state that does not accumulate (a |
| zeroed-every-step moment), a lost lr/bias-correction (never converges), a |
| codebook that collapses v to zero (the update explodes, m/eps).""" |
| hp = dict(lr=5e-2, betas=(0.9, 0.95), weight_decay=0.0) |
| l8, l0 = _quadratic(lambda p: BlockwiseAdamW8bit(p, **hp)) |
| l32, _ = _quadratic(lambda p: torch.optim.AdamW(p, **hp)) |
| assert l32 < 5e-3 * l0, (l32, l0) |
| assert l8 < 5e-3 * l0, (l8, l0) |
| assert abs(l8 - l32) < 1e-3 * l0, (l8, l32, l0) |
|
|
|
|
| def _linreg(opt_fn, steps=300): |
| g = torch.Generator().manual_seed(3) |
| d_in, d_out, n = 128, 64, 256 |
| X = torch.randn(n, d_in, generator=g) |
| W_true = torch.randn(d_out, d_in, generator=g) / math.sqrt(d_in) |
| Y = X @ W_true.T |
| W = torch.nn.Parameter(torch.zeros(d_out, d_in)) |
| bias = torch.nn.Parameter(torch.zeros(d_out)) |
| opt = opt_fn([W, bias]) |
| first = None |
| for _ in range(steps): |
| opt.zero_grad() if hasattr(opt, "zero_grad") else None |
| loss = ((X @ W.T + bias - Y) ** 2).mean() |
| first = loss.item() if first is None else first |
| loss.backward() |
| opt.step() |
| return loss.item(), first |
|
|
|
|
| def test_300_steps_of_tiny_linear_regression_match_torch_adamw(): |
| """Mixed 8-bit (weight) + fp32-fallback (bias) parameters in one |
| optimizer, on a real backward pass.""" |
| hp = dict(lr=1e-2, betas=(0.9, 0.95), weight_decay=0.0) |
| l8, first = _linreg(lambda p: BlockwiseAdamW8bit(p, **hp)) |
| l32, _ = _linreg(lambda p: torch.optim.AdamW(p, **hp)) |
| assert l32 < 1e-3 * first and l8 < 1e-3 * first, (l8, l32, first) |
| assert abs(l8 - l32) < 1e-3 * first, (l8, l32, first) |
|
|
|
|
| def test_first_step_is_exactly_lr_times_sign_because_bias_correction_is_applied(): |
| """After one step, m_hat/(sqrt(v_hat)+eps) == sign(g) (to eps), so every |
| element moves by exactly lr. The update uses the freshly computed fp32 |
| moments, not the requantized ones, so this holds on the 8-bit path to fp32 |
| rounding. Kills: missing bias correction (moves 0.1/sqrt(0.05) = 0.447 lr |
| with beta2=0.95), or an update computed from the dequantized-after-quantize |
| moments (off by the quantization error, ~1e-3 relative).""" |
| p = _param(B * 3, seed=4) |
| p0 = p.detach().clone() |
| p.grad = _randn(B * 3, seed=5) * 10 |
| opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0) |
| assert "m_codes" in opt._state[0] |
| opt.step() |
| moved = (p0 - p.detach()) * p.grad.sign() |
| assert torch.allclose(moved, torch.full_like(moved, 1e-3), rtol=1e-4, atol=1e-7), ( |
| moved.min().item(), |
| moved.max().item(), |
| ) |
|
|
|
|
| def test_weight_decay_is_decoupled_from_the_gradient(): |
| """With a ZERO gradient the moments stay zero, so the only thing that can |
| move a weight is the decay: p <- p * (1 - lr*wd), to fp32 rounding. A |
| COUPLED (L2) implementation adds wd*p to the gradient and Adam then |
| normalises it to a step of ~lr per element -- ~1e3x more movement. Kills: |
| coupled decay, and a dropped weight_decay (p unchanged).""" |
| p = _param(B * 2, seed=6) |
| p0 = p.detach().clone() |
| p.grad = torch.zeros_like(p) |
| opt = BlockwiseAdamW8bit([p], lr=1e-2, betas=(0.9, 0.95), weight_decay=0.5) |
| assert "m_codes" in opt._state[0] |
| opt.step() |
| assert torch.allclose(p.detach(), p0 * (1 - 1e-2 * 0.5), rtol=1e-6, atol=0.0) |
| assert opt._state[0]["m_absmax"].abs().max().item() == 0.0, "decay leaked into the moments" |
|
|
|
|
| def test_chunking_is_invisible_to_the_result(): |
| """Chunks are whole blocks, so processing one block at a time must equal |
| processing the tensor in one pass, BITWISE (params and every code/absmax). |
| Kills: a chunk boundary that is off by a block, a partial-block mishandled |
| at a chunk edge, state written to the wrong row range.""" |
| n = B * 5 + 123 |
| outs = [] |
| for chunk in (B, 2 * B, 10**9): |
| p = _param(n, seed=8) |
| opt = BlockwiseAdamW8bit([p], chunk_elems=chunk, **HP) |
| for step in range(6): |
| _set_grads([p], step) |
| opt.step() |
| outs.append((p.detach().clone(), opt.state_dict()["state"][0])) |
| for pw, st in outs[1:]: |
| assert torch.equal(pw, outs[0][0]) |
| for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): |
| assert torch.equal(st[k], outs[0][1][k]), k |
|
|
|
|
| def test_padding_of_a_partial_last_block_stays_the_zero_code(): |
| """n is not a multiple of the block: the pad tail of the last row must keep |
| quantizing to exact zero forever, or garbage moments would leak into the |
| last block's absmax (and thus its resolution).""" |
| n = 2 * B + 100 |
| p = _param(n, seed=10) |
| opt = BlockwiseAdamW8bit([p], **HP) |
| for step in range(5): |
| _set_grads([p], step) |
| opt.step() |
| st = opt._state[0] |
| assert st["m_codes"].shape == (3, B) |
| assert bool((st["m_codes"][2, 100:] == quant8.zero_code(True)).all()) |
| assert bool((st["v_codes"][2, 100:] == quant8.zero_code(False)).all()) |
|
|
|
|
| def test_param_with_no_grad_is_skipped_and_its_step_count_does_not_advance(): |
| """Bias correction is per tensor: a tensor that sat out a step must not be |
| treated as having taken it. Also: no gradients anywhere is an error, not a |
| silent no-op.""" |
| a, b = _param(B, seed=1), _param(B, seed=2) |
| opt = BlockwiseAdamW8bit([a, b], **HP) |
| a.grad = _randn(B, seed=3) |
| b0 = b.detach().clone() |
| assert opt.step() is True |
| assert opt._state[0]["step"] == 1 and opt._state[1]["step"] == 0 |
| assert torch.equal(b, b0) |
| a.grad = None |
| with pytest.raises(RuntimeError, match="no gradients"): |
| opt.step() |
|
|
|
|
| def test_empty_local_shard_is_tolerated(): |
| """FSDP2 can hand a rank an EMPTY shard (a 1-row param over 2 ranks). It |
| must construct, step and round-trip rather than dividing by zero blocks.""" |
| e, p = torch.nn.Parameter(torch.zeros(0)), _param(B, seed=1) |
| opt = BlockwiseAdamW8bit([e, p], **HP) |
| e.grad, p.grad = torch.zeros(0), _randn(B, seed=2) |
| assert opt.step() is True |
| opt.load_state_dict(opt.state_dict()) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_grad_norm_is_recorded_before_clipping_and_clipping_is_applied(): |
| """THE PPO BLIND-TRIPWIRE INCIDENT's cousin: the norm a tripwire reads must |
| be the PRE-clip norm (post-clip it is <= max by construction and a |
| 10,000x explosion looks like 1.0). And the clip must actually rescale the |
| gradients that enter the update. |
| |
| Adam is scale-invariant on a single step, so this runs TWO steps (a big |
| clipped gradient, then a small one) and compares, bitwise, against an |
| unclipped optimizer fed the pre-clipped gradients by hand. Kills: norm |
| recorded after clipping; clipping not applied (the two trajectories |
| differ by the relative weight of step 1 in the moments); a coefficient |
| other than max/(norm+1e-6).""" |
| ps = [_param(B * 2, seed=1), _param(B, seed=2), _param(7, seed=3)] |
| qs = [torch.nn.Parameter(p.detach().clone()) for p in ps] |
| hp = dict(lr=1e-2, betas=(0.9, 0.95), weight_decay=0.0) |
| clipped = BlockwiseAdamW8bit(ps, max_grad_norm=1.0, **hp) |
| manual = BlockwiseAdamW8bit(qs, max_grad_norm=None, **hp) |
|
|
| for step, scale in enumerate((100.0, 0.01)): |
| _set_grads(ps, step, scale=scale) |
| true_norm = float(torch.sqrt(sum(p.grad.double().pow(2).sum() for p in ps))) |
| for q, p in zip(qs, ps): |
| q.grad = p.grad.clone() |
| if true_norm > 1.0: |
| coef = 1.0 / (true_norm + 1e-6) |
| for q in qs: |
| q.grad = q.grad * coef |
| assert clipped.step() is True |
| manual.step() |
| assert clipped.last_grad_norm == pytest.approx(true_norm, rel=1e-6), ( |
| "last_grad_norm must be the GLOBAL PRE-CLIP norm" |
| ) |
| assert clipped.last_grad_norm < 1.0 |
| for p, q in zip(ps, qs): |
| assert torch.equal(p, q) |
|
|
|
|
| def test_nonfinite_gradient_skips_the_step_and_leaves_every_byte_of_state_alone(): |
| """A NaN/inf discovered after the moments were updated has already poisoned |
| the second moment. The verdict must come BEFORE any state is touched, the |
| norm must read NaN (a 0.0 would read as 'converged'), and a good step after |
| the skip must work. Kills: checking finiteness after the update; counting |
| the skipped step in the bias-correction step number; returning True.""" |
| p = _param(B * 2, seed=1) |
| opt = BlockwiseAdamW8bit([p], **HP) |
| _set_grads([p], 0) |
| opt.step() |
| snap, p0 = opt.state_dict(), p.detach().clone() |
| for bad in (float("nan"), float("inf"), float("-inf")): |
| p.grad = _randn(B * 2, seed=3) |
| p.grad[17] = bad |
| assert opt.step() is False |
| assert math.isnan(opt.last_grad_norm) |
| assert torch.equal(p, p0) |
| assert opt.skipped_steps == 3 |
| now = opt.state_dict() |
| assert now["state"][0]["step"] == snap["state"][0]["step"] == 1 |
| for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): |
| assert torch.equal(now["state"][0][k], snap["state"][0][k]), k |
| _set_grads([p], 1) |
| assert opt.step() is True and not torch.equal(p, p0) |
|
|
|
|
| def test_second_moment_underflow_cannot_produce_an_m_over_eps_update(): |
| """THE ZERO-SECOND-MOMENT CATASTROPHE. In a block with one huge gradient |
| (1.0) the tiny ones (1e-6) have v/absmax ~ 1e-12, below the smallest |
| nonzero code, so their STORED second moment is exactly 0. If the update |
| were computed from stored moments, it would be m / (sqrt(0) + 1e-8) -- 100x |
| lr per step, i.e. the weight explodes. The update must use the FRESH fp32 |
| moments (quantization is a storage format), which bounds the per-step move |
| by lr / sqrt(1 - beta2) (= 4.47 lr here; measured 2.1 lr). |
| |
| Kills: computing the update after requantization; a codebook whose zero code |
| swallows small v.""" |
| p = torch.nn.Parameter(torch.zeros(4096)) |
| opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0) |
| g = torch.full((4096,), 1e-6) |
| g[0] = 1.0 |
| worst = 0.0 |
| for _ in range(20): |
| before = p.detach().clone() |
| p.grad = g.clone() |
| opt.step() |
| worst = max(worst, (p.detach() - before).abs().max().item()) |
| assert (opt._state[0]["v_codes"][0, 1:] == 0).all(), "precondition: v really underflowed" |
| assert worst <= 1e-3 / math.sqrt(1 - 0.95) * 1.05, worst / 1e-3 |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_stochastic_rounding_is_unbiased_and_exact_on_representable_values(): |
| """E[SR(x)] == x, and a value already representable in bf16 is unchanged. |
| Kills: round-down (biased), round-to-nearest, a carry that corrupts the |
| sign/exponent for negative numbers.""" |
| gen = torch.Generator().manual_seed(0) |
| x = torch.full((200000,), 1.0 + 3e-4) |
| out = stochastic_round_to_bf16(x, gen).float() |
| assert set(out.unique().tolist()) <= {1.0, 1.0078125} |
| assert out.mean().item() == pytest.approx(1.0 + 3e-4, abs=3e-5) |
| neg = stochastic_round_to_bf16(-x, gen).float() |
| assert neg.mean().item() == pytest.approx(-(1.0 + 3e-4), abs=3e-5) |
| exact = torch.tensor([0.0, 1.0, -2.5, 0.15625, 3.0e38 / 2**20]).to(torch.bfloat16).float() |
| assert torch.equal(stochastic_round_to_bf16(exact, gen).float(), exact) |
|
|
|
|
| def test_bf16_params_without_a_master_copy_still_move_under_tiny_updates(): |
| """THE NO-MASTER-COPY DEFECT (the silent-no-train class). A bf16 weight of |
| 1.0 has spacing 7.8e-3; an Adam step of lr=1e-4 is 80x below half a ulp, so |
| round-to-nearest returns 1.0 EVERY step and the run trains nothing while |
| its loss and grad-norm logs look healthy. Stochastic rounding makes the |
| parameter move by the right amount in expectation: 200 steps of lr 1e-4 on a |
| constant gradient move the mean by ~0.02. |
| |
| Both directions: nearest rounding really is stuck (the control that makes |
| this test mean something), and SR really moves. Kills: stochastic_rounding |
| ignored / the write-back always rounding to nearest.""" |
| def run(sr): |
| p = torch.nn.Parameter(torch.ones(B * 4, dtype=torch.bfloat16)) |
| opt = BlockwiseAdamW8bit( |
| [p], lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, stochastic_rounding=sr |
| ) |
| for _ in range(200): |
| p.grad = torch.ones_like(p) |
| opt.step() |
| return p.detach().float().mean().item() |
|
|
| assert run(False) == 1.0, "control: round-to-nearest should be stuck" |
| assert run(True) == pytest.approx(1.0 - 200 * 1e-4, abs=2.5e-3) |
|
|
|
|
| def test_bf16_param_with_fp32_grad_is_accepted(): |
| """The fp32-reduce policy hands back fp32 grads for bf16 params.""" |
| p = torch.nn.Parameter(_randn(B * 2, seed=1, dtype=torch.bfloat16)) |
| if hasattr(p, "grad_dtype"): |
| p.grad_dtype = None |
| p.grad = _randn(B * 2, seed=2) |
| p0 = p.detach().clone() |
| BlockwiseAdamW8bit([p], lr=1e-2).step() |
| assert p.dtype == torch.bfloat16 and not torch.equal(p, p0) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _train(opt, params, steps, start=0, dtype=None): |
| for step in range(start, start + steps): |
| _set_grads(params, step, dtype=dtype) |
| opt.step() |
|
|
|
|
| @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| def test_resume_from_state_dict_is_bit_exact(dtype): |
| """Save at step 5, keep training the original; load into a FRESH optimizer |
| over a copy of the step-5 weights and train the same 5 steps: weights and |
| every code/absmax byte must be identical. bf16 covers the stochastic |
| rounding RNG (its generator state is part of the checkpoint). Kills: a |
| state_dict that drops the step count, the absmax, the SR RNG, or the |
| hyper-parameters; a load that re-quantizes instead of copying.""" |
| mk = lambda: [_param(B * 3, seed=1, dtype=dtype), _param(300, seed=2, dtype=dtype), |
| _param(B + 9, seed=3, dtype=dtype)] |
| a = mk() |
| opt_a = BlockwiseAdamW8bit(a, **HP, seed=5) |
| _train(opt_a, a, 5) |
| sd = opt_a.state_dict() |
| b = [torch.nn.Parameter(p.detach().clone()) for p in a] |
| opt_b = BlockwiseAdamW8bit(b, **HP, seed=999) |
| opt_b.load_state_dict(sd) |
| _train(opt_a, a, 5, start=5) |
| _train(opt_b, b, 5, start=5) |
| for pa, pb in zip(a, b): |
| assert torch.equal(pa, pb) |
| sa, sb = opt_a.state_dict(), opt_b.state_dict() |
| for x, y in zip(sa["state"], sb["state"]): |
| assert x.keys() == y.keys() |
| for k in x: |
| assert torch.equal(torch.as_tensor(x[k]), torch.as_tensor(y[k])), k |
|
|
|
|
| def test_state_dict_is_a_cpu_snapshot_not_a_live_view(): |
| """THE SHARED-STATE ALIASING INCIDENT: `.to('cpu')` on a CPU tensor returns |
| the same object, so a 'snapshot' built that way is mutated by the next |
| step -- uint8 codes that no comparison flags as drifting. And loading one |
| dict into two optimizers must not make them share state. Kills: missing |
| copy=True in state_dict; a load that stores references.""" |
| p = _param(B * 2, seed=1) |
| opt = BlockwiseAdamW8bit([p], **HP) |
| _train(opt, [p], 2) |
| sd = opt.state_dict() |
| frozen = copy.deepcopy(sd) |
| for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): |
| assert sd["state"][0][k].device.type == "cpu" |
| assert sd["state"][0][k].data_ptr() != opt._state[0][k].data_ptr() |
| _train(opt, [p], 3, start=2) |
| for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): |
| assert torch.equal(sd["state"][0][k], frozen["state"][0][k]), f"{k}: snapshot drifted" |
| assert sd["state"][0]["step"] == 2 |
|
|
| q1, q2 = _param(B * 2, seed=1), _param(B * 2, seed=1) |
| o1, o2 = BlockwiseAdamW8bit([q1], **HP), BlockwiseAdamW8bit([q2], **HP) |
| o1.load_state_dict(sd) |
| o2.load_state_dict(sd) |
| assert o1._state[0]["m_codes"].data_ptr() != o2._state[0]["m_codes"].data_ptr() |
| assert o1._state[0]["m_codes"].data_ptr() != sd["state"][0]["m_codes"].data_ptr() |
|
|
|
|
| def test_state_dict_survives_a_pickle_roundtrip(): |
| """It rides `gather_object` (pickle) inside the checkpoint: it must hold no |
| live handles (generators, devices) and stay byte-identical.""" |
| import pickle |
|
|
| p = _param(B * 2, seed=1, dtype=torch.bfloat16) |
| opt = BlockwiseAdamW8bit([p], **HP) |
| _train(opt, [p], 3) |
| sd = opt.state_dict() |
| back = pickle.loads(pickle.dumps(sd)) |
| q = torch.nn.Parameter(p.detach().clone()) |
| other = BlockwiseAdamW8bit([q], **HP) |
| other.load_state_dict(back) |
| _train(opt, [p], 2, start=3) |
| _train(other, [q], 2, start=3) |
| assert torch.equal(p, q) |
|
|
|
|
| def test_resume_refuses_a_different_world_size_layout_or_format_and_changes_nothing(): |
| """The state is SHARD-shaped. Loading it under another world size aligns |
| rank r's moments to a different slice of the parameter and training |
| continues plausibly on corrupted state. A refused load must also leave the |
| optimizer untouched (validate everything, then write). Kills: dropping the |
| world_size check; checking numel only after partially copying.""" |
| p = _param(B * 2, seed=1) |
| opt = BlockwiseAdamW8bit([p], **HP) |
| _train(opt, [p], 2) |
| sd = opt.state_dict() |
|
|
| fresh_p = _param(B * 2, seed=1) |
| fresh = BlockwiseAdamW8bit([fresh_p], **HP) |
| before = fresh.state_dict() |
|
|
| with pytest.raises(ValueError, match="world_size"): |
| fresh.load_state_dict(dict(sd, world_size=sd["world_size"] + 4)) |
| wrong_n = copy.deepcopy(sd) |
| wrong_n["state"][0]["numel"] += 1 |
| with pytest.raises(ValueError, match="numel"): |
| fresh.load_state_dict(wrong_n) |
| with pytest.raises(ValueError, match="block_size"): |
| fresh.load_state_dict(dict(sd, block_size=1024)) |
| with pytest.raises(ValueError, match="format"): |
| fresh.load_state_dict({"inner": {}, "state": {}}) |
| bad_shape = copy.deepcopy(sd) |
| bad_shape["state"][0]["m_absmax"] = torch.zeros(5) |
| with pytest.raises(ValueError, match="m_absmax"): |
| fresh.load_state_dict(bad_shape) |
| other = BlockwiseAdamW8bit([_param(B * 2), _param(B * 2)], **HP) |
| with pytest.raises(ValueError, match="tensors"): |
| other.load_state_dict(sd) |
|
|
| after = fresh.state_dict() |
| assert after["state"][0]["step"] == before["state"][0]["step"] == 0 |
| for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): |
| assert torch.equal(after["state"][0][k], before["state"][0][k]) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_state_bytes_match_the_2n_plus_8n_over_block_budget(): |
| """The RSS/VRAM projection must be HONEST: for block-aligned tensors the |
| real footprint is exactly 2*n + 8*n/block bytes; small tensors are 8*n; and |
| `projected_state_bytes` (usable before allocating anything) agrees with |
| what was actually allocated. Kills: an absmax kept per element, fp32 |
| codes, or a projection that forgets the fp32 fallback tensors.""" |
| n_big = 100 * B |
| params = [_param(n_big, seed=1), _param(300, seed=2), _param(5 * B + 17, seed=3)] |
| opt = BlockwiseAdamW8bit(params, **HP) |
| assert opt._state[0]["m_codes"].dtype == torch.uint8 |
| assert opt._state[0]["m_absmax"].dtype == torch.float32 |
| assert opt._state[0]["m_absmax"].numel() == 100 |
| big_only = BlockwiseAdamW8bit([_param(n_big)], **HP).state_bytes() |
| assert big_only == 2 * n_big + 8 * n_big // B |
| got = opt.state_bytes() |
| want = BlockwiseAdamW8bit.projected_state_bytes([p.numel() for p in params]) |
| assert got == want |
| assert got < 0.26 * sum(p.numel() for p in params) * 8, "must be ~4x smaller than fp32 moments" |
| |
| rank_numel = int(7_248_023_552 / 4) |
| proj = BlockwiseAdamW8bit.projected_state_bytes([rank_numel]) |
| assert proj / rank_numel == pytest.approx(2 + 8 / B, rel=1e-3) |
| assert 3.3 < proj / 2**30 < 3.5 |
|
|
|
|
| def test_describe_placement_reports_uint8_codes_on_the_param_device(): |
| opt = BlockwiseAdamW8bit([_param(B * 2)], **HP) |
| ent = {e["key"]: e for e in opt.describe_placement()["state"]} |
| assert ent["m_codes"]["dtype"] == str(torch.uint8) and ent["m_codes"]["device"] == "cpu" |
| assert ent["v_absmax"]["dtype"] == str(torch.float32) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_constructor_refuses_bad_inputs(): |
| p = _param(B) |
| with pytest.raises(ValueError, match="empty"): |
| BlockwiseAdamW8bit([], lr=1e-3) |
| with pytest.raises(ValueError, match="twice"): |
| BlockwiseAdamW8bit([p, p], lr=1e-3) |
| with pytest.raises(ValueError, match="betas"): |
| BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 1.0)) |
| with pytest.raises(ValueError, match="max_grad_norm"): |
| BlockwiseAdamW8bit([p], lr=1e-3, max_grad_norm=0.0) |
| with pytest.raises(ValueError, match="block_size"): |
| BlockwiseAdamW8bit([p], lr=1e-3, block_size=0) |
| with pytest.raises(ValueError, match="floating"): |
| BlockwiseAdamW8bit([torch.zeros(B, dtype=torch.int64)], lr=1e-3) |
| with pytest.raises(ValueError, match="contiguous"): |
| BlockwiseAdamW8bit([torch.nn.Parameter(torch.zeros(64, 64).t())], lr=1e-3) |
| q = _param(B) |
| opt = BlockwiseAdamW8bit([q], lr=1e-3) |
| q.data = torch.zeros(B + 1) |
| q.grad = torch.zeros(B + 1) |
| with pytest.raises(ValueError, match="elements"): |
| opt.step() |
|
|
|
|
| def test_lr_written_through_param_groups_is_what_the_step_uses(): |
| """`Trainer._apply_scheduled_lr` writes the scheduled lr into |
| `param_groups[*]['lr']`; a step that cached the constructor lr would train |
| at a constant rate while logging a decaying one (the scheduler-never-applied |
| incident). Kills: reading self.lr instead of the group.""" |
| p, q = _param(B, seed=1), _param(B, seed=1) |
| a = BlockwiseAdamW8bit([p], lr=1e-3, weight_decay=0.0) |
| b = BlockwiseAdamW8bit([q], lr=7e-3, weight_decay=0.0) |
| a.param_groups[0]["lr"] = 7e-3 |
| for o, x in ((a, p), (b, q)): |
| x.grad = _randn(B, seed=2) |
| o.step() |
| assert torch.equal(p, q) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _tiny_params(): |
| return [_param(64, 64, seed=1), _param(64, seed=2)] |
|
|
|
|
| @pytest.mark.parametrize("name", ["sft.toml", "pretrain.toml"]) |
| def test_make_optimizer_builds_adamw8bit_for_the_shipped_full_ft_configs(name): |
| """THE LAUNCH BLOCKER (review_trainer.md): `load_config(sft.toml) -> |
| make_optimizer(...)` raised NotImplementedError, so no full-FT stage could |
| start. The shipped configs must now construct, tagged for the banner with |
| the 8-bit footprint, and hold no host-side master.""" |
| from align.config import load_config |
| from align.train.offload_adamw import make_optimizer |
|
|
| cfg = load_config(os.path.join(os.path.dirname(__file__), "..", "configs", name)) |
| assert cfg.optimizer.name == "adamw8bit" |
| params = _tiny_params() |
| opt = make_optimizer(cfg, params) |
| assert isinstance(opt, BlockwiseAdamW8bit) |
| assert opt.banner_tag.startswith("adamw8bit") and "8-bit" in opt.banner_tag |
| assert opt.param_groups[0]["lr"] == cfg.optimizer.lr |
| assert opt.param_groups[0]["betas"] == tuple(cfg.optimizer.betas) |
| assert opt.param_groups[0]["weight_decay"] == cfg.optimizer.weight_decay |
| for p in params: |
| p.grad = torch.ones_like(p) |
| assert opt.step() is True and opt.skipped_steps == 0 |
|
|
|
|
| def test_make_optimizer_adamw8bit_is_wired_with_the_config_hyperparameters(): |
| """Kills: a factory that drops betas/eps/weight_decay on the floor (the |
| parity test passes against defaults but not against the run config).""" |
| import types |
|
|
| from align.train.offload_adamw import make_optimizer |
|
|
| o = types.SimpleNamespace( |
| name="AdamW8bit", lr=3e-4, betas=[0.8, 0.9], eps=1e-6, weight_decay=0.25 |
| ) |
| opt = make_optimizer(types.SimpleNamespace(optimizer=o), _tiny_params()) |
| g = opt.param_groups[0] |
| assert (g["lr"], g["betas"], g["eps"], g["weight_decay"]) == (3e-4, (0.8, 0.9), 1e-6, 0.25) |
| assert g["max_grad_norm"] is None |
| o.max_grad_norm = 2.0 |
| assert make_optimizer(types.SimpleNamespace(optimizer=o), _tiny_params()).param_groups[0][ |
| "max_grad_norm" |
| ] == 2.0 |
|
|
|
|
| def test_adamw8bit_never_silently_becomes_offload_adamw_and_typos_still_raise(): |
| import types |
|
|
| from align.train.offload_adamw import OffloadAdamW, make_optimizer |
|
|
| cfg = types.SimpleNamespace( |
| optimizer=types.SimpleNamespace( |
| name="adamw8bit", lr=1e-3, betas=[0.9, 0.95], eps=1e-8, weight_decay=0.0 |
| ) |
| ) |
| assert not isinstance(make_optimizer(cfg, _tiny_params()), OffloadAdamW) |
| cfg.optimizer.name = "adamw8bti" |
| with pytest.raises(ValueError, match="unknown optimizer"): |
| make_optimizer(cfg, _tiny_params()) |
| cfg.optimizer.name = "adafactor" |
| with pytest.raises(ValueError, match="BANNED"): |
| make_optimizer(cfg, _tiny_params()) |
|
|
|
|
| def test_shard_local_shim_constructs_delegates_and_round_trips(): |
| """The legacy class name now builds the real optimizer. Its wrapper |
| state_dict (world_size / n_shards / recursive CPU clone) must round-trip |
| the 8-bit state bit-exactly, and its step must surface the inner skip/norm.""" |
| from align.train.offload_adamw import ShardLocalAdamW8bit |
|
|
| a = [_param(B * 2, seed=1), _param(40, seed=2)] |
| opt = ShardLocalAdamW8bit(a, lr=1e-3, weight_decay=0.1) |
| _train(opt, a, 3) |
| assert opt.last_grad_norm > 0 and opt.skipped_steps == 0 |
| a[0].grad = torch.full_like(a[0], float("nan")) |
| assert opt.step() is False and opt.skipped_steps == 1 and math.isnan(opt.last_grad_norm) |
| opt.zero_grad() |
| sd = opt.state_dict() |
| assert sd["n_shards"] == 2 and sd["skipped_steps"] == 1 |
| assert sd["inner"]["state"][0]["m_codes"].device.type == "cpu" |
|
|
| b = [torch.nn.Parameter(p.detach().clone()) for p in a] |
| other = ShardLocalAdamW8bit(b, lr=1e-3, weight_decay=0.1) |
| other.load_state_dict(sd) |
| _train(opt, a, 2, start=3) |
| _train(other, b, 2, start=3) |
| for pa, pb in zip(a, b): |
| assert torch.equal(pa, pb) |
| with pytest.raises(ValueError, match="world_size"): |
| other.load_state_dict(dict(sd, world_size=sd["world_size"] + 1)) |
| with pytest.raises(ValueError, match="empty"): |
| ShardLocalAdamW8bit([], lr=1e-3) |
|
|
|
|
| |
| |
| |
|
|
| WS = 2 |
| JOIN_TIMEOUT_S = 180 |
|
|
|
|
| def _free_port() -> int: |
| with socket.socket() as s: |
| s.bind(("127.0.0.1", 0)) |
| return s.getsockname()[1] |
|
|
|
|
| def _body_dtensor(rank: int, run_dir: str) -> dict: |
| """Child body. Linear(64, 129): 129 rows over 2 ranks = 65 + 64 rows -> |
| local weights of 4160 (8-bit path) and 4096 (8-bit path, exactly at the |
| threshold); the bias shards (65/64 elements) take the fp32 fallback. So one |
| optimizer exercises UNEVEN shards, both state paths, and DTensor grads.""" |
| import torch.distributed as dist |
| from torch.distributed.device_mesh import init_device_mesh |
| from torch.distributed.fsdp import fully_shard |
|
|
| torch.manual_seed(0) |
| ref = torch.nn.Linear(64, 129) |
| sharded = copy.deepcopy(ref) |
| fully_shard(sharded, mesh=init_device_mesh("cpu", (WS,))) |
| g = torch.Generator().manual_seed(3) |
| x, y = torch.randn(256, 64, generator=g), torch.randn(256, 129, generator=g) |
|
|
| hp = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0) |
| ref_opt = torch.optim.AdamW(ref.parameters(), **hp) |
| opt = BlockwiseAdamW8bit(sharded.parameters(), **hp) |
| |
| |
| plain = [torch.nn.Parameter(p.to_local().detach().clone()) for p in sharded.parameters()] |
| plain_opt = BlockwiseAdamW8bit(plain, **hp) |
|
|
| def full(m): |
| return {n: p.full_tensor().detach().clone() for n, p in m.named_parameters()} |
|
|
| def loss(m): |
| return ((m(x) - y) ** 2).mean() |
|
|
| norms, ref_norms, mean_diff, plain_equal = [], [], 0.0, True |
| for _ in range(6): |
| loss(ref).backward() |
| ref_norms.append(float(torch.sqrt(sum(p.grad.double().pow(2).sum() for p in ref.parameters())))) |
| ref_opt.step() |
| ref_opt.zero_grad() |
|
|
| loss(sharded).backward() |
| for pl, p in zip(plain, sharded.parameters()): |
| pl.grad = p.grad.to_local().detach().clone() |
| assert opt.step() is True |
| assert plain_opt.step() is True |
| opt.zero_grad() |
| norms.append(opt.last_grad_norm) |
| plain_equal &= all( |
| torch.equal(pl.detach(), p.to_local()) for pl, p in zip(plain, sharded.parameters()) |
| ) |
| got = full(sharded) |
| mean_diff = max( |
| mean_diff, |
| *[float((p.detach() - q).abs().mean()) for (_, p), q in zip(ref.named_parameters(), got.values())], |
| ) |
|
|
| |
| pre = full(sharded) |
| sharded(x).pow(2).mean().backward() |
| if rank == 1: |
| next(iter(sharded.parameters())).grad.to_local().fill_(float("nan")) |
| stepped = opt.step() |
| opt.zero_grad() |
| post = full(sharded) |
| skip_diff = max(float((pre[k] - post[k]).abs().max()) for k in pre) |
|
|
| |
| |
| sd = opt.state_dict() |
| gathered = [None] * WS |
| dist.all_gather_object(gathered, norms) |
| return { |
| "norms": norms, |
| "ref_norms": ref_norms, |
| "mean_param_diff": mean_diff, |
| "plain_equal": plain_equal, |
| "all_rank_norms": gathered, |
| "local_numel": [int(p.to_local().numel()) for p in sharded.parameters()], |
| "quantized": [bool(s["quantized"]) for s in sd["state"]], |
| "stepped_on_nan": stepped, |
| "skip_diff": skip_diff, |
| "skipped": opt.skipped_steps, |
| "sd_world_size": sd["world_size"], |
| } |
|
|
|
|
| _BODIES = {"dtensor": _body_dtensor} |
|
|
|
|
| def _child(rank: int, port: int, mode: str, run_dir: str) -> None: |
| import torch.distributed as dist |
|
|
| os.environ.update( |
| RANK=str(rank), |
| WORLD_SIZE=str(WS), |
| LOCAL_RANK=str(rank), |
| MASTER_ADDR="127.0.0.1", |
| MASTER_PORT=str(port), |
| ) |
| dist.init_process_group("gloo", rank=rank, world_size=WS) |
| code = 0 |
| try: |
| result = _BODIES[mode](rank, run_dir) |
| with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f: |
| json.dump(result, f) |
| except BaseException: |
| with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f: |
| f.write(traceback.format_exc()) |
| code = 1 |
| finally: |
| try: |
| dist.destroy_process_group() |
| except Exception: |
| pass |
| if code: |
| sys.exit(code) |
|
|
|
|
| def _run(mode: str, run_dir: str, monkeypatch) -> list[dict]: |
| monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") |
| ctx = mp.get_context("spawn") |
| port = _free_port() |
| procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)] |
| for p in procs: |
| p.start() |
| for p in procs: |
| p.join(JOIN_TIMEOUT_S) |
| hung = [p for p in procs if p.is_alive()] |
| for p in hung: |
| p.kill() |
| p.join() |
| errs = [ |
| f"rank {r}:\n{open(os.path.join(run_dir, f'{mode}_{r}.err')).read()}" |
| for r in range(WS) |
| if os.path.isfile(os.path.join(run_dir, f"{mode}_{r}.err")) |
| ] |
| assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs) |
| assert not errs, "\n".join(errs) |
| return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)] |
|
|
|
|
| def test_adamw8bit_steps_fsdp2_dtensors_with_a_global_norm_and_a_global_skip(tmp_workdir, monkeypatch): |
| """THE FSDP2 DTENSOR DEFECT + THE SHARD-LOCAL NORM DEFECT, for the 8-bit |
| optimizer (bnb's fused kernel died on DTensors: 'mixed torch.Tensor and |
| DTensor'; OffloadAdamW died the same way until `_localize`). |
| |
| (1) Stepping the LOCAL shards of an uneven FSDP2 weight is BIT-IDENTICAL to |
| the same optimizer over plain clones of those shards (DTensor handling |
| changes nothing), and tracks a single-process torch AdamW on the unsharded |
| model on average (mean abs diff < 2e-4 (measured 6e-5) against ~6e-3 of travel; the max is |
| not bounded tightly, see the cold-start note in adamw8bit.py). |
| (2) `last_grad_norm` is the GLOBAL norm and the SAME number on every rank: |
| a rank-local slice norm lets one rank cross a halt tripwire while its peer |
| waits in a collective. (3) A NaN on one rank's shard makes both ranks skip. |
| (4) State really is shard-shaped and quantized on at least one rank.""" |
| rows = _run("dtensor", tmp_workdir, monkeypatch) |
| for r in rows: |
| assert r["plain_equal"], "stepping DTensor shards != stepping plain clones of them" |
| assert r["mean_param_diff"] < 2e-4, r["mean_param_diff"] |
| |
| |
| assert r["norms"][0] == pytest.approx(r["ref_norms"][0], rel=1e-5) |
| assert r["norms"] == pytest.approx(r["ref_norms"], rel=2e-3), (r["norms"], r["ref_norms"]) |
| assert r["all_rank_norms"][0] == r["all_rank_norms"][1], r["all_rank_norms"] |
| assert r["stepped_on_nan"] is False and r["skipped"] == 1, r |
| assert r["skip_diff"] == 0.0, "a rank stepped although its peer's gradient was non-finite" |
| assert r["sd_world_size"] == WS |
| assert r["quantized"][0] is True, "the weight shard must take the 8-bit path" |
| assert rows[0]["local_numel"] != rows[1]["local_numel"], "129 rows must shard unevenly" |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _shipped_configs_naming(name): |
| import glob |
| import os |
|
|
| from align.config import load_config |
| from tests.conftest import REPO_ROOT |
|
|
| out = [] |
| for path in sorted(glob.glob(os.path.join(REPO_ROOT, "configs", "*.toml"))): |
| try: |
| cfg = load_config(path) |
| except Exception: |
| continue |
| if getattr(getattr(cfg, "optimizer", None), "name", None) == name: |
| out.append((os.path.basename(path), cfg)) |
| return out |
|
|
|
|
| def test_every_shipped_adamw8bit_config_constructs_the_real_optimizer_and_steps(): |
| """THE LAUNCH BLOCKER (review_trainer.md): every shipped full-FT config names adamw8bit and |
| make_optimizer used to raise NotImplementedError for it, so no stage could construct an |
| optimizer. Build from the REAL shipped config on a tiny stub model, and take a step. Mutant: |
| restore the raise (or return OffloadAdamW) -> the isinstance check fails.""" |
| from align.train.adamw8bit import BlockwiseAdamW8bit |
| from align.train.offload_adamw import make_optimizer |
|
|
| found = _shipped_configs_naming("adamw8bit") |
| assert found, "no shipped config names adamw8bit; this test would be vacuous" |
| assert "sft.toml" in [n for n, _ in found] |
| for fname, cfg in found: |
| params = [torch.nn.Parameter(torch.randn(n).bfloat16()) for n in (4096 * 3, 100)] |
| opt = make_optimizer(cfg, params) |
| assert isinstance(opt, BlockwiseAdamW8bit), (fname, type(opt)) |
| assert "adamw8bit" in opt.banner_tag and "zero host RSS" in opt.banner_tag |
| before = [p.detach().clone() for p in params] |
| for p in params: |
| p.grad = torch.randn_like(p) * 1e-2 |
| assert opt.step() is True |
| assert any(not torch.equal(a, b.detach()) for a, b in zip(before, params)), fname |
|
|
|
|
| def test_optimizer_config_max_grad_norm_reaches_the_optimizer_and_is_recorded_pre_clip(): |
| """`OptimizerConfig.max_grad_norm` (1.0 by default since G0; None disables clipping explicitly) |
| must reach the optimizer, which clips by the GLOBAL norm and records the PRE-clip |
| norm. Mutant: make_optimizer drops the field -> the param-group assertion fails (the clip |
| arithmetic itself is pinned bitwise by the grad-norm test above).""" |
| from align.config import OptimizerConfig, RunConfig |
| from align.train.offload_adamw import make_optimizer |
|
|
| assert OptimizerConfig().max_grad_norm == 1.0 |
|
|
| def build(clip): |
| cfg = RunConfig() |
| cfg.optimizer.name = "adamw8bit" |
| cfg.optimizer.max_grad_norm = clip |
| torch.manual_seed(0) |
| p = torch.nn.Parameter(torch.randn(8192).bfloat16()) |
| return p, make_optimizer(cfg, [p]) |
|
|
| p_clip, o_clip = build(1.0) |
| p_free, o_free = build(None) |
| g = torch.randn(8192) * 100.0 |
| for p in (p_clip, p_free): |
| p.grad = g.bfloat16().clone() |
| o_clip.step() |
| o_free.step() |
| assert o_clip.last_grad_norm == pytest.approx(o_free.last_grad_norm, rel=1e-6) |
| assert o_clip.last_grad_norm > 1000.0, "must record the PRE-clip norm, not <= max_grad_norm" |
| |
| assert o_clip.param_groups[0]["max_grad_norm"] == 1.0 |
| assert o_free.param_groups[0]["max_grad_norm"] is None |
|
|
|
|
| def test_train_config_per_rank_optimizer_state_is_settable_from_struct_mode_toml(tmp_path): |
| """The Trainer read this flag with getattr because the field did not exist, so struct-mode TOML |
| (unknown keys raise) could not turn it on. It must parse, default OFF, and round-trip.""" |
| from align.config import TrainConfig, load_config |
|
|
| assert TrainConfig().per_rank_optimizer_state is False |
| f = tmp_path / "c.toml" |
| f.write_text("[train]\nper_rank_optimizer_state = true\n") |
| assert load_config(str(f)).train.per_rank_optimizer_state is True |
|
|
|
|
| def test_train_module_default_extra_state_is_an_empty_noop(): |
| """Subclasses (GRPO tripwire) override the pair; the base must exist so the Trainer's resume.pt |
| always carries the key and `load_extra_state` never needs a getattr guard to be safe.""" |
| from align.train.train_module import TrainModule |
|
|
| m = TrainModule.__new__(TrainModule) |
| assert m.extra_state() == {} |
| assert m.load_extra_state({}) is None |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_second_moment_overflow_cannot_poison_the_codes_with_nan(): |
| """THE V-OVERFLOW POISON. The finiteness gate tests the float64 NORM, which stays finite for |
| a gradient whose fp32 SQUARE is inf: one element of 1e21 gives v = (1-beta2) * g^2 = inf. |
| torch.optim.AdamW survives (m / (inf + eps) = 0); here inf became the block's absmax, the |
| requantize divided inf/inf = NaN into the uint8 codes, and the block's weights were NaN one |
| step LATER with the gate still seeing only finite gradients. |
| |
| Kills: removing the `v_full.clamp_(max=fp32_max)` before requantization.""" |
| p = torch.nn.Parameter(_randn(B * 4, seed=1)) |
| ref = torch.nn.Parameter(p.detach().clone()) |
| opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95)) |
| ropt = torch.optim.AdamW([ref], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0) |
| for step in range(3): |
| g = _randn(B * 4, seed=10 + step) |
| if step == 0: |
| g[5] = 1e21 |
| p.grad, ref.grad = g.clone(), g.clone() |
| assert opt.step() is True |
| ropt.step() |
| st = opt._state[0] |
| assert torch.isfinite(st["v_absmax"]).all() and torch.isfinite(st["m_absmax"]).all(), step |
| assert torch.isfinite(p).all(), f"weights went non-finite at step {step}" |
| |
| |
| assert abs(p[5].item() - ref[5].item()) < 50 * 1e-3 |
|
|
|
|
| def test_load_state_dict_rewinds_the_rounding_rng_of_an_optimizer_that_already_stepped(): |
| """`load_state_dict` promises a byte-for-byte restore. The stochastic-rounding generator is |
| lazily created, so a state saved BEFORE the first bf16 step carries no `sr_rng`; loading it |
| into an optimizer that has since stepped left that optimizer's generator advanced, and the |
| 'rewound' run rounded differently from a never-stepped one. |
| |
| Kills: not clearing `_generators` on load.""" |
| torch.manual_seed(0) |
| base = (torch.randn(B * 3) * 0.05).to(torch.bfloat16) |
|
|
| def mk(): |
| return torch.nn.Parameter(base.clone()) |
|
|
| fresh_p = mk() |
| sd0 = BlockwiseAdamW8bit([fresh_p], lr=1e-3).state_dict() |
| assert sd0["sr_rng"] == {}, "precondition: nothing stepped, so no generator to save" |
|
|
| used_p, ref_p = mk(), mk() |
| used = BlockwiseAdamW8bit([used_p], lr=1e-3) |
| ref = BlockwiseAdamW8bit([ref_p], lr=1e-3) |
| for s in range(3): |
| _set_grads([used_p], s) |
| used.step() |
| used_p.data.copy_(base) |
| used.load_state_dict(sd0) |
| for s in range(4): |
| _set_grads([used_p], s) |
| _set_grads([ref_p], s) |
| used.step() |
| ref.step() |
| assert torch.equal(used_p, ref_p) |
|
|
|
|
| def test_persistent_tiny_gradient_overstep_stays_inside_the_documented_bound(): |
| """The second-moment underflow is PERMANENT for an element whose gradient stays below |
| ~5e-4 of its block's largest (v is re-floored to 0 every step), not the transient the module |
| docstring used to claim: measured against torch.optim.AdamW with constant g = 1e-4 of the |
| block max, 300 steps, the 8-bit travel is 3.9x at beta2 = 0.95 and 9.9x at 0.999. What must |
| hold regardless of how the owner resolves that is the per-step ceiling lr / sqrt(1 - beta2): |
| the update comes from the FRESH moment, so a long run of underflowed steps can never turn |
| into the m / (0 + eps) = 100 lr step. |
| |
| Kills: computing the update from the requantized (stored, zero) v, at ANY step; a floor |
| that is applied to the fresh moment.""" |
| lr, beta2 = 1e-3, 0.95 |
| p = torch.nn.Parameter(torch.zeros(B * 2)) |
| opt = BlockwiseAdamW8bit([p], lr=lr, betas=(0.9, beta2), weight_decay=0.0) |
| g = torch.full((B * 2,), 1e-4) |
| g[0] = 1.0 |
| worst = 0.0 |
| for _ in range(200): |
| before = p.detach().clone() |
| p.grad = g.clone() |
| opt.step() |
| worst = max(worst, (p.detach() - before)[1:].abs().max().item()) |
| assert (opt._state[0]["v_codes"][0, 1:] == 0).all(), "precondition: v really is underflowed" |
| assert worst <= lr / math.sqrt(1 - beta2) * 1.05, worst / lr |
|
|
|
|
| def _body_dtensor_clip(rank: int, run_dir: str) -> dict: |
| """Child body: a 1-row layer (rank 1 holds an EMPTY shard), an uneven 129-row layer, in-optimizer |
| global clipping and weight decay, against torch clip_grad_norm_ + AdamW on the unsharded model.""" |
| from torch.distributed.device_mesh import init_device_mesh |
| from torch.distributed.fsdp import fully_shard |
|
|
| class Net(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.a = torch.nn.Linear(64, 1) |
| self.b = torch.nn.Linear(64, 129) |
|
|
| def forward(self, x): |
| return torch.cat([self.a(x), self.b(x)], 1) |
|
|
| torch.manual_seed(0) |
| ref = Net() |
| start = copy.deepcopy(ref) |
| sharded = copy.deepcopy(ref) |
| mesh = init_device_mesh("cpu", (WS,)) |
| fully_shard(sharded.a, mesh=mesh) |
| fully_shard(sharded.b, mesh=mesh) |
| fully_shard(sharded, mesh=mesh) |
| g = torch.Generator().manual_seed(3) |
| x, y = torch.randn(128, 64, generator=g), torch.randn(128, 130, generator=g) * 10 |
|
|
| hp = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1) |
| ref_opt = torch.optim.AdamW(ref.parameters(), **hp) |
| opt = BlockwiseAdamW8bit(sharded.parameters(), max_grad_norm=0.5, **hp) |
| rel_norm_err, clipped = [], 0 |
| for _ in range(8): |
| ((ref(x) - y) ** 2).mean().backward() |
| total = float(torch.nn.utils.clip_grad_norm_(ref.parameters(), 0.5)) |
| clipped += total > 0.5 |
| ref_opt.step() |
| ref_opt.zero_grad() |
| ((sharded(x) - y) ** 2).mean().backward() |
| assert opt.step() is True |
| opt.zero_grad() |
| rel_norm_err.append(abs(opt.last_grad_norm - total) / total) |
| with torch.no_grad(): |
| diff = max( |
| float((p.full_tensor() - r).abs().mean()) for p, r in zip(sharded.parameters(), ref.parameters()) |
| ) |
| travel = max(float((r - r0).abs().mean()) for r, r0 in zip(ref.parameters(), start.parameters())) |
| sd = opt.state_dict() |
| opt.load_state_dict(sd) |
| return { |
| "local_numel": [int(p.to_local().numel()) for p in sharded.parameters()], |
| "rel_norm_err": rel_norm_err, |
| "clipped_steps": int(clipped), |
| "mean_diff": diff, |
| "travel": travel, |
| } |
|
|
|
|
| _BODIES["dtensor_clip"] = _body_dtensor_clip |
|
|
|
|
| def test_adamw8bit_global_clip_and_empty_shard_match_torch_on_two_ranks(tmp_workdir, monkeypatch): |
| """The in-optimizer clip must use the GLOBAL norm: a shard-local norm is ~1/sqrt(world) of the |
| true one, so the clip would fire late (or never) and the two ranks would scale differently. A |
| rank whose shard of a 1-row layer is EMPTY (0 local elements) must still join every |
| collective, skip its own update, and round-trip its state. |
| |
| Kills: `coef` computed from a rank-local norm; skipping `_grad_norm`'s all-reduce; treating a |
| zero-element shard as 'no gradient' (the 'no gradients on any parameter' raise would hang the |
| peer in its all-reduce).""" |
| rows = _run("dtensor_clip", tmp_workdir, monkeypatch) |
| assert rows[1]["local_numel"][:2] == [0, 0], "rank 1 must hold the empty shard of the 1-row layer" |
| for r in rows: |
| assert r["clipped_steps"] >= 6, "the clip must actually be active in this test" |
| assert max(r["rel_norm_err"]) < 1e-3, r["rel_norm_err"] |
| assert r["mean_diff"] < 0.02 * r["travel"], (r["mean_diff"], r["travel"]) |
|
|
