PyTorch 2.6 のセキュリティ: モデルチェックポイントロードの安全対策

PyTorch 2.6 のセキュリティ: モデルチェックポイントロードの安全対策

こんにちは!

今日は、魅力的なPyTorchチェックポイントが配られているからと言って無邪気に使っちゃうと、超ヤバイよ、というお話になります。

みなさまモデルチェックポイントをロードする際のセキュリティリスクについて考えたことはありますでしょうか。実はモデルの重みファイルだとばかり思っていたチェックポイントが、思わぬセキュリティホールになる可能性があります。

本記事では、PyTorchのtorch.load関数の安全性と、モデルチェックポイントを適切に扱うための実践的なガイドラインを紹介します。

モデルチェックポイントの隠れた危険性

PyTorchのtorch.load関数は非常に便利な一方で、セキュリティ上の重大なリスクを含んでいます。

その理由は、

  • チェックポイント単なるパラメータだけではないよ!
    チェックポイントファイルには、モデルの重み(weights)だけでなく、クラスや関数など任意のPythonコードを含めることが可能です。
  • 実行可能なコードが入ってるよ!
    これは、チェックポイントが単なる「データファイル」ではなく、Pythonコードを実行できる可能性を持つファイルであることを意味します。
  • ってことで潜在的な脆弱性があるよ!
    悪意あるチェックポイントを不用意にロードすると、ファイルをロードした瞬間に任意のコードが実行される可能性があります。

そうです、チェックポイントは重みデータだけではないんです。

チェックポイントファイルに入れられるもの

ということで、PyTorchのチェックポイントファイル(.ptや.pth,.ckptファイル)には、モデルの重み(パラメータ)だけでなく、実際にPythonのコードも含めることができてしまいます。

これは、PyTorchが内部的にPythonのpickleシリアライゼーション形式を使用しているためです。

pickleは、Pythonオブジェクトの状態を保存・復元するための仕組みですが、クラス定義や関数など、実行可能なコードも含めて保存できる特性があります。

具体的には、チェックポイントには以下のような要素が含まれてる可能性があります。

  • モデルの重みパラメータ
  • モデルのアーキテクチャ情報
  • オプティマイザの状態
  • 学習率スケジューラの状態
  • カスタムクラス・関数の定義(←ここが危ない)
  • その他のメタデータ

「カスタムクラス・関数の定義」が、セキュリティ上の懸念となります。悪意のある人が作成したチェックポイントには、ファイル削除やシステムコマンドの実行など、危険なコードが含まれている可能性があります。

だから、魅力的なチェックポイントが配られているからと言って無邪気に使っちゃうと、超ヤバイんです。

具体的な攻撃シナリオ

悪意あるチェックポイントが引き起こす可能性のある問題は次のようなものがあります

  1. (あなたが)不正なファイルをダウンロードする
  2. (あなたが)それをweights_only=Falseでロードする
  3. (わるい奴が仕込んだ)チェックポイントに含まれた悪意あるコードが即座に実行される
  4. (最悪の場合は)攻撃者はPC上で任意の操作(ファイル削除、情報漏洩、マルウェアのインストールなど)が可能になる

PyTorch 2.6でのセキュリティ強化

PyTorch 2.6以降では、セキュリティ対策が強化されました。具体的な変更点としては、

  • torch.load()weights_only引数がデフォルトでTrueに設定されるようになりました
  • これにより、デフォルトでモデルのパラメータ(重み)のみが安全にロードされ、潜在的に危険なコードの実行が防止されます

この変更により、以下のようなエラーが発生したり

