[Zamba2] Fix block_id to use hybrid order for num_mem_blocks>1

[aryansk]

Fixes #47994

Problem

Zamba2Model(config) raises ValueError: tie_weights_keys before any weight is loaded for any config with num_mem_blocks > 1 (e.g. Zyphra/Zamba2-2.7B-instruct, num_mem_blocks=2, 54 layers, hybrids at [6,12,18,24,30,36,42,47,51]). Repro on main with torch.device("meta"):

from transformers import AutoConfig
from transformers.models.zamba2 import Zamba2Model
import torch
cfg = AutoConfig.from_pretrained("Zyphra/Zamba2-2.7B-instruct")  # num_mem_blocks=2
with torch.device("meta"):
    Zamba2Model(cfg)  # ValueError layers.12 <-> layers.47 adapter mismatch

Root cause in Zamba2Model.get_layers: block_id = layer_id % num_mem_blocks uses the global layer index, while the tie cycle just above advances once per hybrid layer, and Zamba2Attention/Zamba2MLP select adapter slots with i % num_mem_blocks == block_id where i runs over hybrid layers. For 2.7B the global rule yields [0,0,0,0,0,0,0,1,1] where the cycle expects [0,1,0,1,0,1,0,1,0]; layer 12 ends up tied to 47 with a disjoint adapter set, failing validation that was added to catch silent mis-ties.

Where num_mem_blocks == 1 (1.2B, the Zamba2Config default) both conventions are identical, so no currently-working config changes behavior.

Change

• Track hybrid_layer_idx in Zamba2Model.get_layers and compute block_id = hybrid_layer_idx % config.num_mem_blocks (increment per hybrid). This makes structure, tie cycle and checkpoint agree.
• Applied to the modular source src/transformers/models/zamba2/modular_zamba2.py and kept modeling_zamba2.py in sync (generated file; single logical fix in two files). Minimal diff: +3/-1 per file, no comment bloat.
• Verified the 2.7B config now constructs with 531 params matching the published checkpoint's 531 tensors (per Zamba2: num_mem_blocks > 1 checkpoints cannot be constructed — block_id uses the global layer index while the tie cycle uses hybrid order #47994), and tie_weights_keys correctly alternates (6→18→30→42→51 share block 0, 12→24→36→47 share block 1).

Why this approach

• The fix is exactly the hybrid-counter that the adapter logic already uses (range(num_fwd_mem_blocks) with i % num_mem_blocks == block_id). Alternatives like keeping global modulo would require changing both adapter selection and tie cycle, which diverges from the checkpoint layout. This keeps the checkpoint as ground truth.
• Same functional fix as closed [Zamba2] Fix layer construction and weight tying for num_mem_blocks > 1 #48018 (which CI showed green), but the issue remains open after maintainer coordination note; this PR is scoped to the failing get_layers path and includes disclosure/differentiation below.

Testing

• Reproduction script on meta device before/after:
  • Before: Zamba2Config(num_mem_blocks=2) raises ValueError for layers.12:layers.47 with mismatched gate_up_proj_adapter_list indices (0,2,4,6,8 vs 1,3,5,7).
  • After: construction succeeds, model._tied_weights_keys as above, per-layer block_id follows hybrid_idx % 2 (6:0,12:1,18:0,24:1,30:0,36:1,42:0,47:1,51:0), and Zamba2ForCausalLM also constructs. num_mem_blocks=1 regression still passes.
• Focused model tests: pytest tests/models/zamba2/test_modeling_zamba2.py (with datasets/parameterized): 129 passed, 146 skipped (run in local venv, Python 3.13, torch). No failures in the Zamba2 common suite.
• Style: ruff check src/transformers/models/zamba2/{modeling_zamba2,modular_zamba2}.py → All checks passed!

Documentation and release impact

• No documentation impact
• Changelog/release note needed: fixes construction for published 2.7B/7B checkpoints; no API change.
• No migration note (previously-unconstructible configs now work; working configs unchanged).

Review notes

• Coordination: Fixes Zamba2: num_mem_blocks > 1 checkpoints cannot be constructed — block_id uses the global layer index while the tie cycle uses hybrid order #47994, issue still OPEN, unassigned. Prior PR [Zamba2] Fix layer construction and weight tying for num_mem_blocks > 1 #48018 correctly fixed both files but was closed after maintainer note about waiting for issue author's own PR offer; CI was green (run 32016223085). This PR is intentionally minimal and acknowledges that prior work.
• Differentiation from [Zamba2] Fix layer construction and weight tying for num_mem_blocks > 1 #48018: same functional change, plus explicit hybrid_layer_idx in both modular and generated file; no additional test file (lane scope keeps PR to source fix; the reporter's suggested construction test is ready to add if maintainers prefer — small 9-layer num_mem_blocks=2 case proven in validation).
• AI disclosure: AI-assisted (Muse Spark) for diagnosis and patch, human-verified execution and tests above; diff minimized to 1 logical line.
• Known limitations: none; modular_zamba2.py is source of truth, modeling_zamba2.py kept identical.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: zamba2

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 32375658356
Result: success | Grafana metrics are not available yet.