File size: 3,494 Bytes
16533aa | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, AutoModelForCausalLM
from .configuration_parallel_mlp import ParallelMLPConfig
class ParallelMLPAdapter(nn.Module):
"""Simple bottleneck MLP (norm → up → SiLU → down) parallel to a transformer block.
down_proj is zero-initialized so the adapter is a strict pass-through at the
start of training, avoiding representation shock.
"""
def __init__(self, hidden_size: int, intermediate_size: int, rms_norm_eps: float = 1e-5):
super().__init__()
self.norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
nn.init.zeros_(self.down_proj.weight)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x + self.down_proj(F.silu(self.up_proj(self.norm(x))))
class ParallelMLPBlockWrapper(nn.Module):
"""Wraps a transformer decoder layer, adding a parallel MLP adapter.
The adapter sees the same pre-block hidden states as the original block.
Its delta (adapter output minus input) is added to the original block's output,
so neither path depends on the other — true parallel execution.
"""
def __init__(self, original_block: nn.Module, mlp_adapter: ParallelMLPAdapter):
super().__init__()
self.original_block = original_block
self.mlp_adapter = mlp_adapter # name used by freeze filter: "mlp_adapter"
def __getattr__(self, name: str):
try:
return super().__getattr__(name)
except AttributeError:
return getattr(self.original_block, name)
def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
block_outputs = self.original_block(hidden_states, *args, **kwargs)
out_hidden = block_outputs[0] if isinstance(block_outputs, tuple) else block_outputs
# adapter(x) returns x + delta; extract only the delta to add to block output
adapter_delta = self.mlp_adapter(hidden_states) - hidden_states
out_hidden = out_hidden + adapter_delta
if isinstance(block_outputs, tuple):
return (out_hidden,) + block_outputs[1:]
return out_hidden
class UnifiedParallelMLPForCausalLM(PreTrainedModel):
config_class = ParallelMLPConfig
def __init__(self, config: ParallelMLPConfig, **kwargs):
super().__init__(config)
self.backbone = AutoModelForCausalLM.from_pretrained(
config.base_model_name_or_path,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
attn_implementation="flash_attention_2",
)
backbone_dtype = next(self.backbone.parameters()).dtype
layers = self.backbone.model.layers
for pos in config.mlp_positions:
if 0 < pos <= len(layers):
original_layer = layers[pos - 1]
adapter = ParallelMLPAdapter(
config.hidden_size,
config.mlp_intermediate_size,
config.rms_norm_eps,
).to(dtype=backbone_dtype)
layers[pos - 1] = ParallelMLPBlockWrapper(original_layer, adapter)
def forward(self, *args, **kwargs):
return self.backbone(*args, **kwargs)
def generate(self, *args, **kwargs):
return self.backbone.generate(*args, **kwargs)
|