From 8458ad23e518291f1fcef028af187e4838dc8a99 Mon Sep 17 00:00:00 2001 From: dlyr3 <152425764+dlyr3@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:58:32 +0200 Subject: [PATCH 1/6] test & benchmark for reproducer --- benchmarks/bench_psgd_fp64.py | 105 ++++++++++++++++++++++++++++++++++ test/test_psgd_lb_dtype.py | 84 +++++++++++++++++++++++++++ 2 files changed, 189 insertions(+) create mode 100644 benchmarks/bench_psgd_fp64.py create mode 100644 test/test_psgd_lb_dtype.py diff --git a/benchmarks/bench_psgd_fp64.py b/benchmarks/bench_psgd_fp64.py new file mode 100644 index 0000000..25f1c38 --- /dev/null +++ b/benchmarks/bench_psgd_fp64.py @@ -0,0 +1,105 @@ +"""Reproducer for FP64 promotion in PSGD preconditioner lower bound. + +Run BEFORE and AFTER the fix to compare: + python benchmarks/bench_psgd_fp64.py --tag before + python benchmarks/bench_psgd_fp64.py --tag after + +Outputs: + bench_psgd_{tag}.json Perfetto trace (open in chrome://tracing or ui.perfetto.dev) + bench_psgd_{tag}.txt Profiler key_averages table +""" + +import argparse +import os +import time + +import torch +from torch import nn +from torch.profiler import ProfilerActivity, profile, schedule + +import heavyball + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--tag", default="run", help="Tag for output files (e.g. 'before' or 'after')") + parser.add_argument("--dim", type=int, default=2048, help="Linear layer dimension") + parser.add_argument("--warmup", type=int, default=10, help="Warmup steps (not profiled)") + parser.add_argument("--steps", type=int, default=20, help="Profiled steps") + args = parser.parse_args() + + out_dir = os.path.dirname(os.path.abspath(__file__)) + json_path = os.path.join(out_dir, f"bench_psgd_{args.tag}.json") + txt_path = os.path.join(out_dir, f"bench_psgd_{args.tag}.txt") + + torch.manual_seed(42) + model = nn.Linear(args.dim, args.dim, bias=False, device="cuda", dtype=torch.bfloat16) + opt = heavyball.PSGDPRO(model.parameters(), lr=1e-3) + data = torch.randn(64, args.dim, device="cuda", dtype=torch.bfloat16) + target = torch.randn(64, args.dim, device="cuda", dtype=torch.bfloat16) + + for _ in range(args.warmup): + loss = ((model(data) - target) ** 2).mean() + loss.backward() + opt.step() + opt.zero_grad() + + torch.cuda.synchronize() + + wall_times = [] + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=schedule(wait=0, warmup=0, active=args.steps), + record_shapes=True, + with_flops=True, + ) as prof: + for _ in range(args.steps): + t0 = time.perf_counter() + loss = ((model(data) - target) ** 2).mean() + loss.backward() + opt.step() + opt.zero_grad() + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + prof.step() + + prof.export_chrome_trace(json_path) + + table = prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=30) + with open(txt_path, "w") as f: + f.write(table) + + kernel_events = [e for e in prof.key_averages() if e.device_type == torch.autograd.DeviceType.CUDA] + fp64_kernels = [e for e in kernel_events if "d884" in e.key] + total_cuda = sum(e.self_cuda_time_total for e in kernel_events) + + print(f"\n{'=' * 60}") + print(f"PSGD FP64 Benchmark [{args.tag}]") + print(f"{'=' * 60}") + print(f"Model: Linear({args.dim}, {args.dim}) bfloat16") + print(f"Steps: {args.steps} (after {args.warmup} warmup)") + print(f"Wall time/step: {sum(wall_times) / len(wall_times) * 1000:.2f}ms " + f"(+/- {torch.tensor(wall_times).std().item() * 1000:.2f}ms)") + print(f"Total CUDA: {total_cuda / 1000:.2f}ms") + print() + + if fp64_kernels: + fp64_time = sum(e.self_cuda_time_total for e in fp64_kernels) + print(f"FP64 GEMM (d884) FOUND: {len(fp64_kernels)} kernel type(s), " + f"{fp64_time / 1000:.2f}ms ({fp64_time / total_cuda * 100:.1f}% of GPU time)") + else: + print("FP64 GEMM (d884): NONE -- fix is working") + + print(f"\nTop 5 GPU kernels by self CUDA time:") + sorted_kernels = sorted(kernel_events, key=lambda e: e.self_cuda_time_total, reverse=True) + for e in sorted_kernels[:5]: + pct = e.self_cuda_time_total / total_cuda * 100 if total_cuda else 0 + print(f" {pct:5.1f}% | {e.self_cuda_time_total / 1000:8.2f}ms | {e.count:4}x | {e.key[:90]}") + + print(f"\nTraces saved to:") + print(f" {json_path}") + print(f" {txt_path}") + + +if __name__ == "__main__": + main() diff --git a/test/test_psgd_lb_dtype.py b/test/test_psgd_lb_dtype.py new file mode 100644 index 0000000..7fab5ab --- /dev/null +++ b/test/test_psgd_lb_dtype.py @@ -0,0 +1,84 @@ +"""Test that PSGD preconditioner lower bound uses fp32, not fp64. + +Verifies: +1. Convergence is not degraded with fp32 lower bound +2. No fp64 tensors leak into optimizer state +""" + +import pytest +import torch +from torch import nn + +import heavyball +from heavyball.utils import clean, set_torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +_PSGD_OPTIMIZERS = [ + (heavyball.PSGDPRO, 1e-3, {}), + (heavyball.PSGDKron, 1e-3, {}), +] + + +def _train(model, opt, data, target, steps): + losses = [] + for _ in range(steps): + p = next(model.parameters()) + d = data.to(p.dtype) if p.dtype != data.dtype else data + loss = ((model(d) - target.to(d.dtype)) ** 2).mean().float() + loss.backward() + opt.step() + opt.zero_grad() + losses.append(loss.item()) + return losses + + +@pytest.mark.parametrize("opt_cls,lr,extra_kw", _PSGD_OPTIMIZERS, ids=[t[0].__name__ for t in _PSGD_OPTIMIZERS]) +def test_psgd_convergence_fp32_lb(opt_cls, lr, extra_kw): + set_torch() + torch.manual_seed(42) + data = torch.randn(128, 64, device="cuda") + target = torch.randn(128, 32, device="cuda") + + torch.manual_seed(0) + model = nn.Linear(64, 32, bias=False, device="cuda") + opt = opt_cls(model.parameters(), lr=lr, **extra_kw) + + losses = _train(model, opt, data, target, 200) + assert losses[-1] < losses[0] * 0.5, f"Loss did not converge: {losses[0]:.4f} -> {losses[-1]:.4f}" + + del model, opt + clean() + + +@pytest.mark.parametrize("opt_cls,lr,extra_kw", _PSGD_OPTIMIZERS, ids=[t[0].__name__ for t in _PSGD_OPTIMIZERS]) +def test_psgd_no_fp64_in_state(opt_cls, lr, extra_kw): + set_torch() + torch.manual_seed(42) + data = torch.randn(32, 64, device="cuda") + target = torch.randn(32, 32, device="cuda") + + torch.manual_seed(0) + model = nn.Linear(64, 32, bias=False, device="cuda") + opt = opt_cls(model.parameters(), lr=lr, **extra_kw) + + _train(model, opt, data, target, 5) + + for param in model.parameters(): + state = opt.state[param] + for key, val in state.items(): + if isinstance(val, torch.Tensor) and val.is_floating_point(): + if key == "step": + continue + assert val.dtype != torch.float64, ( + f"fp64 tensor found in optimizer state: key={key!r}, dtype={val.dtype}" + ) + if isinstance(val, list): + for i, v in enumerate(val): + if isinstance(v, torch.Tensor) and v.is_floating_point(): + assert v.dtype != torch.float64, ( + f"fp64 tensor found in optimizer state: key={key!r}[{i}], dtype={v.dtype}" + ) + + del model, opt + clean() From df06e1391e9f2f2e3bcccfa9649efdad4b8eb1b6 Mon Sep 17 00:00:00 2001 From: dlyr3 <152425764+dlyr3@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:46:47 +0200 Subject: [PATCH 2/6] fix --- heavyball/chainable.py | 6 +++--- heavyball/utils.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/heavyball/chainable.py b/heavyball/chainable.py index b3c13ba..da77234 100644 --- a/heavyball/chainable.py +++ b/heavyball/chainable.py @@ -1038,7 +1038,7 @@ def _init_psgd_kron(state, group, update, grad, param, cached: bool = False, pro dtype=getattr(torch, group["q_dtype"]), ) state["Q"] = utils.triu_to_line(Q) if group["store_triu_as_line"] else Q - state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float64) for q in Q] + state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float32) for q in Q] state["step"] = torch.zeros((), device=param.device, dtype=torch.float64) if not cached: return @@ -1060,7 +1060,7 @@ def _init_psgd_eigen_kron(state, group, update, grad, param, prob: Optional[call tmp.get("vector"), dtype=getattr(torch, group["q_dtype"]), ) - state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float64) for q in Q] + state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float32) for q in Q] state["step"] = torch.zeros((), device=param.device, dtype=torch.float64) _update_psgd_precond( @@ -1093,7 +1093,7 @@ def _init_psgd_pro_kron(state, group, update, grad, param, cached: bool = False, dtype=getattr(torch, group["q_dtype"]), ) state["Q"] = Q - state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float64) for q in Q] + state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float32) for q in Q] state["step"] = torch.zeros((), device=param.device, dtype=torch.float64) if not cached: return diff --git a/heavyball/utils.py b/heavyball/utils.py index a49297a..e71ef45 100644 --- a/heavyball/utils.py +++ b/heavyball/utils.py @@ -3093,6 +3093,7 @@ def _update_lb(ell: Tensor, lb_state: Tensor, beta: Tensor) -> Tensor: ell = promote(ell) ell = ell.maximum(promote(lb_state) + (ell - promote(lb_state)) * (1 - beta)) copy_stochastic_(lb_state, ell) + ell = ell.float() return ell From 03a3bfea4c3060d2d7022cf42374f7238689ce0c Mon Sep 17 00:00:00 2001 From: dlyr3 <152425764+dlyr3@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:01:54 +0200 Subject: [PATCH 3/6] fix benchmark and test for profiler API compatibility --- benchmarks/bench_psgd_fp64.py | 24 +++++++++++------------- test/test_psgd_lb_dtype.py | 28 ++++++++++++++-------------- 2 files changed, 25 insertions(+), 27 deletions(-) diff --git a/benchmarks/bench_psgd_fp64.py b/benchmarks/bench_psgd_fp64.py index 25f1c38..1fb7c92 100644 --- a/benchmarks/bench_psgd_fp64.py +++ b/benchmarks/bench_psgd_fp64.py @@ -15,7 +15,7 @@ import torch from torch import nn -from torch.profiler import ProfilerActivity, profile, schedule +from torch.profiler import ProfilerActivity, profile import heavyball @@ -34,7 +34,7 @@ def main(): torch.manual_seed(42) model = nn.Linear(args.dim, args.dim, bias=False, device="cuda", dtype=torch.bfloat16) - opt = heavyball.PSGDPRO(model.parameters(), lr=1e-3) + opt = heavyball.PSGDPRO(model.parameters(), lr=1e-3, compile_step=False) data = torch.randn(64, args.dim, device="cuda", dtype=torch.bfloat16) target = torch.randn(64, args.dim, device="cuda", dtype=torch.bfloat16) @@ -49,7 +49,6 @@ def main(): wall_times = [] with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - schedule=schedule(wait=0, warmup=0, active=args.steps), record_shapes=True, with_flops=True, ) as prof: @@ -61,17 +60,16 @@ def main(): opt.zero_grad() torch.cuda.synchronize() wall_times.append(time.perf_counter() - t0) - prof.step() prof.export_chrome_trace(json_path) - table = prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=30) + table = prof.key_averages().table(sort_by="self_device_time_total", row_limit=30) with open(txt_path, "w") as f: f.write(table) - kernel_events = [e for e in prof.key_averages() if e.device_type == torch.autograd.DeviceType.CUDA] + kernel_events = [e for e in prof.key_averages() if e.self_device_time_total > 0] fp64_kernels = [e for e in kernel_events if "d884" in e.key] - total_cuda = sum(e.self_cuda_time_total for e in kernel_events) + total_cuda = sum(e.self_device_time_total for e in kernel_events) print(f"\n{'=' * 60}") print(f"PSGD FP64 Benchmark [{args.tag}]") @@ -80,21 +78,21 @@ def main(): print(f"Steps: {args.steps} (after {args.warmup} warmup)") print(f"Wall time/step: {sum(wall_times) / len(wall_times) * 1000:.2f}ms " f"(+/- {torch.tensor(wall_times).std().item() * 1000:.2f}ms)") - print(f"Total CUDA: {total_cuda / 1000:.2f}ms") + print(f"Total GPU: {total_cuda / 1000:.2f}ms") print() if fp64_kernels: - fp64_time = sum(e.self_cuda_time_total for e in fp64_kernels) + fp64_time = sum(e.self_device_time_total for e in fp64_kernels) print(f"FP64 GEMM (d884) FOUND: {len(fp64_kernels)} kernel type(s), " f"{fp64_time / 1000:.2f}ms ({fp64_time / total_cuda * 100:.1f}% of GPU time)") else: print("FP64 GEMM (d884): NONE -- fix is working") - print(f"\nTop 5 GPU kernels by self CUDA time:") - sorted_kernels = sorted(kernel_events, key=lambda e: e.self_cuda_time_total, reverse=True) + print(f"\nTop 5 GPU kernels by self device time:") + sorted_kernels = sorted(kernel_events, key=lambda e: e.self_device_time_total, reverse=True) for e in sorted_kernels[:5]: - pct = e.self_cuda_time_total / total_cuda * 100 if total_cuda else 0 - print(f" {pct:5.1f}% | {e.self_cuda_time_total / 1000:8.2f}ms | {e.count:4}x | {e.key[:90]}") + pct = e.self_device_time_total / total_cuda * 100 if total_cuda else 0 + print(f" {pct:5.1f}% | {e.self_device_time_total / 1000:8.2f}ms | {e.count:4}x | {e.key[:90]}") print(f"\nTraces saved to:") print(f" {json_path}") diff --git a/test/test_psgd_lb_dtype.py b/test/test_psgd_lb_dtype.py index 7fab5ab..1643fa4 100644 --- a/test/test_psgd_lb_dtype.py +++ b/test/test_psgd_lb_dtype.py @@ -51,6 +51,19 @@ def test_psgd_convergence_fp32_lb(opt_cls, lr, extra_kw): clean() +def _check_no_fp64(obj, path=""): + if isinstance(obj, torch.Tensor) and obj.is_floating_point(): + if "step" in path: + return + assert obj.dtype != torch.float64, f"fp64 tensor at {path}: dtype={obj.dtype}" + elif isinstance(obj, dict): + for k, v in obj.items(): + _check_no_fp64(v, f"{path}.{k}" if path else str(k)) + elif isinstance(obj, (list, tuple)): + for i, v in enumerate(obj): + _check_no_fp64(v, f"{path}[{i}]") + + @pytest.mark.parametrize("opt_cls,lr,extra_kw", _PSGD_OPTIMIZERS, ids=[t[0].__name__ for t in _PSGD_OPTIMIZERS]) def test_psgd_no_fp64_in_state(opt_cls, lr, extra_kw): set_torch() @@ -65,20 +78,7 @@ def test_psgd_no_fp64_in_state(opt_cls, lr, extra_kw): _train(model, opt, data, target, 5) for param in model.parameters(): - state = opt.state[param] - for key, val in state.items(): - if isinstance(val, torch.Tensor) and val.is_floating_point(): - if key == "step": - continue - assert val.dtype != torch.float64, ( - f"fp64 tensor found in optimizer state: key={key!r}, dtype={val.dtype}" - ) - if isinstance(val, list): - for i, v in enumerate(val): - if isinstance(v, torch.Tensor) and v.is_floating_point(): - assert v.dtype != torch.float64, ( - f"fp64 tensor found in optimizer state: key={key!r}[{i}], dtype={v.dtype}" - ) + _check_no_fp64(opt.state[param]) del model, opt clean() From e2ae8d736c3d1f11677606fd39a1f55bba9adb0a Mon Sep 17 00:00:00 2001 From: dlyr3 <152425764+dlyr3@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:21:55 +0200 Subject: [PATCH 4/6] prod bench + convergence tests with plot gen --- benchmarks/bench_psgd_prod.py | 135 +++++++++++++++++++++++++++++++++ benchmarks/convergence_test.py | 70 +++++++++++++++++ benchmarks/plot_convergence.py | 61 +++++++++++++++ 3 files changed, 266 insertions(+) create mode 100644 benchmarks/bench_psgd_prod.py create mode 100644 benchmarks/convergence_test.py create mode 100644 benchmarks/plot_convergence.py diff --git a/benchmarks/bench_psgd_prod.py b/benchmarks/bench_psgd_prod.py new file mode 100644 index 0000000..48ccf62 --- /dev/null +++ b/benchmarks/bench_psgd_prod.py @@ -0,0 +1,135 @@ +"""Benchmark PSGD FP64 lower-bound fix on a production-like model. + +An alternative to the toy benchmark bench_psgd_fp64.py + - 2-layer MLP: Linear(512,512) -> ReLU -> Linear(512,512) + - bfloat16, batch=64, bias=True + - PSGDPRO with default settings (compile_step default, max-autotune) + - MSE loss + +Run BEFORE and AFTER the fix: + python benchmarks/bench_psgd_prod.py --tag before + python benchmarks/bench_psgd_prod.py --tag after +""" + +import argparse +import json +import os +import time + +import torch +from torch import nn +from torch.profiler import ProfilerActivity, profile, record_function, schedule + +import heavyball + +def make_model(dim, device): + return nn.Sequential( + nn.Linear(dim, dim, bias=True, device=device, dtype=torch.bfloat16), + nn.ReLU(), + nn.Linear(dim, dim, bias=True, device=device, dtype=torch.bfloat16), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--tag", default="run") + parser.add_argument("--dim", type=int, default=512) + parser.add_argument("--batch", type=int, default=64) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--steps", type=int, default=5) + args = parser.parse_args() + + out_dir = os.path.dirname(os.path.abspath(__file__)) + json_path = os.path.join(out_dir, f"prod_{args.tag}.json") + txt_path = os.path.join(out_dir, f"prod_{args.tag}.txt") + loss_path = os.path.join(out_dir, f"prod_{args.tag}_losses.json") + + torch.manual_seed(42) + model = make_model(args.dim, "cuda") + opt = heavyball.PSGDPRO(model.parameters(), lr=1e-3) + data = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) + target = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) + + for _ in range(args.warmup): + with record_function("fwdbwd"): + with record_function("forward"): + out = model(data) + loss = nn.functional.mse_loss(out, target) + with record_function("backward"): + loss.backward() + with record_function("optimizer_step"): + with record_function("psgd_pro_step"): + opt.step() + opt.zero_grad() + + torch.cuda.synchronize() + + wall_times = [] + losses = [] + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=schedule(wait=1, warmup=1, active=args.steps), + record_shapes=True, + with_flops=True, + on_trace_ready=lambda p: p.export_chrome_trace(json_path), + ) as prof: + for _ in range(1 + 1 + args.steps): + t0 = time.perf_counter() + with record_function("fwdbwd"): + with record_function("forward"): + out = model(data) + loss = nn.functional.mse_loss(out, target) + with record_function("backward"): + loss.backward() + with record_function("optimizer_step"): + with record_function("psgd_pro_step"): + opt.step() + opt.zero_grad() + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + losses.append(loss.item()) + prof.step() + + active_times = wall_times[2:] + active_losses = losses[2:] + + table = prof.key_averages().table(sort_by="self_device_time_total", row_limit=30) + with open(txt_path, "w") as f: + f.write(table) + + with open(loss_path, "w") as f: + json.dump({"wall_times": active_times, "losses": active_losses, "tag": args.tag}, f, indent=2) + + kernel_events = [e for e in prof.key_averages() if e.self_device_time_total > 0] + fp64_kernels = [e for e in kernel_events if "d884" in e.key] + total_cuda = sum(e.self_device_time_total for e in kernel_events) + + print(f"\n{'=' * 60}") + print(f"PSGD Prod Benchmark [{args.tag}]") + print(f"{'=' * 60}") + print(f"Model: MLP(Linear({args.dim},{args.dim}) -> ReLU -> Linear({args.dim},{args.dim})) bf16") + print(f"Batch: {args.batch}") + print(f"Steps: {args.steps} (after {args.warmup} warmup)") + print(f"Wall time/step: {sum(active_times) / len(active_times) * 1000:.2f}ms") + print(f"Total GPU: {total_cuda / 1000:.2f}ms") + + if fp64_kernels: + fp64_time = sum(e.self_device_time_total for e in fp64_kernels) + print(f"\nFP64 GEMM (d884): {len(fp64_kernels)} type(s), " + f"{fp64_time / 1000:.2f}ms ({fp64_time / total_cuda * 100:.1f}% of GPU time)") + else: + print("\nFP64 GEMM (d884): NONE") + + print(f"\nTop 10 GPU kernels:") + sorted_kernels = sorted(kernel_events, key=lambda e: e.self_device_time_total, reverse=True) + for e in sorted_kernels[:10]: + pct = e.self_device_time_total / total_cuda * 100 if total_cuda else 0 + print(f" {pct:5.1f}% | {e.self_device_time_total / 1000:8.2f}ms | {e.count:4}x | {e.key[:80]}") + + print(f"\nOutputs: {json_path}") + print(f" {txt_path}") + print(f" {loss_path}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/convergence_test.py b/benchmarks/convergence_test.py new file mode 100644 index 0000000..8287302 --- /dev/null +++ b/benchmarks/convergence_test.py @@ -0,0 +1,70 @@ +"""Convergence test for PSGD FP64 lower-bound fix. + +Trains a model for N steps and logs per-step loss to JSON. +Run before and after the fix with the same seed to produce +comparable convergence curves. + + python benchmarks/convergence_test.py --tag before --steps 500 + python benchmarks/convergence_test.py --tag after --steps 500 +""" + +import argparse +import json +import os + +import torch +from torch import nn + +import heavyball + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--tag", default="run") + parser.add_argument("--dim", type=int, default=512) + parser.add_argument("--batch", type=int, default=64) + parser.add_argument("--steps", type=int, default=500) + parser.add_argument("--lr", type=float, default=1e-3) + args = parser.parse_args() + + out_dir = os.path.dirname(os.path.abspath(__file__)) + out_path = os.path.join(out_dir, f"convergence_{args.tag}.json") + + torch.manual_seed(42) + model = nn.Sequential( + nn.Linear(args.dim, args.dim, bias=True, device="cuda", dtype=torch.bfloat16), + nn.ReLU(), + nn.Linear(args.dim, args.dim, bias=True, device="cuda", dtype=torch.bfloat16), + ) + opt = heavyball.PSGDPRO(model.parameters(), lr=args.lr) + data = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) + target = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) + + losses = [] + for step in range(args.steps): + out = model(data) + loss = nn.functional.mse_loss(out, target).float() + loss.backward() + opt.step() + opt.zero_grad() + losses.append(loss.item()) + if step % 50 == 0 or step == args.steps - 1: + print(f" step {step:4d}/{args.steps} loss={losses[-1]:.6e}") + + result = { + "tag": args.tag, + "dim": args.dim, + "batch": args.batch, + "steps": args.steps, + "lr": args.lr, + "losses": losses, + } + with open(out_path, "w") as f: + json.dump(result, f, indent=2) + + print(f"\nFinal loss: {losses[-1]:.6e}") + print(f"Saved to: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/plot_convergence.py b/benchmarks/plot_convergence.py new file mode 100644 index 0000000..5cc08e2 --- /dev/null +++ b/benchmarks/plot_convergence.py @@ -0,0 +1,61 @@ +"""Generate convergence plot from before/after JSON files. + + python benchmarks/plot_convergence.py --before convergence_before.json --after convergence_after.json --out convergence.png +""" + +import argparse +import json +import os + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--before", required=True) + parser.add_argument("--after", required=True) + parser.add_argument("--out", default="convergence.png") + args = parser.parse_args() + + with open(args.before) as f: + before = json.load(f) + with open(args.after) as f: + after = json.load(f) + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) + + ax1.plot(before["losses"], label="before (fp64 lower bound)", + color="#1f77b4", alpha=0.9, linewidth=2.5) + ax1.plot(after["losses"], label="after (fp32 lower bound)", + color="#d62728", alpha=0.7, linewidth=1.5, linestyle="--") + ax1.set_xlabel("Step") + ax1.set_ylabel("Loss") + ax1.set_title("Convergence: linear scale") + ax1.legend() + ax1.grid(True, alpha=0.3) + + ax2.semilogy(before["losses"], label="before (fp64 lower bound)", + color="#1f77b4", alpha=0.9, linewidth=2.5) + ax2.semilogy(after["losses"], label="after (fp32 lower bound)", + color="#d62728", alpha=0.7, linewidth=1.5, linestyle="--") + ax2.set_xlabel("Step") + ax2.set_ylabel("Loss (log scale)") + ax2.set_title("Convergence: log scale") + ax2.legend() + ax2.grid(True, alpha=0.3) + + fig.suptitle( + f"PSGD FP64 Fix - Convergence Test\n" + f"MLP(Linear({before['dim']},{before['dim']}) -> ReLU -> Linear({before['dim']},{before['dim']})), " + f"bf16, batch={before['batch']}, lr={before['lr']}", + fontsize=10, + ) + fig.tight_layout() + fig.savefig(args.out, dpi=150, bbox_inches="tight") + print(f"Saved: {args.out}") + + +if __name__ == "__main__": + main() From 601cbc71f2daea9b35ca62e0d11f90cd6efec50f Mon Sep 17 00:00:00 2001 From: dlyr3 <152425764+dlyr3@users.noreply.github.com> Date: Sat, 11 Jul 2026 17:40:26 +0200 Subject: [PATCH 5/6] removed reproducer --- benchmarks/bench_psgd_fp64.py | 103 -------------------------- benchmarks/bench_psgd_prod.py | 135 ---------------------------------- 2 files changed, 238 deletions(-) delete mode 100644 benchmarks/bench_psgd_fp64.py delete mode 100644 benchmarks/bench_psgd_prod.py diff --git a/benchmarks/bench_psgd_fp64.py b/benchmarks/bench_psgd_fp64.py deleted file mode 100644 index 1fb7c92..0000000 --- a/benchmarks/bench_psgd_fp64.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Reproducer for FP64 promotion in PSGD preconditioner lower bound. - -Run BEFORE and AFTER the fix to compare: - python benchmarks/bench_psgd_fp64.py --tag before - python benchmarks/bench_psgd_fp64.py --tag after - -Outputs: - bench_psgd_{tag}.json Perfetto trace (open in chrome://tracing or ui.perfetto.dev) - bench_psgd_{tag}.txt Profiler key_averages table -""" - -import argparse -import os -import time - -import torch -from torch import nn -from torch.profiler import ProfilerActivity, profile - -import heavyball - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--tag", default="run", help="Tag for output files (e.g. 'before' or 'after')") - parser.add_argument("--dim", type=int, default=2048, help="Linear layer dimension") - parser.add_argument("--warmup", type=int, default=10, help="Warmup steps (not profiled)") - parser.add_argument("--steps", type=int, default=20, help="Profiled steps") - args = parser.parse_args() - - out_dir = os.path.dirname(os.path.abspath(__file__)) - json_path = os.path.join(out_dir, f"bench_psgd_{args.tag}.json") - txt_path = os.path.join(out_dir, f"bench_psgd_{args.tag}.txt") - - torch.manual_seed(42) - model = nn.Linear(args.dim, args.dim, bias=False, device="cuda", dtype=torch.bfloat16) - opt = heavyball.PSGDPRO(model.parameters(), lr=1e-3, compile_step=False) - data = torch.randn(64, args.dim, device="cuda", dtype=torch.bfloat16) - target = torch.randn(64, args.dim, device="cuda", dtype=torch.bfloat16) - - for _ in range(args.warmup): - loss = ((model(data) - target) ** 2).mean() - loss.backward() - opt.step() - opt.zero_grad() - - torch.cuda.synchronize() - - wall_times = [] - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=True, - with_flops=True, - ) as prof: - for _ in range(args.steps): - t0 = time.perf_counter() - loss = ((model(data) - target) ** 2).mean() - loss.backward() - opt.step() - opt.zero_grad() - torch.cuda.synchronize() - wall_times.append(time.perf_counter() - t0) - - prof.export_chrome_trace(json_path) - - table = prof.key_averages().table(sort_by="self_device_time_total", row_limit=30) - with open(txt_path, "w") as f: - f.write(table) - - kernel_events = [e for e in prof.key_averages() if e.self_device_time_total > 0] - fp64_kernels = [e for e in kernel_events if "d884" in e.key] - total_cuda = sum(e.self_device_time_total for e in kernel_events) - - print(f"\n{'=' * 60}") - print(f"PSGD FP64 Benchmark [{args.tag}]") - print(f"{'=' * 60}") - print(f"Model: Linear({args.dim}, {args.dim}) bfloat16") - print(f"Steps: {args.steps} (after {args.warmup} warmup)") - print(f"Wall time/step: {sum(wall_times) / len(wall_times) * 1000:.2f}ms " - f"(+/- {torch.tensor(wall_times).std().item() * 1000:.2f}ms)") - print(f"Total GPU: {total_cuda / 1000:.2f}ms") - print() - - if fp64_kernels: - fp64_time = sum(e.self_device_time_total for e in fp64_kernels) - print(f"FP64 GEMM (d884) FOUND: {len(fp64_kernels)} kernel type(s), " - f"{fp64_time / 1000:.2f}ms ({fp64_time / total_cuda * 100:.1f}% of GPU time)") - else: - print("FP64 GEMM (d884): NONE -- fix is working") - - print(f"\nTop 5 GPU kernels by self device time:") - sorted_kernels = sorted(kernel_events, key=lambda e: e.self_device_time_total, reverse=True) - for e in sorted_kernels[:5]: - pct = e.self_device_time_total / total_cuda * 100 if total_cuda else 0 - print(f" {pct:5.1f}% | {e.self_device_time_total / 1000:8.2f}ms | {e.count:4}x | {e.key[:90]}") - - print(f"\nTraces saved to:") - print(f" {json_path}") - print(f" {txt_path}") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/bench_psgd_prod.py b/benchmarks/bench_psgd_prod.py deleted file mode 100644 index 48ccf62..0000000 --- a/benchmarks/bench_psgd_prod.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Benchmark PSGD FP64 lower-bound fix on a production-like model. - -An alternative to the toy benchmark bench_psgd_fp64.py - - 2-layer MLP: Linear(512,512) -> ReLU -> Linear(512,512) - - bfloat16, batch=64, bias=True - - PSGDPRO with default settings (compile_step default, max-autotune) - - MSE loss - -Run BEFORE and AFTER the fix: - python benchmarks/bench_psgd_prod.py --tag before - python benchmarks/bench_psgd_prod.py --tag after -""" - -import argparse -import json -import os -import time - -import torch -from torch import nn -from torch.profiler import ProfilerActivity, profile, record_function, schedule - -import heavyball - -def make_model(dim, device): - return nn.Sequential( - nn.Linear(dim, dim, bias=True, device=device, dtype=torch.bfloat16), - nn.ReLU(), - nn.Linear(dim, dim, bias=True, device=device, dtype=torch.bfloat16), - ) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--tag", default="run") - parser.add_argument("--dim", type=int, default=512) - parser.add_argument("--batch", type=int, default=64) - parser.add_argument("--warmup", type=int, default=5) - parser.add_argument("--steps", type=int, default=5) - args = parser.parse_args() - - out_dir = os.path.dirname(os.path.abspath(__file__)) - json_path = os.path.join(out_dir, f"prod_{args.tag}.json") - txt_path = os.path.join(out_dir, f"prod_{args.tag}.txt") - loss_path = os.path.join(out_dir, f"prod_{args.tag}_losses.json") - - torch.manual_seed(42) - model = make_model(args.dim, "cuda") - opt = heavyball.PSGDPRO(model.parameters(), lr=1e-3) - data = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) - target = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) - - for _ in range(args.warmup): - with record_function("fwdbwd"): - with record_function("forward"): - out = model(data) - loss = nn.functional.mse_loss(out, target) - with record_function("backward"): - loss.backward() - with record_function("optimizer_step"): - with record_function("psgd_pro_step"): - opt.step() - opt.zero_grad() - - torch.cuda.synchronize() - - wall_times = [] - losses = [] - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - schedule=schedule(wait=1, warmup=1, active=args.steps), - record_shapes=True, - with_flops=True, - on_trace_ready=lambda p: p.export_chrome_trace(json_path), - ) as prof: - for _ in range(1 + 1 + args.steps): - t0 = time.perf_counter() - with record_function("fwdbwd"): - with record_function("forward"): - out = model(data) - loss = nn.functional.mse_loss(out, target) - with record_function("backward"): - loss.backward() - with record_function("optimizer_step"): - with record_function("psgd_pro_step"): - opt.step() - opt.zero_grad() - torch.cuda.synchronize() - wall_times.append(time.perf_counter() - t0) - losses.append(loss.item()) - prof.step() - - active_times = wall_times[2:] - active_losses = losses[2:] - - table = prof.key_averages().table(sort_by="self_device_time_total", row_limit=30) - with open(txt_path, "w") as f: - f.write(table) - - with open(loss_path, "w") as f: - json.dump({"wall_times": active_times, "losses": active_losses, "tag": args.tag}, f, indent=2) - - kernel_events = [e for e in prof.key_averages() if e.self_device_time_total > 0] - fp64_kernels = [e for e in kernel_events if "d884" in e.key] - total_cuda = sum(e.self_device_time_total for e in kernel_events) - - print(f"\n{'=' * 60}") - print(f"PSGD Prod Benchmark [{args.tag}]") - print(f"{'=' * 60}") - print(f"Model: MLP(Linear({args.dim},{args.dim}) -> ReLU -> Linear({args.dim},{args.dim})) bf16") - print(f"Batch: {args.batch}") - print(f"Steps: {args.steps} (after {args.warmup} warmup)") - print(f"Wall time/step: {sum(active_times) / len(active_times) * 1000:.2f}ms") - print(f"Total GPU: {total_cuda / 1000:.2f}ms") - - if fp64_kernels: - fp64_time = sum(e.self_device_time_total for e in fp64_kernels) - print(f"\nFP64 GEMM (d884): {len(fp64_kernels)} type(s), " - f"{fp64_time / 1000:.2f}ms ({fp64_time / total_cuda * 100:.1f}% of GPU time)") - else: - print("\nFP64 GEMM (d884): NONE") - - print(f"\nTop 10 GPU kernels:") - sorted_kernels = sorted(kernel_events, key=lambda e: e.self_device_time_total, reverse=True) - for e in sorted_kernels[:10]: - pct = e.self_device_time_total / total_cuda * 100 if total_cuda else 0 - print(f" {pct:5.1f}% | {e.self_device_time_total / 1000:8.2f}ms | {e.count:4}x | {e.key[:80]}") - - print(f"\nOutputs: {json_path}") - print(f" {txt_path}") - print(f" {loss_path}") - - -if __name__ == "__main__": - main() From 64438321ca8cbbffcdff83711b5dda4b893b559e Mon Sep 17 00:00:00 2001 From: dlyr3 <152425764+dlyr3@users.noreply.github.com> Date: Sat, 11 Jul 2026 18:57:58 +0200 Subject: [PATCH 6/6] fix --- benchmarks/convergence_test.py | 70 ---------------------------------- benchmarks/plot_convergence.py | 61 ----------------------------- heavyball/chainable.py | 6 +-- heavyball/utils.py | 4 +- test/test_psgd_lb_dtype.py | 21 +++++----- 5 files changed, 15 insertions(+), 147 deletions(-) delete mode 100644 benchmarks/convergence_test.py delete mode 100644 benchmarks/plot_convergence.py diff --git a/benchmarks/convergence_test.py b/benchmarks/convergence_test.py deleted file mode 100644 index 8287302..0000000 --- a/benchmarks/convergence_test.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Convergence test for PSGD FP64 lower-bound fix. - -Trains a model for N steps and logs per-step loss to JSON. -Run before and after the fix with the same seed to produce -comparable convergence curves. - - python benchmarks/convergence_test.py --tag before --steps 500 - python benchmarks/convergence_test.py --tag after --steps 500 -""" - -import argparse -import json -import os - -import torch -from torch import nn - -import heavyball - - -def main(): - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--tag", default="run") - parser.add_argument("--dim", type=int, default=512) - parser.add_argument("--batch", type=int, default=64) - parser.add_argument("--steps", type=int, default=500) - parser.add_argument("--lr", type=float, default=1e-3) - args = parser.parse_args() - - out_dir = os.path.dirname(os.path.abspath(__file__)) - out_path = os.path.join(out_dir, f"convergence_{args.tag}.json") - - torch.manual_seed(42) - model = nn.Sequential( - nn.Linear(args.dim, args.dim, bias=True, device="cuda", dtype=torch.bfloat16), - nn.ReLU(), - nn.Linear(args.dim, args.dim, bias=True, device="cuda", dtype=torch.bfloat16), - ) - opt = heavyball.PSGDPRO(model.parameters(), lr=args.lr) - data = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) - target = torch.randn(args.batch, args.dim, device="cuda", dtype=torch.bfloat16) - - losses = [] - for step in range(args.steps): - out = model(data) - loss = nn.functional.mse_loss(out, target).float() - loss.backward() - opt.step() - opt.zero_grad() - losses.append(loss.item()) - if step % 50 == 0 or step == args.steps - 1: - print(f" step {step:4d}/{args.steps} loss={losses[-1]:.6e}") - - result = { - "tag": args.tag, - "dim": args.dim, - "batch": args.batch, - "steps": args.steps, - "lr": args.lr, - "losses": losses, - } - with open(out_path, "w") as f: - json.dump(result, f, indent=2) - - print(f"\nFinal loss: {losses[-1]:.6e}") - print(f"Saved to: {out_path}") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/plot_convergence.py b/benchmarks/plot_convergence.py deleted file mode 100644 index 5cc08e2..0000000 --- a/benchmarks/plot_convergence.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Generate convergence plot from before/after JSON files. - - python benchmarks/plot_convergence.py --before convergence_before.json --after convergence_after.json --out convergence.png -""" - -import argparse -import json -import os - -import matplotlib -matplotlib.use("Agg") -import matplotlib.pyplot as plt - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--before", required=True) - parser.add_argument("--after", required=True) - parser.add_argument("--out", default="convergence.png") - args = parser.parse_args() - - with open(args.before) as f: - before = json.load(f) - with open(args.after) as f: - after = json.load(f) - - fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) - - ax1.plot(before["losses"], label="before (fp64 lower bound)", - color="#1f77b4", alpha=0.9, linewidth=2.5) - ax1.plot(after["losses"], label="after (fp32 lower bound)", - color="#d62728", alpha=0.7, linewidth=1.5, linestyle="--") - ax1.set_xlabel("Step") - ax1.set_ylabel("Loss") - ax1.set_title("Convergence: linear scale") - ax1.legend() - ax1.grid(True, alpha=0.3) - - ax2.semilogy(before["losses"], label="before (fp64 lower bound)", - color="#1f77b4", alpha=0.9, linewidth=2.5) - ax2.semilogy(after["losses"], label="after (fp32 lower bound)", - color="#d62728", alpha=0.7, linewidth=1.5, linestyle="--") - ax2.set_xlabel("Step") - ax2.set_ylabel("Loss (log scale)") - ax2.set_title("Convergence: log scale") - ax2.legend() - ax2.grid(True, alpha=0.3) - - fig.suptitle( - f"PSGD FP64 Fix - Convergence Test\n" - f"MLP(Linear({before['dim']},{before['dim']}) -> ReLU -> Linear({before['dim']},{before['dim']})), " - f"bf16, batch={before['batch']}, lr={before['lr']}", - fontsize=10, - ) - fig.tight_layout() - fig.savefig(args.out, dpi=150, bbox_inches="tight") - print(f"Saved: {args.out}") - - -if __name__ == "__main__": - main() diff --git a/heavyball/chainable.py b/heavyball/chainable.py index da77234..b3c13ba 100644 --- a/heavyball/chainable.py +++ b/heavyball/chainable.py @@ -1038,7 +1038,7 @@ def _init_psgd_kron(state, group, update, grad, param, cached: bool = False, pro dtype=getattr(torch, group["q_dtype"]), ) state["Q"] = utils.triu_to_line(Q) if group["store_triu_as_line"] else Q - state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float32) for q in Q] + state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float64) for q in Q] state["step"] = torch.zeros((), device=param.device, dtype=torch.float64) if not cached: return @@ -1060,7 +1060,7 @@ def _init_psgd_eigen_kron(state, group, update, grad, param, prob: Optional[call tmp.get("vector"), dtype=getattr(torch, group["q_dtype"]), ) - state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float32) for q in Q] + state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float64) for q in Q] state["step"] = torch.zeros((), device=param.device, dtype=torch.float64) _update_psgd_precond( @@ -1093,7 +1093,7 @@ def _init_psgd_pro_kron(state, group, update, grad, param, cached: bool = False, dtype=getattr(torch, group["q_dtype"]), ) state["Q"] = Q - state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float32) for q in Q] + state["running_lower_bound"] = [torch.zeros((grad.shape[0],), device=q.device, dtype=torch.float64) for q in Q] state["step"] = torch.zeros((), device=param.device, dtype=torch.float64) if not cached: return diff --git a/heavyball/utils.py b/heavyball/utils.py index e71ef45..ee74b77 100644 --- a/heavyball/utils.py +++ b/heavyball/utils.py @@ -3090,11 +3090,11 @@ def calcG_expr(q_dim, g_dim): def _update_lb(ell: Tensor, lb_state: Tensor, beta: Tensor) -> Tensor: + orig_dtype = ell.dtype ell = promote(ell) ell = ell.maximum(promote(lb_state) + (ell - promote(lb_state)) * (1 - beta)) copy_stochastic_(lb_state, ell) - ell = ell.float() - return ell + return ell.to(orig_dtype) @decorator_no_fullgraph diff --git a/test/test_psgd_lb_dtype.py b/test/test_psgd_lb_dtype.py index 1643fa4..14d489a 100644 --- a/test/test_psgd_lb_dtype.py +++ b/test/test_psgd_lb_dtype.py @@ -1,8 +1,8 @@ -"""Test that PSGD preconditioner lower bound uses fp32, not fp64. +"""Test that PSGD preconditioner does not promote matrices to fp64. Verifies: -1. Convergence is not degraded with fp32 lower bound -2. No fp64 tensors leak into optimizer state +1. Convergence is not degraded +2. No fp64 matrices leak into optimizer state (scalar/vector fp64 is intentional) """ import pytest @@ -51,21 +51,20 @@ def test_psgd_convergence_fp32_lb(opt_cls, lr, extra_kw): clean() -def _check_no_fp64(obj, path=""): +def _check_no_fp64_matrices(obj, path=""): if isinstance(obj, torch.Tensor) and obj.is_floating_point(): - if "step" in path: - return - assert obj.dtype != torch.float64, f"fp64 tensor at {path}: dtype={obj.dtype}" + if obj.ndim >= 2 and obj.dtype == torch.float64: + assert False, f"fp64 matrix at {path}: shape={obj.shape}, dtype={obj.dtype}" elif isinstance(obj, dict): for k, v in obj.items(): - _check_no_fp64(v, f"{path}.{k}" if path else str(k)) + _check_no_fp64_matrices(v, f"{path}.{k}" if path else str(k)) elif isinstance(obj, (list, tuple)): for i, v in enumerate(obj): - _check_no_fp64(v, f"{path}[{i}]") + _check_no_fp64_matrices(v, f"{path}[{i}]") @pytest.mark.parametrize("opt_cls,lr,extra_kw", _PSGD_OPTIMIZERS, ids=[t[0].__name__ for t in _PSGD_OPTIMIZERS]) -def test_psgd_no_fp64_in_state(opt_cls, lr, extra_kw): +def test_psgd_no_fp64_matrices_in_state(opt_cls, lr, extra_kw): set_torch() torch.manual_seed(42) data = torch.randn(32, 64, device="cuda") @@ -78,7 +77,7 @@ def test_psgd_no_fp64_in_state(opt_cls, lr, extra_kw): _train(model, opt, data, target, 5) for param in model.parameters(): - _check_no_fp64(opt.state[param]) + _check_no_fp64_matrices(opt.state[param]) del model, opt clean()