@@ -55,8 +55,8 @@ model_id = "tiiuae/falcon-7b"
5555tokenizer = AutoTokenizer.from_pretrained(model_id)
5656
5757model = AutoModelForCausalLM.from_pretrained(
58- model_id,
59- torch_dtype = torch.bfloat16,
58+ model_id,
59+ torch_dtype = torch.bfloat16,
6060 attn_implementation = " flash_attention_2" ,
6161)
6262```
@@ -112,7 +112,7 @@ model_id = "tiiuae/falcon-7b"
112112tokenizer = AutoTokenizer.from_pretrained(model_id)
113113
114114model = AutoModelForCausalLM.from_pretrained(
115- model_id,
115+ model_id,
116116 load_in_8bit = True ,
117117 attn_implementation = " flash_attention_2" ,
118118)
@@ -130,7 +130,7 @@ model_id = "tiiuae/falcon-7b"
130130tokenizer = AutoTokenizer.from_pretrained(model_id)
131131
132132model = AutoModelForCausalLM.from_pretrained(
133- model_id,
133+ model_id,
134134 load_in_4bit = True ,
135135 attn_implementation = " flash_attention_2" ,
136136)
@@ -149,7 +149,7 @@ model_id = "tiiuae/falcon-7b"
149149tokenizer = AutoTokenizer.from_pretrained(model_id)
150150
151151model = AutoModelForCausalLM.from_pretrained(
152- model_id,
152+ model_id,
153153 load_in_4bit = True ,
154154 attn_implementation = " flash_attention_2" ,
155155)
@@ -173,7 +173,7 @@ BetterTransformerは、テキスト、画像、およびオーディオモデル
173173<Tip >
174174
175175Flash Attentionは、fp16またはbf16のdtypeを使用するモデルにのみ使用できます。BetterTransformerを使用する前に、モデルを適切なdtypeにキャストしてください。
176-
176+
177177</Tip >
178178
179179### Encoder models
@@ -214,11 +214,12 @@ model.to_bettertransformer()
214214# Use it for training or inference
215215```
216216
217- SDPAは、ハードウェアや問題のサイズに応じて[ Flash Attention] ( https://arxiv.org/abs/2205.14135 ) カーネルを使用することもできます。Flash Attentionを有効にするか、特定の設定(ハードウェア、問題サイズ)で使用可能かどうかを確認するには、[ ` torch.backends.cuda.sdp_kernel ` ] ( https://pytorch.org/docs/master/backends.html# torch.backends.cuda.sdp_kernel ) をコンテキストマネージャとして使用します。
217+ SDPAは、ハードウェアや問題のサイズに応じて[ Flash Attention] ( https://arxiv.org/abs/2205.14135 ) カーネルを使用することもできます。Flash Attentionを有効にするか、特定の設定(ハードウェア、問題サイズ)で使用可能かどうかを確認するには、[ ` torch.nn.attention.sdpa_kernel ` ] ( https://pytorch.org/docs/stable/generated/ torch.nn.attention.sdpa_kernel.html ) をコンテキストマネージャとして使用します。
218218
219219
220220``` diff
221221import torch
222+ + from torch.nn.attention import SDPBackend, sdpa_kernel
222223from transformers import AutoModelForCausalLM, AutoTokenizer
223224
224225tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
@@ -229,7 +230,7 @@ model.to_bettertransformer()
229230input_text = "Hello my dog is cute and"
230231inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
231232
232- + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False ):
233+ + with sdpa_kernel(SDPBackend.FLASH_ATTENTION ):
233234 outputs = model.generate(**inputs)
234235
235236print(tokenizer.decode(outputs[0], skip_special_tokens=True))
@@ -421,6 +422,7 @@ In this example, the first GPU will use 1GB of memory and the second 2GB.
421422
422423``` py
423424import torch
425+ from torch.nn.attention import SDPBackend, sdpa_kernel
424426from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
425427
426428quantization_config = BitsAndBytesConfig(
@@ -434,7 +436,7 @@ model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", quantization_c
434436input_text = " Hello my dog is cute and"
435437inputs = tokenizer(input_text, return_tensors = " pt" ).to(" cuda" )
436438
437- with torch.backends.cuda.sdp_kernel( enable_flash = True , enable_math = False , enable_mem_efficient = False ):
439+ with sdpa_kernel(SDPBackend. FLASH_ATTENTION ):
438440 outputs = model.generate(** inputs)
439441
440442print (tokenizer.decode(outputs[0 ], skip_special_tokens = True ))
0 commit comments