Prism/Shampoo-Prism Sweep Part 1: Hyperparameter search

[leloykun]

Part 1 of #2695 , focusing on hyperparameter search

Vibecoding warning: lib/levanter/tests/test_optim_util.py is entirely vibecoded. More eyes may be needed there, but it seems to work pretty well.

[Copilot]

Pull request overview

Adds initial infrastructure for Prism/Shampoo-Prism hyperparameter sweeps by extending an optimizer utility to map across multiple PyTrees (flattening Linear layers consistently), adding a focused regression test, and introducing a new PrismH sweep experiment script for Qwen3 scaling runs.

Changes:

• Extend map_flattened_linear_layers to accept additional PyTrees (*rest) and map over them after flattening Linear leaves.
• Add a unit test ensuring additional PyTrees are flattened consistently when mapping.
• Add an experiment script to sweep PrismH hyperparameters (gamma_L, gamma_R, bidirectional) for Qwen3 speedruns.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
lib/levanter/src/levanter/optim/util.py	Generalizes map_flattened_linear_layers to support zipping/mapping over multiple parameter/update trees.
lib/levanter/tests/test_optim_util.py	Adds coverage for the new *rest flatten+map behavior.
experiments/speedrun/prism_qwen3_scaling/prism_sweep_1_hp_sweep.py	Introduces PrismH hyperparameter sweep experiment and a PrismH optimizer implementation used by the sweep.

[leloykun]

test errors are all huggingface-related