fix(optim): honor user-provided max_unorm in LAMB (+ 8-bit arg guards)

[egeozkoc]

Problem

The LAMB, LAMB8bit, and LAMB32bit constructors accept a max_unorm argument but pass a hardcoded 1.0 to the base optimizer:

super().__init__("lamb", ..., min_8bit_size, max_unorm=1.0)  # ignores the argument

On the 32-bit LAMB path (LAMB with the default optim_bits=32, and LAMB32bit), max_unorm drives the trust-ratio clipping of the update (optimizer_update_32bit → _compute_update_norm_and_scale). So the setting was silently ignored: a tighter max_unorm had no effect, and it could not be changed from 1.0.

Two related inconsistencies in the 8-bit optimizers, in the same spirit as the Adam8bit/AdamW8bit guards that already exist (relates to #1261):

• LAMB8bit accepts amsgrad, which the base optimizer never uses.
• Adagrad8bit accepts optim_bits, but always runs 8-bit.

Fix

• Thread max_unorm through in all three LAMB classes. Default behavior is unchanged (the default is 1.0, which equals the previously hardcoded value).
• LAMB8bit: raise on amsgrad=True; Adagrad8bit: raise on optim_bits != 8; add docstring notes. Mirrors the existing Adam8bit/AdamW8bit guards.

Notes

• The 8-bit blockwise update kernel does not apply update-norm clipping, so max_unorm has no numerical effect on LAMB8bit. The value is now stored consistently and this limitation is documented in the docstring.
• No public API changes; the guards only reject values that were already silently ignored, and no code in the repo constructs these classes with the rejected arguments.

Tests (tests/test_optim.py)

• test_lamb_max_unorm_changes_update — behavioral: on the 32-bit path a tight vs loose max_unorm now yields different updates (fails on main, where the value was ignored). Runs on CPU.
• test_lamb_max_unorm_threaded_to_config — the value reaches optimizer.args for all three LAMB classes.
• test_lamb8bit_rejects_amsgrad, test_adagrad8bit_rejects_non_8_optim_bits — the guards.

pre-commit run --all-files passes; verified the new assertions fail on main and pass here.

Relates to #1261

🤖 Generated with Claude Code

[matthewdouglas]

Thanks for this. The 32-bit LAMB fix is the real win — max_unorm was stuck at 1.0 and couldn't be changed, and this fixes it. Keep the fix, the docstring correction, the test, and the guards.

The one thing I'd change is LAMB8bit. The 8-bit path has no update-norm clipping — the kernel just doesn't implement it — so max_unorm does nothing there no matter what. Threading it through doesn't change any 8-bit behavior: a user who passes max_unorm=0.1 to LAMB8bit still gets a silent no-op with no error, and reasonably assumes it's doing something.

I'd rather make that explicit — raise (or warn) when a non-default max_unorm is passed to LAMB8bit, the same way you're already guarding amsgrad. As written, the diff is a little inconsistent: it rejects one unsupported argument (amsgrad) but silently accepts another (max_unorm). Rejecting both is cleaner.

Implementing the clip in the 8-bit kernel would be the proper fix, but that's a bigger change and out of scope here.

So: keep everything on the 32-bit side, and just have LAMB8bit reject a non-default max_unorm instead of quietly accepting it.

🤖 Generated with Claude Code

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[egeozkoc]

Done, that makes sense. LAMB8bit now raises on any non-default max_unorm, matching the amsgrad guard and the existing optim_bits pattern (the default 1.0 is allowed for signature compatibility). Tests updated to match: the config-threading test covers LAMB and LAMB32bit only, and there's a new guard test for LAMB8bit. Agreed that implementing the clip in the 8-bit kernel would be the proper fix and out of scope here.

[matthewdouglas]

LGTM, thanks!