Fix(test): match the current scan_layers mismatch error in test_scan_layers_mismatch_tpu

[ecnal-cienet]

Description

tests/integration/checkpointing_test.py::test_scan_layers_mismatch_tpu fails on scheduled CI (run 29234138806). The behavior under test is correct — restoring a scan_layers=True checkpoint into a scan_layers=False model still raises ValueError, and the message still names scan_layers. Only the wording the test matches on has moved, so this is a test-only change; nothing in src/ is touched.

The mismatch is now reported by the weight-mismatch check in _raise_on_weight_mismatch, added in #4343 (fbb35123a), which lists each weight the checkpoint didn't carry:

ValueError: Checkpoint does not match the model:
  - 'decoder/layers/0/mlp/wi_0/kernel': missing (model expects (128, 128) float32)
  ...
Verify the checkpoint matches the model architecture (emb_dim, mlp_dim, num layers, scan_layers).

The test instead expected one of two older strings, and neither can match:

• "Failed to restore checkpoint" is raised nowhere in the repository — grepping the tree, the only occurrence is the assertion itself.
• "Configuration mismatch" comes from verify_and_sync_scan_layers, which returns early unless load_parameters_path is set. This test resumes from base_output_directory and never sets that flag, so the path is unreachable from here. It is already covered directly by tests/unit/model_creation_utils_test.py.

So the assertion is rewritten against the message the restore path actually produces, rather than OR-ing in strings that cannot fire.

The test carries scheduled_only, so it does not run on presubmit — which is why #4343's own CI was green and the staleness only surfaced on the next scheduled run.

The "scan_layers" in message assertion is kept but is no longer sufficient on its own: that substring now comes from the error's generic trailing hint (...num layers, scan_layers)), which is printed for any architecture mismatch — a wrong emb_dim would satisfy it too. The new "decoder/layers/0/" assertion pins the failure to the missing unscanned per-layer weight, which is what a scan/unscan mismatch specifically produces.

Tests

tests/integration/checkpointing_test.py::test_scan_layers_mismatch_tpu passes on a v6e-8 (it fails on main at the same assertion reported by CI, and passes with this change).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!