fix(exporters): apply fine-tune weights to text embedder + decoder prefill

[rylinjames]

Bug

When a user exports a fine-tuned SmolVLA checkpoint via export_vlm_prefix, the fine-tuned state_dict was applied to the vision encoder only. export_text_embedder and export_decoder_prefill each called AutoModel.from_pretrained independently and applied no fine-tune weights, silently shipping base-model text embeddings and decoder KV prefix.

A user deploying a fine-tuned SmolVLA via the decomposed path got:

• vision_encoder.onnx -- fine-tuned (correct)
• text_embedder.onnx -- base model (wrong; embed_tokens table was never updated)
• decoder_prefill.onnx -- base model (wrong; all 16 decoder layer k/v projections were never updated)

Key-namespace evidence (no silent no-op)

AutoModel.from_pretrained(SmolVLM2) returns SmolVLMModel with state_dict keys:

text_model.embed_tokens.weight
text_model.layers.0.self_attn.k_proj.weight
text_model.layers.0.self_attn.v_proj.weight
...
vision_model.embeddings.patch_embedding.weight
connector.weight

After stripping model.vlm_with_expert.vlm. from the checkpoint:

text_model.embed_tokens.weight           <- matches
text_model.layers.0.self_attn.k_proj.weight  <- matches
...
connector.weight                         <- matches

These are identical namespaces. Verified with a synthetic _SmolVLMModel module: load_state_dict(strict=False) returns 0 missing, 0 unexpected -- applied == total_model_keys.

Applied-count logging fix

_apply_checkpoint_vlm_weights previously returned len(rebased) (number of checkpoint keys) instead of the actual applied count (total_model_keys - missing). Now it correctly computes and logs:

[vlm-weights/text_embedder] applied 247/247 fine-tune weights (0 missing, 0 unexpected)

Zero-applied raises a WARNING so silent no-ops are visible in future.

Changes

• _apply_checkpoint_vlm_weights: add tag param for per-sub-model log attribution; fix return value to actual applied count; add zero-applied WARNING.
• export_text_embedder: add checkpoint_state_dict param; apply weights to full AutoModel before extracting embed_tokens (Python reference semantics ensure the extracted submodule already carries fine-tuned weights).
• export_decoder_prefill: add checkpoint_state_dict param; apply weights to full AutoModel before extracting text_model.
• export_vlm_prefix: thread checkpoint_state_dict=state_dict to both call sites; also fix two pre-existing F841 ruff violations.

Tests (no HF download)

9 new unit tests in TestApplyCheckpointVLMWeights using a synthetic _SmolVLMModel-shaped nn.Module:

1. test_non_zero_keys_applied -- guard against silent no-op: applied > 0
2. test_applied_count_equals_total_model_keys -- full coverage checkpoint applies all keys
3. test_finetune_values_actually_loaded -- param values equal fill value after apply
4. test_none_checkpoint_is_a_noop -- None checkpoint leaves model unchanged
5. test_embed_tokens_reference_updated_after_apply -- text embedder pattern: apply to full model, extract embed_tokens, values updated
6. test_text_model_reference_updated_after_apply -- decoder prefill pattern: apply to full model, extract text_model, k_proj/v_proj updated
7. test_no_prefix_match_returns_zero -- bad checkpoint prefix returns 0, model unchanged
8. test_partial_checkpoint_partial_apply -- text_model-only checkpoint updates only text_model keys
9. test_tag_does_not_affect_result -- tag is cosmetic only

All 34 tests pass (1 integration test skipped); ruff clean; py_compile clean.

Follow-up (noted, not blocked)

Tokenizer: export_text_embedder loads the model via AutoModel, not AutoTokenizer, so any fine-tuned vocabulary additions in the tokenizer are not applied here. The tokenizer is loaded separately in VLMPrefixOrchestrator._load_tokenizer_and_processor from checkpoint_path_or_id -- this is likely correct already, but worth an explicit integration test once TETHER_INTEGRATION=1 infra is wired.

Generated with Claude Code