model

.gitignore

@@ -31,6 +31,15 @@ outputs/
 logs/
 wandb/
 
+# Test cache
+.pytest_cache/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Log files
+*.log
+
 # OS
 .DS_Store
 Thumbs.db

src/easytranslate/model/attention.py

@@ -57,8 +57,15 @@ class MultiHeadAttention(nn.Module):
         self.d_model = d_model
         self.nhead = nhead
         self.d_k = d_model // nhead
+        self.use_rotary_embedding = use_rotary_embedding
 
-        raise NotImplementedError("TODO: Person B 实现 MultiHeadAttention.__init__")
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+        self.dropout = nn.Dropout(p=dropout)
+        self.rope: Optional[nn.Module] = None
 
     def forward(
         self,
@@ -67,8 +74,61 @@ class MultiHeadAttention(nn.Module):
         value: torch.Tensor,   # [B, L_v, D]
         key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, L_k]
         attn_mask: Optional[torch.Tensor] = None,             # [L_q, L_k]
+        is_causal: bool = False,
     ) -> torch.Tensor:
-        raise NotImplementedError("TODO: Person B 实现 MultiHeadAttention.forward")
+        B, L_q, _ = query.size()
+        L_k = key.size(1)
+        L_v = value.size(1)
+
+        # 1. 线性投影
+        Q = self.q_proj(query)  # [B, L_q, D]
+        K = self.k_proj(key)    # [B, L_k, D]
+        V = self.v_proj(value)  # [B, L_v, D]
+
+        # 2. reshape 为 [B, nhead, L, d_k]
+        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
+        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
+        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]
+
+        # 3. (可选) 应用 RoPE
+        if self.use_rotary_embedding and self.rope is not None:
+            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
+
+        # 4. 计算 attention scores: QK^T / sqrt(d_k)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, H, L_q, L_k]
+
+        # 5. 应用 masks
+        if key_padding_mask is not None:
+            # key_padding_mask: [B, L_k] -> [B, 1, 1, L_k]
+            scores = scores.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
+            )
+        if is_causal:
+            # 生成 causal mask
+            L_q, L_k_local = scores.size(-2), scores.size(-1)
+            causal_mask = torch.triu(
+                torch.ones(L_q, L_k_local, device=scores.device), diagonal=1
+            ).bool()
+            scores = scores.masked_fill(
+                causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")
+            )
+        if attn_mask is not None:
+            # attn_mask: [L_q, L_k] -> [1, 1, L_q, L_k]
+            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+        # 6. Softmax + Dropout
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        # 7. 加权求和 V
+        attn_output = torch.matmul(attn_weights, V)  # [B, H, L_q, d_k]
+
+        # 8. reshape 回 [B, L_q, d_model]
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
+
+        # 9. 输出投影
+        output = self.out_proj(attn_output)
+        return output
 
 
 class FlashMultiHeadAttention(nn.Module):
@@ -93,7 +153,19 @@ class FlashMultiHeadAttention(nn.Module):
         use_rotary_embedding: bool = False,
     ):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 FlashMultiHeadAttention.__init__")
+        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
+        self.d_model = d_model
+        self.nhead = nhead
+        self.d_k = d_model // nhead
+        self.use_rotary_embedding = use_rotary_embedding
+        self.dropout_p = dropout
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+        self.rope: Optional[nn.Module] = None
 
     def forward(
         self,
@@ -103,4 +175,57 @@ class FlashMultiHeadAttention(nn.Module):
         key_padding_mask: Optional[torch.BoolTensor] = None,
         is_causal: bool = False,
     ) -> torch.Tensor:
