Fix Zamba2 construction for num_mem_blocks > 1 checkpoints (#48325)
* Fix Zamba2 block_id for num_mem_blocks > 1: derive it from hybrid-layer order
Zamba2Model.get_layers derived block_id from the global layer index while
the weight-tie cycle and the adapter slots inside each shared block follow
hybrid-layer order. Every published num_mem_blocks=2 checkpoint (the
Zamba2-2.7B and Zamba2-7B families) has hybrid layers whose indices are
not all congruent modulo num_mem_blocks, so construction paired
structurally different blocks and raised before any weight was read.
Derive block_id from a hybrid-layer counter instead — the convention the
tie cycle and the published checkpoints already use. For num_mem_blocks=1
(the Zamba2-1.2B family, and the config default) the two rules agree, so
nothing changes there.
* Fold the Zamba2 num_mem_blocks regression into one slow generation test
Review feedback: the slow test alone is enough for this bug, and it should
greedily generate a small text and pin the output with `Expectations()`.
Drop `test_num_mem_blocks_mixed_residue_hybrid_layout`, and rewrite
`test_num_mem_blocks_2_official_checkpoint` to greedily generate 20 tokens
through the chat template and compare against an `Expectations()` entry
instead of only checking that the logits are finite.
The checkpoint keeps its mixed-residue hybrid layout -- hybrid layers
`[6, 12, 18, 24, 30, 36, 42, 47, 51]` are not all congruent modulo
`num_mem_blocks`. Against the pre-fix source the test fails at construction
with `We found layers.12.shared_transformer... to tie into
layers.47.shared_transformer...`, while an evenly spaced layout would
construct fine either way, so a synthetic config cannot replace it.
The expected string was measured on CPU; float32 and bfloat16 produce
identical token ids.