TGI and turn off Flash Attention v2
We are running our own TGI container and trying to boot Mistral Instruct. It’s dieing trying to utilize Flash Attention 2. I know this is because I am using a T4 GPU, but for the life of me I can’t figure out how to tell TGI not to use  Flash Attention 2.Somehow, when we deploy it through HuggingFace on an AWS T4, it knows. Looking at the logs for HF deployment I see:2024-08-01T01:48:41.507+00:00 {“timestamp”:“2024-08-01T01:48:41.507739Z”,“level”:“INFO”,“fields”:{“message”:“Successfully downloaded weights.”},“target”:“text_generation_launcher”,“span”:{“name”:“download”},“spans”:[{“name”:“download”}]}2024-08-01T01:48:41.507+00:00 {“timestamp”:“2024-08-01T01:48:41.507904Z”,“level”:“INFO”,“fields”:{“message”:“Starting shard”},“target”:“text_generation_launcher”,“span”:{“rank”:0,“name”:“shard-manager”},“spans”:[{“rank”:0,“name”:“shard-manager”}]}2024-08-01T01:48:41.507+00:00 {“timestamp”:“2024-08-01T01:48:41.507914Z”,“level”:“INFO”,“fields”:{“message”:“Starting shard”},“target”:“text_generation_launcher”,“span”:{“rank”:1,“name”:“shard-manager”},“spans”:[{“rank”:1,“name”:“shard-manager”}]}2024-08-01T01:48:41.507+00:00 {“timestamp”:“2024-08-01T01:48:41.507952Z”,“level”:“INFO”,“fields”:{“message”:“Starting shard”},“target”:“text_generation_launcher”,“span”:{“rank”:2,“name”:“shard-manager”},“spans”:[{“rank”:2,“name”:“shard-manager”}]}2024-08-01T01:48:41.508+00:00 {“timestamp”:“2024-08-01T01:48:41.507977Z”,“level”:“INFO”,“fields”:{“message”:“Starting shard”},“target”:“text_generation_launcher”,“span”:{“rank”:3,“name”:“shard-manager”},“spans”:[{“rank”:3,“name”:“shard-manager”}]}2024-08-01T01:48:46.187+00:00 {“timestamp”:“2024-08-01T01:48:46.187088Z”,“level”:“WARN”,“fields”:{“message”:“Unable to use Flash Attention V2: GPU with CUDA capability 7 5 is not supported for Flash Attention V2\n”},“target”:“text_generation_launcher”}2024-08-01T01:48:46.204+00:00 {“timestamp”:“2024-08-01T01:48:46.204816Z”,“level”:“WARN”,“fields”:{“message”:“Unable to use Flash Attention V2: GPU with CUDA capability 7 5 is not supported for Flash Attention V2\n”},“target”:“text_generation_launcher”}2024-08-01T01:48:46.210+00:00 {“timestamp”:“2024-08-01T01:48:46.210281Z”,“level”:“WARN”,“fields”:{“message”:“Unable to use Flash Attention V2: GPU with CUDA capability 7 5 is not supported for Flash Attention V2\n”},“target”:“text_generation_launcher”}2024-08-01T01:48:46.210+00:00 {“timestamp”:“2024-08-01T01:48:46.210527Z”,“level”:“WARN”,“fields”:{“message”:“Unable to use Flash Attention V2: GPU with CUDA capability 7 5 is not supported for Flash Attention V2\n”},“target”:“text_generation_launcher”}2024-08-01T01:48:51.516+00:00 {“timestamp”:“2024-08-01T01:48:51.516157Z”,“level”:“INFO”,“fields”:{“message”:“Waiting for shard to be ready…”},“target”:“text_generation_launcher”,“span”:{“rank”:3,“name”:“shard-manager”},“spans”:[{“rank”:3,“name”:“shard-manager”}]}2024-08-01T01:48:51.516+00:00 {“timestamp”:“2024-08-01T01:48:51.516157Z”,“level”:“INFO”,“fields”:{“message”:“Waiting for shard to be ready…”},“target”:“text_generation_launcher”,“span”:{“rank”:2,“name”:“shard-manager”},“spans”:[{“rank”:2,“name”:“shard-manager”}]}2024-08-01T01:48:51.516+00:00 {“timestamp”:“2024-08-01T01:48:51.516157Z”,“level”:“INFO”,“fields”:{“message”:“Waiting for shard to be ready…”},“target”:“text_generation_launcher”,“span”:{“rank”:0,“name”:“shard-manager”},“spans”:[{“rank”:0,“name”:“shard-manager”}]}2024-08-01T01:48:51.516+00:00 {“timestamp”:“2024-08-01T01:48:51.516157Z”,“level”:“INFO”,“fields”:{“message”:“Waiting for shard to be ready…”},“target”:“text_generation_launcher”,“span”:{“rank”:1,“name”:“shard-manager”},“spans”:[{“rank”:1,“name”:“shard-manager”}]}Notice it correctly says it can’t use Flash Attention V2.When we run this on our own T4 in our AWS, we don’t see this output.  Instead, it bombs out trying to use Flash Attention v2 with:Traceback (most recent call last):File “/opt/conda/bin/text-generation-server”, line 8, insys.exit(app())File “/opt/conda/lib/python3.10/site-packages/typer/main.py”, line 311, incallreturn get_command(self)(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/click/core.py”, line 1157, incallreturn self.main(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/typer/core.py”, line 778, in mainreturn _main(File “/opt/conda/lib/python3.10/site-packages/typer/core.py”, line 216, in _mainrv = self.invoke(ctx)File “/opt/conda/lib/python3.10/site-packages/click/core.py”, line 1688, in invokereturn _process_result(sub_ctx.command.invoke(sub_ctx))File “/opt/conda/lib/python3.10/site-packages/click/core.py”, line 1434, in invokereturn ctx.invoke(self.callback, **ctx.params)File “/opt/conda/lib/python3.10/site-packages/click/core.py”, line 783, in invokereturn __callback(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/typer/main.py”, line 683, in wrapperreturn callback(**use_params)  # type: ignoreFile “/opt/conda/lib/python3.10/site-packages/text_generation_server/cli.py”, line 109, in serveserver.serve(File “/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py”, line 274, in serveasyncio.run(File “/opt/conda/lib/python3.10/asyncio/runners.py”, line 44, in runreturn loop.run_until_complete(main)File “/opt/conda/lib/python3.10/asyncio/base_events.py”, line 636, in run_until_completeself.run_forever()File “/opt/conda/lib/python3.10/asyncio/base_events.py”, line 603, in run_foreverself._run_once()File “/opt/conda/lib/python3.10/asyncio/base_events.py”, line 1909, in _run_oncehandle._run()File “/opt/conda/lib/python3.10/asyncio/events.py”, line 80, in _runself._context.run(self._callback, *self._args)File “/opt/conda/lib/python3.10/site-packages/grpc_interceptor/server.py”, line 165, in invoke_intercept_methodreturn await self.intercept(File “/opt/conda/lib/python3.10/site-packages/text_generation_server/interceptor.py”, line 21, in interceptreturn await responseFile “/opt/conda/lib/python3.10/site-packages/opentelemetry/instrumentation/grpc/_aio_server.py”, line 120, in _unary_interceptorraise errorFile “/opt/conda/lib/python3.10/site-packages/opentelemetry/instrumentation/grpc/_aio_server.py”, line 111, in _unary_interceptorreturn await behavior(request_or_iterator, context)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py”, line 123, in Warmupmax_supported_total_tokens = self.model.warmup(batch)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_causal_lm.py”, line 1099, in warmup_, batch, _ = self.generate_token(batch)File “/opt/conda/lib/python3.10/contextlib.py”, line 79, in innerreturn func(*args, **kwds)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_causal_lm.py”, line 1374, in generate_tokenout, speculative_logits = self.forward(batch, adapter_data)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/flash_causal_lm.py”, line 1299, in forwardlogits, speculative_logits = self.model.forward(File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/flash_mistral_modeling.py”, line 516, in forwardhidden_states = self.model(File “/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1553, in _wrapped_call_implreturn self._call_impl(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1562, in _call_implreturn forward_call(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/flash_mistral_modeling.py”, line 440, in forwardhidden_states, residual = layer(File “/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1553, in _wrapped_call_implreturn self._call_impl(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1562, in _call_implreturn forward_call(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/flash_mistral_modeling.py”, line 365, in forwardattn_output = self.self_attn(File “/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1553, in _wrapped_call_implreturn self._call_impl(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1562, in _call_implreturn forward_call(*args, **kwargs)File “/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/flash_mistral_modeling.py”, line 218, in forwardattn_output = attention(File “/opt/conda/lib/python3.10/site-packages/text_generation_server/layers/attention/cuda.py”, line 221, in attentionreturn flash_attn_2_cuda.varlen_fwd(RuntimeError: FlashAttention only supports Ampere GPUs or newer.Does anyone know how I am supposed to tell TGI not to use Flash Attention v2, or how HF TGI deployment is getting just a warning and moving on?
2024-08-06T05:01:59Z
I am also desperate for this. Hello fellow (likely) govcloud user.
2024-08-13T21:43:17Z
FYI, they fixed this in the latest version of the image:GitHubBuild software better, togetherGitHub is where people build software. More than 100 million people use GitHub to discover, fork, and contribute to over 420 million projects.
2024-08-22T21:42:22Z
Hi@rkettrey, just to clarify your fix- the model you were trying to use actually didn’t need FlashAttention 2? And once you used this new version of the TGI, the SageMaker deployment script from the HuggingFace website no longer tried to utilize FlashAttention 2?
2024-08-23T11:54:36Z
The model we were trying to use was using FlashAttention 2, but our chip we were booting on (T4) is not supported for FA2. Therefore, we were expecting the boot process to NOT try to use FA2 and it was trying to.The updated TGI now properly detects and does not try to use FA2.