-        raise NotImplementedError("TODO: Person B 实现 FlashMultiHeadAttention.forward")
+        B, L_q, _ = query.size()
+        L_k = key.size(1)
+        L_v = value.size(1)
+
+        # 1. 线性投影
+        Q = self.q_proj(query)  # [B, L_q, D]
+        K = self.k_proj(key)    # [B, L_k, D]
+        V = self.v_proj(value)  # [B, L_v, D]
+
+        # 2. reshape 为 [B, nhead, L, d_k]
+        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
+        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
+        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]
+
+        # 3. (可选) 应用 RoPE
+        if self.use_rotary_embedding and self.rope is not None:
+            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
+
+        # 4. 构建 attn_mask 以适配 scaled_dot_product_attention
+        # PyTorch >= 2.0 支持 [B, nhead, L, d_k] 的 4D 输入
+        # 注意: scaled_dot_product_attention 不允许同时设置 attn_mask 和 is_causal=True
+        attn_mask: Optional[torch.Tensor] = None
+        if is_causal or key_padding_mask is not None:
+            attn_mask = torch.zeros(
+                B, self.nhead, L_q, L_k, dtype=Q.dtype, device=Q.device
+            )
+            if is_causal:
+                # 生成 causal mask (上三角为 -inf)
+                causal_mask = torch.triu(
+                    torch.ones(L_q, L_k, device=Q.device), diagonal=1
+                ).bool()
+                attn_mask = attn_mask.masked_fill(
+                    causal_mask[None, None, :, :], float("-inf")
+                )
+            if key_padding_mask is not None:
+                # key_padding_mask: True = padding (忽略)
+                _bool_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
+                _bool_mask = _bool_mask.expand(B, self.nhead, L_q, L_k)
+                attn_mask = attn_mask.masked_fill(_bool_mask, float("-inf"))
+
+        # 5. Flash Attention (PyTorch 原生)
+        attn_output = F.scaled_dot_product_attention(
+            Q, K, V,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=False,  # 已通过 attn_mask 处理
+        )  # [B, H, L_q, d_k]
+
+        # 6. reshape 回 [B, L_q, d_model]
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
+
+        # 7. 输出投影
+        output = self.out_proj(attn_output)
+        return output

src/easytranslate/model/decoder.py

@@ -13,21 +13,23 @@ Transformer Decoder 模块 — Person B 负责实现
 
 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。
 
-    TODO [Person B]: 实现以下组件:
-    1. Masked Self-Attention (因果掩码，防止看到未来)
-    2. Cross-Attention (decoder 查询 encoder 输出)
-    3. Feed-Forward Network
-    4. 三个 LayerNorm
-    5. Residual connections + Dropout
+    架构 (Pre-LayerNorm):
+        x → LN → Masked Self-Attention → Residual
+          → LN → Cross-Attention → Residual
+          → LN → FFN → Residual
     """
 
     def __init__(
@@ -42,7 +44,42 @@ class TransformerDecoderLayer(nn.Module):
         pre_norm: bool = True,
     ):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 TransformerDecoderLayer.__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,
@@ -52,28 +89,71 @@ class TransformerDecoderLayer(nn.Module):
         memory_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
         tgt_key_padding_mask: Optional[torch.BoolTensor] = None,     # [B, T]
     ) -> torch.Tensor:
-        """
-        TODO [Person B]: Pre-LayerNorm 前向传播:
-        1. Masked Self-Attention with causal mask
-        2. Cross-Attention with encoder output
-        3. FFN
-        每步都有 residual connection 和 dropout
-        """
-        raise NotImplementedError("TODO: Person B 实现 TransformerDecoderLayer.forward")
+        """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。
-
-    TODO [Person B]:
-    1. 堆叠 N 个 TransformerDecoderLayer
-    2. 最终加一个 LayerNorm
     """
 
     def __init__(self, decoder_layer: TransformerDecoderLayer, num_layers: int):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 TransformerDecoder.__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,
@@ -83,4 +163,14 @@ class TransformerDecoder(nn.Module):
         memory_key_padding_mask: Optional[torch.BoolTensor] = None,
         tgt_key_padding_mask: Optional[torch.BoolTensor] = None,
     ) -> torch.Tensor:
-        raise NotImplementedError("TODO: Person B 实现 TransformerDecoder.forward")
+        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

src/easytranslate/model/encoder.py

@@ -11,23 +11,21 @@ Transformer Encoder 模块 — Person B 负责实现
 
 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 TransformerEncoderLayer(nn.Module):
     """
     单层 Transformer Encoder。
 
