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

発話音声からリアルなリップシンクを生成する技術 第2回:AIを使ったドリフト補正

発話音声からリアルなリップシンクを生成する技術 第2回:AIを使ったドリフト補正

こんにちは! 前回の記事では、当社のMotionVoxで使用している「リップシンク」技術について、wav2vecを用いた音声特徴量抽出の仕組みを解説しました。音声から正確な口の動きを予測するための基礎技術について理解いただけたかと思います。 今回は、その続編として、リップシンク制作における重要な技術的課題である「累積ドリフト」に焦点を当てます。wav2vecで高精度な音素認識ができても、実際の動画制作では複数の音声セグメントを時系列に配置する際、わずかなタイミング誤差が蓄積して最終的に大きなずれとなる現象が発生します。 本記事では、この累積ドリフトのメカニズムと、機械学習を活用した最新の補正技術について、実際の測定データを交えながら詳しく解説していきます。前回のwav2vecによる特徴抽出と今回のドリフト補正技術を組み合わせることで、MotionVoxがどのように高品質なリップシンクを実現しているのか、その全体像が見えてくるはずです。 累積ドリフトとは何か 基本概念 累積ドリフトとは、個々の音声セグメントが持つ微小なタイミング誤差が、時間の経過とともに蓄積していく現象で

By Qualiteg 研究部
AIエージェント時代の新たな番人「ガーディアンエージェント」とは?

AIエージェント時代の新たな番人「ガーディアンエージェント」とは?

こんにちは!今日は先日ガートナーが発表したガーディアンエージェントについて解説します ガートナーの公式定義 ハイプカーブで有名なガートナーは2025年6月に、ガーディアンエージェントに関する見解を発表しました。ガーディアン・エージェントとは、AIとの安全で信頼できるやりとりを支援するために設計されたAIベースのテクノロジです。 ざっくりいうと、 「AIエージェントが来るよ」と予言したガートナー社は、次は、「ガーディアンエージェントが来るよ」と予言しました。なぜガーディアンエージェントが来るのでしょうか?本稿では、そのあたりを考察していきたいと思います。 なぜ今、AIの「監視役」が必要なのか 2025年、私たちは本格的なAIエージェント時代の入り口に立っています。AIが単なるツールから、自律的に判断し行動する「エージェント」へと進化する中で、新たな課題が浮上しています。 従来のAIとエージェント型AIの違い さて、ガーディアンエージェントが必要になる理由として、生成AI(以後AIと呼びます)の急速な進化があげられます。従来のAIとエージェント型AIの違いを思い出

By Qualiteg コンサルティング
LLM推論基盤プロビジョニング講座 第4回 推論エンジンの選定

LLM推論基盤プロビジョニング講座 第4回 推論エンジンの選定

こんにちは!前回までの講座では、LLMサービス構築に必要なリクエスト数の見積もりや、使用モデルの推論時消費メモリ計算について詳しく解説してきました。今回は7ステッププロセスの4番目、「推論エンジンの選定」について詳しく掘り下げていきます。 推論エンジンとは何か 推論エンジンとは、GPU上でLLMモデルの推論計算(テキスト生成)を効率的に行うために設計された専用のソフトウェアプログラムです。一般的なディープラーニングフレームワーク(PyTorch、TensorFlowなど)でも推論は可能ですが、実運用環境では専用の推論エンジンを使用することで、大幅なパフォーマンス向上とリソース効率化が期待できます。 推論エンジンは単なる実行環境ではなく、様々な最適化技術を実装しています。特定のモデルアーキテクチャに特化した最適化機能を実装したものや、推論速度の高速化に特化したもの、前回解説したKVキャッシュのメモリ効率化機能を備えたものなど、それぞれ特徴が異なります。そのため、自社で採用したLLMモデルや運用環境、要件に合致した推論エンジンを選定することが重要です。 推論エンジン選定のアプロ

By Qualiteg コンサルティング
発話音声からリアルなリップシンクを生成する技術 第1回:音素とwav2vec

発話音声からリアルなリップシンクを生成する技術 第1回:音素とwav2vec

こんにちは! 今日は当社のMotionVox でも実際に使っている「リップシンク」技術について総合的に解説してみたいとおもいます。 音声に合わせて自然な口の動きを生成するリップシンク技術は、AIアバターや3Dアニメーション制作においても重要な技術です。 本記事では、最新のディープラーニング技術を活用したリップシンク学習の基礎から実装まで、技術的な観点から詳しく解説します。 1. リップシンク学習の基礎概念 1.1 問題設定 リップシンク学習とは、音声データから対応する口の動きを予測する回帰問題ととらえることができます f: 音声特徴量(t) → 口の動きパラメータ(t) この問題のコアは 音韻(音の特徴)と視素(視覚的な口の形)の対応関係を学習する ことにあります。 1.2 音韻-視素マッピングの複雑性 ただし! 人間の発話における音と口の形の関係は、単純な1対1マッピングではないんです。 同じ音でも文脈で変化 「あ」の発音でも: - 「か」の後の「あ」→ 口がやや狭めから開く - 「ん」の後の「あ」→ 口が閉じた状態から大きく開く 調音結合

By Qualiteg 研究部, Qualiteg コンサルティング