File size: 2,310 Bytes
932bc69 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """Regression: BF16 counters above 256 can stop incrementing after resume."""
import json
import torch
from ngram_train_resume_safe import restore_adamw_steps
parameter = torch.nn.Parameter(torch.ones(3))
optimizer = torch.optim.AdamW([parameter], lr=0.01)
parameter.sum().backward()
optimizer.step()
counter = torch.tensor(1001.0, dtype=torch.bfloat16)
before = counter.item()
counter.add_(1)
assert counter.item() == before == 1000.0
optimizer.state[parameter]["step"] = counter
assert restore_adamw_steps([optimizer], 1001) == 1
assert optimizer.state[parameter]["step"].dtype == torch.float32
optimizer.zero_grad()
parameter.sum().backward()
optimizer.step()
assert optimizer.state[parameter]["step"].item() == 1002
optimizer.state[parameter]["step"] = torch.tensor(900.0)
try:
restore_adamw_steps([optimizer], 1001)
except ValueError:
pass
else:
raise AssertionError("Mismatched counters must not be silently overwritten")
# The real checkpointer saves BF16 then promotes all floating states to the
# model dtype on load. Large, non-BF16-exact training positions must still pass.
checkpoint_step = 25680
saved = optimizer.state_dict()
saved["state"][0]["step"] = torch.tensor(checkpoint_step, dtype=torch.bfloat16).float()
optimizer.load_state_dict(saved)
assert optimizer.state[parameter]["step"].dtype == torch.float32
assert optimizer.state[parameter]["step"].item() == 25728
moments = {
name: optimizer.state[parameter][name].clone()
for name in ("exp_avg", "exp_avg_sq")
}
assert restore_adamw_steps([optimizer], checkpoint_step) == 1
assert optimizer.state[parameter]["step"].item() == checkpoint_step
for name, value in moments.items():
torch.testing.assert_close(optimizer.state[parameter][name], value, rtol=0, atol=0)
optimizer.zero_grad()
parameter.sum().backward()
optimizer.step()
assert optimizer.state[parameter]["step"].item() == checkpoint_step + 1
print(
json.dumps(
{
"passed": True,
"legacy_bf16_counter_stalled_at": before,
"repaired_fp32_counter_after_update": 1002,
"mismatched_counter_rejected": True,
"bf16_promoted_to_fp32_roundtrip_passed": True,
"moments_unchanged_by_repair": True,
"large_counter_after_update": checkpoint_step + 1,
}
)
)
|