-    TODO [Person B]: 实现以下组件:
-    1. Self-Attention: MultiHeadAttention (支持 Flash Attention)
-    2. Feed-Forward Network: Linear → Activation → Dropout → Linear
-    3. 两个 LayerNorm
-    4. Residual connections
-    5. Dropout
-
-    注意: 使用 Pre-LayerNorm 架构 (先 norm 再 attention/ffn)
+    架构 (Pre-LayerNorm):
+        x → LayerNorm → MultiHeadAttention → Residual → LayerNorm → FFN → Residual
     """
 
     def __init__(
@@ -42,43 +40,88 @@ class TransformerEncoderLayer(nn.Module):
         pre_norm: bool = True,
     ):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 TransformerEncoderLayer.__init__")
+        self.pre_norm = pre_norm
+
+        # Self-Attention
+        attn_cls = FlashMultiHeadAttention if use_flash_attention else MultiHeadAttention
+        self.self_attn = attn_cls(
+            d_model=d_model,
+            nhead=nhead,
+            dropout=dropout,
+            use_rotary_embedding=use_rotary_embedding,
+        )
+
+        # 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)
+
+        # LayerNorm
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # Dropout for residual
+        self.dropout1 = nn.Dropout(p=dropout)
+        self.dropout2 = nn.Dropout(p=dropout)
 
     def forward(
         self,
         src: torch.Tensor,                             # [B, S, D]
         src_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
     ) -> torch.Tensor:
-        """
-        TODO [Person B]: Pre-LayerNorm 前向传播:
-        1. residual = src
-        2. src = layer_norm_1(src)
-        3. src = self_attention(src, src, src, key_padding_mask=src_key_padding_mask)
-        4. src = residual + dropout(src)
-        5. residual = src
-        6. src = layer_norm_2(src)
-        7. src = ffn(src)
-        8. src = residual + dropout(src)
-        """
-        raise NotImplementedError("TODO: Person B 实现 TransformerEncoderLayer.forward")
+        """Pre-LayerNorm 前向传播。"""
+        if self.pre_norm:
+            # 1. Self-Attention sublayer
+            residual = src
+            src = self.norm1(src)
+            src = self.self_attn(
+                src, src, src,
+                key_padding_mask=src_key_padding_mask,
+            )
+            src = residual + self.dropout1(src)
+
+            # 2. FFN sublayer
+            residual = src
+            src = self.norm2(src)
+            src = self.linear2(self.dropout(self.activation(self.linear1(src))))
+            src = residual + self.dropout2(src)
+        else:
+            # Post-LayerNorm (备用)
+            residual = src
+            src = self.self_attn(
+                src, src, src,
+                key_padding_mask=src_key_padding_mask,
+            )
+            src = self.norm1(residual + self.dropout1(src))
+
+            residual = src
+            src = self.linear2(self.dropout(self.activation(self.linear1(src))))
+            src = self.norm2(residual + self.dropout2(src))
+
+        return src
 
 
 class TransformerEncoder(nn.Module):
     """
     多层 Transformer Encoder。
-
-    TODO [Person B]:
-    1. 堆叠 N 个 TransformerEncoderLayer
-    2. 最终加一个 LayerNorm (Pre-Norm 架构需要)
     """
 
     def __init__(self, encoder_layer: TransformerEncoderLayer, num_layers: int):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 TransformerEncoder.__init__")
+        self.layers = nn.ModuleList(
+            [copy.deepcopy(encoder_layer) for _ in range(num_layers)]
+        )
+        self.num_layers = num_layers
+        self.norm = nn.LayerNorm(encoder_layer.self_attn.d_model)
 
     def forward(
         self,
         src: torch.Tensor,
         src_key_padding_mask: Optional[torch.BoolTensor] = None,
     ) -> torch.Tensor:
-        raise NotImplementedError("TODO: Person B 实现 TransformerEncoder.forward")
+        output = src
+        for layer in self.layers:
+            output = layer(output, src_key_padding_mask=src_key_padding_mask)
+        output = self.norm(output)
+        return output

src/easytranslate/model/finetune.py

@@ -50,7 +50,23 @@ def load_pretrained_model(
     Returns:
         (model, tokenizer)
     """
