Skip to content

Load the SAM 3.1 multiplex checkpoint once, without spurious key warnings (fixes #506) - #643

Open
YoussefMadkour wants to merge 1 commit into
facebookresearch:mainfrom
YoussefMadkour:fix-506-multiplex-double-load
Open

YoussefMadkour wants to merge 1 commit into
facebookresearch:mainfrom
YoussefMadkour:fix-506-multiplex-double-load

Conversation

@YoussefMadkour

Copy link
Copy Markdown

Fixes #506

Problem

Building the SAM 3.1 predictor with build_sam3_multiplex_video_predictor prints hundreds of Missing keys / Unexpected keys warnings, which makes it look like the checkpoint failed to load. There are two causes:

  1. The merged checkpoint is loaded twice. The predictor passes checkpoint_path to build_sam3_multiplex_video_model, which loads the merged checkpoint into the bare tracker. The merged checkpoint's keys are prefixed tracker.model.* / detector.*, while the tracker's own keys are not, so this first load matches nothing. It only prints warnings and reads the multi-GB file an extra time. The same file is then loaded again, correctly, into the full demo model a few lines later.
  2. 64 RoPE buffers are reported missing. With use_rope_real=True (the predictor's default), each ViT attention block registers freqs_cis_real / freqs_cis_imag. These are derived from freqs_cis at init and aren't in the checkpoint, which accounts for the Missing keys (64) in the issue.

Fix

  • Build the tracker with checkpoint_path=None inside build_sam3_multiplex_video_predictor. Its weights still come from the single load into the demo model.
  • Register freqs_cis_real / freqs_cis_imag with persistent=False. They're still created and moved with the module as before, but they're no longer expected in, or saved to, state dicts. freqs_cis itself is unchanged.

The weights the model ends up with don't change: I built the tracker on its own and confirmed none of its 931 state-dict keys start with tracker. or detector., so the first load never loaded anything.

Tests

Added test/test_multiplex_checkpoint_loading.py (no checkpoint or GPU needed):

  • build_sam3_multiplex_video_predictor no longer passes the checkpoint to the tracker build (it stops the build at that call with a mock).
  • A small ViT with use_rope_real=True keeps the real/imag buffers on the module, leaves them out of state_dict(), and strictly loads a state dict saved without them.

Both tests fail on main and pass with this change.

…ings (facebookresearch#506)

build_sam3_multiplex_video_predictor passed the merged checkpoint to
build_sam3_multiplex_video_model, which loaded it into the bare tracker.
None of the merged keys (tracker.model.*, detector.*) match the tracker,
so nothing was loaded and hundreds of missing/unexpected keys were
printed; the checkpoint was then read again and loaded correctly into the
full demo model. Skip the first load.

Also register the real-valued RoPE buffers (freqs_cis_real/imag) as
non-persistent: they are derived from freqs_cis at init and are not in
the checkpoint, which caused the remaining 64 missing-key warnings.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Checkpoint key mismatch when loading pretrained SAM3.1 weights

1 participant