File size: 3,885 Bytes
8b099fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Deeper audio projector for the HF Gemma4 unified model.
The stock model maps audio encoder features into LM space with a single
`nn.Linear(1536 -> 5376, bias=False)` at `model.model.embed_audio.embedding_projection`.
`DeepAudioProjector` is a drop-in replacement that keeps that linear as a
warm-started backbone and adds a *zero-initialised* residual MLP, so at init
the output is bit-identical to the original projector (we reuse the learned
phase-1 weights exactly), and SFT learns the deeper correction on top.
It exposes a `.weight` property because the parent `embed_audio.forward` reads
`self.embedding_projection.weight.dtype` to decide the input cast dtype.
"""
from __future__ import annotations
import torch
import torch.nn as nn
class DeepAudioProjector(nn.Module):
def __init__(
self,
in_dim: int = 1536,
out_dim: int = 5376,
hidden: int = 4096,
n_hidden_layers: int = 2,
dropout: float = 0.0,
out_dtype: torch.dtype = torch.bfloat16,
):
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.out_dtype = out_dtype
# Warm-started backbone (copied from the trained phase-1 Linear).
self.proj = nn.Linear(in_dim, out_dim, bias=False)
# Residual branch (starts at exactly 0 -> identity behaviour at init).
self.ln = nn.LayerNorm(in_dim)
layers: list[nn.Module] = []
d = in_dim
for _ in range(n_hidden_layers):
layers += [nn.Linear(d, hidden), nn.GELU()]
if dropout > 0:
layers.append(nn.Dropout(dropout))
d = hidden
self.mlp = nn.Sequential(*layers)
self.out = nn.Linear(d, out_dim)
nn.init.zeros_(self.out.weight)
nn.init.zeros_(self.out.bias)
@property
def weight(self) -> torch.Tensor:
# Parent module reads `.weight.dtype`; expose the backbone weight.
return self.proj.weight
def forward(self, x: torch.Tensor) -> torch.Tensor:
wdt = self.proj.weight.dtype
x = x.to(wdt)
base = self.proj(x)
res = self.out(self.mlp(self.ln(x)))
return (base + res).to(self.out_dtype)
@classmethod
def from_linear(cls, linear: nn.Linear, **kwargs) -> "DeepAudioProjector":
out_dim, in_dim = linear.weight.shape
m = cls(in_dim=in_dim, out_dim=out_dim, **kwargs)
with torch.no_grad():
m.proj.weight.copy_(linear.weight)
return m
def find_audio_projection_parent(model):
"""Return (parent_module, attr_name) for the audio embedding_projection."""
inner = getattr(model, "model", model)
embed_audio = getattr(inner, "embed_audio", None)
if embed_audio is None:
raise AttributeError("Could not find model.model.embed_audio")
if not hasattr(embed_audio, "embedding_projection"):
raise AttributeError("embed_audio has no embedding_projection")
return embed_audio, "embedding_projection"
def install_deep_projector(
model,
hidden: int = 4096,
n_hidden_layers: int = 2,
dropout: float = 0.0,
param_dtype: torch.dtype = torch.float32,
):
"""Replace the audio embedding_projection with a warm-started DeepAudioProjector.
Returns the new projector module (params left in `param_dtype`, e.g. fp32 for
stable optimisation; forward output is cast back to the LM dtype).
"""
parent, attr = find_audio_projection_parent(model)
old = getattr(parent, attr)
assert isinstance(old, nn.Linear), f"expected nn.Linear, got {type(old)}"
out_dtype = old.weight.dtype
deep = DeepAudioProjector.from_linear(
old, hidden=hidden, n_hidden_layers=n_hidden_layers,
dropout=dropout, out_dtype=out_dtype,
)
deep = deep.to(device=old.weight.device, dtype=param_dtype)
setattr(parent, attr, deep)
return deep
|