-    raise NotImplementedError("TODO: Person B 实现 load_pretrained_model")
+    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    # 设置语言对
+    if hasattr(tokenizer, "lang_code_to_id"):
+        tokenizer.src_lang = src_lang
+        tokenizer.tgt_lang = tgt_lang
+    if hasattr(model.config, "forced_bos_token_id"):
+        model.config.forced_bos_token_id = tokenizer.lang_code_to_id.get(
+            tgt_lang, tokenizer.bos_token_id
+        )
+
+    model = model.to(device)
+    logger.info(f"Loaded pretrained model: {model_name}")
+    return model, tokenizer
 
 
 def setup_lora(
@@ -63,33 +79,38 @@ def setup_lora(
     """
     为模型配置 LoRA 微调。
 
-    TODO [Person B]: 实现以下逻辑:
-    1. 定义 LoraConfig:
-       - r: LoRA 秩 (低秩分解维度)
-       - lora_alpha: 缩放因子
-       - lora_dropout: LoRA dropout
-       - target_modules: 需要加 LoRA 的模块 (如 q_proj, v_proj)
-       - task_type: SEQ_2_SEQ_LM
-    2. 使用 get_peft_model(model, config) 包装模型
-    3. 打印可训练参数数量和比例
-    4. 返回 LoRA 模型
-
-    参考: https://huggingface.co/docs/peft
-
     Returns:
         peft_model: LoRA 包装后的模型
     """
+    from peft import LoraConfig, get_peft_model, TaskType
+
     if target_modules is None:
         target_modules = ["q_proj", "v_proj"]
 
-    raise NotImplementedError("TODO: Person B 实现 setup_lora")
+    config = LoraConfig(
+        r=r,
+        lora_alpha=alpha,
+        lora_dropout=dropout,
+        target_modules=target_modules,
+        task_type=TaskType.SEQ_2_SEQ_LM,
+    )
 
+    model = get_peft_model(model, config)
 
-def freeze_model_except_lora(model: nn.Module):
-    """
-    冻结模型所有参数，只保留 LoRA 参数可训练。
+    # 打印可训练参数信息
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    ratio = 100 * trainable_params / total_params if total_params > 0 else 0
+    logger.info(
+        f"LoRA setup: trainable params={trainable_params:,} "
+        f"/ total={total_params:,} ({ratio:.4f}%)"
+    )
 
-    TODO [Person B]: 遍历 model.named_parameters()，
-    如果参数名中不包含 "lora"，则设置 requires_grad = False。
-    """
-    raise NotImplementedError("TODO: Person B 实现 freeze_model_except_lora")
+    return model
+
+
+def freeze_model_except_lora(model: nn.Module):
+    """冻结模型所有参数，只保留 LoRA 参数可训练。"""
+    for name, param in model.named_parameters():
+        if "lora" not in name.lower():
+            param.requires_grad = False

src/easytranslate/model/positional.py

@@ -32,7 +32,17 @@ class SinusoidalPositionalEncoding(nn.Module):
 
     def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 SinusoidalPositionalEncoding.__init__")
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_seq_len, d_model)
+        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]
+        self.register_buffer("pe", pe)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -41,7 +51,8 @@ class SinusoidalPositionalEncoding(nn.Module):
         Returns:
             x + positional_encoding: [B, L, D]
         """
-        raise NotImplementedError("TODO: Person B 实现 SinusoidalPositionalEncoding.forward")
+        x = x + self.pe[:, : x.size(1), :]
+        return self.dropout(x)
 
 
 class RotaryPositionalEmbedding(nn.Module):
@@ -72,10 +83,29 @@ class RotaryPositionalEmbedding(nn.Module):
 
     def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
         super().__init__()
-        raise NotImplementedError("TODO: Person B 实现 RotaryPositionalEmbedding.__init__")
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2).float() / dim)
+        )
+        self.register_buffer("inv_freq", inv_freq)
+
+        # 预缓存 cos/sin
+        self._cached_seq_len = 0
+        self._cached_cos: torch.Tensor | None = None
+        self._cached_sin: torch.Tensor | None = None
 
     def _compute_rope(self, seq_len: int, device: torch.device):
-        raise NotImplementedError("TODO: Person B 实现 _compute_rope")
+        if seq_len > self._cached_seq_len or self._cached_cos is None:
+            t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+            freqs = torch.outer(t, self.inv_freq)  # [seq_len, dim//2]
+            emb = torch.cat((freqs, freqs), dim=-1)  # [seq_len, dim]
+            self._cached_cos = emb.cos()[None, None, :, :]  # [1, 1, seq_len, dim]
+            self._cached_sin = emb.sin()[None, None, :, :]  # [1, 1, seq_len, dim]
+            self._cached_seq_len = seq_len
+        return self._cached_cos, self._cached_sin
 
     @staticmethod
     def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -94,4 +124,11 @@ class RotaryPositionalEmbedding(nn.Module):
         Returns:
             (q_rotated, k_rotated)
         """
