AdEMAMix32bit 和 PagedAdEMAMix32bit 分配了单个大小的 state1,并在第一步出错

bnb.optim.AdEMAMix32bit and bnb.optim.PagedAdEMAMix32bit subclass Optimizer2State directly (bitsandbytes/optim/ademamix.py:355 and :385 on 8336490), so they never run AdEMAMix.init_state. That override allocates state1 as (2, *p.shape) to hold m1 and m2. The base Optimizer2State.init_state allocates p.shape. Every ademamix kernel then reads m2 from the second half of state1:

• CPU backend (_optimizer_update_32bit_cpu in bitsandbytes/backends/cpu/ops.py) indexes state1[0] and state1[1], so the first step() raises RuntimeError: output with shape [] doesn't match the broadcast shape [4096].
• Default backend (_optimizer_update_32bit in bitsandbytes/backends/default/ops.py, used on MPS) does the same indexing and raises the same error.
• CUDA (kOptimizer32bit2State in csrc/kernels.cu, line 674) loads m2 with Load(&(state1[n + i]), s3_vals, ...) on a buffer of n floats. I do not have a CUDA device to run this. From the source, it reads and later writes past the end of state1.

The scheduler is lost as well. Optimizer2State.update_step never applies t_alpha or t_beta3, so AdEMAMix32bit(t_alpha=100, t_beta3=100) runs unscheduled without any warning.