1. 概要

昨今社会を激変させ、これからも一変させるであろうGPT の中身、気になりませんか？
本記事では、比較的小さいGPT モデルであるnanoGPT を基に、GPT の中身をコードベースで解説します。

GitHub に書かれている、nanoGPT の説明は以下の通りです。

最もシンプルで、速い中サイズGPTのトレーニング/ファインチューニング用リポジトリです。これは教育よりも効率を優先するminGPTの書き換えです。まだ積極的な開発中ですが、現在のところtrain.pyファイルは、単一の8XA100 40GBノード上でOpenWebText上でGPT-2 (124M)を約4日間のトレーニングで再現します。コード自体は明快で読みやすいものです：train.pyは約300行の標準的なトレーニングループ、model.pyは約300行のGPTモデル定義で、必要に応じてOpenAIのGPT-2の重みを読み込むことができます。

nanoGPT のGitHub にはいくつかのPython ファイルがあります。本記事では、特にモデル構造が書かれているmain.py の中身について見ていきます。

以下の画像はLLM Visualization から引用させていただいた、nanoGPT の構造を可視化した3 次元図です。

また、本記事では参考にしていませんが、YouTube にnanoGPT の解説動画があります。丁寧に解説されており、見やすいです。キャッチアップには最適かと思います。
YouTube: Let's build GPT: from scratch, in code, spelled out.

目次

nanoGPT のモデルの中身を見てみよう！

2. 各クラスの簡単なまとめ

main.py に入っているクラスについて以下で簡単にまとめました。

LayerNorm Class
Layer Normalization（ニューラルネットワークの各層における入力の各サンプルに対して、平均と分散を計算し、それを用いて正規化をする）を行うクラス。
CausalSelfAttention Class
Causal Self-Attention（因果的自己注意）を実装するクラス。この機構により、モデルは各入力位置がそれ自身とそれより前の位置の情報のみに「注意（Attention）」を払うようになる。「自己」とは、コンテキスト内の各要素が自分自身のコンテキスト内でどのように位置付けられているかを自己評価することを指す。「因果的」とは、ある位置の出力がその位置より前の入力にのみ依存するという意味を指す。
Block Class
Transformerモデルの一つのブロックが書いてあるクラス。2 つのLayerNormレイヤー、1 つの自己注意レイヤー、そしてMLPレイヤーが含まれている。
GPTConfig Class
GPTモデルの構成設定を管理するクラス。ブロックサイズ、語彙サイズ、レイヤーの数、アテンションヘッドの数、埋め込み次元の数、ドロップアウト率、バイアスの有無などの設定が含まれている。
GPT Class
上記クラスを使用し、GPTモデルを定義するクラス。

3. GPT Class のメソッドの簡単なまとめ

上記のクラス内の関数は基本的に__init__ とforward 関数のみです。その中で、GPT Class 関数はそれら以外の関数が多く含まれるので、それらの関数について簡単に解説します。

GPT Class
- get_num_param
  モデル内のパラメータ数を返すクラスメソッド。
- _init_weights
  線形層と埋め込み層の重みを正規分布で初期化し、バイアスがある場合は0で初期化するクラスメソッド。
- forward
  GPTモデルのフォワードパスを定義するクラスメソッド。入力トークンのインデックスから、トークンと位置の埋め込みを生成し、これをドロップアウトレイヤーを通して、各トランスフォーマーブロックに順に適用する。目標となるトークンが与えられている場合は損失を計算し、そうでない場合は最終層のロジットのみを返す。
- crop_block_size
  モデルのブロックサイズを必要に応じて小さく調整するために使用されるクラスメソッド。
- from_pretrained
  事前訓練済みのモデルをロードして初期化するクラスメソッド。
- configure_optimizers
  モデルのためのオプティマイザ（最適化アルゴリズム）を設定するクラスメソッド。
- estimate_mfu
  モデルのFLOPS（浮動小数点演算数）利用率を推定するクラスメソッド。
- generate
  与えられた条件付けシーケンスのインデックスから新しいトークンを生成する機能を提供するクラスメソッド。

4. クラスと関数の具体的な中身

nanoGPT のmain.py に入っているクラスとその中身について、コードベースで見ていきます。クラス内に日本語での説明を付与いたしました。

4.1. LayerNorm Class

class LayerNorm(nn.Module):
    """LayerNormの変種で、オプションとしてバイアスを持つ。PyTorchは単純にbias=Falseをサポートしない。
    このクラスは、Layer Normalizationを実装しており、オプションとしてバイアスの有無を指定できます。__init__メソッドで重みとバイアスのパラメータを初期化し、forwardメソッドで入力テンソルに対して正規化処理を行います。

    Args:
        ndim (int): 正規化される次元数。
        bias (bool): バイアスを使用するかどうか。

    """

    def __init__(self, ndim, bias):
        """初期化関数。
        
        Args:
            ndim (int): 正規化される次元数。
            bias (bool): バイアスを使用するかどうか。
        """
        super().__init__()
        # 重みを初期化。サイズはndimと一致
        self.weight = nn.Parameter(torch.ones(ndim))
        # バイアスを使用する場合、それを初期化。さもなければNone
        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self, input):
        """フォワードパスを定義。
        この関数は、自己注意機構を用いて入力テンソルxを処理します。最初に、クエリ、キー、バリューを計算し、それらを適切な形状に変換します。その後、Flash Attentionを使用するか、もしくは標準の自己注意計算を行います。最後に、すべてのヘッド出力を再構成し、出力プロジェクションを適用して戻り値としています。

        Args:
            input (Tensor): 入力テンソル。

        Returns:
            Tensor: 正規化されたテンソル。
        """
        # layer_norm関数を用いて正規化を実行
        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

4.2. CausalSelfAttention Class

class CausalSelfAttention(nn.Module):
    """因果的自己注意（Causal Self-Attention）を実装するクラス。

    PyTorchのモジュールを継承し、因果関係を考慮した自己注意機構を提供します。

    Args:
        config (object): 設定オブジェクト。モデルの各種パラメータを含みます。

    """

    def __init__(self, config):
        """初期化関数。
        
        Args:
            config (object): 設定オブジェクト。
        """
        super().__init__()
        assert config.n_embd % config.n_head == 0  # 埋め込み次元がヘッド数で割り切れることを確認
        # キー、クエリ、値の投影を全てのヘッドに対してバッチ処理として行う
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # 出力の投影
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # 正則化
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        # 設定からヘッド数と埋め込み次元を取得
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        # flash attentionのサポートチェック（PyTorch >= 2.0が必要）
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            # フラッシュアテンションが使えない場合の警告
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # 入力シーケンスの左側にのみ注意を適用するように因果マスクを作成
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                        .view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        """フォワードパスを定義。

        Args:
            x (Tensor): 入力テンソル。形状は [バッチサイズ, シーケンス長, 埋め込み次元数]。

        Returns:
            Tensor: 自己注意を適用した後のテンソル。
        """
        # 入力テンソルのサイズを取得（バッチサイズ、シーケンス長、埋め込み次元数）
        B, T, C = x.size()

        # クエリ、キー、バリューを計算し、ヘッド次元をバッチ次元へ移動
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # 自己注意の計算
        if self.flash:
            # Flash Attentionを使用した効率的な注意計算
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # 手動での注意計算
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)

        # ヘッド出力を再構成
        y = y.transpose(1, 2).contiguous().view(B, T, C)

        # 出力のプロジェクション
        y = self.resid_dropout(self.c_proj(y))
        return y

4.3. Block Class

class Block(nn.Module):
    """Transformerモデルのブロッククラス。
    このクラスはTransformerモデルの一つのブロックを表し、二つのLayerNormレイヤー、一つの自己注意レイヤー、そしてMLPレイヤーが含まれています。

    Args:
        config (オブジェクト): モデル構成情報を含むオブジェクト。

    """

    def __init__(self, config):
        """初期化関数。

        Args:
            config (オブジェクト): モデル構成情報を含むオブジェクト。
        """
        super().__init__()
        # LayerNormレイヤーを初期化（第一層）
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        # 自己注意レイヤーを初期化
        self.attn = CausalSelfAttention(config)
        # LayerNormレイヤーを初期化（第二層）
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        # MLPレイヤーを初期化
        self.mlp = MLP(config)

    def forward(self, x):
        """フォワードパスを定義。

        Args:
            x (Tensor): 入力テンソル。

        Returns:
            Tensor: 出力テンソル。
        """
        # LayerNormと自己注意を適用し、元の入力と加算
        x = x + self.attn(self.ln_1(x))
        # LayerNormとMLPを適用し、再度元の入力と加算
        x = x + self.mlp(self.ln_2(x))
        return x

4.4. GPTConfig Class

@dataclass
class GPTConfig:
    """GPTモデルの構成設定を定義するデータクラス。
    このクラスは、GPTモデルの構成設定を管理するためのものです。ブロックサイズ、語彙サイズ、レイヤーの数、アテンションヘッドの数、埋め込み次元の数、ドロップアウト率、バイアスの有無などの設定が含まれています。これらの設定は、モデルのトレーニングや推論時に使用されます

    Attributes:
        block_size (int): ブロックサイズ。デフォルトは1024。
        vocab_size (int): 語彙サイズ。GPT-2の語彙サイズは50257で、効率のため64の倍数にパディングされて50304になっている。
        n_layer (int): レイヤーの数。デフォルトは12。
        n_head (int): アテンションヘッドの数。デフォルトは12。
        n_embd (int): 埋め込み次元の数。デフォルトは768。
        dropout (float): ドロップアウト率。デフォルトは0.0。
        bias (bool): 線形層とLayerNormにバイアスを使用するか。Trueの場合はGPT-2のようにバイアスを使用。Falseの場合は少し良くて高速。
    """
    block_size: int = 1024
    vocab_size: int = 50304
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = True

4.5. GPT Class

class GPT(nn.Module):
    """GPTモデルの定義。

    Args:
        config (GPTConfig): GPTモデルの設定。

    """

    def __init__(self, config):
        """モデルの初期化。
        モデルの各部分（埋め込み、トランスフォーマーブロック、最終LayerNorm、言語モデルヘッド）を初期化し、重みを適切に設定します。

        Args:
            config (GPTConfig): GPTモデルの設定。
        """
        super().__init__()
        # 設定値の確認
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config

        # トランスフォーマーの各コンポーネントを初期化
        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),  # トークンの埋め込み
            wpe=nn.Embedding(config.block_size, config.n_embd),  # 位置の埋め込み
            drop=nn.Dropout(config.dropout),  # ドロップアウト層
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),  # トランスフォーマーブロックのリスト
            ln_f=LayerNorm(config.n_embd, bias=config.bias),  # 最終的なLayerNorm
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)  # 言語モデルのヘッド

        # 重みのタイイング（言語モデルのヘッドとトークン埋め込みの重みを共有）
        self.transformer.wte.weight = self.lm_head.weight

        # 重みの初期化
        self.apply(self._init_weights)

        # 特殊なスケーリングで残差プロジェクションの重みを初期化（GPT-2論文に従って）
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

        # パラメータ数の報告
        print(f"number of parameters: {self.get_num_params()/1e6:.2f}M")

    def get_num_params(self, non_embedding=True):
        """モデル内のパラメータ数を返す。
        モデル内の全パラメータ数を計算し、必要に応じて位置埋め込みのパラメータを除外します。

        Args:
            non_embedding (bool): 埋め込みを除外したパラメータ数をカウントする場合はTrue。デフォルトはTrue。

        Returns:
            int: モデル内のパラメータ数。
        """
        # モデルの全パラメータ数を計算
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding:
            # 位置埋め込みのパラメータ数を除外
            n_params -= self.transformer.wpe.weight.numel()
        return n_params

    def _init_weights(self, module):
        """モジュールの重みを初期化する。
        線形層と埋め込み層の重みを正規分布で初期化し、バイアスがある場合は0で初期化します。

        Args:
            module (nn.Module): 初期化するモジュール。
        """
        if isinstance(module, nn.Linear):
            # 線形層の重みを正規分布で初期化
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                # バイアスが存在する場合、0で初期化
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            # 埋め込み層の重みを正規分布で初期化
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        """GPTモデルのフォワードパス。
        この関数は、GPTモデルのフォワードパスを実装します。入力トークンのインデックスから、トークンと位置の埋め込みを生成し、これをドロップアウトレイヤーを通して、各トランスフォーマーブロックに順に適用します。目標となるトークンが与えられている場合は損失を計算し、そうでない場合は最終層のロジットのみを返します。

        Args:
            idx (Tensor): 入力トークンのインデックス。
            targets (Tensor, optional): 目標となるトークンのインデックス。デフォルトはNone。

        Returns:
            Tuple[Tensor, Tensor]: ロジットと、もしtargetsが指定されていれば損失値。
        """
        # 入力インデックスからデバイス情報を取得し、サイズを確認
        device = idx.device
        b, t = idx.size()
        # 入力サイズがブロックサイズを超えないことを確認
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        # 位置インデックスを生成
        pos = torch.arange(0, t, dtype=torch.long, device=device) # 形状 (t)

        # GPTモデル自体のフォワードパス
        tok_emb = self.transformer.wte(idx) # トークン埋め込み、形状 (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # 位置埋め込み、形状 (t, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb) # 埋め込みの加算後、ドロップアウト適用
        for block in self.transformer.h:
            x = block(x) # 各トランスフォーマーブロックを順に適用
        x = self.transformer.ln_f(x) # 最終的なLayerNorm適用

        if targets is not None:
            # 目標値が指定されている場合、損失も計算
            logits = self.lm_head(x) # 言語モデルヘッドを適用しロジットを計算
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) # クロスエントロピー損失計算
        else:
            # 推論時の最適化: 言語モデルヘッドを最後の位置にのみ適用
            logits = self.lm_head(x[:, [-1], :]) # 注: リスト[-1]を使用して時系列次元を保持
            loss = None

        return logits, loss

    def crop_block_size(self, block_size):
        """モデルのブロックサイズを変更する。
        この関数は、モデルのブロックサイズを必要に応じて小さく調整するために使用されます。新しいブロックサイズが現在のブロックサイズを超えていないことを確認した後、モデルの設定と位置埋め込みの重みを更新します。また、各トランスフォーマーブロックの注意機構のバイアスも新しいブロックサイズに合わせて調整されます。これにより、モデルを特定のニーズに合わせてカスタマイズすることが可能になります。

        Args:
            block_size (int): 新しいブロックサイズ。

        """
        # 新しいブロックサイズが現在のブロックサイズ以下であることを確認
        assert block_size <= self.config.block_size
        # モデル設定のブロックサイズを更新
        self.config.block_size = block_size
        # 位置埋め込みの重みを新しいブロックサイズに合わせて調整
        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
        # 各トランスフォーマーブロックの注意機構のバイアスを新しいブロックサイズに合わせて調整
        for block in self.transformer.h:
            if hasattr(block.attn, 'bias'):
                block.attn.bias = block.attn.bias[:, :, :block_size, :block_size]

    @classmethod
    def from_pretrained(cls, model_type, override_args=None):
        """事前訓練済みのモデルをロードして初期化するクラスメソッド。
        指定されたタイプの事前訓練済みGPTモデルをロードし、必要に応じてドロップアウト率をオーバーライドします。モデルの各レイヤー数、ヘッド数、埋め込み次元数は、モデルタイプに基づいて設定されます。

        Args:
            model_type (str): モデルのタイプ（'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'のいずれか）。
            override_args (dict, optional): ドロップアウト率などのオーバーライド引数。デフォルトはNone。

        Returns:
            GPT: 初期化されたGPTモデル。
        """
        # モデルタイプのバリデーション
        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
        override_args = override_args or {}

        # ドロップアウト以外の引数のオーバーライドを防止
        assert all(k == 'dropout' for k in override_args)

        # HuggingFaceの事前訓練済みモデルをロード
        print(f"loading weights from pretrained gpt: {model_type}")

        # モデル設定を定義
        config_args = {
            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),
            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024),
            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280),
            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600),
        }[model_type]

        # 常に固定された語彙サイズ、ブロックサイズ、バイアス設定
        config_args['vocab_size'] = 50257
        config_args['block_size'] = 1024
        config_args['bias'] = True

        # ドロップアウト率をオーバーライドする場合
        if 'dropout' in override_args:
            print(f"overriding dropout rate to {override_args['dropout']}")
            config_args['dropout'] = override_args['dropout']

        # GPTモデルを初期化
        config = GPTConfig(**config_args)
        model = GPT(config)

        # モデルの状態を取得し、特定のパラメータを除外
        sd = model.state_dict()
        sd_keys = [k for k in sd.keys() if not k.endswith('.attn.bias')]

        # HuggingFaceモデルの状態をロード
        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
        sd_hf = model_hf.state_dict()

        # パラメータを確認し、必要に応じて転置してコピー
        for k in sd_keys:
            if k in sd_hf:
                # 転置が必要なパラメータに対して特別な処理
                if k in ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']:
                    assert sd_hf[k].shape[::-1] == sd[k].shape
                    with torch.no_grad():
                        sd[k].copy_(sd_hf[k].t())
                else:
                    # 他のパラメータをそのままコピー
                    assert sd_hf[k].shape == sd[k].shape
                    with torch.no_grad():
                        sd[k].copy_(sd_hf[k])

        return model

    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
        """オプティマイザを設定する。
        この関数は、モデルのためのオプティマイザ（最適化アルゴリズム）を設定します。全てのパラメータを取得し、重み減衰を適用すべきパラメータとそうでないパラメータを分けてグループ化します。その後、AdamWオプティマイザを作成し、CUDAで利用可能な場合はfusedバージョンを使用します。最終的にこのオプティマイザを返します。

        Args:
            weight_decay (float): 重み減衰の値。
            learning_rate (float): 学習率。
            betas (Tuple[float, float]): AdamWオプティマイザのベータ値。
            device_type (str): 使用するデバイスのタイプ ('cuda' または 'cpu')。

        Returns:
            torch.optim.Optimizer: 設定されたオプティマイザ。
        """
        # 全パラメータを取得
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # 勾配が必要なパラメータのみにフィルタリング
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        # 重み減衰を適用するパラメータグループとしないグループを作成
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
        optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0}
        ]
        # AdamWオプティマイザの作成。CUDAで利用可能な場合はfusedバージョンを使用
        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_available and device_type == 'cuda'
        extra_args = dict(fused=True) if use_fused else dict()
        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)

        return optimizer

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        """モデルのFLOPS利用率（MFU）を推定する。
        この関数は、モデルのFLOPS（浮動小数点演算数）利用率を推定します。モデルのパラメータ数と設定から、トークンごとのFLOPSを計算し、それをイテレーションごとのFLOPSに拡張します。最終的に、このFLOPSのスループットをA100 GPUのピークFLOPSと比較して、モデルのFLOPS利用率を求めます。

        Args:
            fwdbwd_per_iter (int): 各イテレーションでの前方および後方パスの回数。
            dt (float): 測定期間（秒）。

        Returns:
            float: モデルのFLOPS利用率（MFU）。

        """
        # 各イテレーションでのFLOPS（浮動小数点演算数）の推定
        N = self.get_num_params()  # モデルのパラメータ数
        cfg = self.config  # モデルの設定
        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size
        flops_per_token = 6 * N + 12 * L * H * Q * T  # トークンごとのFLOPS
        flops_per_fwdbwd = flops_per_token * T  # 前方および後方パスでのFLOPS
        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter  # イテレーションごとのFLOPS

        # FLOPSのスループットをA100 bfloat16ピークFLOPSと比較
        flops_achieved = flops_per_iter * (1.0 / dt)  # 秒間FLOPS
        flops_promised = 312e12  # A100 GPU bfloat16ピークFLOPSは312 TFLOPS
        mfu = flops_achieved / flops_promised  # FLOPS利用率

        return mfu

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """与えられたインデックスシーケンスから新しいトークンを生成する。
        この関数は、与えられた条件付けシーケンスのインデックスから新しいトークンを生成する機能を提供します。各イテレーションで、モデルを使用してシーケンスの次のトークンのロジットを計算し、特定の温度パラメータで調整します。必要に応じて、上位のロジットを選択し、ソフトマックスを適用して確率に変換します。その後、確率分布から次のトークンのインデックスをサンプリングし、シーケンスに追加します。

        Args:
            idx (LongTensor): 条件付けシーケンスのインデックス（形状は(b,t)）。
            max_new_tokens (int): 生成する新しいトークンの最大数。
            temperature (float): 温度パラメータ。デフォルトは1.0。
            top_k (int, optional): トップkロジットを選択するためのパラメータ。デフォルトはNone。

        Returns:
            LongTensor: 生成されたインデックスシーケンス。
        """
        for _ in range(max_new_tokens):
            # シーケンスが長すぎる場合はブロックサイズで切り取る
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            # モデルを用いてシーケンスのロジットを取得
            logits, _ = self(idx_cond)
            # 最後のステップのロジットを取得し、温度パラメータで調整
            logits = logits[:, -1, :] / temperature
            # top_kが指定されている場合、上位のロジットのみを選択
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            # ソフトマックスを適用し、確率に変換
            probs = F.softmax(logits, dim=-1)
            # 確率分布から次のインデックスをサンプリング
            idx_next = torch.multinomial(probs, num_samples=1)
            # サンプリングされたインデックスをシーケンスに追加
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

コードを見て分かるように、GPT 自体は単純な構造ですね。私はOpenAI が凄さはRLHF を含めたチューニングの技術と考えています。モデルの構築、学習、RLFH を含めたパラメータチューニング一通りやってみたいです。

簡単な解説でしたが本記事は以上です。少しでもGPT を詳しく勉強するきっかけに少しでもなれば幸いです！