-        raise NotImplementedError("TODO: Person B 实现 apply_rotary_pos_emb")
+        cos, sin = self._compute_rope(q.size(2), q.device)
+        q_embed = (q * cos[:, :, : q.size(2), :]) + (
+            self._rotate_half(q) * sin[:, :, : q.size(2), :]
+        )
+        k_embed = (k * cos[:, :, : k.size(2), :]) + (
+            self._rotate_half(k) * sin[:, :, : k.size(2), :]
+        )
+        return q_embed, k_embed

src/easytranslate/model/transformer.py

@@ -20,8 +20,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from easytranslate.model.encoder import TransformerEncoder
-from easytranslate.model.decoder import TransformerDecoder
+from easytranslate.model.encoder import TransformerEncoder, TransformerEncoderLayer
+from easytranslate.model.decoder import TransformerDecoder, TransformerDecoderLayer
 from easytranslate.model.positional import SinusoidalPositionalEncoding, RotaryPositionalEmbedding
 
 logger = logging.getLogger(__name__)
@@ -78,33 +78,94 @@ class TransformerTranslationModel(nn.Module):
         share_embedding: bool = False,
     ):
         super().__init__()
-        # TODO [Person B]: 实现模型初始化
-        # 保存超参数
         self.d_model = d_model
         self.pad_id = pad_id
+        self.use_rotary_embedding = use_rotary_embedding
+        self.max_seq_len = max_seq_len
 
-        raise NotImplementedError("TODO: Person B 实现 TransformerTranslationModel.__init__")
+        # Embeddings
+        self.src_embed = nn.Embedding(src_vocab_size, d_model)
+        self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)
+        self.embed_scale = math.sqrt(d_model)
 
-    def _init_weights(self):
-        """
-        参数初始化。
+        # Positional encoding
+        if use_rotary_embedding:
+            # RoPE 在 attention 内部应用到 Q/K，不需要额外的位置编码层
+            self.pos_encoding: Optional[nn.Module] = None
+            rope = RotaryPositionalEmbedding(
+                dim=d_model // nhead,
+                max_seq_len=max_seq_len,
+            )
+        else:
+            self.pos_encoding = SinusoidalPositionalEncoding(
+                d_model=d_model,
+                max_seq_len=max_seq_len,
+                dropout=dropout,
+            )
+            rope = None
 
-        TODO [Person B]: 实现 Xavier/Kaiming 初始化:
-        - Embedding: normal_(0, d_model^-0.5)
-        - Linear: xavier_uniform_
-        - LayerNorm: ones_ / zeros_
-        """
-        raise NotImplementedError("TODO: Person B 实现 _init_weights")
+        # Encoder
+        encoder_layer = TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            use_flash_attention=use_flash_attention,
+            use_rotary_embedding=use_rotary_embedding,
+            pre_norm=pre_norm,
+        )
+        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers)
 
