Preserve input dtype in GeneralizedComponent

[koriyoshi2041]

Description

GeneralizedComponent.forward currently infers a compute dtype from the first floating-point parameter in the wrapped module. That parameter can be an FP32 auxiliary value or FP8 quantization storage, so the wrapper can change a valid mixed-precision activation before the original component sees it.

This removes the generic parameter-based cast and preserves the dtype produced by hook_in. Model-specific conversions can remain explicit in their adapters. The regression covers both positional and keyword input paths with a BF16 activation and an FP32 auxiliary parameter.

Fixes #1712

Type of change

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

Checklist

• The generic hook wrapper preserves the original component boundary dtype
• Tests cover positional and keyword input paths
• No public API or documentation changes are required

Validation

• uv run pytest tests/unit/model_bridge/generalized_components/test_base.py -q — 29 passed
• uv run black --check transformer_lens/model_bridge/generalized_components/base.py tests/unit/model_bridge/generalized_components/test_base.py
• uv run isort --check-only transformer_lens/model_bridge/generalized_components/base.py tests/unit/model_bridge/generalized_components/test_base.py
• uv run mypy transformer_lens/model_bridge/generalized_components/base.py
• git diff --check

[jlarson4]

@koriyoshi2041 Thanks for tackling this issue! The fix looks great, there are just a couple other locations that we will want to also incorporate this.

Identical logic still runs in:

• generalized_components/attention.py:800-835 (AttentionBridge.forward, three call sites)
• generalized_components/position_embeddings_attention.py:355-363 (PositionEmbeddingsAttentionBridge.forward)

A worse variant also still runs in:

• generalized_components/moe.py:284-307
• generalized_components/unembedding.py:89-101

Can we extend the fix to these locations as well? As a follow up PR, if you're interested, it may be worth while to build out a helper function for handling this specific problem, so it does not need to be mirror fixed across 5 locations in the future.

[jlarson4]

The test pins the FP32-auxiliary case, but the failure reported in the issue was Promotion for Float8 Types is not supported, from an FP8 weight and an FP8 scale. Nothing here would catch a future change that re-introduced inference for FP8 specifically.

Let's add a parametrization with a torch.float8_e4m3fn parameter asserting the BF16 input is unchanged.

[koriyoshi2041]

Thanks for the pointers — I extended the fix at the current head to all four locations you listed. I kept the helper extraction out of this PR, and added direct mixed-dtype regression coverage for the specialized attention, position-embedding attention, MoE, and unembedding forward paths.

Local verification: 57 focused tests, Black/isort, targeted mypy, and git diff --check all pass. Hosted compatibility/format/type/docstring/benchmark/notebook checks are green; full coverage is still running.

[jlarson4]

Hi @koriyoshi2041, just pushed a CI update to your branch that should help stabilize Full Code Coverage, there was a flaw in dev that was causing failures. Everything looks good, once this run passes I'll merge, thanks for your patience!