| 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 |
|
|
| 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_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) |
|
|
