safetensors: AIモデルの新しい標準ファイルフォーマットの仕組みと利点

機械学習モデルの保存と読み込みを安全かつ高速に行うsafetensorsの仕組みと利点について解説します。セキュリティ監査を通過し、主要なAI企業が採用を進める背景も紹介します。

機械学習モデルの大規模化と共に、モデルの保存と読み込みの効率性とセキュリティの重要性が増しています。safetensorsは、これらの課題に対応するためにHugging Faceが開発した新しいファイルフォーマットです。2024年初頭にTrail of Bitsによるセキュリティ監査を通過し、EleutherAIStability AIなどの主要なAI企業が標準フォーマットとしての採用を進めています。

本記事では、safetensorsの仕組みと利点、実際の使用方法について解説します。

safetensorsとは

safetensorsは、深層学習モデルで使用されるテンソル(多次元配列データ)を効率的に保存・読み込むためのファイルフォーマットです。PyTorchが従来使用していたpickleフォーマットの安全性の課題を解決し、さらに読み込み速度を大幅に改善しています。

Rust言語で実装されており、主要な機械学習フレームワーク(PyTorch、TensorFlow、JAX、PaddlePaddle、NumPy)をサポートしています。特に最近では、Stable Diffusion XLなどの大規模モデルで標準的に使用されるようになっています。

safetensorsの主な利点

safetensorsの採用が進んでいる背景には、以下のような重要な利点があります:

  1. セキュリティ:外部セキュリティ監査により安全性が確認されています。モデルファイルに悪意のあるコードが含まれる可能性を排除し、安全な実行を保証します。

  2. 高速な読み込み:CPU上での読み込みが従来のpickleフォーマットと比べて約100倍高速です。これは効率的なシリアライゼーションと圧縮アルゴリズムの採用によるものです。

  3. 遅延読み込み(Lazy Loading):大規模言語モデル(LLM)などで重要な、必要な部分のみを効率的に読み込む機能をサポートしています。

  4. フレームワーク非依存:同じファイルをPyTorchやTensorFlowなど、異なるフレームワークで読み込むことができます。

セキュリティの重要性

従来のPyTorchモデルで使用されるpickleフォーマットには重大なセキュリティ上の課題がありました。pickleを使用すると、モデルファイルに任意のコードを埋め込むことが可能で、ユーザーの知らないうちにマルウェアを実行される危険性がありました。

2024年初頭、Trail of Bitsによる外部セキュリティ監査により、safetensorsには重大なセキュリティ上の欠陥がないことが確認されました。主な監査結果は以下の通りです:

  • 任意のコード実行につながる重大な脆弱性は発見されず
  • 仕様の不明確な部分が特定され修正
  • polyglotファイルの問題が修正
  • テストスイートの大幅な改善

また、Rust言語の持つメモリ安全性の保証も、追加のセキュリティ層として機能しています。これらの特徴により、safetensorsはオープンソースのAIモデル共有において、重要な役割を果たすことが期待されています。

safetensorsファイルの構造

safetensorsファイルは、バイナリファイルフォーマットでありながら、シンプルで効率的な構造を持っています。ファイルは大きく分けて「ヘッダー部」と「データ部」の2つの領域で構成されています。ヘッダー部はJSONで記述され、各テンソルのメタデータ(データ型、形状、位置など)を保持します。データ部は実際のテンソル値をバイナリ形式で効率的に格納します。

safetensorsファイルの特徴的な設計ポイントは以下の通りです:

  1. ヘッダーサイズの明示:ファイルの先頭8バイトにヘッダーの正確なサイズを格納することで、効率的なファイル解析を可能にしています。

  2. JSON形式のヘッダー:人間が読める形式でメタデータを保存し、簡単な検証と拡張性を提供します。ヘッダーには各テンソルの名前、データ型、形状、およびデータバッファ内の位置情報が含まれます。

  3. メタデータのサポート:__metadata__キーを使用して、文字列ベースの任意のメタデータを保存できます。

  4. 効率的なデータレイアウト:

    • リトルエンディアン形式でデータを格納
    • 行優先(C-order)でテンソルを配置
    • パディングや余分なデータを含まない連続したバイト配列
  5. セキュリティ対策:

    • 重複キーの禁止
    • ヘッダーサイズの制限(100MB以下)
    • データ領域の重複禁止(メモリ攻撃の防止)

