前言

這篇筆記整理 torch.nn.Transformer 中一個常見限制：d_model 必須能被 nhead 整除。

從理論上看，Multi-Head Attention 可以想成對同一份 input 做多組 Self-Attention，再把多個 head 的輸出接起來。因此直覺上會覺得 feature size 與 head 數量不一定要整除。但 PyTorch 的實作為了效率，會把 embedding dimension 平均切給每個 head，這就是限制的來源。

Multi-Head Transformer 理論

Multi-Head Transformer 的概念是：對同一個 input 做多個 Self-Attention，將多次輸出 concat 後，再透過一個矩陣投影回原本大小。具體流程如下：

在這個抽象描述裡，輸入 x 的 feature size 和 head 數量看起來沒有硬性關係。也就是說，如果只看理論流程，不論 feature size 與 head 數量是多少，似乎都可以訓練。

為什麼 d_model 需要被 nhead 整除

如果對 nn.Transformer 填入任意的 feature size 與 head 數量，可能會遇到錯誤訊息，提示 embed_dim 必須能被 num_heads 整除。

原因可以從 PyTorch source code 看出來。首先看 nn.Transformer，其中與 nhead、d_model 相關的部分會進到 TransformerEncoderLayer。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

class Transformer(Module):
    #......
    def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
                 custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None,
                 layer_norm_eps: float = 1e-5, batch_first: bool = False, norm_first: bool = False,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        torch._C._log_api_usage_once(f"torch.nn.modules.{self.__class__.__name__}")

        if custom_encoder is not None:
            self.encoder = custom_encoder
        else:
            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout,
                                                    activation, layer_norm_eps, batch_first, norm_first,
                                                    **factory_kwargs)
            encoder_norm = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
        ### ......

接著看 TransformerEncoderLayer，可以發現真正處理 attention 的類別是 MultiheadAttention。

1
2
3
4
5
6
7
8
9
10
11
12

class TransformerEncoderLayer(Module):
    #......
    def __init__(self, d_model: int, nhead: int, dim_feedforward: int = 2048, dropout: float = 0.1,
                 activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
                 layer_norm_eps: float = 1e-5, batch_first: bool = False, norm_first: bool = False,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
                                            **factory_kwargs)
        #......
    #......

繼續追 MultiheadAttention，會看到它把 forward 的核心邏輯交給 F.multi_head_attention_forward。

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

class MultiheadAttention(Module):
    #.......
    def forward(
            self,
            query: Tensor,
            key: Tensor,
            value: Tensor,
            key_padding_mask: Optional[Tensor] = None,
            need_weights: bool = True,
            attn_mask: Optional[Tensor] = None,
            average_attn_weights: bool = True,
            is_causal : bool = False) -> Tuple[Tensor, Optional[Tensor]]:
        #.......
        if not self._qkv_same_embed_dim:
            attn_output, attn_output_weights = F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask,
                use_separate_proj_weight=True,
                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
                v_proj_weight=self.v_proj_weight,
                average_attn_weights=average_attn_weights,
                is_causal=is_causal)
        else:
            attn_output, attn_output_weights = F.multi_head_attention_forward(
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj.weight, self.out_proj.bias,
                training=self.training,
                key_padding_mask=key_padding_mask,
                need_weights=need_weights,
                attn_mask=attn_mask,
                average_attn_weights=average_attn_weights,
                is_causal=is_causal)
        if self.batch_first and is_batched:
            return attn_output.transpose(1, 0), attn_output_weights
        else:
            return attn_output, attn_output_weights
    #......

最後看 F.multi_head_attention_forward。關鍵在這段：PyTorch 會先用 embed_dim // num_heads 算出每個 head 分到的維度，並 assert 這個拆分必須剛好整除。

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

def multi_head_attention_forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    embed_dim_to_check: int,
    num_heads: int,
    #......
    )
    #......
    else:
        head_dim = embed_dim // num_heads
    assert head_dim * num_heads == embed_dim, f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
    #......
    q = q.view(bsz, num_heads, tgt_len, head_dim)
        k = k.view(bsz, num_heads, src_len, head_dim)
        v = v.view(bsz, num_heads, src_len, head_dim)

        attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
        attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)

        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
        if not is_batched:
            # squeeze the output if input was unbatched
            attn_output = attn_output.squeeze(1)
        return attn_output, None

