File size: 5,617 Bytes
c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 4d62693 | 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | """
Transformer Decoder 模块 — Person B 负责实现
包含:
- TransformerDecoderLayer: 单层解码器
- TransformerDecoder: 多层解码器堆叠
架构 (Pre-LayerNorm):
x → LN → Masked Self-Attention → Residual
→ LN → Cross-Attention → Residual
→ LN → FFN → Residual
"""
from __future__ import annotations
import copy
import torch
import torch.nn as nn
from typing import Optional
from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
class TransformerDecoderLayer(nn.Module):
"""
单层 Transformer Decoder。
架构 (Pre-LayerNorm):
x → LN → Masked Self-Attention → Residual
→ LN → Cross-Attention → Residual
→ LN → FFN → Residual
"""
def __init__(
self,
d_model: int = 512,
nhead: int = 8,
dim_feedforward: int = 2048,
dropout: float = 0.1,
activation: str = "gelu",
use_flash_attention: bool = True,
use_rotary_embedding: bool = True,
pre_norm: bool = True,
):
super().__init__()
self.pre_norm = pre_norm
self.d_model = d_model
attn_cls = FlashMultiHeadAttention if use_flash_attention else MultiHeadAttention
# 1. Masked Self-Attention
self.self_attn = attn_cls(
d_model=d_model,
nhead=nhead,
dropout=dropout,
use_rotary_embedding=use_rotary_embedding,
)
# 2. Cross-Attention (decoder queries encoder memory)
self.multihead_attn = attn_cls(
d_model=d_model,
nhead=nhead,
dropout=dropout,
use_rotary_embedding=False, # cross-attention 不使用 RoPE
)
# 3. Feed-Forward Network
self.linear1 = nn.Linear(d_model, dim_feedforward)
self.activation = nn.GELU() if activation == "gelu" else nn.ReLU()
self.dropout = nn.Dropout(p=dropout)
self.linear2 = nn.Linear(dim_feedforward, d_model)
# 4. LayerNorms
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
# 5. Dropouts for residuals
self.dropout1 = nn.Dropout(p=dropout)
self.dropout2 = nn.Dropout(p=dropout)
self.dropout3 = nn.Dropout(p=dropout)
def forward(
self,
tgt: torch.Tensor, # [B, T, D]
memory: torch.Tensor, # [B, S, D] (encoder output)
tgt_mask: Optional[torch.Tensor] = None, # [T, T] causal mask
memory_key_padding_mask: Optional[torch.BoolTensor] = None, # [B, S]
tgt_key_padding_mask: Optional[torch.BoolTensor] = None, # [B, T]
) -> torch.Tensor:
"""Pre-LayerNorm 前向传播。"""
# 适配 Flash Attention: 使用 is_causal 替代显式 causal mask
is_causal = tgt_mask is not None
if self.pre_norm:
# 1. Masked Self-Attention
residual = tgt
tgt = self.norm1(tgt)
tgt = self.self_attn(
tgt, tgt, tgt,
key_padding_mask=tgt_key_padding_mask,
is_causal=is_causal,
)
tgt = residual + self.dropout1(tgt)
# 2. Cross-Attention
residual = tgt
tgt = self.norm2(tgt)
tgt = self.multihead_attn(
tgt, memory, memory,
key_padding_mask=memory_key_padding_mask,
)
tgt = residual + self.dropout2(tgt)
# 3. FFN
residual = tgt
tgt = self.norm3(tgt)
tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
tgt = residual + self.dropout3(tgt)
else:
# Post-LayerNorm (备用)
residual = tgt
tgt = self.self_attn(
tgt, tgt, tgt,
key_padding_mask=tgt_key_padding_mask,
is_causal=is_causal,
)
tgt = self.norm1(residual + self.dropout1(tgt))
residual = tgt
tgt = self.multihead_attn(
tgt, memory, memory,
key_padding_mask=memory_key_padding_mask,
)
tgt = self.norm2(residual + self.dropout2(tgt))
residual = tgt
tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
tgt = self.norm3(residual + self.dropout3(tgt))
return tgt
class TransformerDecoder(nn.Module):
"""
多层 Transformer Decoder。
"""
def __init__(self, decoder_layer: TransformerDecoderLayer, num_layers: int):
super().__init__()
self.layers = nn.ModuleList(
[copy.deepcopy(decoder_layer) for _ in range(num_layers)]
)
self.num_layers = num_layers
self.norm = nn.LayerNorm(decoder_layer.d_model)
def forward(
self,
tgt: torch.Tensor,
memory: torch.Tensor,
tgt_mask: Optional[torch.Tensor] = None,
memory_key_padding_mask: Optional[torch.BoolTensor] = None,
tgt_key_padding_mask: Optional[torch.BoolTensor] = None,
) -> torch.Tensor:
output = tgt
for layer in self.layers:
output = layer(
output,
memory,
tgt_mask=tgt_mask,
memory_key_padding_mask=memory_key_padding_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
)
output = self.norm(output)
return output
|