Posts Tags Categories

Contents

WSL (Ubuntu) + uv 環境で Attention Kernel (flash-attn 等) をインストールする際の注意点と手順

   539 words   3 minutes 

最近、Windows Subsystem for Linux (WSL) を使って Linux 環境で開発を行うことが増えてきました。特に GPU を活用した機械学習の分野では、WSL 2 が NVIDIA CUDA をサポートしたことで、Windows 上でも Linux ベースの開発環境を構築しやすくなりましたね。

さらに、Python のパッケージ管理ツールとして、Rust 製の超高速ツール uv が登場し、pip や venv の代替として注目を集めています。

しかし、この便利な WSL + uv 環境でも、特定のライブラリ、特に GPU アクセラレーションに深く依存するパッケージ (例えば flash-attn のような Attention 実装) のインストールでつまずくことがあります。

この記事では、私が実際に WSL (Ubuntu) + uv 環境で flash-attn をインストールしようとして遭遇したエラーとその解決策、そして関連する Attention Kernel (xformers, sage-attention) のインストール方法、最後に uv コマンドの使い分けについて、備忘録も兼ねてまとめます。

目標

WSL (Ubuntu) 環境で xformers, flash-attn, sage-attention を uv を使ってインストールする。

特に flash-attn のインストール時に発生しやすい問題を解決する。

環境

OS: Windows 11 + WSL 2 (Ubuntu 22.04 LTS)

GPU: NVIDIA GeForce RTX シリーズ

Windows側 NVIDIA Driver: 最新版 (Game Ready または Studio)

パッケージマネージャー: uv (0.1.x 以降)

Python環境: uv で作成・管理される仮想環境 (.venv)

Attention Kernel のインストール手順

ソフトウェアによっては、計算効率やメモリ効率を改善するために、標準の PyTorch Attention 以外に xformers, flash-attn, sage-attention といった最適化された Attention Kernel を利用できる場合があります。

ここでは、uv を使ってこれらをインストールする方法を見ていきましょう。

  1. Sage Attention
    比較的簡単にインストールできます。バージョン指定が推奨されている場合もあります。
uv pip install sageattention
  1. Xformers
uv pip install ninja # 先に入れておくのがおすすめ
uv pip install xformers
  1. Flash Attention (問題発生ポイント)
    flash-attn は非常に高速ですが、インストール時に CUDA Toolkit を使ったコンパイルが必要となり、ここで問題が発生しやすいです。

まずは通常通りインストールを試みます。

uv pip install flash-attn

問題 1: ModuleNotFoundError: No module named ’torch'

  × Failed to build `flash-attn==2.7.4.post1`
  (...)
  ModuleNotFoundError: No module named 'torch'

  hint: This error likely indicates that `flash-attn@2.7.4.post1` depends on `torch`,
  but doesn't declare it as a build dependency. (...)
  Otherwise, `uv pip install torch` into the environment and re-run with `--no-build-isolation`.

これは、flash-attn のビルドプロセスが torch を必要とするのに、ビルド依存関係として宣言されていないため、uv が作成する隔離されたビルド環境に torch が含まれずエラーになる、という問題です。

解決策 1: ヒントにある通り、–no-build-isolation オプションを付けて再試行します。これは、隔離環境を使わず、現在アクティブな Python 環境のパッケージを使ってビルドするオプションです。 この際、uv add ではなく uv pip install を使うのがポイントです (後述)。

# まず PyTorch (CUDA版) が環境にあることを確認
# (なければ uv pip install torch --extra-index-url ... でインストール)

# --no-build-isolation を付けて再実行
uv pip install flash-attn --no-build-isolation

すると、次の問題が発生しました。

問題 2: ModuleNotFoundError: No module named ‘setuptools’

  × Failed to build `flash-attn==2.7.4.post1`
  (...)
  ModuleNotFoundError: No module named 'setuptools'