也就是說，PyTorch 的實作會把 Q、K、V reshape 成 4 維：

• batch size
• number of head
• target/source length
• head dimension

原本的 feature size 會被拆成 number of head * head dimension。如果 embed_dim 不能被 num_heads 整除，就無法平均 reshape，因此會直接 assert。

簡單來說，torch.nn.Transformer 的 multi-head 實作不是把完整 input 重複餵給每個 head，而是將 feature size 平均拆成多份，每份交給不同 head 計算，最後再接回來。這種做法可以大幅節省運算量與記憶體使用，但代價就是 d_model 必須被 nhead 整除。

小結

這個例子展示了理論描述與工程實作之間的差異。理論上 Multi-Head Attention 可以用比較抽象的方式理解，但在框架實作中，為了讓張量 reshape、batch 運算與 GPU 加速更有效率，會加入更明確的維度限制。

這也是讀 framework source code 很有價值的地方：除了理解模型，也能看到實作用哪些假設換取效能與穩定性。

Introduction to Hardware

DNN 常見的運算包含 Convolution 與 FC Layer，這兩類 Layer 的核心多半是大量點積運算，而點積又會仰賴 MAC（Multiply-Accumulate，乘法累加器）。在硬體裡，MAC 通常會放在 ALU 中，因此一塊運算單元會配置不少 ALU 來支撐高吞吐量。

這篇 survey 將 DNN 運算硬體分成兩大類：

• Temporal Architecture：代表是 CPU、GPU。
• Spatial Architecture：代表是 ASIC、FPGA 這類 DNN accelerator。

Temporal Architecture and Spatial Architecture

Temporal Architecture 主要包含 CPU 與 GPU。這類架構通常透過 SIMD、SIMT 與演算法最佳化來縮短運算時間。

Spatial Architecture 則比較接近加速器的設計思路，常見實作是 ASIC 或 FPGA。它會將 Register File 與 Control Logic 放在 ALU 附近，藉此減少資料搬移，降低能耗。

Temporal Architecture

Temporal Architecture 的重點比較偏向「如何把 DNN 運算轉成適合 CPU/GPU 執行的形式」。

FC Layer and Convolutional Layer

FC Layer 本質上就是 2 維 Input 與 Filter 的內積。如果 Input 原本不是 2 維，例如包含 Channel 或 Batch Size，就會先 flatten 成 2 維矩陣，其中一維通常是 Batch Size。轉成矩陣乘法後，CPU/GPU 就可以用 SIMD 或 SIMT 來加速。

Convolutional Layer 也可以用類似方式處理：把 convolution kernel 攤成 1 維陣列，再把 Input 攤成對應的 2 維陣列，最後轉成矩陣內積。缺點是轉換後的 Input feature map 會變大，尤其在多個 Filter 同時計算時，Input feature map 會被重複展開。

Convolution Optimization

作者提到兩種常見的矩陣內積最佳化方法：

• FFT：適合把 convolution 轉到頻域計算。
• Strassen Algorithm：適合加速矩陣乘法。

FFT 的複雜度可從 O(No^2 * Nf^2) 降到 O(No^2 * log No)，其中 Nf 是 Filter 的長寬，No 是 Output 的長寬。不過 Nf^2 不一定比 log No 大，因此作者建議 Filter 較大時，例如 filter > 5 * 5，再考慮 FFT。

Strassen Algorithm 則是把矩陣乘法從 O(N^3) 降到 O(N^(log7/log2))。由於 convolution 轉矩陣乘法時，Filter 越大，展開後的 Input 越容易膨脹，所以 Strassen 比較適合 Filter 較小的情境，例如 filter < 3 * 3。

Spatial Architecture

MAC 運作時會從 DRAM 讀取資料、進行乘法與加法，最後再把結果寫回 DRAM。從能耗比較可以看到，DRAM 存取的成本是 ALU 運算的上百倍。

因此 Spatial Architecture 的目標是減少資料搬移。它會把 Register File 與 Control Logic 放在 ALU 旁邊，組成 Processing Engine（PE），讓資料能在更靠近運算單元的位置被重複使用。

Data Reuse

PE 中的 Register File 離 ALU 很近，因此設計重點會變成如何提高 Data Reuse。作者將 Data Reuse 分成三種：

• Convolutional Reuse：Convolution 運算中，Filter 與 Input 都會被重複使用。
• Fmap Reuse：一筆 Feature Map 資料可能被多個 Filter 使用。
• Filter Reuse：多筆資料可能共用同一個 Filter，常見於 Batch Size > 1。

Term Explanation

在介紹 AI 加速器的種類前，先整理幾個常用術語。

一般電腦程式會先透過 Compiler 編譯成 Binary code，再交給 Processor 執行；Processor 的結構通常稱為 Architecture。

AI 加速器的流程很像：DNN 會先透過 Mapper 轉成對應的 Mapping，再交給 DNN Accelerator 執行；DNN Accelerator 的資料流設計則稱為 Dataflow。

Dataflow

AI 加速器的 Dataflow 主要可以分成 4 種：

Weight Stationary（WS）

WS 的目標是最小化讀取 Weights 的能耗，也就是盡量讓 Weight 留在 Register File 中重複使用。

Output Stationary（OS）

OS 的目標是最小化讀寫 Partial Sum 的能耗，也就是盡量讓 Partial Sum 留在 Register File 中。

OS Dataflow 又可以依照 Channel 與 Activation 的數量分成 OSA、OSB、OSC。OSA 主要處理 Convolution，OSC 主要處理 FC Layer，OSB 則介於兩者之間。

No Local Reuse（NLR）

Register File 可以減少能耗，但會增加面積。NLR 則反過來弱化 local reuse，目標是最大化 Global Buffer 的儲存能力，並最小化 Off-Chip Memory Bandwidth。

Row Stationary（RS）

RS 的目標是最大化所有類型資料在 Register File 中的重複使用機會。以 1 維 Convolution 為例，PE 中的 Filter 幾乎固定不動，只位移 Input 與 Partial Sum 來完成運算。

2 維陣列的概念也類似，只是會使用多個 PE。Filter 仍然維持相對固定，每次計算 Row 時主要位移 Input 與 Partial Sum。

如果有多個 Channel 或 Batch Size 大於 1，則可以透過連接與交錯的方式得到對應輸出。

Dataflow Example

作者使用 Eyeriss DNN Accelerator 作為範例，它的 PE Array 大小為 12 * 14。

這時會遇到兩個問題：

1. PE Array 和 Layer 大小不同：Layer 較小時，可以一次塞多張 Layer 到 PE Array；Layer 較大時，則可以透過裁切或 folding 塞進 PE Array。

2. 資料不知道要傳到哪個 PE：可以使用 Multicast Network 解決。最簡單的做法是廣播資料，再由 PE Array 中的 Control Logic 篩選每個 PE 需要的資料。

Dataflow Comparision

接著比較 WS、OS、NLR 與 RS 的能耗表現。

在 Convolutional Layer 中，RS 因為最大化 Register File 中的資料重複使用，所以 RF 能耗較高，但整體能耗最低。NLR 因為沒有 RF，資料都放在 Buffer，因此 Buffer 能耗最高。OSA 專門處理 Convolution，所以能耗也比 OSC 更低。

從另一個角度看，WS 因為最大化 Weight reuse，所以 Weight 能耗最低；OS 因為最大化 Partial Sum reuse，所以 Partial Sum 能耗最低。

在 FC Layer 中，OSC 因為更適合 FC Layer，所以能耗比 OSA 小。

最後一張圖是使用 RS Dataflow 跑 AlexNet 的能耗分析。L1 ~ L5 多為 Convolutional Layer，RF 能耗較高；L6 ~ L8 多為 FC Layer，DRAM 能耗較高。整體來看，L1 ~ L5 消耗了大部分能量，而後續神經網路也越來越偏向大量使用 Convolution，因此改善 Convolution 的資料搬移與重複使用會非常重要。

Reference

• Efficient Processing of Deep Neural Networks: A Tutorial and Survey, Part V
• FFT Convolution 參考
• Strassen Algorithm 參考