Repository navigation
Conversation
_init_psgd_kron_states built Q with torch.eye and L with torch.ones without a dtype, so they were always float32. For a float64 or bfloat16 param the first step then failed inside the compiled apply_preconditioner, because tensordot got a float32 Q and a momentum buffer in the param dtype. Create both in grad.dtype, the same way exp_avg already follows the param through zeros_like. Signed-off-by: Dmitry Voropaev <dy.voropaev@gmail.com>
 [Medium impact] The change reaches training runs that use PSGDPro. Restore factor precision after checkpoint loading before merging; resumed Findings
SummaryUpdates
Checkpoint loading still needs to restore the factors' intended dtype. Reviews (2) · Last reviewed commit: "Keep PSGD Kronecker factors in fp32 for ..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| self.assertEqual(param.dtype, dtype) | ||
| self.assertTrue(torch.isfinite(param).all()) | ||
| for q in optimizer.state[param]["Q"]: | ||
| self.assertEqual(q.dtype, dtype) |
There was a problem hiding this comment.
The new test checks only that three steps finish with finite parameters and the expected dtype. This change also makes procrustes_step run in bfloat16, although its comment says it relies on float32 for numerical stability. Finite results alone do not check that the factor updates remain accurate.
Add a numerical check against a float32 reference, using a non-symmetric factor so the rotation actually runs.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
Good catch on procrustes_step. I changed it so bf16 params keep the factors in fp32 (float64 params get float64) and the momentum is cast to that before it meets Q. The bf16 test now runs the same problem in fp32 alongside and compares the factors and params. With bf16 factors the factors were off by about 1e-2 after five steps, with fp32 factors it's about 1e-3, and the test checks 5e-3.
There was a problem hiding this comment.
That addresses the concern. Keeping bf16 factors in fp32 avoids running procrustes_step in reduced precision, while retaining float64 factors for float64 parameters. Casting exp_avg before the preconditioner path also keeps operand dtypes consistent.
The matched fp32 reference, deterministic damping noise, factor comparisons, and parameter comparison provide the numerical coverage that was missing. The tighter factor error with fp32 factors is consistent with the intended stability improvement.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
procrustes_step is written for fp32, so creating Q in bf16 for a bf16 param ran it in bf16. Use promote_types(param dtype, float32) for the factors instead, which is fp32 for bf16 and fp16 and float64 for float64 params, and cast the momentum to that dtype before it meets Q. float32 params are unchanged. The bf16 test now runs the same problem in fp32 next to it and compares the factors and params. With bf16 factors the factors drift by about 1e-2 after five steps, with fp32 factors by about 1e-3. Signed-off-by: Dmitry Voropaev <dy.voropaev@gmail.com>
v0ropaev
changed the title
| exp_avg = state["exp_avg"] | ||
| exp_avg.lerp_(grad, 1 - group["momentum"]) | ||
| # The Kronecker factors can be kept in higher precision than the param, see _kron_factor_dtype. | ||
| exp_avg = exp_avg.to(_kron_factor_dtype(grad.dtype)) |
There was a problem hiding this comment.
After loading a bfloat16 optimizer checkpoint, the next step() fails. PyTorch's inherited load_state_dict() casts floating-point state, including Q, to the parameter dtype. This line casts exp_avg to float32, but existing factors skip initialization and remain in bfloat16. The next apply_preconditioner() therefore calls torch.tensordot() with different dtypes and raises an error.
Restore Q and L to _kron_factor_dtype() after loading, and add a checkpoint round-trip test.
Knowledge Base Used: PSGD preconditioner
There was a problem hiding this comment.
Right, that's the same load_state_dict cast I opened #316 about. It hits every optimizer here that keeps fp32 state for bf16 params (Iso, SOAP, Shampoo and so on), and PSGD now joins that list. I'd rather fix it once for all of them there than add a PSGD-only workaround in this PR. On main bf16 PSGD doesn't step at all, so this isn't a regression.
There was a problem hiding this comment.
That makes sense. Since the checkpoint-casting issue affects every optimizer with higher-precision state and is already tracked in #316, a PSGD-only restoration would duplicate the eventual shared fix. I agree this is not a regression introduced by this PR—the new bf16 path cannot step on main at all—so I’m withdrawing this finding for the PR and leaving checkpoint compatibility to the cross-optimizer fix.
PSGDPro only works with float32 params right now. With a float64 or bfloat16 param the first
step()fails withTorchRuntimeError: RuntimeError when making fake tensor call, which comes from tensordot in the compiledapply_preconditionergetting operands of different dtypes. Unlike Muon there is no explicit fp32 check, so you just get that dynamo error._init_psgd_kron_statescreates Q withtorch.eyeand L withtorch.oneswithout a dtype, so both end up float32, whileexp_avgcomes fromzeros_like(grad)and follows the param. Since procrustes_step is written for fp32, the factors now usetorch.promote_types(param dtype, float32), so fp32 for bf16 and fp16 params and float64 for float64 params, and the momentum is cast to that dtype before it goes into the preconditioner update andapply_preconditioner. The momentum buffer itself stays in the param dtype as before, and float32 params go through exactly the same path as today.Tests in tests/test_psgd_pro.py:
test_step_with_float64_paramruns three steps and checks the param stays finite and Q is float64.test_bfloat16_param_matches_float32_referenceruns the same five steps on a bf16 param and on an fp32 copy with the same grads and noise seed, and checks that Q is fp32 and matches the fp32 run within 5e-3, and the params within 2e-2. With the factors kept in bf16 instead, Q drifts by about 1e-2 and the test fails.On main both tests error out with the dynamo error above. With the change, test_psgd_pro.py, test_psgd_utils.py, test_psgd_contractions.py and test_registry.py pass with
--device=cpuwith no seed,--seed=42and--seed=3735928559on torch 2.14.1 CPU on macOS. I did not run the CUDA suite. ruff 0.9.9 check and format are clean.