96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
# coding=utf-8
|
|
# Copyright 2023 The HuggingFace Inc. team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Utility that checks the supports of 3rd party libraries are listed in the documentation file. Currently, this includes:
|
|
- flash attention support
|
|
- SDPA support
|
|
|
|
Use from the root of the repo with (as used in `make repo-consistency`):
|
|
|
|
```bash
|
|
python utils/check_support_list.py
|
|
```
|
|
|
|
It has no auto-fix mode.
|
|
"""
|
|
import os
|
|
from glob import glob
|
|
|
|
|
|
# All paths are set with the intent you should run this script from the root of the repo with the command
|
|
# python utils/check_doctest_list.py
|
|
REPO_PATH = "."
|
|
|
|
|
|
def check_flash_support_list():
|
|
with open(os.path.join(REPO_PATH, "docs/source/en/perf_infer_gpu_one.md"), "r") as f:
|
|
doctext = f.read()
|
|
|
|
doctext = doctext.split("FlashAttention-2 is currently supported for the following architectures:")[1]
|
|
doctext = doctext.split("You can request to add FlashAttention-2 support")[0]
|
|
|
|
patterns = glob(os.path.join(REPO_PATH, "src/transformers/models/**/modeling_*.py"))
|
|
patterns_tf = glob(os.path.join(REPO_PATH, "src/transformers/models/**/modeling_tf_*.py"))
|
|
patterns_flax = glob(os.path.join(REPO_PATH, "src/transformers/models/**/modeling_flax_*.py"))
|
|
patterns = list(set(patterns) - set(patterns_tf) - set(patterns_flax))
|
|
archs_supporting_fa2 = []
|
|
for filename in patterns:
|
|
with open(filename, "r") as f:
|
|
text = f.read()
|
|
|
|
if "_supports_flash_attn_2 = True" in text:
|
|
model_name = os.path.basename(filename).replace(".py", "").replace("modeling_", "")
|
|
archs_supporting_fa2.append(model_name)
|
|
|
|
for arch in archs_supporting_fa2:
|
|
if arch not in doctext:
|
|
raise ValueError(
|
|
f"{arch} should be in listed in the flash attention documentation but is not. Please update the documentation."
|
|
)
|
|
|
|
|
|
def check_sdpa_support_list():
|
|
with open(os.path.join(REPO_PATH, "docs/source/en/perf_infer_gpu_one.md"), "r") as f:
|
|
doctext = f.read()
|
|
|
|
doctext = doctext.split(
|
|
"For now, Transformers supports SDPA inference and training for the following architectures:"
|
|
)[1]
|
|
doctext = doctext.split("Note that FlashAttention can only be used for models using the")[0]
|
|
|
|
patterns = glob(os.path.join(REPO_PATH, "src/transformers/models/**/modeling_*.py"))
|
|
patterns_tf = glob(os.path.join(REPO_PATH, "src/transformers/models/**/modeling_tf_*.py"))
|
|
patterns_flax = glob(os.path.join(REPO_PATH, "src/transformers/models/**/modeling_flax_*.py"))
|
|
patterns = list(set(patterns) - set(patterns_tf) - set(patterns_flax))
|
|
archs_supporting_sdpa = []
|
|
for filename in patterns:
|
|
with open(filename, "r") as f:
|
|
text = f.read()
|
|
|
|
if "_supports_sdpa = True" in text:
|
|
model_name = os.path.basename(filename).replace(".py", "").replace("modeling_", "")
|
|
archs_supporting_sdpa.append(model_name)
|
|
|
|
for arch in archs_supporting_sdpa:
|
|
if arch not in doctext:
|
|
raise ValueError(
|
|
f"{arch} should be in listed in the SDPA documentation but is not. Please update the documentation."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
check_flash_support_list()
|
|
check_sdpa_support_list()
|