Skip to content

[T7] decompose.prepare_for_export mis-handles Gemma-family RMSNorm #290

Description

@DsThakurRawat

Labels: bug, export, parity, priority:medium
File: src/tether/decompose.py:212-226 (decompose_rmsnorm), :257 (prepare_for_export); correct dispatcher unused at :56 (swap_rmsnorm_variants) / :36 (DecomposedGemmaRMSNorm)

Description. The public prepare_for_export/decompose_rmsnorm apply Llama-style DecomposedRMSNorm to every RMSNorm, ignoring Gemma's (1 + weight) parameterization that DecomposedGemmaRMSNorm exists to handle. Production exporters hand-pick the right class, but the public API and its tests bless the wrong behavior.

Why it matters. Anyone calling prepare_for_export(model) on a PaliGemma backbone (pi0/pi05) gets silently wrong normalization → a runnable-but-incorrect ONNX. A latent trap for contributors following the public API.

Tasks.

  • Route decompose_rmsnorm/prepare_for_export through the same class-name dispatch as swap_rmsnorm_variants (Gemma → DecomposedGemmaRMSNorm), or delete the legacy functions if dead.

Dedup. #21 is a pi0 attention patch — different. Not covered.

CI / TOOLING / PACKAGING / DOCS

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions