LLM【事前学習Ⅱ：基礎理論 B-3】

事前学習Ⅱ

B-3：自動微分と最適化を実装する

前回B-2では、長さの違う文章をバッチへまとめ、paddingをAttentionと損失計算から適切に除外し、RoPEで位置情報を与えるところまで実装しました。これで、実際のテキストデータをTransformerへ渡し、次トークン予測の損失を計算できるようになりました。

けれど、損失が計算できることと、モデルが学習することは別です。loss.backward()を呼んでも重みは更新されませんし、勾配を消すタイミングを間違えると、意図しないデータの影響が次の更新へ混ざります。学習ループは数行でも、その数行が何を担当しているかを理解する必要があります。

今回は、B-2のモデルとDataLoaderに、勾配計算と重み更新を加えます。対象は、モデルの順方向計算を実装し、次に「正しく学習しているか」を確かめたい方です。目的は、最適化手法の名前を覚えることではありません。損失から勾配、勾配から重みの更新へ進む流れを理解し、学習が進まないときに確認する場所を判断できるようになることです。

  1. 1. 今回、どこまで実装するのか？
  2. 2. 自動微分は、何を計算しているのか？
    1. 2.1 計算グラフを逆にたどり、各重みの勾配を求める
    2. 2.2 勾配計算と重みの更新は、別の仕事である
  3. 3. 学習の1ステップは、どの順番で進めるのか？
    1. 3.1 勾配を消してから、新しい勾配を計算する
    2. 3.2 B-2の入力形式で、最小の更新を書く
  4. 4. AdamWは、何を調整しているのか？
    1. 4.1 過去の勾配を使い、パラメータごとの更新量を調整する
    2. 4.2 どのパラメータへweight decayを適用するか決める
  5. 5. 勾配クリッピングは、何を防ぐのか？
    1. 5.1 全体のノルムが大きいときに、勾配を縮める
    2. 5.2 クリッピングだけでは、発散の原因は解消しない
  6. 6. 学習率は、何を基準に変えるのか？
    1. 6.1 ウォームアップとコサイン減衰を組み合わせる
    2. 6.2 バッチを読んだ回数と、重みを更新した回数を分ける
  7. 7. 勾配累積では、何を平均する必要があるのか？
    1. 7.1 小さなバッチを複数回処理し、一度だけ更新する
    2. 7.2 可変長バッチでは、有効トークン数で重みを付ける
    3. 7.3 最後の端数も更新し、平均化してからクリップする
  8. 8. B-2のモデルへ、学習ループを接続する
    1. 8.1 損失の集計と、重みの更新をまとめる
    2. 8.2 小さな固定データで、損失が下がるか確かめる
  9. 9. 学習が進まないとき、どこを見るのか？
    1. 9.1 損失・勾配ノルム・学習率を一緒に見る
    2. 9.2 よくある実装ミスを、処理の役割から見直す
  10. 10. まとめ：学習ループを、説明できる処理へ変える
  11. 11. 今回のブログの考察：B-4への橋渡し
  12. 参考文献
  13. このシリーズの案内

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

B-2から引き継ぐ処理と、今回追加する処理を整理します。

段階 処理 今回の確認ポイント
入力 可変長バッチとvalid_maskを用意する B-2の入力形式を保つ
順方向計算 次トークン予測のlogitsを返す 学習時は勾配追跡を有効にする
損失計算 paddingを除いて予測の誤差を集約する 有効な正解トークン数を基準にする
逆伝播 backward()で勾配を計算する 勾配は加算されることを理解する
重み更新 クリッピング後、AdamWで更新する 勾配計算と更新を分けて考える
学習管理 学習率と記録を更新する バッチ数と重みの更新回数を区別する

まずCPU・単精度で仕組みを確認します。混合精度はB-4で追加し、今回は単一デバイスでの学習に範囲を絞ります。

コードでは、B-1のTransformerConfigとFeedForward、B-2のRoPELanguageModel、TokenizedDataset、DynamicPaddingCollatorを再利用します。これらを定義した環境で、以下のコードを順に実行する想定です。B-2のラベルはすでに一つずらしてあるため、今回の学習ループでは再度シフトしません。

2. 自動微分は、何を計算しているのか？

2.1 計算グラフを逆にたどり、各重みの勾配を求める

自動微分は、順方向計算で実行した演算を記録し、連鎖律を使って微分を計算する仕組みです。損失を少し変えるには、各重みをどちらへ動かせばよいか。その局所的な手掛かりになるのが勾配です。

たとえば、損失を $L=x^2+1$ とすると、微分は $dL/dx=2x$ です。$x=2$ なら、勾配は4になります。

import torch
import torch.nn as nn
import torch.nn.functional as F

x = torch.tensor(2.0, requires_grad=True)
loss = x.square() + 1
loss.backward()
assert x.grad.item() == 4.0
assert x.item() == 2.0  # backwardだけでは値は変わらない

# 新しく順方向計算しても、既存の勾配へ加算される
(x.square() + 1).backward()
assert x.grad.item() == 8.0
x.grad = None

Transformerでも原理は同じです。Embedding、Attention、FFNなどの演算を通って得た損失から、学習対象の各パラメータへ勾配が伝わります。通常の学習で整数のトークンID自体を微分する必要はありません。更新したいのは、IDに対応するEmbeddingなどの重みです。

PyTorchの通常の実行では、計算グラフは順方向計算のたびに作られます。逆伝播に必要な中間値は通常のbackward()後に解放されるため、次のバッチでは新たに順方向計算します。Autogradの公式解説

2.2 勾配計算と重みの更新は、別の仕事である

loss.backward()は、学習対象のパラメータの.gradへ勾配を蓄積します。重みを実際に変更するのは、後で呼ぶoptimizer.step()です。backwardの公式ドキュメント

この分担を押さえると、「勾配はあるのに学習しない」という状況を調べやすくなります。step()を呼び忘れていないか、Optimizerに意図したパラメータを渡しているか、学習率が小さすぎないかを順に確認できます。

また、model.train()はDropoutなどを学習時の動作へ切り替える命令であり、勾配追跡を有効にする命令ではありません。B-2の確認で使ったtorch.no_grad()の内側に学習処理を置くと、必要な計算グラフを記録できません。

3. 学習の1ステップは、どの順番で進めるのか？

3.1 勾配を消してから、新しい勾配を計算する

勾配累積を使わない基本の順番は、次の通りです。

  1. 前回の勾配を消す
  2. 順方向計算と損失計算を行う
  3. 逆伝播で勾配を計算する
  4. 必要に応じて勾配をクリッピングする
  5. Optimizerで重みを更新する
  6. 次の更新に使う学習率を設定する

optimizer.zero_grad(set_to_none=True)は、Optimizerが管理するパラメータの勾配をNoneへ戻します。重みやAdamWの内部状態を初期化する処理ではありません。zero_gradの公式ドキュメント

3.2 B-2の入力形式で、最小の更新を書く

まずはスケジューラや勾配累積を入れず、一度の更新を確認します。以下では、B-2のmodelとbatchをCPU上で使います。

model = model.cpu()
model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
optimizer.zero_grad(set_to_none=True)

logits = model(batch["input_ids"], batch["valid_mask"])
loss = F.cross_entropy(
    logits.reshape(-1, logits.size(-1)),
    batch["labels"].reshape(-1),
    ignore_index=-100,
)
if not torch.isfinite(loss).item():
    raise RuntimeError("損失にNaNまたはInfがあります")

loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(), max_norm=1.0, error_if_nonfinite=True,
)
optimizer.step()
print(f"loss={loss.item():.4f}, grad_norm={grad_norm.item():.4f}")

ここで表示する損失は、更新前の予測に対する値です。更新後の損失を知りたい場合は、再び順方向計算します。また、一度更新すれば必ず損失が下がるとは限りません。まず確認したいのは、損失と勾配が有限で、意図した重みが更新されていることです。

4. AdamWは、何を調整しているのか？

4.1 過去の勾配を使い、パラメータごとの更新量を調整する

単純な勾配降下法は、勾配に学習率を掛け、その分だけ重みを動かします。Adamは、勾配の一次モーメントと二次モーメントの移動平均を使い、パラメータごとの更新量を調整します。

AdamWでは、この勾配に基づく更新と、重みを縮めるweight decayを分離します。Adamのような適応的な更新では、損失へL2正則化項を加えることと、重みに直接減衰を適用することが同じにならないためです。AdamWの原論文

設定 主な役割 調整時に見ること
lr 更新全体の大きさを決める 損失の低下、発散、更新の停滞
betas 勾配と勾配の二乗の履歴を平滑化する 変化への追従性、学習の安定性
eps 更新式の分母を安定させる 数値精度との相性
weight_decay 勾配による更新と分離して重みを縮める 学習損失と検証損失の関係

AdamWを選んでも、学習率の調整が不要になるわけではありません。3e-4や0.01といった値は、ここでは動作確認の出発点です。データ、モデル規模、バッチ構成に対する最適値を示してはいません。

4.2 どのパラメータへweight decayを適用するか決める

すべてのパラメータへ同じ減衰を掛ける構成も作れますが、今回は行列の重みと、bias・LayerNormの係数を分けます。B-2のモデルでは、次元数を使って簡潔に分類できます。

def build_optimizer(model, lr=3e-4, weight_decay=0.01):
    decay, no_decay = [], []
    for parameter in model.parameters():
        if not parameter.requires_grad:
            continue
        if parameter.ndim >= 2:
            decay.append(parameter)
        else:
            no_decay.append(parameter)

    return torch.optim.AdamW(
        [
            {"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0},
        ],
        lr=lr,
        betas=(0.9, 0.999),
        eps=1e-8,
    )

この分類ではEmbeddingにもweight decayが掛かります。次元数による分類は今回のモデル向けの方針であり、すべての構成に共通する規則ではありません。モデルを拡張したら、各グループに含まれるパラメータを確認します。グループごとに減衰係数を指定できることは、AdamWの公式ドキュメントでも確認できます。

5. 勾配クリッピングは、何を防ぐのか？

5.1 全体のノルムが大きいときに、勾配を縮める

勾配が急に大きくなると、重みの更新が不安定になる場合があります。その対策の一つが、勾配全体のノルムへ上限を設けるクリッピングです。

全パラメータの勾配を一つにつないだベクトルを $\mathbf{g}$、上限を $c$ とすると、考え方は次の式で表せます。ノルムが0の場合は変更しません。

$$ \tilde{\mathbf{g}} = \mathbf{g}\min\left(1,\frac{c}{|\mathbf{g}|_2}\right) $$

たとえば、ノルムが10で上限が1なら、全体をおよそ10分の1に縮めます。各要素を個別に一定範囲へ切り詰める方法とは異なります。

clip_grad_norm_は勾配をその場で変更し、クリッピング前の全体ノルムを返します。第3章のログが1を超えていても、それだけではクリッピングの失敗を意味しません。クリッピングの公式ドキュメント

5.2 クリッピングだけでは、発散の原因は解消しない

上限1.0は今回の実装例の値です。常に最適とは限りません。また、NaNになった勾配をクリッピングで正常な値へ戻すことはできません。

毎回のように強くクリップされる場合は、上限だけでなく、学習率、損失の集約方法、異常な入力、マスクを確認します。Transformerの勾配が大きくなる理由を、系列長に沿った指数的な増大だけで説明することもできません。層構成、データ、数値精度などを含めて調べる必要があります。

AdamWでは過去の勾配やweight decayも更新に関わるため、勾配ノルムを1に抑えても、パラメータの更新量そのものに同じ上限が付くわけではありません。

6. 学習率は、何を基準に変えるのか？

6.1 ウォームアップとコサイン減衰を組み合わせる

学習初期は、重みもOptimizerの内部状態も十分に調整されていません。小さな学習率から始めて段階的に上げるウォームアップを使い、その後は学習率を下げる構成が考えられます。

今回は、線形ウォームアップとコサイン減衰を組み合わせます。ウォームアップの長さは実験条件として指定し、常に全体の10%が適切だとは仮定しません。

import math
from torch.optim.lr_scheduler import LambdaLR

def build_scheduler(optimizer, total_updates, warmup_updates):
    if total_updates < 1 or not 0 <= warmup_updates < total_updates:
        raise ValueError("0 <= warmup_updates < total_updatesが必要です")

    def lr_factor(update_index):
        if update_index < warmup_updates:
            return (update_index + 1) / warmup_updates
        progress = (update_index - warmup_updates) / (
            total_updates - warmup_updates
        )
        progress = min(max(progress, 0.0), 1.0)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    return LambdaLR(optimizer, lr_lambda=lr_factor)

LambdaLRは、基準学習率へ掛ける係数を返す関数を受け取ります。この例では、ウォームアップを10更新にすると、最初の更新は基準学習率の10分の1で始まります。LambdaLRの公式ドキュメント

6.2 バッチを読んだ回数と、重みを更新した回数を分ける

ここではscheduler.step()をoptimizer.step()の後に呼び、次の更新で使う学習率を設定します。この呼び出し順序は、Optimizerとスケジューラの公式解説に沿っています。

重要なのは、何を1ステップと数えるかです。勾配累積で4バッチごとに重みを更新する場合、スケジューラも4バッチごとに進めます。バッチを読むたびに進めると、想定の4倍の速さで学習率が変わってしまいます。

7. 勾配累積では、何を平均する必要があるのか？

7.1 小さなバッチを複数回処理し、一度だけ更新する

勾配累積は、各マイクロバッチで逆伝播を行い、複数回分の勾配を足してから重みを更新する方法です。マイクロバッチとは、一度にメモリへ載せて処理する小さなバッチを指します。

たとえば、毎回2文書ずつ処理し、4バッチ分を累積して更新すれば、通常の更新では8文書分の勾配を使えます。各バッチで逆伝播を済ませるため、8文書分の計算グラフを同時に保持する必要はありません。ただし、モデルの重み、勾配、Optimizerの内部状態に必要なメモリは別途残ります。

同じデータ量を学習するときに、所要時間が必ず累積回数倍になるわけではありません。一括処理との速度差はGPUの使用効率などに依存します。また、Dropoutや浮動小数点の丸めもあるため、大きなバッチとの結果の完全一致を前提にはしません。

7.2 可変長バッチでは、有効トークン数で重みを付ける

各バッチの平均損失を単に累積回数で割る方法は、今回の可変長データでは注意が必要です。

マイクロバッチ 有効な正解トークン数 平均損失 損失の合計
A 100 2.0 200
B 300 4.0 1,200

バッチ平均を均等に平均すると3.0ですが、全400トークンの平均損失は、$(200+1200)/400=3.5$です。B-2と同じ「各正解トークンを等しい重みで扱う」目的を保つには、後者に合わせます。

そこで、今回は各バッチでreduction="sum"の損失を逆伝播し、更新直前に累積した勾配を有効トークン総数で割ります。重みを変えずに累積している間は、損失の和の微分が勾配の和になるためです。ignore_index=-100で除外した位置は、分母にも含めません。CrossEntropyLossの公式ドキュメント

7.3 最後の端数も更新し、平均化してからクリップする

DataLoaderが10バッチで累積回数が4なら、更新単位は4・4・2バッチです。最後の2バッチを捨てたり、4バッチあるものとして扱ったりしないようにします。

また、クリッピングは累積した勾配をトークン数で平均化した後に一度だけ行います。各バッチを別々にクリップしてから足すと、最後に全体をクリップする場合と更新の意味が変わります。

8. B-2のモデルへ、学習ループを接続する

8.1 損失の集計と、重みの更新をまとめる

以下は、有限の長さを持つ通常のDataLoaderを対象にした、単一デバイス・単精度の1エポック分の処理です。勾配累積を使わない場合はaccumulation_steps=1にします。

def train_epoch(
    model, loader, optimizer, scheduler=None,
    accumulation_steps=1, max_grad_norm=1.0, device="cpu",
):
    if accumulation_steps < 1 or len(loader) == 0:
        raise ValueError("累積回数は1以上、loaderは空でない必要があります")

    model.train()
    optimizer.zero_grad(set_to_none=True)
    epoch_loss_sum = 0.0
    epoch_tokens = 0
    window_tokens = 0
    updates = 0
    last_grad_norm = None
    last_used_lr = None

    for batch_index, batch in enumerate(loader):
        input_ids = batch["input_ids"].to(device)
        labels = batch["labels"].to(device)
        valid_mask = batch["valid_mask"].to(device)
        num_tokens = int((labels != -100).sum().item())
        if num_tokens == 0:
            raise ValueError("有効な正解トークンがないバッチです")

        logits = model(input_ids, valid_mask)
        loss_sum = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            labels.reshape(-1),
            ignore_index=-100,
            reduction="sum",
        )
        if not torch.isfinite(loss_sum).item():
            raise RuntimeError("損失にNaNまたはInfがあります")
        loss_sum.backward()

        epoch_loss_sum += loss_sum.detach().item()
        epoch_tokens += num_tokens
        window_tokens += num_tokens
        is_boundary = (batch_index + 1) % accumulation_steps == 0
        is_last = batch_index + 1 == len(loader)
        if is_boundary or is_last:
            # 有効トークン当たりの平均勾配へ変換する
            for parameter in model.parameters():
                if parameter.grad is not None:
                    parameter.grad.div_(window_tokens)

            grad_norm = torch.nn.utils.clip_grad_norm_(
                model.parameters(), max_norm=max_grad_norm,
                error_if_nonfinite=True,
            )
            # schedulerを進める前に、今回使う学習率を記録する
            last_used_lr = optimizer.param_groups[0]["lr"]
            optimizer.step()
            if scheduler is not None:
                scheduler.step()
            optimizer.zero_grad(set_to_none=True)
            last_grad_norm = grad_norm.item()
            window_tokens = 0
            updates += 1

    return {
        "loss": epoch_loss_sum / epoch_tokens,
        "tokens": epoch_tokens,
        "updates": updates,
        "last_grad_norm": last_grad_norm,
        "last_used_lr": last_used_lr,
    }

返すlossは、エポック中に観測した損失を有効トークン数で平均した値です。更新したバッチの損失だけを集めたり、文書数で割ったりしていません。一方、last_grad_normは最後の更新直前の値であり、エポック全体の平均ではありません。

この例は、モデルの全学習パラメータをOptimizerが管理する前提です。GPUを使う場合は、Optimizerを作る前にモデルをデバイスへ移動し、関数のdeviceにも同じデバイスを渡します。Optimizerとスケジューラはエポックの外で一度だけ作り、履歴を保ちます。

8.2 小さな固定データで、損失が下がるか確かめる

最初の確認には、B-2で使った説明用の2文書を利用できます。これは実装の動作確認であり、自然言語の品質や汎化性能を評価する実験ではありません。

from torch.utils.data import DataLoader

torch.manual_seed(0)
config = TransformerConfig(
    vocab_size=32, block_size=8, d_model=32,
    n_heads=4, n_layers=1, d_ff=64, dropout=0.0,
)
model = RoPELanguageModel(config).to("cpu")
dataset = TokenizedDataset([[11, 12, 13, 14], [21, 22]], block_size=8)
loader = DataLoader(
    dataset, batch_size=1, shuffle=False, num_workers=0,
    collate_fn=DynamicPaddingCollator(pad_token_id=0),
)

num_epochs = 60
accumulation_steps = 2
updates_per_epoch = math.ceil(len(loader) / accumulation_steps)
optimizer = build_optimizer(model, lr=3e-3, weight_decay=0.0)
scheduler = build_scheduler(
    optimizer,
    total_updates=num_epochs * updates_per_epoch,
    warmup_updates=5,
)

for epoch in range(num_epochs):
    stats = train_epoch(
        model, loader, optimizer, scheduler,
        accumulation_steps=accumulation_steps, device="cpu",
    )
    if epoch == 0 or (epoch + 1) % 10 == 0:
        print(
            f"epoch={epoch + 1}, loss={stats['loss']:.4f}, "
            f"grad_norm={stats['last_grad_norm']:.4f}, "
            f"lr={stats['last_used_lr']:.2e}"
        )

この例では、正解トークンが3個の文書と1個の文書を別々に処理し、4トークン分を平均して一度更新します。Dropoutとweight decayは確認を単純にするため無効にしています。ここでの学習率も、小さなモデルでの動作確認用です。

編集時に上のコードをCPU・PyTorch 2.9.0で実行したところ、平均損失は1エポック目の約3.128から、60エポック目の約0.060へ下がりました。これは説明用の2文書を覚えられたという動作確認であり、未学習データでの性能を示す結果ではありません。

まず、初期と終盤の損失を比べます。下がらない場合は、学習時間を延ばす前に、ラベルの対応、勾配の有無、更新の実行、学習率を確認します。固定した少量データを覚えられたら、次に実データを学習用と検証用へ分け、未学習の文章でも損失が改善するかを調べます。

9. 学習が進まないとき、どこを見るのか？

9.1 損失・勾配ノルム・学習率を一緒に見る

一つの数値だけでは原因は決まりません。少なくとも次の情報を同じ更新単位で残すと、問題を追いやすくなります。

観測 考えられる要因 最初に確認すること
損失が変わらない 未更新、学習率過小、ラベル不整合 step()の実行、使用学習率、少量データでの過学習
勾配がNoneのまま 損失までの経路がない、追跡が無効 requires_grad、no_grad()、detach()の位置
損失や勾配がNaN・Inf 数値不安定、不正な入力やマスク 異常が初めて出たバッチと演算
大きな勾配が続く 学習率、損失の集約、入力の異常 平均化の分母、データ、クリップ前ノルム
学習損失だけが下がる 過学習、データ分布や前処理の差 検証データ、分割、重複、正則化

Noneと、値が0の勾配は区別します。0は計算した結果として起こることがありますし、一部のパラメータがそのバッチで使われなければ、Noneも直ちに異常とは限りません。

学習損失が下がっても、入力と正解が同じトークンになっていたり、未来のトークンが見えていたりすれば、次トークン予測が正しく学べたとは言えません。B-1・B-2で確認したCausal Maskとラベルのシフトは、最適化を加えてからも前提です。

9.2 よくある実装ミスを、処理の役割から見直す

特に混同しやすいのは、次の四つです。

  1. 勾配を消し忘れる：独立した更新のはずが、前回の勾配まで加算される。更新単位の境界でzero_grad()を呼ぶ。
  2. 累積の途中で更新する：小バッチごとにstep()を呼ぶと、複数バッチ分を同じ重みで累積する目的から外れる。
  3. 不要なretain_graph=Trueを付ける：通常の勾配累積は別々の順方向計算を使うため、過去の計算グラフを保持する必要はない。
  4. ログへTensorを保存し続ける：損失を記録するならloss.item()などで数値へ変換し、計算グラフを参照するTensorをリストへため続けない。

Optimizerを比較するときも、単一の「最終精度」の表だけでは判断できません。同じデータ、モデル初期値、学習トークン数、評価方法をそろえ、各手法の学習率などを調整したうえで、検証損失と所要時間を比べます。今回の動作確認だけで、AdamWが他の手法より常に優れているとは結論付けられません。

10. まとめ：学習ループを、説明できる処理へ変える

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

  • backward()は勾配を計算・蓄積し、optimizer.step()が重みを更新する
  • 勾配を消すタイミングは、重みを更新する単位に合わせる
  • AdamWは適応的な勾配更新とweight decayを分離するが、学習率の検証は必要である
  • 勾配クリッピングは、平均化した勾配へ更新直前に適用し、クリップ前ノルムを記録する
  • 可変長データの勾配累積では、有効な正解トークン数で平均し、最後の端数も更新する
  • 学習率スケジュールは重みの更新回数で進め、まず小さな固定データで学習を確認する

学習ループで大切なのは、処理を増やすことだけではありません。どのデータから計算した勾配を、どの大きさへ整え、いつ重みに反映するかを一貫させることです。この対応が説明できれば、損失が下がらないときも、設定をやみくもに変える前に確認する順番が見えてきます。

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

B-1で予測を計算するモデルを作り、B-2で可変長データと位置情報を接続し、B-3でその予測から重みを更新する流れを組み立てました。ここまでで、入力、損失、勾配、更新を小さな実験で確かめるための土台ができました。

次に考えたいのは、この学習を限られたGPUメモリと時間の中でどう進めるかです。次回B-4では、単精度の学習ループを基準に、Autocastによる混合精度演算と、必要に応じたGradScalerを導入します。

今回確認した順序は、そのまま役立ちます。勾配をスケーリングする構成では、クリッピング前にスケールを戻す必要があります。AMPの公式実装例を踏まえ、B-4では「どの演算を低精度にし、どこで数値の安定性を確認するか」を考えます。まず単精度で学習が進む基準を作っておけば、混合精度を加えた影響も比較しやすくなります。

参考文献

  1. PyTorch. Autograd mechanics、Tensor.backward、Optimizer.zero_grad.
  2. Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019.
  3. PyTorch. AdamW、torch.optim、LambdaLR.
  4. PyTorch. clip_gradnorm、CrossEntropyLoss.
  5. PyTorch. Automatic Mixed Precision examples.

このシリーズの案内

ブログB「実装詳細」では、Transformerの実装から学習の安定化、計算・メモリ効率の改善へと進みます。B-1〜B-7の流れの中で、今回は第3回です。

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

コメント

タイトルとURLをコピーしました