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じゃなくて勝手にCPUで実行されてるかも  ~ONNX RuntimeのcuDNN 警告と対策~

その処理、GPUじゃなくて勝手にCPUで実行されてるかも ~ONNX RuntimeのcuDNN 警告と対策~

こんにちは! 本日は、ONNX RuntimeでGPU推論時の「libcudnn.so.9: cannot open shared object file」エラーの解決方法についての内容となります。 ONNX Runtimeを使用してGPU推論を行う際、CUDAプロバイダの初期化エラーに遭遇することがありますので、このエラーの原因と解決方法を解説いたします。 エラーメッセージの詳細 [E:onnxruntime:Default, provider_bridge_ort.cc:2195 TryGetProviderInfo_CUDA] /onnxruntime_src/onnxruntime/core/session/provider_bridge_ort.cc:1778 onnxruntime::Provider& onnxruntime::ProviderLibrary::Get() [ONNXRuntimeError] : 1 : FAIL : Failed to load

By Qualiteg プロダクト開発部
大企業のAIセキュリティを支える基盤技術 - 今こそ理解するActive Directory 第3回 クライアントとサーバーのドメイン参加

大企業のAIセキュリティを支える基盤技術 - 今こそ理解するActive Directory 第3回 クライアントとサーバーのドメイン参加

こんにちは、今回はシリーズ第3回クライアントとサーバーのドメイン参加について解説いたします! はじめに こんにちは!シリーズ第3回「クライアントとサーバーのドメイン参加」へようこそ。 前回(第2回)では、Active Directoryドメイン環境の構築手順について、ドメインコントローラーのセットアップからDNS設定まで詳しく解説しました。ドメイン環境の「土台」が整ったところで、今回はいよいよ実際にコンピューターをドメインに参加させる手順に進みます。 「ドメインユーザーアカウントを作ったのに、なぜかログインできない」「新しいPCを追加したけど、ドメイン認証が使えない」といった経験はありませんか?実は、Active Directoryの世界では、ユーザーアカウントを作成しただけでは不十分で、そのユーザーが使用するコンピューター自体もドメインに「参加」させる必要があるのです。 本記事では、このドメイン参加について、単なる手順の説明にとどまらず、「なぜドメイン参加が必要なのか」「裏側で何が起きているのか」という本質的な仕組みまで、初心者の方にも分かりやすく解説していきます。Win

By Qualiteg コンサルティング
使い捨てソフトウェア時代の幕開け ― 市場構造の根本的変革と日本企業

使い捨てソフトウェア時代の幕開け ― 市場構造の根本的変革と日本企業

こんにちは、株式会社Qualiteg コンサルティング部門です。 昨今、生成AIの急速な進化により、ソフトウェア開発の在り方が根本から変わりつつあります。2024年にはClaude、GPT-4、Geminiなどの大規模言語モデルがコード生成能力を飛躍的に向上させ、GitHub CopilotやCursor、Windsurf等の開発支援ツールが実際の開発現場で広く活用されるようになりました。さらに、Devin、OpenAI Canvas、Anthropic Claude Codingといった、より高度な自律的コーディング機能を持つAIエージェントも登場しています。 このような技術革新を背景に、当部門では今後のソフトウェア産業の構造変化について詳細な分析を行いました。本シリーズでは、特に注目すべき変化として、従来1000人月規模を要していた企業向けSaaSプラットフォームや、基幹システムが、AIエージェントを効果的に活用することで、わずか2-3名のチームが数日から数週間で実装可能になるという、開発生産性の劇的な向上について考察してまいります。 これは単なる効率化ではなく、ソフトウェア

By Qualiteg コンサルティング
NVIDIA GeForce RTX 50xx with CUDA capability sm_120 is not compatible with the current PyTorch installation. が発生したとき

NVIDIA GeForce RTX 50xx with CUDA capability sm_120 is not compatible with the current PyTorch installation. が発生したとき

こんにちは、PyTorch 2.6.0 環境で以下のような問題が発生したときの対処方法について解説いたします。 NVIDIA GeForce RTX 5090 with CUDA capability sm_120 is not compatible with the current PyTorch installation. The current PyTorch install supports CUDA capabilities sm_50 sm_60 sm_70 sm_75 sm_80 sm_86 sm_90. 他のBlackwell GeForce の場合は以下のようなメッセージとなります。 NVIDIA GeForce RTX

By Qualiteg プロダクト開発部