# coding=utf-8 # Copyright 2023 The HuggingFace Inc. team. All rights reserved. # # 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. """Tools that fills the neuron cache with common models for the supported tasks.""" import os import sys import time from argparse import ArgumentParser from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Dict, List, Union from optimum.neuron.utils.cache_utils import get_num_neuron_cores # Important to do it before importing the tests. os.environ["RUN_SLOW"] = "1" path_tests = Path(__file__).parent.parent / "tests" sys.path.insert(0, str(path_tests)) if TYPE_CHECKING: from test_examples import ExampleTesterBase # noqa: E402 from test_examples import ( # noqa: E402 ExampleTestMeta, ImageClassificationExampleTester, MultipleChoiceExampleTester, QuestionAnsweringExampleTester, SummarizationExampleTester, TextClassificationExampleTester, TokenClassificationExampleTester, TranslationExampleTester, ) TESTER_CLASSES = { "sequence-classification": TextClassificationExampleTester, "token-classification": TokenClassificationExampleTester, "multiple-choice": MultipleChoiceExampleTester, "question-answering": QuestionAnsweringExampleTester, "summarization": SummarizationExampleTester, "translation": TranslationExampleTester, "image-classification": ImageClassificationExampleTester, } ARCHITECTURES_TO_COMMON_PRETRAINED_WEIGHTS = { "bart": { "facebook/bart-base": { "translation": {"batch_size": 8, "source_sequence_length": 512, "target_sequence_length": 512}, "summarization": {"batch_size": 8, "source_sequence_length": 200, "target_sequence_length": 1024}, }, "facebook/bart-large": { "translation": {"batch_size": 8, "source_sequence_length": 512, "target_sequence_length": 512}, "summarization": {"batch_size": 8, "source_sequence_length": 200, "target_sequence_length": 1024}, }, }, "bert": { "bert-base-uncased": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, "bert-large-uncased": { "default": {"batch_size": 8, "sequence_length": 128}, "token-classification": {"batch_size": 4, "sequence_length": 512}, "multiple-choice": {"batch_size": 4, "sequence_length": 512}, }, }, "camembert": { "camembert-base": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, "camembert/camembert-large": { "default": {"batch_size": 8, "sequence_length": 128}, "token-classification": {"batch_size": 4, "sequence_length": 512}, "multiple-choice": {"batch_size": 4, "sequence_length": 512}, }, }, "distilbert": { "distilbert-base-uncased": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, }, "electra": { "google/electra-small-discriminator": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, "google/electra-base-discriminator": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, "google/electra-large-discriminator": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 4, "sequence_length": 512}, "multiple-choice": {"batch_size": 4, "sequence_length": 512}, }, }, "gpt2": { "gpt2": { "default": {"batch_size": 16, "sequence_length": 128}, }, }, "marian": { "Helsinki-NLP/opus-mt-en-es": { "translation": {"batch_size": 4, "source_sequence_length": 512, "target_sequence_length": 512}, }, "Helsinki-NLP/opus-mt-en-hi": { "translation": {"batch_size": 4, "source_sequence_length": 512, "target_sequence_length": 512}, }, "Helsinki-NLP/opus-mt-es-en": { "translation": {"batch_size": 4, "source_sequence_length": 512, "target_sequence_length": 512}, }, }, "roberta": { "roberta-base": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, "roberta-large": { "default": {"batch_size": 8, "sequence_length": 128}, "token-classification": {"batch_size": 4, "sequence_length": 512}, "multiple-choice": {"batch_size": 4, "sequence_length": 512}, }, }, "t5": { "t5-small": { "translation": {"batch_size": 8, "source_sequence_length": 512, "target_sequence_length": 512}, "summarization": {"batch_size": 8, "source_sequence_length": 200, "target_sequence_length": 512}, }, "t5-base": { "translation": {"batch_size": 4, "source_sequence_length": 512, "target_sequence_length": 512}, "summarization": {"batch_size": 4, "source_sequence_length": 200, "target_sequence_length": 512}, }, }, "vit": { "google/vit-base-patch16-224": {"default": {"batch_size": 16}}, "google/vit-base-patch16-224-in21k": {"default": {"batch_size": 16}}, "google/vit-large-patch16-224-in21k": {"default": {"batch_size": 8}}, }, "xlm-roberta": { "xlm-roberta-base": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, "xlm-roberta-large": { "default": {"batch_size": 16, "sequence_length": 128}, "token-classification": {"batch_size": 8, "sequence_length": 512}, "multiple-choice": {"batch_size": 8, "sequence_length": 512}, }, }, } def get_testers_for_model_type(model_type: str) -> List["ExampleTesterBase"]: testers = [] for task, cls in TESTER_CLASSES.items(): for attr_name in dir(cls): if attr_name.startswith("test") and model_type in attr_name: testers.append((task, cls(), attr_name)) return testers def remove_extra_command_line_argument( command_prefix: str, extra_command_line_arguments: List[Union[str, Dict[str, str]]] ): argument_idx = None for idx, cmd_line_argument in enumerate(extra_command_line_arguments): if isinstance(cmd_line_argument, dict): starts_with_command_prefix = any(v.startswith(command_prefix) for v in cmd_line_argument.values()) else: starts_with_command_prefix = cmd_line_argument.startswith(command_prefix) if starts_with_command_prefix: argument_idx = idx if argument_idx is not None: extra_command_line_arguments.pop(argument_idx) def run_auto_fill_cache_for_model_name( model_type: str, model_name: str, shape_values_for_task: Dict[str, int], tester: "ExampleTesterBase", method_name: str, neuron_cache: str, num_neuron_cores: int, bf16: bool, ): batch_size = shape_values_for_task.get("batch_size") sequence_length = shape_values_for_task.get("sequence_length") source_sequence_length = shape_values_for_task.get("source_sequence_length") target_sequence_length = shape_values_for_task.get("target_sequence_length") extra_command_line_arguments = tester.EXTRA_COMMAND_LINE_ARGUMENTS if extra_command_line_arguments is None: extra_command_line_arguments = [] if batch_size is not None: if not bf16: batch_size = min(int(batch_size / 2), 1) tester.TRAIN_BATCH_SIZE = batch_size tester.EVAL_BATCH_SIZE = batch_size if sequence_length is not None: remove_extra_command_line_argument("--max_seq_length", extra_command_line_arguments) extra_command_line_arguments.append(f"--max_seq_length {sequence_length}") if source_sequence_length is not None: remove_extra_command_line_argument("--max_source_length", extra_command_line_arguments) extra_command_line_arguments.append(f"--max_source_length {source_sequence_length}") if target_sequence_length is not None: remove_extra_command_line_argument("--max_target_length", extra_command_line_arguments) extra_command_line_arguments.append(f"--max_target_length {target_sequence_length}") tester.EXTRA_COMMAND_LINE_ARGUMENTS = extra_command_line_arguments tester.NEURON_CACHE = neuron_cache tester.DO_PRECOMPILATION = False tester.MAX_STEPS = 200 tester.SAVE_STEPS = 100 tester.NPROC_PER_NODE = num_neuron_cores tester.BF16 = bf16 tester.EVAL_IS_SUPPORTED = False tester.GRADIENT_ACCUMULATION_STEPS = 1 setattr(tester, method_name, ExampleTestMeta._create_test(model_type, model_name)) getattr(tester, method_name)(tester) def parse_args(): parser = ArgumentParser(description="Tool that runs precompilation to fill the Neuron Cache.") parser.add_argument( "--cache-path", type=Path, default="neuron-cache", help="The directory in which all the precompiled neff files will be stored.", ) parser.add_argument("--models", type=str, default="all", nargs="+", help="The models to precompile.") parser.add_argument("--tasks", type=str, default="all", nargs="+", help="The tasks to precompile.") return parser.parse_args() def open_and_append_to_file(filename: str, new_line_to_append: str): with open(filename, "w+") as fp: content = fp.read() new_content = f"{content}\n{new_line_to_append}" fp.write(new_content) def main(): args = parse_args() now = datetime.now() now_str = now.strftime("%d%m%Y_%H_%M") success_filename = f"success_{now_str}.txt" failure_filename = f"failure_{now_str}.txt" if args.models == "all": models = list(ARCHITECTURES_TO_COMMON_PRETRAINED_WEIGHTS.keys()) else: models = args.models # Test examples are slow, this allows us to run them. os.environ["RUN_SLOW"] = "1" total_num_cores = get_num_neuron_cores() num_cores_configurations = [1, 2, 8, total_num_cores] num_cores_configurations = sorted(set(num_cores_configurations)) for model_type in models: testers = get_testers_for_model_type(model_type) for task, tester, method_name in testers: if args.tasks != "all" and task not in args.tasks: continue for model_name, shape_values in ARCHITECTURES_TO_COMMON_PRETRAINED_WEIGHTS[model_type].items(): print(f"Running precompilation for {model_name} on {task}...") for num_cores in num_cores_configurations: print(f"For {num_cores} neuron cores") if num_cores == 1: tester.MULTI_PROC = "false" else: tester.MULTI_PROC = "true" start = time.time() shape_values_for_task = shape_values.get(task) if shape_values_for_task is None: shape_values_for_task = shape_values["default"] example = ["*" * 20] example.append(f"Model:\t{model_name}") example.append("Shapes:") for name, value in shape_values_for_task.items(): example.append(f"\t{name} = {value}") example.append(f"Num cores:\t{num_cores}") example_str = "" try: start = time.time() example.append("Precision:\tBF16") example.append("*" * 20) example_str = "\n".join(example) run_auto_fill_cache_for_model_name( model_type, model_name, shape_values_for_task, tester, method_name, args.cache_path, num_cores, True, ) except Exception as e: print(e) open_and_append_to_file(failure_filename, example_str) else: open_and_append_to_file(success_filename, example_str) end = time.time() print(f"Done! Duration: {end - start:.3f}.") try: start = time.time() example.append("Precision:\tFull-precision") example.append("*" * 20) example_str = "\n".join(example) run_auto_fill_cache_for_model_name( model_type, model_name, shape_values_for_task, tester, method_name, args.cache_path, num_cores, False, ) except Exception: open_and_append_to_file(failure_filename, example_str) else: open_and_append_to_file(success_filename, example_str) end = time.time() print(f"Done! Duration: {end - start:.3f}.") os.environ["RUN_SLOW"] = "0" if __name__ == "__main__": main()