fix: raise on stop_at_layer when no 'blocks' stack is registered - #1789
Conversation
Replace the stop_at_layer reject-list with an allowlist guard and check self._modules instead of hasattr for both stop_at_layer and start_at_layer. Addresses TransformerLensOrg#1769
|
Retargeted this to dev-4.x (it was accidentally opened against main, which pulled in 92 unrelated commits). The new test_stop_at_layer_guard.py passes. The remaining failures (test_generate_stream_batch_decode.py, test_hf_generate_batch_padding.py with AttributeError: ... no attribute '_driver') are in generate_stream, which this PR doesn't touch, so they seem to come from dev-4.x itself. Happy to look into them separately if useful. cc @jlarson4 |
|
@Aurnawr I pulled in the latest |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Clean implementation of both halves of the approach from the issue, and the regression test is solid. Two comments on the diff's leftover scratch comments and the guard's test coverage, plus one on the docstring.
test: cover a wrapped model exposing its own .blocks Addresses TransformerLensOrg#1769
|
This looks great! Thank you for quickly getting to those, I will merge this and make sure it gets into today's release |
2 participants
Description
TransformerBridge.forward(..., stop_at_layer=N)silently did nothing on architectures that don't register ablocksstack (e.g. Raven/Huginn, whose layers areprelude/core_block/coda). The old guard only rejected four known names (L_blocks,H_blocks,encoder_blocks,decoder_blocks) and only set_stop_at_layer_idxwhenhasattr(self, "blocks"), so Raven fell through both checks and the forward pass ran to completion.Following @jlarson4's suggested approach in the issue:
stop_at_layernow raisesNotImplementedError("stop_at_layer requires a 'blocks' stack; ...")whenever noblocksstack is registered. This covers Raven and any future non-standard adapter, and removes the four-name reject-list._modulesinstead ofhasattr: added a private helper_has_registered_blocks()that checks"blocks"in the bridge's registered submodules.hasattr(self, "blocks")can returnTruevia__getattr__falling through to the wrapped HF model. The existingstart_at_layerguard now uses the same helper.tests/unit/model_bridge/test_stop_at_layer_guard.pywith a regression test that builds a bare bridge with Raven-styleprelude/core_block/codalists (no checkpoint download) and assertsstop_at_layer=0raises.No new dependencies.
Addresses #1769. Per the maintainer's note, the issue should stay open to track actual
stop_at_layersupport for Raven via top-level wiring.Type of change
Note: architectures with
L_blocks/H_blocks/encoder_blocks/decoder_blocksstill raiseNotImplementedError, but the message text changed.Checklist:
Tests run locally:
uv run pytest tests/unit/model_bridge/test_stop_at_layer_guard.py tests/unit/model_bridge/test_audio_start_at_layer_guard.py -v: all pass; the new test fails on the pre-fix code.black --check/isort --check-onlyon the changed files: clean.