RuntimeError: ('Attempted to deserialize object from torch.nn.Module that contains non-parameter/buffer types, which could potentially lead to security vulnerabilities. ...

 _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
         (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
         (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
         WeightsUnpickler error: Unsupported global:...

このエラーは、チェックポイントファイルに単なるモデルの重みだけでなく、クラスや関数などのコード情報も含まれている場合に発生します。例えば、fairseq.data.dictionary.Dictionaryのようなクラスが含まれていると、デフォルト設定では未許可のクラスとして拒否されます。

安全の基準:何が「安全」で何が「安全でない」か

🟢 「安全」と判断できるファイル

次のような場合はほぼ「安全」と考えることができます

  • 自分自身で作成したモデルのチェックポイント(内容を完全に把握しているため)
  • 信頼できる公式ソースから直接取得したモデルファイル(有名な大学、著名な企業、公式レポジトリなど)
  • 公式のGitHubリポジトリや公式ウェブサイトから取得したモデルで、提供者が広く知られており、モデルのチェックポイントの内容が明示的に示されている場合(例:Hugging Faceの公式リポジトリ、Facebookのfairseq公式リポジトリなど)

(とはいえ、注意深くセキュリティ情報は把握しておく必要はありますし、基本は weights_only =True でロードです。)

🔴 「安全でない」可能性のあるファイル

次のような場合は「安全でない」と疑うべきです:

  • 出所不明のサイトや第三者が匿名で共有したモデル
  • SNSや非公式フォーラムなど、素性がわからない人物から直接ダウンロードしたモデルファイル
  • モデルファイルの配布者が特定できず、かつモデルの構造やコードについて十分な情報が公開されていない場合

このようなファイルは絶対にweights_only=Falseでロードしてはいけません。

🔐 安全なモデルロードのためのコード例と対処法

エラーが発生した場合の解決方法

ここではfacebookの fairseqを使っていて実際に発生したエラーと対処法についてみてみます

 _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
         (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
         (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
         WeightsUnpickler error: Unsupported global: GLOBAL fairseq.data.dictionary.Dictionary was not an allowed global by default. Please use torch.serialization.add_safe_globals([Dictionary]) or the torch.serialization.safe_globals([Dictionary]) context manager to allowlist this global if you trust this class/function.

このエラーは、PyTorchのtorch.load()でモデルのチェックポイントを読み込む際、セキュリティのためデフォルト設定が厳しくなったために発生したものです。

具体的には、以下の原因によります

  • PyTorch 2.6以降で、torch.load()weights_only引数がデフォルトでTrueとなり、モデルのパラメータ(重み)のみを安全にロードしようとします。
  • このエラーが発生するチェックポイントファイルには、単なるモデルの重みだけでなく、クラスや関数などのコード情報も含まれています。そのため、weights_only=Trueのままロードしようとすると、未許可のクラスや関数が含まれているとして拒否されます。
  • 今回のエラーの場合、未許可として拒否されたクラスは fairseq.data.dictionary.Dictionary となっていることがこのメッセージから読み取れますね。

① 推奨の解決方法(安全性高)

信頼できるクラスを「許可リスト」に登録してから読み込みます。

import torch
from fairseq.data.dictionary import Dictionary

# このクラスが安全と分かっている場合のみ実行
torch.serialization.add_safe_globals([Dictionary])

# その後にチェックポイントをロード
checkpoint = torch.load('checkpoint.pt')

もしくは、一時的に許可する場合はコンテキストマネージャを利用します。

import torch
from fairseq.data.dictionary import Dictionary

with torch.serialization.safe_globals([Dictionary]):
    checkpoint = torch.load('checkpoint.pt')

② 旧来の方法(安全性低)

weights_only=Falseを指定してロードします。ただし、この方法は安全ではないため、信頼できるソースから取得したファイルにのみ使用してください。

import torch

checkpoint = torch.load('checkpoint.pt', weights_only=False)

実装例:安全にチェックポイントをロードする関数

以下は、安全性を考慮したチェックポイントロード関数の例です

import torch
from typing import List, Type, Optional

def load_checkpoint_safely(
    checkpoint_path: str,
    trusted_classes: Optional[List[Type]] = None,
    force_weights_only: bool = False
) -> dict:
    """
    安全にチェックポイントをロードする関数
    
    Args:
        checkpoint_path: チェックポイントファイルのパス
        trusted_classes: 信頼できるクラスのリスト(None の場合は空リスト)
        force_weights_only: True の場合、weights_only=True で強制的にロード
        
    Returns:
        チェックポイントの内容
        
    Raises:
        RuntimeError: ロードに失敗した場合
    """
    if force_weights_only:
        return torch.load(checkpoint_path, weights_only=True)
        
    trusted_classes = trusted_classes or []
    
    try:
        # まず安全にロードを試みる
        return torch.load(checkpoint_path)
    except RuntimeError as e:
        if "non-parameter/buffer types" in str(e):
            # 信頼できるクラスがある場合はそれらを使用
            if trusted_classes:
                with torch.serialization.safe_globals(trusted_classes):
                    return torch.load(checkpoint_path)
            else:
                raise RuntimeError(
                    f"チェックポイントに未許可のクラスが含まれています。"
                    f"信頼できるソースの場合は、trusted_classes パラメータを使用してください。"
                    f"原エラー: {str(e)}"
                )
        else:
            # その他のエラーはそのまま再発生
            raise

さて、最後に今回の事例についてベストプラクティスをまとめておきましょう

ベストプラクティスのまとめ

  1. デフォルトは安全第一
    • PyTorch 2.6以降のデフォルト設定(weights_only=True)を活用する
    • 信頼性が不明なファイルは常にデフォルト設定のままロードを試みる
  2. 明示的な許可リストの使用
    • 信頼できるクラスはadd_safe_globalsまたはsafe_globalsコンテキストマネージャで明示的に許可する
    • これにより、必要最小限のクラスのみを許可し、安全性を確保できる
  3. weights_only=Falseは最終手段です
    • 信頼できるファイルの場合のみ、最終手段として使用する
    • 不明なソースのチェックポイントには絶対に使用しない

まとめ:セキュリティ意識の向上を

PyTorch 2.6以降のセキュリティ強化は、深層学習モデルの安全な共有と利用を前進させました。

基本的な安全の基準は非常にシンプルで

  • 出所が明確で信頼(安全のエビデンスがある)できる → 安全
  • 出所不明、匿名、不審な提供元 → 安全でない

当社でもプロダクトのベースを PyTorch 2.6 にしたら、いままで動いていたコードが動かなくなり、多くの見直しも発生しました。が、そのおかげで、モデル・セキュリティの意識と知識を一歩深めることができたと感じております。

ディープラーニングの適用例やプレイヤーが増えるにつれ多くの有用なモデルが配布されるようになりましたが、プレイヤーが増えると、このようによからぬことを考えるプレイヤーも増えるという事なので、今後もチェックポイントのロードには細心の注意を払っていきたいと思います。旧バージョンのPyTorchを使ってるプロダクト開発する必要がある場合は特にきをつけるべきですね。

それでは、また次回お会いしましょう!

Read more

GPUサーバーの最適容量計算: キューイング理論と実践的モデル

GPUサーバーの最適容量計算: キューイング理論と実践的モデル

最大ユーザーサポート数計算ツール 同時に1件のみ処理できるGPU変換サーバーの最大ユーザーサポート数を計算します 処理時間 (t_p) 分/件 1件の変換処理にかかる時間 目標システム利用率 (ρ) 0 〜 1 安定稼働のための目標稼働率(推奨: 0.7〜0.8) ピーク係数 (P_c) 倍 最も混雑する時間帯の平均アクセス倍率 稼働時間 (H) 時間/日 システムが1日に稼働している総時間 アクセス確率 (P_a) 0 〜 1 1人のユーザーが1日にシステムを利用する確率 1ユーザーあたりの変換回数 (F) 回/日 利用する日の平均変換処理回数 計算過程を表示 計算結果 サポート可能な総ユーザー数: 人 計算式: N = (ρ × μ × H) ÷ (P_a

By Qualiteg プロダクト開発部
[AI新規事業創出]Qualitegオリジナル、アイデア評価、事業アイデア選定方法

[AI新規事業創出]Qualitegオリジナル、アイデア評価、事業アイデア選定方法

Qualiteg blogを訪問してくださった皆様、こんにちは。Micheleです。AIを活用した新規事業やマーケティングを手がけている私には、クライアントからよく寄せられる質問があります。AIを用いた事業展開を検討されている方々が共通して直面するであろう課題に対して、このブログを通じて私なりの解答をご提供したいと思います。 AIを活用した事業アイデア評価と選定方法 | Qualitegオリジナルアプローチ 新規事業の立ち上げは、アイデアの創出から始まりますが、その後の評価と選定プロセスこそが成功の鍵を握ります。Qualitegでは、AIを積極的に活用した独自の評価・選定メソッドを開発し、より客観的かつ多角的な視点でビジネスアイデアを検証しています。今回は、私たちの実践的なアプローチをご紹介します。 AIを活用したアイデア評価の基本フレームワーク 当社のアイデア評価プロセスは、以下の2段階で構成しております。 1. 多次元評価マトリックスによる定量分析 まず、出てきたアイデアについて、ChatGPTなどの大規模言語モデル(LLM)を活用し、以下の8つの評価軸でアイデア

By Join us, Michele on Qualiteg's adventure to innovation
Windows Terminal で「無効な "icon" を持つプロファイルが見つかりました。既定では、そのプロファイルにアイコンはありません。」が出たときの対処法

Windows Terminal で「無効な "icon" を持つプロファイルが見つかりました。既定では、そのプロファイルにアイコンはありません。」が出たときの対処法

何度か、WSL にいろんなバージョンのLinux を入れたり消したりしたときに遭遇した現象です ユーザー設定の読み込み中にエラーが発生しました 無効な "icon" を持つプロファイルが見つかりました。既定では、そのプロファイルにアイコンはありません。"icon" を設定するときは、値が画像への有効なファイルパスとなっていることをご確認ください。 が発生するときの原因と対象法のレポートです 原因 使われなくなったゾンビ・プロファイルがWindows Terminal (のキャッシュ)に残り続ける 対処法 このメッセージを解消するには、いったん、プロファイルをリセットする必要がありました。 ※既存プロファイル設定が消える場合があるので留意すること Step1 Windows Terminal を落とす Windows Terminal をいったんすべて落とす Step2 settings.json を消す エクスプローラーで settings.json のあるフォルダに移動しファイルを削除する %LOCALAPPDATA%\Packages\Micros

By Qualiteg プロダクト開発部
本番運用におけるPyTorch+CUDAサーバーでの「Unknown Error」問題とその対策

本番運用におけるPyTorch+CUDAサーバーでの「Unknown Error」問題とその対策

こんにちは!Qualitegプロダクト開発部です。 今日は、GPUをつかった商用サービスにて悩ましい、テストは全部通るけど、長時間運用をしていると急に起こる「CUDA error: unknown error」についての内容です。 これ、出会うと残念な気持ちになりますが、けっこうGPU商用サービス界隈では「あるある」なんです。 原因を真面目に探るには CUDAバージョン、PyTorchバージョンの調合具合、実際のアプリケーションコードまですべてソースまで追う必要があるのですが、多くの場合、運用でカバーします。 なぜなら仮に1つ原因をみつけて対処できたとしても、CUDAバージョンはしょっちゅうあがりますし、PyTorchもそれに追従して頻繁に更新されます。さらにやっかいなことに、1日、2日、いや1週間くらいは安定的に動作しているようにみえて、数週間後にとつぜんエラーが出るといった具合なので、修正確認の難易度が高いんです。 そこで本日は「開発環境や実験環境」ではなく「本番環境」で発生しがちなこのCUDA Unknown Error について問題の原因と実践的な対策につい

By Qualiteg プロダクト開発部