hint: This usually indicates a problem with the package or the build environment.
–no-build-isolation を指定したことで現在の環境でビルドしようとしましたが、今度はその環境にビルドツールである setuptools がなかったようです。

解決策 2: 現在の環境にビルドツールをインストールします。

uv pip install --upgrade setuptools wheel

そして、再度インストールを試みます。

uv pip install flash-attn --no-build-isolation

これでうまくいくかと思いきや、最後の関門が待っていました。

問題 3: OSError: CUDA_HOME environment variable is not set / nvcc was not found

  × Failed to build `flash-attn==2.7.4.post1`
  (...)
  <string>:106: UserWarning: flash_attn was requested, but nvcc was not found. (...)
  (...)
  OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.

これは、flash-attn の CUDA コードをコンパイルするために必要な NVIDIA CUDA Toolkit (特にコンパイラ nvcc) が WSL 環境内に見つからない、というエラーです。WSL で CUDA を使うには、Windows 側のドライバだけでなく、WSL 内にも CUDA Toolkit をインストール し、そのパスを環境変数で設定する必要があります。

解決策 3: NVIDIA の公式ドキュメントに従い、WSL 内に CUDA Toolkit をインストールし、環境変数を設定します。

CUDA Toolkit のインストール (WSL 内):

NVIDIA CUDA Toolkit ダウンロードページ へアクセスします。

Linux -> x86_64 -> WSL-Ubuntu -> (バージョン選択) -> deb (network または local) を選択します。

表示される手順に従って、リポジトリ設定とインストールを行います。 重要: sudo apt-get install cuda-toolkit-XX-Y のように、ツールキットのみを指定します。cuda や cuda-drivers をインストールすると、不要な Linux ドライバがインストールされてしまうため避けてください。

(例: CUDA 12.5 の場合)

# PIN ファイル設定 (サイトの指示に従う)
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
# リポジトリ DEB 設定 (サイトの指示に従う、バージョン注意)
wget https://developer.download.nvidia.com/compute/cuda/12.5.0/local_installers/cuda-repo-wsl-ubuntu-12-5-local_12.5.0-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-12-5-local_12.5.0-1_amd64.deb
# GPGキーコピー
sudo cp /var/cuda-repo-wsl-ubuntu-12-5-local/cuda-*-keyring.gpg /usr/share/keyrings/
# 更新
sudo apt-get update
# ツールキットのみインストール！
sudo apt-get -y install cuda-toolkit-12-5

環境変数の設定:
~/.bashrc や ~/.zshrc に以下を追加し、source ~/.bashrc などで反映させます。パスのバージョン (cuda-12.5) はインストールしたものに合わせてください。

export CUDA_HOME=/usr/local/cuda-12.5
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

確認:
ターミナルで nvcc –version を実行し、CUDA のバージョンが表示されれば OK です。

❯ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Tue_Feb_27_16:19:40_PST_2024
Cuda compilation tools, release 12.4, V12.4.99
Build cuda_12.4.r12.4/compiler.33961267_0
flash-attn インストール再々試行:

今度こそ、インストールが成功するはずです。念のため ninja も入れておきましょう。

uv pip install ninja
uv pip install flash-attn --no-build-isolation

コンパイルには少し時間がかかりますが、エラーが出なければ成功です！

まとめ

WSL (Ubuntu) + uv 環境は非常に強力ですが、GPU 依存ライブラリのインストール、特にコンパイルが必要なもの (flash-attn など) では、いくつかの落とし穴があります。

  • ビルド隔離 (–no-build-isolation で回避)

  • ビルドツールの不足 (setuptools, wheel のインストール)

  • WSL 環境内への CUDA Toolkit の欠如 (cuda-toolkit-X-Y パッケージのインストールと環境変数設定)

これらの問題を一つずつ解決し、uv コマンド (add と pip install) を適切に使い分けることで、WSL 環境でも快適に最新の機械学習ライブラリを活用した開発を進めることができます。

この記事が、同様の問題に直面している方の一助となれば幸いです。

Updated on 2025-04-19
Back | Home