-    def _generate_square_subsequent_mask(self, sz: int, device: torch.device) -> torch.Tensor:
-        """
-        生成因果注意力掩码 (causal mask)。
+        # Decoder
+        decoder_layer = TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            use_flash_attention=use_flash_attention,
+            use_rotary_embedding=use_rotary_embedding,
+            pre_norm=pre_norm,
+        )
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
 
-        TODO [Person B]:
-        返回上三角矩阵 mask，shape [sz, sz]，
-        mask[i][j] = -inf if j > i else 0
-        """
-        raise NotImplementedError("TODO: Person B 实现 _generate_square_subsequent_mask")
+        # 将 RoPE 注入到 encoder/decoder 的 attention 模块中
+        if rope is not None:
+            for layer in self.encoder.layers:
+                layer.self_attn.rope = rope
+            for layer in self.decoder.layers:
+                layer.self_attn.rope = rope
+                # cross-attention 不使用 RoPE
+                layer.multihead_attn.rope = None
+
+        # Output projection
+        self.output_projection = nn.Linear(d_model, tgt_vocab_size)
+
+        # 可选: 共享目标语言 embedding 和输出投影权重
+        self.share_embedding = share_embedding
+        if share_embedding:
+            self.output_projection.weight = self.tgt_embed.weight
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """参数初始化。"""
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+        for module in self.modules():
+            if isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0, std=self.d_model ** -0.5)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def _generate_square_subsequent_mask(self, sz: int, device: torch.device) -> torch.Tensor:
+        """生成因果注意力掩码 (causal mask)。"""
+        mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1)
+        mask = mask.masked_fill(mask == 1, float("-inf"))
+        return mask
 
     def forward(
         self,
@@ -119,19 +180,50 @@ class TransformerTranslationModel(nn.Module):
         Returns:
             logits: [B, T, tgt_vocab_size]
         """
-        raise NotImplementedError("TODO: Person B 实现 forward")
+        # 1. Embedding
+        src_emb = self.src_embed(src_ids) * self.embed_scale  # [B, S, D]
+        tgt_emb = self.tgt_embed(tgt_input_ids) * self.embed_scale  # [B, T, D]
+
+        # 2. Positional encoding (如果不使用 RoPE)
+        if self.pos_encoding is not None:
+            src_emb = self.pos_encoding(src_emb)
+            tgt_emb = self.pos_encoding(tgt_emb)
+
+        # 3. Masks
+        if src_padding_mask is None:
+            src_padding_mask = src_ids.eq(self.pad_id)
+        if tgt_padding_mask is None:
+            tgt_padding_mask = tgt_input_ids.eq(self.pad_id)
+
+        tgt_seq_len = tgt_input_ids.size(1)
+        tgt_mask = self._generate_square_subsequent_mask(tgt_seq_len, tgt_input_ids.device)
+
+        # 4. Encoder
+        encoder_output = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
+
+        # 5. Decoder
+        decoder_output = self.decoder(
+            tgt_emb,
+            encoder_output,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=src_padding_mask,
+            tgt_key_padding_mask=tgt_padding_mask,
+        )
+
+        # 6. Output projection
+        logits = self.output_projection(decoder_output)
+        return logits
 
     @torch.no_grad()
     def encode(self, src_ids: torch.Tensor, src_padding_mask: Optional[torch.BoolTensor] = None) -> torch.Tensor:
-        """
-        仅编码（用于推理时复用 encoder 输出）。
-
-        TODO [Person B]:
-        1. src embedding + positional encoding
-        2. encoder forward
-        3. 返回 encoder_output
-        """
-        raise NotImplementedError("TODO: Person B 实现 encode")
+        """仅编码（用于推理时复用 encoder 输出）。"""
+        src_emb = self.src_embed(src_ids) * self.embed_scale
+        if self.pos_encoding is not None:
+            src_emb = self.pos_encoding(src_emb)
+        if src_padding_mask is None:
+            src_padding_mask = src_ids.eq(self.pad_id)
+        encoder_output = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
+        return encoder_output
 
     @torch.no_grad()
     def decode_step(
@@ -140,17 +232,24 @@ class TransformerTranslationModel(nn.Module):
         encoder_output: torch.Tensor,
         src_padding_mask: Optional[torch.BoolTensor] = None,
     ) -> torch.Tensor:
-        """
-        解码一步（用于自回归推理）。
-
-        TODO [Person B]:
-        1. tgt embedding + positional encoding
-        2. 生成 causal mask
-        3. decoder forward
-        4. 取最后一个 token 的 logits
-        5. 返回 logits [B, vocab_size]
-        """
-        raise NotImplementedError("TODO: Person B 实现 decode_step")
+        """解码一步（用于自回归推理）。"""
+        tgt_emb = self.tgt_embed(tgt_input_ids) * self.embed_scale
+        if self.pos_encoding is not None:
+            tgt_emb = self.pos_encoding(tgt_emb)
+
+        tgt_seq_len = tgt_input_ids.size(1)
+        tgt_mask = self._generate_square_subsequent_mask(tgt_seq_len, tgt_input_ids.device)
+
+        decoder_output = self.decoder(
+            tgt_emb,
+            encoder_output,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=src_padding_mask,
+        )
+
+        # 取最后一个 token 的 logits
+        logits = self.output_projection(decoder_output[:, -1, :])
+        return logits
 
     def count_parameters(self) -> int:
         """返回可训练参数数量。"""