Fix FP64 promotion in PSGD running_lower_bound - #95
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e2ae8d736c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ell = promote(ell) | ||
| ell = ell.maximum(promote(lb_state) + (ell - promote(lb_state)) * (1 - beta)) | ||
| copy_stochastic_(lb_state, ell) | ||
| ell = ell.float() |
There was a problem hiding this comment.
When PSGD is configured with q_dtype="float64"/double gradients, this unconditional cast forces the lower-bound denominator back to fp32 right before it is used in update / ell. For double-precision runs whose preconditioner eigenvalue scale is outside fp32 range, ell.float() can underflow or overflow to 0/inf, producing infinite or zero preconditioner updates even though the Q factors and covariance were computed in fp64. Cast to the preconditioner/update dtype, or only force fp32 for the default non-fp64 Q/legacy-state path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Casting is in the function rather than in the callable to avoid redundancy and potential oversight.
Also i didn't pass q_ dtype to not mix responsabilities and also because ell was promoted on its own, without regards to q_
|
that's a huge finding, thank you! cc @opooladz: PSGD speedup incoming |
| ell = promote(ell) | ||
| ell = ell.maximum(promote(lb_state) + (ell - promote(lb_state)) * (1 - beta)) | ||
| copy_stochastic_(lb_state, ell) | ||
| ell = ell.float() |
| @@ -0,0 +1,103 @@ | |||
| """Reproducer for FP64 promotion in PSGD preconditioner lower bound. | |||
There was a problem hiding this comment.
we generally don't need issue reproduction scripts in the repository, but thank you for sharing it!
| import heavyball | ||
|
|
||
| def make_model(dim, device): | ||
| return nn.Sequential( |
There was a problem hiding this comment.
unrelated to heavyball, but bias is true by default and .to(device) and .to(dtype) are usually done outside of the construction
this specific file should not be in upstream either
| ) | ||
| 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] |
There was a problem hiding this comment.
rlb was set to fp64 so that we have higher precision accumulation. it's a scalar, so it's ~free!
| @@ -0,0 +1,84 @@ | |||
| """Test that PSGD preconditioner lower bound uses fp32, not fp64. | |||
|
|
|||
There was a problem hiding this comment.
The scalar lower bound arithmethic should be fp64, only matmul and matrix-level elementwise computation should be standard precision
|
thanks for your review ! please let me know if anything more could be improved The test was ran and successful on my end |
|
lgtm! |
Problem
running_lower_boundwas initialized as float64.This propagated through
_update_lb()into preconditioner matmuls (R @ q_,R @ RQ), causing CUTLASS to dispatch d884gemm (FP64 GEMM kernels).On consumer GPUs (1:64 FP64:FP32 throughput) this consumed ~80% of GPU time.
Fix
Changed
running_lower_boundinit dtype from float64 to float32 (chainable.py lines 1041/1063/1096).Added defensive
.float()cast in_update_lb().Results (RTX 3090, MLP 512x512 bf16, PSGDPRO defaults)
Before:
After:
Convergence:
