PyTorch FSDP で optim_state_dict を使ってオプティマイザーの状態を保存・復元

PyTorch: Fully Sharded Data Parallel - torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict() 解説

torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict() は、PyTorch の Fully Sharded Data Parallel (FSDP) で使用される関数です。FSDP は、大規模なモデルを複数の GPU に分散させて効率的にトレーニングするための技術です。この関数は、FSDP で使用されるオプティマイザーの状態辞書を取得するために使用されます。

詳細

optim_state_dict() 関数は、以下の引数を受け取ります。

• model: FSDP でラップされたモデル
• optim: モデルの訓練に使用されるオプティマイザー
• optim_state_dict: オプティマイザーの状態辞書
• is_named_optimizer: オプティマイザーが NamedOptimizer または KeyedOptimizer であるかどうか
• load_directly: True の場合、この関数は optim.load_state_dict(result) を呼び出して、結果をオプティマイザーに直接読み込みます。
• group: モデルのパラメータがシャードされているプロセスグループ

戻り値

この関数は、以下の情報を格納した辞書を返します。

• state: オプティマイザーの状態
• param_groups: パラメータグループの情報

コード例

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

optim_state_dict = fsdp_model.optim_state_dict(optim)

# オプティマイザーの状態を保存
torch.save(optim_state_dict, "optim_state_dict.pt")

# オプティマイザーの状態をロード
optim_state_dict = torch.load("optim_state_dict.pt")

# オプティマイザーに状態を復元
optim.load_state_dict(optim_state_dict)

補足

• FSDP は、PyTorch 1.9 以降で利用可能です。
• FSDP を使用するには、torch.distributed モジュールをインストールする必要があります。

PyTorch FSDP optim_state_dict サンプルコード集

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

# オプティマイザーの状態を取得
optim_state_dict = fsdp_model.optim_state_dict(optim)

# オプティマイザーの状態を保存
torch.save(optim_state_dict, "optim_state_dict.pt")

# オプティマイザーの状態をロード
optim_state_dict = torch.load("optim_state_dict.pt")

# オプティマイザーに状態を復元
optim.load_state_dict(optim_state_dict)

NamedOptimizer/KeyedOptimizer の使用

from torch.optim import Optimizer

class NamedOptimizer(Optimizer):
    def __init__(self, params, lr=0.01):
        super().__init__(params, lr)
        self.param_groups = [{"params": params, "lr": lr}]

class KeyedOptimizer(Optimizer):
    def __init__(self, params, lr=0.01):
        super().__init__(params, lr)
        self.param_groups = [{"params": params, "lr": lr, "key": "param_group_1"}]

model = torch.nn.Linear(10, 1)

# NamedOptimizer の場合
optim_named = NamedOptimizer(model.parameters())

# KeyedOptimizer の場合
optim_keyed = KeyedOptimizer(model.parameters())

fsdp_model = FSDP(model)

# NamedOptimizer の場合
optim_state_dict_named = fsdp_model.optim_state_dict(optim_named, is_named_optimizer=True)

# KeyedOptimizer の場合
optim_state_dict_keyed = fsdp_model.optim_state_dict(optim_keyed, is_named_optimizer=True)

# ...

load_directly オプションの使用

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

# オプティマイザーの状態を取得
optim_state_dict = fsdp_model.optim_state_dict(optim)

# オプティマイザーに状態を直接読み込む
fsdp_model.optim_state_dict(optim, load_directly=True)

特定のパラメータグループのみの状態を取得

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

# 特定のパラメータグループのみの状態を取得
optim_state_dict = fsdp_model.optim_state_dict(
    optim, param_groups=["param_group_1"]
)

特定のプロセスグループのみの状態を取得

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

# プロセスグループを作成
group = torch.distributed.new_group(ranks=[0, 1])

fsdp_model = FSDP(model, process_group=group)

# 特定のプロセスグループのみの状態を取得
optim_state_dict = fsdp_model.optim_state_dict(
    optim, group=group
)

これらのサンプルコードは、PyTorch FSDP optim_state_dict 関数の使用方法を理解するのに役立ちます。

補足

• FSDP を使用するには、`

PyTorch FSDP optim_state_dict のその他の使用方法

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

# パラメータの名前を取得
param_names = [name for name, _ in model.named_parameters()]

# パラメータの名前と状態を格納する辞書を作成
optim_state_dict = {}
for name in param_names:
    optim_state_dict[name] = fsdp_model.optim_state_dict(optim, param_groups=[name])

# ...

optim.state_dict() を直接使用

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

# オプティマイザーの状態を取得
optim_state_dict = optim.state_dict()

# ...

カスタムな方法で状態を取得

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

def custom_optim_state_dict(model, optim):
    # カスタムな方法で状態を取得
    # ...

model = torch.nn.Linear(10, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)

fsdp_model = FSDP(model)

# カスタムな方法で状態を取得
optim_state_dict = custom_optim_state_dict(fsdp_model, optim)

# ...

これらの方法は、optim_state_dict 関数のデフォルトの動作を変更したい場合に役立ちます。