safetensorsの使用方法

safetensorsは、シンプルなAPIを通じて簡単に使用することができます。各主要フレームワークでの基本的な使用方法から、より高度な使用例まで見ていきましょう。

基本的な使用例

basic_usage.py
1import torch
2from safetensors.torch import save_file, load_file
3
4# テンソルの作成と保存
5tensors = {
6    "weight": torch.randn(10, 10),
7    "bias": torch.zeros(10)
8}
9save_file(tensors, "model.safetensors")
10
11# テンソルのロード
12loaded = load_file("model.safetensors")
13print(loaded["weight"].shape)  # torch.Size([10, 10])

Transformersライブラリでの使用

transformers_example.py
1from transformers import AutoModel
2
3# モデルの読み込み - safetensorsファイルが自動的に使用される
4model = AutoModel.from_pretrained("bert-base-uncased")
5print(model.config.model_type)  # 'bert'
6
7# モデルの保存 - safetensorsフォーマットで保存
8model.save_pretrained("my-model", safe_serialization=True)

遅延読み込みの活用

lazy_loading.py
1from safetensors import safe_open
2
3# ファイルを開いてテンソル情報を確認
4with safe_open("large_model.safetensors", framework="pt") as f:
5    # 利用可能なテンソルの確認
6    tensor_names = list(f.keys())
7    print("Available tensors:", tensor_names)
8    
9    # 特定のレイヤーのみを読み込み
10    if "encoder.layer.0.attention.self.query.weight" in tensor_names:
11        query_weight = f.get_tensor("encoder.layer.0.attention.self.query.weight")
12        print("Query weight shape:", query_weight.shape)

その他のフレームワークでの使用

safetensorsは、主要な機械学習フレームワークをすべてサポートしています。フレームワーク固有のモジュールを使用することで、それぞれのフレームワークに最適化された方法でテンソルを扱えます:

  • TensorFlow: safetensors.tensorflowモジュール
  • JAX: safetensors.jaxモジュール
  • NumPy: safetensors.numpyモジュール
  • PaddlePaddle: safetensors.paddleモジュール

フレームワーク間での相互運用性も確保されており、あるフレームワークで保存したモデルを別のフレームワークで読み込むことも可能です。詳しい使用方法は公式ドキュメントを参照してください。

PyTorchを使用した基本的な例から見ていきます。テンソルの保存と読み込みは、直感的なAPIを通じて簡単に行えます:

basic_usage.py
1import torch
2from safetensors.torch import save_file, load_file
3
4# テンソルの作成と保存
5tensors = {
6    "weight": torch.randn(10, 10),
7    "bias": torch.zeros(10)
8}
9save_file(tensors, "model.safetensors")
10
11# テンソルのロード
12loaded = load_file("model.safetensors")
13print(loaded["weight"].shape)  # torch.Size([10, 10])

Hugging FaceTransformersライブラリでは、safetensorsが標準のファイルフォーマットとして採用されつつあります。ライブラリをインストールするだけで、自動的にsafetensorsファイルが優先的に使用されます:

transformers_example.py
1from transformers import AutoModel
2
3# モデルの読み込み - safetensorsファイルが自動的に使用される
4model = AutoModel.from_pretrained("bert-base-uncased")
5print(model.config.model_type)  # 'bert'
6
7# モデルの保存 - safetensorsフォーマットで保存
8model.save_pretrained("my-model", safe_serialization=True)

大規模モデルを扱う場合、safetensorsの遅延読み込み機能を活用することで、メモリ使用量を最適化できます。必要な部分のみを効率的に読み込むことが可能です:

lazy_loading.py
1from safetensors import safe_open
2
3# ファイルを開いてテンソル情報を確認
4with safe_open("large_model.safetensors", framework="pt") as f:
5    # 利用可能なテンソルの確認
6    tensor_names = list(f.keys())
7    print("Available tensors:", tensor_names)
8    
9    # 特定のレイヤーのみを読み込み
10    if "encoder.layer.0.attention.self.query.weight" in tensor_names:
11        query_weight = f.get_tensor("encoder.layer.0.attention.self.query.weight")
12        print("Query weight shape:", query_weight.shape)

