transformers
323f2434 - Fix Zamba2 construction for num_mem_blocks > 1 checkpoints (#48325)

Commit
54 days ago
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.
Author
Parents
Loading