ColaboratoryでFlashAttentionをインストールするとめちゃめちゃ時間がかかる
kun432とあるノートブック試したくて、そのノートブック内でFlashAttentionがインストールされている。
pip install flash-attn --no-build-isolation
ただこれめちゃめちゃ時間がかかる。ノートブックではT4を使用しているようだが、試した限り少なくとも1時間では終わらなかった。L4でも試したが同じ。（ちなみにFlashAttention2はT4のような新しいアーキテクチャには対応していないとレポジトリには書いてあるようなのだが、はて？）
自分は待てなかったので止めた。
でこういうポストを見かけた
どうやらFlashAttention2のレポジトリで事前ビルド済みのものが提供されているみたい
さらにpip install flash-attn --no-build-isolationは、この事前ビルド済みのものをダウンロードしてきてインストールするのだが、この時、環境内のPython/CUDA/PyTorchなどのバージョンをチェックして選択している模様。ただ、Colaboratoryのバージョン似合ったものが用意されていないので、結果ビルドすることになり時間がかかっていると思われる。
その辺の経緯が以下にまとまっている
で、自分もColab環境で各種バージョンを確認してやってみたのだが、ピンポイントでマッチするものはなく、近しいバージョンと思われるものをインストールしてみたのだが、やはりうまくいかなかった。
kun432こういうポストを見かけた
レポジトリはここ。GitHub Actionsで複数のパターンでビルドされている。
リリースに事前ビルド済みパッケージがある
ただColaboratoryだとPython3.10向けが必要になるが、上記のレポジトリではPython-3.10は定義されていない。フォークさせていただき、バージョン定義を修正して、ビルドさせた（めちゃめちゃ時間がかかる）
ColabでQwen-VL2-2Bを動かしてみた。ランタイムはL4。
!pip install accelerate qwen-vl-utils torchvision torch datasets huggingface_hub[hf_transfer] polars ninja
!pip install git+https://github.com/huggingface/transformers.git
ColabだとPython-3.10/CUDA-12.1/PyTorch-2.5.0っぽいので、それにあったパッケージをインストール。
!wget https://github.com/kun432/flash-attention-prebuild-wheels/releases/download/v0.0.0-test/flash_attn-2.6.3+cu121torch2.5-cp310-cp310-linux_x86_64.whl
!pip install --no-dependencies --upgrade flash_attn-2.6.3+cu121torch2.5-cp310-cp310-linux_x86_64.whl
ロード
import flash_attn
flash_attn.__version__
ここはOK
2.6.3
お試しでQwen2-VL-2B-Instructを動かしてみる。"flash_attention_2"を有効化。
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
model = Qwen2VLForConditionalGeneration.from_pretrained(
"Qwen/Qwen2-VL-2B-Instruct",
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
device_map="auto",
)
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
messages = [
{
"role": "user",
"content": [
{
"type": "image",
"image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
},
{"type": "text", "text": "この画像について説明して。"},
],
}
]
text = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
)
inputs = inputs.to("cuda")
generated_ids = model.generate(**inputs, max_new_tokens=200)
generated_ids_trimmed = [
out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
output_text
['この画像は、ビーチで犬と人との間の楽しい瞬間を描いたものです。ビーチの砂浜に座っている女性は、犬と一緒に手を合わせ、笑顔で写っています。犬は、背中を向けてカメラに向かっており、その上に手を伸ばしています。ビーチの背景には、海が広がり、日が沈むような夕日が照らしています。このシーンは、人間と動物の間の親しみや愛を表現しています。']
行けてそう。無駄な待ち時間とColaboratoryのコンピュータユニットを節約できてめっちゃ助かる。元レポの方に感謝したい。
kun432T4だと動かない（FlashAttentionのレポジトリにも新しすぎるアーキテクチャにはまだ対応できていないので、1.xを使えとある）ので、1.xのパッケージをビルドすればいけそう（試していない）