safetensorsフォーマットの機械学習モデルは、主に以下のプラットフォームから入手できます。多くの主要なAIモデルが、セキュリティと効率性の観点からsafetensorsフォーマットでの配布を採用しています。

主な入手先

  1. Hugging Face Hub:

  2. Civitai:

  3. ModelScope:

既存モデルの変換

他のフォーマット(.ckpt, .bin等)からsafetensorsへの変換には以下の方法があります:

  1. Hugging Face Convert Space:

    • Webインターフェースで簡単に変換
    • アップロード済みモデルの安全な変換
    • セキュリティスキャン付きで安全
  2. Safe and Stable Ckpt2Safetensors:

    • GUIツールでローカル変換
    • バッチ処理対応
    • Windows/Linux対応
  3. Python API:

    import torch
    from safetensors.torch import save_file
    
    # .ckptや.binファイルを読み込み
    state_dict = torch.load("model.ckpt")["state_dict"]
    
    # safetensorsとして保存
    save_file(state_dict, "model.safetensors")
    

モデルの配置場所

各ツールでの標準的なモデル配置場所は以下の通りです:

  1. Automatic1111 Web UI:

    stable-diffusion-webui/
    ├── models/
    │   ├── Stable-diffusion/  # ベースモデル(.safetensors)
    │   ├── Lora/             # LoRAモデル
    │   └── VAE/              # VAEモデル
    └── embeddings/           # Textual Inversion
    
  2. ComfyUI:

    ComfyUI/
    └── models/
        ├── checkpoints/  # ベースモデル
        ├── loras/        # LoRAモデル
        └── vae/          # VAEモデル
    
  3. InvokeAI:

    invokeai/
    └── models/
        └── sd/          # すべてのモデル
    

いずれのツールでも、設定ファイルでカスタムパスの指定が可能です。また、多くの場合シンボリックリンクの使用もサポートされています。

Stable Diffusionでの実践的な使用

Stable Diffusionでsafetensorsモデルを使用する方法と、主要なツールの特徴を見ていきましょう。特に2023年後半以降、ほとんどのモデルがsafetensorsフォーマットで配布されるようになっています。

主要なStable Diffusionツール

  1. Automatic1111 Web UI

  2. ComfyUI

  3. InvokeAI

  4. Fooocus

    • Midjourneyライクな簡単操作
    • SDXL特化の最適化
    • 最小限の設定で高品質な結果

基本的なワークフロー

効率的なワークフローのステップは以下の通りです:

  1. 準備

  2. モデルの準備

  3. モデルの配置

    • 各ツールの指定ディレクトリにファイルを配置
    • 必要に応じてサブフォルダを作成
    • ファイル名の特殊文字を確認
  4. 動作確認と最適化

    • ツールの起動と初期テスト生成
    • VAEの設定確認
    • 必要に応じてxFormersの有効化

実践的なTips

より効率的な使用のためのTipsをご紹介します:

  1. メモリ管理

  2. パフォーマンス最適化

  3. トラブルシューティング

    • TypeError: __init__() got an unexpected keyword argument 'ensure_ascii' → safetensorsのバージョンを確認・更新
    • RuntimeError: Couldn't load safetensors → ファイル名の特殊文字を確認
    • VAE関連のエラー → 対応するVAEファイルの配置を確認
  4. バージョン管理

今後の展望

safetensorsは、機械学習モデルの保存と読み込みにおける新しい標準として急速に普及しています。Hugging FaceEleutherAIStability AIなどの主要企業が採用を進めており、特に大規模言語モデル(LLM)のトレーニングや推論において、その重要性は増しています。

主要な展望として以下が挙げられます:

  1. 標準化の進展

  2. 機能の拡張

    • LLMトレーニング特有の課題に対応する新機能
    • 分散学習環境での最適化
    • より効率的なメモリ管理
  3. エコシステムの発展

    • 変換ツールの充実
    • より多くのフレームワークとの互換性
    • サードパーティツールのサポート拡大

参考リソース

より詳しい情報は以下のリソースを参照してください:

公式ドキュメント・リポジトリ:

セキュリティ関連:

フレームワークごとのガイド:

実践的なガイド:

safetensorsに関する最新の情報は、Hugging Face BlogGitHub Discussionsで確認できます。また、実装の詳細についてはAPIリファレンスも参照してください。