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

AGI時代に向けたプログラマーの未来:役割変化とキャリア戦略

AGI時代に向けたプログラマーの未来:役割変化とキャリア戦略

はじめに 私がはじめてコードを書いたのは1989年です。 当時NECのPC88というパソコンを中古でかってもらい N-88 Basic というBASIC言語のコードをみようみまねで書いて動かしたあの日から何年経つのでしょうか。 当時、電波新聞社のマイコンBASICマガジンという雑誌があり、ベーマガにはいろんなパソコン向けのプログラムコードが掲載されていました。 そんなわけでもう35年以上趣味や仕事でプログラミングに従事していますが、開発環境、情報流通の仕組みには革命といっていいほどの変化、進化がおこりました。 しかしながら、そんな中でも、あくまでコードを書くのは「私」という生身の人間でした。 そうしたある種の古き良き時代は、いよいよ本格的に終わりを告げようとしています。 2023年ごろからのLLM技術の飛躍的進歩により、プログラミング業界は大きな転換期を迎えています。 特に、OpenAI o3,o1やClaude 3.5、Gemini2.0などの大規模言語モデル(LLM)の進化や、その先にある将来的な汎用人工知能(AGI)の出現は、プログラマーやAIエンジニアの役割に根

By Tomonori Misawa / CEO
PythonとWSL開発のトラブルシューティング: PyCharmとCondaの環境不一致問題

PythonとWSL開発のトラブルシューティング: PyCharmとCondaの環境不一致問題

こんにちは! 今回は、WSL上のConda環境をPyCharmから利用する際に発生した「同じ環境なのにパッケージリストが一致しない」という問題に遭遇したため、その原因と対策について書いてみたいとおもいます 問題の状況 開発の流れは以下のようなものでした 1. WSL環境でConda仮想環境を作成 2. その環境をPyCharmのプロジェクトインタプリタとして設定 3. 開発を進める中で奇妙な現象に気づく 具体的には、次のような不一致が発生していました * PyCharmのプロジェクト設定で表示されるpipパッケージのリスト * WSLでConda環境をアクティベートした後にpip listコマンドで表示されるパッケージのリスト これらが一致せず、「WSL側のシェルから直接インストールしたパッケージがPyCharmで認識されない」という問題が生じていました。 この手の問題でよくある原因は、PyCharm側がWSL側の更新を得るのに少し時間がかかったり、 Indexing が遅れているなどなのですが、今回はそれが原因ではありませんでした。 危険な「静かな

By Qualiteg プロダクト開発部
人気ゲーム「ヒット&ブロー」で学ぶ情報理論

人気ゲーム「ヒット&ブロー」で学ぶ情報理論

こんにちは! Qualiteg研究部です! 今日はAIにおいても非常に重要な情報理論について、Nintendo Switchの人気ゲーム「世界のアソビ大全51」にも収録されている「ヒット&ブロー」というゲームを題材に解説いたします! はじめに 論理的思考力を鍛える定番パズルゲームとして長年親しまれている「ヒット&ブロー」(海外では「Mastermind」として知られています)。 このゲームは一見シンプルながらも、その攻略には深い論理的アプローチが必要とされております。 本稿では、このゲームについて情報理論という数学的概念を用いてゲームの素性を分析する方法について掘り下げてみたいとおもいます。 さらに、この情報理論が現代の人工知能(AI)技術においてどのように活用されているかについても触れていきます。 ヒット&ブローのルール説明 ヒット&ブローは、相手が秘密に設定した色や数字の組み合わせを推測するゲームです。日本では主に数字を使った「数当てゲーム」として親しまれていますが、本記事では色を使ったバージョン(マスターマインド)に焦点を当てます。 Nintendo Sw

By Qualiteg 研究部
Model Context Protocol(MCP)入門:いよいよセマンティックWebの世界へ

Model Context Protocol(MCP)入門:いよいよセマンティックWebの世界へ

こんにちは! きょうは話題のMCPについて解説いたします! はじめに 「AIが便利なのはわかるけど、自分のデータにアクセスさせたり、他のアプリと連携させたりするのは難しそう...」 このような悩みを持っている方は多いのではないでしょうか。 実際、従来のAIには大きな壁がありました。トレーニングデータの範囲でしか回答できない、リアルタイム情報にアクセスできない、外部アプリケーションを操作できないなどの制約です。 トレーニングデータの外側にあるデータをうまく検索する技術としてLLM黎明期からRAGとよばれる技術が発展してきました。 データ検索だけではなく、あらゆる分野でAIが半ば自動で連携してくれる技術が登場しました。 それが「Model Context Protocol(MCP)」です。 本記事では、AIと外部ツールの連携を革新的に簡単にするMCPについて、基本から実用まで詳しく解説します。 MCPの本質:AIのための標準インターフェース MCPは、AIモデルと外部ツール・アプリケーションの間の通信を標準化するプロトコルです。これはインターネットの世界でいえば、

By Qualiteg プロダクト開発部