Skip to content

fix(bridge): restore native stop_at_layer and input_to_embed - #1633

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:dev-4.xfrom
emerardd:fix/native-stop-at-layer
Aug 10, 2026
Merged

fix(bridge): restore native stop_at_layer and input_to_embed#1633
jlarson4 merged 1 commit into
TransformerLensOrg:dev-4.xfrom
emerardd:fix/native-stop-at-layer

Conversation

@emerardd

@emerardd emerardd commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

Description

Fixes #1632.

Native TransformerBridge blocks are named layers.N, but the block layer-index parser only recognized layers.N when it had a preceding dot. As a result, direct forward(..., stop_at_layer=...) calls never stopped and input_to_embed() returned final vocabulary logits instead of the residual stream entering block 0.

This change:

  • makes the layer-name parser boundary-aware for blocks.N, h.N, and layers.N, including top-level native names;
  • normalizes negative stop_at_layer indices before passing the stop point to blocks;
  • lets the native model accept inputs_embeds, which is required to resume a residual with start_at_layer;
  • adds direct stopping and input_to_embed() round-trip regression coverage for an offline boot_native() model.

No new dependencies are required.

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Screenshots

Not applicable.

Verification

  • Before the fix, the four new regression cases failed: direct stops returned logits with the d_vocab dimension, and residual re-entry raised because the native model required input_ids.
  • uv run pytest tests/unit/model_bridge/test_boot_native.py — 35 passed
  • Focused direct-stop and round-trip regression selection — 4 passed, 31 deselected
  • uv run pytest tests/integration/model_bridge/test_bridge_input_to_embed.py tests/integration/model_bridge/test_bridge_stop_at_layer.py::test_manual_hooks_stop_at_layer_compat_with_processing — 5 passed
  • uv run mypy . — success, 424 source files
  • pycln, isort, and black checks on the four changed files — passed
  • git diff --check — passed

The complete unit-test suite was not rerun; only the affected unit and integration surface was run locally.

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation (not applicable; behavior now matches the documented API contract)
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes (only the affected test surface was run)
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4

Copy link
Copy Markdown
Collaborator

Thanks for filing and resolving this issue @emerardd, looks great, merging as is.

@jlarson4
jlarson4 merged commit 0d1259a into TransformerLensOrg:dev-4.x Aug 10, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants