B-2：バッチ処理とRoPEを実装する

前回B-1では、Attention、FFN、残差接続、正規化を組み合わせ、固定長のトークン列から次トークン予測のlogitsを返す最小のTransformerを実装しました。入力から出力までの流れは確認できましたが、そのままでは長さの違う文章を同じバッチへ入れられません。

たとえば、100トークンの文章と250トークンの文章をまとめて処理するとき、短い文章の空いた部分をどう埋めるのでしょうか。その埋めた部分をAttentionや損失計算が本物の文章として扱ってしまうと、形状は合っていても、学習する内容が変わってしまいます。

今回は、可変長のデータをバッチへまとめる処理と、A-5で学んだRoPEをQ・Kへ適用する処理を実装します。対象は、B-1のコードを読み、次に実際のテキストデータへ接続したい方です。目的は機能を増やすことだけではありません。バッチへまとめても、各文章の有効な位置では同じ計算になることを、小さな例で確かめられるようになることです。

1. 今回、どこまで実装するのか？

B-1から変更する箇所を、先に整理します。

箇所	B-1の構成	B-2での変更
入力データ	同じ長さのトークン列	長さの違うトークン列
バッチ生成	同じ形状の入力を用意	バッチ内の最大長へ右側padding
Attentionのマスク	Causal Mask	Causal Maskとpadding用マスクの併用
位置情報	学習可能な絶対位置Embedding	Q・KにRoPEを適用
損失計算	全位置を対象	正解のない位置を除外

以下では、1サンプルを1文書のトークン列として扱います。トークン化と文書末尾のEOS付与は前処理で済ませ、入力にはpaddingが含まれていないものとします。学習用の右側paddingに範囲を絞り、左側padding、複数文書のpacking、生成時のKVキャッシュは扱いません。

コードは上から順に実行する想定です。第5章では、B-1のTransformerConfigとFeedForwardを再利用します。まずCPU・単精度で接続を確認し、パラメータ更新は次回B-3で加えます。

2. 可変長データでは、何をそろえる必要があるのか？

2.1 paddingで形状をそろえる

通常の密なテンソルとしてバッチを作るには、系列長をそろえる必要があります。長さ100・250・80の3サンプルを最大長250へそろえると、次のようになります。

サンプル	実トークン数	追加するpadding数	処理する長さ
文書A	100	150	250
文書B	250	0	250
文書C	80	170	250
合計	430	320	750

この例では、入力位置の約43%がpaddingです。全サンプルを常に512へそろえるより、バッチ内の最大長へそろえるほうが、今回の例では不要な位置を減らせます。この方法を動的paddingと呼びます。

ただし、位置数が43%減れば実行時間も43%減る、という意味ではありません。Attentionの計算量は系列長の影響を強く受け、実際の速度はカーネルやハードウェアにも依存します。

2.2 Attentionと損失計算では、除外する目的が違う

paddingを追加するだけでは不十分です。二つの処理を分けて考えます。

処理	防ぎたいこと	今回の方法
Attention	実トークンがpaddingのKeyを参照する	有効なKeyの位置をマスクへ反映する
損失計算	paddingを正解トークンとして学習する	対象外ラベルを-100にする

さらに、Decoder-onlyモデルでは未来のトークンを参照させないCausal Maskも必要です。padding用のマスクを加えたからといって、因果制約が不要になるわけではありません。

本記事ではvalid_mask=Trueを実トークンと定義します。後で使うscaled_dot_product_attentionも、真偽値のattn_maskではTrueが参照可能な位置です。マスクの意味はAPIごとに異なるため、変数名と説明をそろえておきましょう。PyTorch公式ドキュメント

3. DataLoaderで可変長バッチを作る

3.1 入力と正解を一つずらす

次トークン予測では、入力が[11, 12, 13]なら正解は[12, 13, 14]です。元の列[11, 12, 13, 14]から、最後を除いた入力と、先頭を除いたラベルを作ります。

今回は、このシフトをDatasetで一度だけ行います。損失計算時にもう一度ずらすと、二つ先のトークンを予測するタスクになってしまいます。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader

class TokenizedDataset(Dataset):
    def __init__(self, documents, block_size=64):
        if block_size < 1:
            raise ValueError("block_sizeは1以上にしてください")
        # 次トークンの正解を作れない文書は対象外にする
        self.documents = [list(ids) for ids in documents if len(ids) >= 2]
        self.block_size = block_size

    def __len__(self):
        return len(self.documents)

    def __getitem__(self, index):
        tokens = torch.tensor(
            self.documents[index][:self.block_size + 1], dtype=torch.long
        )
        return {"input_ids": tokens[:-1], "labels": tokens[1:]}

class DynamicPaddingCollator:
    def __init__(self, pad_token_id):
        self.pad_token_id = pad_token_id

    def __call__(self, samples):
        inputs = [sample["input_ids"] for sample in samples]
        labels = [sample["labels"] for sample in samples]
        lengths = torch.tensor([len(ids) for ids in inputs])
        input_ids = pad_sequence(
            inputs, batch_first=True, padding_value=self.pad_token_id
        )
        labels = pad_sequence(labels, batch_first=True, padding_value=-100)
        positions = torch.arange(input_ids.size(1))
        valid_mask = positions[None, :] < lengths[:, None]
        return {
            "input_ids": input_ids,
            "labels": labels,
            "valid_mask": valid_mask,
        }

# 説明用のトークンID。0をpadding専用として予約する
sample_tokens = [[11, 12, 13, 14], [21, 22]]
dataset = TokenizedDataset(sample_tokens, block_size=64)
loader = DataLoader(
    dataset, batch_size=2, shuffle=False, num_workers=0,
    collate_fn=DynamicPaddingCollator(pad_token_id=0),
)
batch = next(iter(loader))
assert batch["input_ids"].tolist() == [[11, 12, 13], [21, 0, 0]]
assert batch["labels"].tolist() == [[12, 13, 14], [22, -100, -100]]
assert batch["valid_mask"].tolist() == [[True, True, True], [True, False, False]]

3.2 paddingのIDと、損失から除外する値を分ける

input_idsはEmbeddingへ渡すため、語彙の範囲内のIDが必要です。一方、labelsの-100は損失から除外するための値で、Embeddingへは渡しません。cross_entropyではignore_index=-100を指定します。CrossEntropyLoss公式ドキュメント

実際のpad_token_idは使用するトークナイザで確認してください。全モデルで0がpaddingとは限りません。また、EOSとpaddingに同じIDを使う場合、IDの一致だけでマスクを作ると本物のEOSまで除外してしまいます。上のコードでは元の長さからマスクを作るため、この混同を避けられます。

なお、このDatasetは説明を簡単にするため長い文書の先頭だけを使います。そのまま全量の事前学習へ使うと後半のデータを捨てるため、実運用では分割やpackingを別途設計します。

4. RoPEでは、位置情報をどこへ入れるのか？

4.1 Embeddingへの加算から、Q・Kの回転へ変える

B-1では、Token Embeddingへ学習可能な位置Embeddingを加えました。今回のRoPEでは、この加算を外し、Attention内のQueryとKeyを位置に応じて回転させます。Valueには適用しません。

1ヘッド内の隣り合う2次元を一組とし、位置$m$で角度$m\theta_j$だけ回転させます。

$$ \theta_j = \mathrm{base}^{-2j/d},\qquad j=0,\ldots,d/2-1 $$

$$ \begin{pmatrix}x’{2j}\x’{2j+1}\end

\begin{pmatrix} \cos(m\theta_j)&-\sin(m\theta_j)\ \sin(m\theta_j)&\cos(m\thetaj) \end{pmatrix} \begin{pmatrix}x{2j}\x_{2j+1}\end{pmatrix} $$

$d$はヘッド次元です。QとKの内積には二つの位置の差に依存する回転が現れ、相対位置をAttentionへ組み込めます。RoPEの数式上の背景は、A-5とRoFormer原論文で確認できます。

4.2 baseの値だけで長文性能は決まらない

ここではbase=10000を使用します。$j=0$では角周波数が1で、$j$が大きい組ほど角周波数が小さくなります。つまり、同じ位置の変化に対して速く回る組とゆっくり回る組を持たせます。単位は時間に対するHzではなく、トークン位置あたりの回転角です。

baseを大きくすれば一部の組の回転は緩やかになりますが、それだけで学習時より長い文脈へ対応できるとは断定できません。計算できる長さと、品質を保って利用できる長さは分けて評価する必要があります。

4.3 回転処理を実装する

複素数でも表現できますが、今回は2次元の回転が見える実数のコードにします。学習可能なパラメータではない角周波数はbufferとして登録し、モデルのデバイス移動へ追従させます。

class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, base=10000.0):
        super().__init__()
        if head_dim % 2 != 0:
            raise ValueError("RoPEのヘッド次元は偶数にしてください")
        inv_freq = base ** (-torch.arange(0, head_dim, 2).float() / head_dim)
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x):
        # x: (B, H, T, D)。今回は右側paddingなので位置は0〜T-1
        positions = torch.arange(x.size(-2), device=x.device).float()
        angles = positions[:, None] * self.inv_freq.float()[None, :]
        cos = angles.cos()[None, None, :, :]
        sin = angles.sin()[None, None, :, :]
        even, odd = x.float()[..., 0::2], x.float()[..., 1::2]
        rotated = torch.stack(
            (even * cos - odd * sin, even * sin + odd * cos), dim=-1
        )
        return rotated.flatten(-2).to(x.dtype)

隣接する次元を組にする実装と、ベクトルの前半・後半を組にする実装では、重みの配置との対応が異なります。既存モデルの重みを読み込む場合は、そのモデルの規約に合わせてください。今回のコードはB-1と同様、最初から学習する小さなモデルを対象にしています。

5. RoPEと二つのマスクをAttentionへ組み込む

5.1 参照できる位置を一つのマスクへまとめる

位置$q$から位置$k$を参照できる条件は、「$k \leq q$であり、かつ$k$が実トークンであること」です。下三角のCausal Maskと、Keyの有効位置をANDで結合します。

ここでは結合済みのマスクをattn_maskへ渡し、is_causalはデフォルトのFalseにします。DropoutもB-1と同様、評価時には明示的に0へ切り替えます。SDPA公式ドキュメント

class RoPECausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.d_model % config.n_heads == 0
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
        self.out_proj = nn.Linear(config.d_model, config.d_model)
        self.dropout = config.dropout
        self.rope = RotaryEmbedding(self.head_dim)

    def forward(self, x, valid_mask):
        batch_size, seq_len, d_model = x.shape
        qkv = self.qkv(x).view(
            batch_size, seq_len, 3, self.n_heads, self.head_dim
        ).permute(2, 0, 3, 1, 4)
        query, key, value = qkv.unbind(dim=0)
        query, key = self.rope(query), self.rope(key)

        causal = torch.ones(
            seq_len, seq_len, device=x.device, dtype=torch.bool
        ).tril()
        # (B, 1, T, T)。TrueのKeyだけを参照できる
        allowed = causal[None, None, :, :] & valid_mask[:, None, None, :]
        attended = F.scaled_dot_product_attention(
            query, key, value, attn_mask=allowed,
            dropout_p=self.dropout if self.training else 0.0,
        )
        attended = attended.transpose(1, 2).contiguous()
        attended = attended.view(batch_size, seq_len, d_model)
        return self.out_proj(attended)

class RoPETransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm1 = nn.LayerNorm(config.d_model)
        self.attention = RoPECausalSelfAttention(config)
        self.norm2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)  # B-1の定義を再利用

    def forward(self, x, valid_mask):
        x = x + self.attention(self.norm1(x), valid_mask)
        return x + self.ffn(self.norm2(x))

class RoPELanguageModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList(
            [RoPETransformerBlock(config) for _ in range(config.n_layers)]
        )
        self.final_norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

    def forward(self, input_ids, valid_mask):
        if input_ids.size(1) > self.config.block_size:
            raise ValueError("入力長がblock_sizeを超えています")
        x = self.token_embedding(input_ids)  # 絶対位置Embeddingの加算を外す
        for block in self.blocks:
            x = block(x, valid_mask)
        return self.lm_head(self.final_norm(x))

5.2 padding位置の出力は、ゼロになる必要があるのか？

この実装ではpaddingのQueryも過去の実トークンを参照するため、その出力はゼロとは限りません。それでも、実トークンからpaddingのKeyへの参照は全層で禁止され、padding位置は損失からも除外されます。今回の構成では、有効な位置の予測へ混入しないようにすることが重要です。

各入力には最低1個の実トークンがあり、右側paddingなので、全てのQuery行に少なくとも一つ参照先があります。空の入力や左側paddingまで受け付ける場合は、この前提を見直す必要があります。

6. バッチ化しても計算が変わらないことを確認する

6.1 出力形状と有限なLossを確認する

B-1のTransformerConfigを定義したうえで、ここまでのコードを順に実行します。

torch.manual_seed(0)
config = TransformerConfig()
model = RoPELanguageModel(config)
model.eval()  # 比較時はDropoutを無効にする

with torch.no_grad():
    logits = model(batch["input_ids"], batch["valid_mask"])
    loss = F.cross_entropy(
        logits.reshape(-1, config.vocab_size),
        batch["labels"].reshape(-1),
        ignore_index=-100,
    )
assert logits.shape == (2, 3, config.vocab_size)
assert torch.isfinite(loss)

ラベルはDatasetですでにシフト済みなので、ここでは全位置のlogitsとそのまま対応させます。このLossは有効な正解トークンについての平均であり、文書ごとのLossを等しい重みで平均したものではありません。

6.2 短い文章の単独入力と、padding付き入力を比べる

形状が合っていても、マスクが正しいとは限りません。短い文書を単独で入力した結果と、長い文書と同じバッチへ入れた結果を比較します。さらに、未来のトークンを書き換えても過去の出力が変わらないことを確かめます。

with torch.no_grad():
    single_ids = batch["input_ids"][1:2, :1]
    single = model(single_ids, torch.ones_like(single_ids, dtype=torch.bool))
    torch.testing.assert_close(
        single[0, 0], logits[1, 0], rtol=1e-4, atol=1e-5
    )

    changed_ids = batch["input_ids"].clone()
    changed_ids[0, 2] = 99  # 位置2だけ変更する
    changed_logits = model(changed_ids, batch["valid_mask"])
    torch.testing.assert_close(
        changed_logits[0, :2], logits[0, :2], rtol=1e-4, atol=1e-5
    )

前者はバッチ化による影響、後者は未来の情報の漏れを調べています。計算順序の違いによる浮動小数点誤差を考慮し、完全一致ではなく許容誤差付きで比較します。これらはモデル品質の評価ではなく、今回加えた処理の基本的な接続確認です。

7. DataLoaderは、何を見て速くするのか？

7.1 まず単一プロセスで確認する

最初はnum_workers=0で動かすと、Datasetやcollatorの例外を追いやすくなります。動作を確認してから、データ読み込みが学習を待たせているかを測り、ワーカー数を増やします。メモリ上の小さなデータでは、プロセス間通信の負担で遅くなる場合もあります。

設定	役割	判断するときの観点
num_workers	データ読み込みを複数プロセスへ分担	0・2・4などを比較して待ち時間を測る
pin_memory	CUDA転送に使うホストメモリをページ固定する	CUDA利用時に転送時間とメモリ使用量を確認する
prefetch_factor	ワーカーごとに先読みするバッチ数	ワーカー利用時に待ち時間とCPUメモリを比較する
persistent_workers	エポック終了後もワーカーを維持する	ワーカー利用時の再起動コストを確認する

複数ワーカーを使うスクリプトでは、Datasetとcollatorをトップレベルに定義し、DataLoaderの作成・実行をif __name__ == "__main__":の内側へ置きます。特にmacOSやWindowsでは、子プロセスによるスクリプト再実行を考慮する必要があります。DataLoader公式ドキュメント

7.2 サンプル数だけでなく、有効トークン数を見る

長さの違う文章では、同じ32サンプルでも処理量が大きく変わります。速度比較には、有効トークン数/秒とpadding率も使いましょう。今回のバッチなら、有効位置はvalid_mask.sum()、総位置数はvalid_mask.numel()から求められます。

長さの近いサンプルを同じバッチへ集める長さ別バケット化も、paddingを減らす選択肢です。ただし、並べ方によってデータ順序が偏らないよう、シャッフルとの組み合わせを設計します。複数文書を連結するpackingはさらに別の方法であり、文書境界でAttentionを許可するかも含めて考える必要があります。

8. よくある実装ミスは何か？

8.1 paddingだけを隠し、未来のトークンを見せてしまう

たとえば正解のトークンが次の入力位置に存在する場合、Causal Maskがなければそれを参照できてしまいます。Lossが下がることだけでは正しさを判断できません。第6章のように、未来の入力を変更して過去の出力を比べます。

8.2 ラベルを二重にずらす

B-1では損失計算の直前に[:, :-1]と[:, 1:]で対応を作りました。B-2ではDatasetがその役割を担います。どちらも使うとずれが重なるため、シフトを行う場所を一つに決めてください。

8.3 padding用のIDを正解として残す

たとえば短いサンプルを大量のpaddingで埋め、その位置まで損失へ含めると、モデルはpaddingを予測することにも最適化されます。ラベルの対象外位置を-100にし、有効トークンだけでLossを計算します。

8.4 RoPEを加えても絶対位置Embeddingを残す

二つを併用する設計も考えられますが、今回の目的はB-1の位置表現をRoPEへ置き換えることです。絶対位置Embeddingの加算を外し、Q・Kへの回転へ変更した箇所を確認します。

8.5 ワーカー数を増やせば必ず速くなると考える

CPU数、ストレージ、前処理の重さ、メモリ容量によって適した設定は変わります。GPUの台数だけでワーカー数を決めず、バッチ待ち時間と有効トークン数/秒を比較して判断しましょう。

9. まとめ：形状をそろえ、意味を保つ

今回の要点を整理します。

• 動的paddingは、バッチ内の最大長へ入力をそろえる方法である
• Attentionでpaddingを参照させない処理と、損失から除外する処理は別に必要になる
• Causal Maskと有効なKeyのマスクを結合し、未来の情報を隠す
• 次トークン予測用のシフトは一度だけ行い、対象外ラベルを-100にする
• RoPEはQ・Kを位置に応じて回転させ、相対位置をAttentionへ組み込む
• 単独入力との比較と未来の入力の書き換えで、マスクの基本動作を確認できる

バッチ処理は、データを同じ形状へ詰めるだけの作業ではありません。どのトークンを参照し、どの位置を正解として学習するかを決める処理です。形状と意味の両方がそろって初めて、次の学習ステップへ進めます。

10. 今回のブログの考察：B-3への橋渡し

B-1でモデルの骨格を作り、B-2では長さの違うデータを渡して、位置情報と損失計算まで接続しました。ここまでの確認により、Lossが想定外の動きをしたときに、入力、マスク、ラベルのどこから調べるかが明確になります。

次に必要なのは、このLossから勾配を求め、パラメータを実際に更新することです。次回B-3では、PyTorchの自動微分、backward()、AdamW、勾配クリッピングを扱います。今回の有効トークンに対するLossを出発点に、B-1で紹介した少数データへの過学習テストを、更新処理まで含めて確かめていきましょう。

---

参考文献

1. Su, J., et al. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv:2104.09864.
2. PyTorch. scaled_dot_product_attention.
3. PyTorch. CrossEntropyLoss.
4. PyTorch. torch.utils.data.

---

このシリーズの案内

ブログB「Transformer実装」では、Aシリーズで整理した設計を、PyTorchの実装・最適化・運用上の工夫へ段階的に落とし込みます。

#	テーマ	内容
B-1	Transformerブロック実装	Causal Attention、FFN、Pre-Norm、次トークン予測
B-2	バッチ処理と位置エンコーディング	padding、DataLoader、RoPE（今回）
B-3	自動微分と最適化	backward、AdamW、勾配クリッピング
B-4	混合精度演算	AMP、損失スケーリング、数値安定性
B-5	FlashAttention	Attentionの高速化とメモリ効率
B-6	分散学習とDDP	複数GPUでの学習
B-7	メモリ最適化	活性化・KVキャッシュ・チェックポイント

• 前の記事：B-1：TransformerブロックをPyTorchで実装する
• 次に読む：B-3：自動微分とバックプロパゲーション最適化
• 理論を確認する：A-5：位置エンコーディングの進化