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