# 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. """Defines the command line related to dealing with the Neuron cache repo.""" from typing import TYPE_CHECKING from ...neuron.cache import get_hub_cached_entries, synchronize_hub_cache from ...neuron.utils.cache_utils import ( CACHE_REPO_NAME, HF_HOME_CACHE_REPO_FILE, create_custom_cache_repo, set_custom_cache_repo_name_in_hf_home, ) from ...neuron.utils.require_utils import requires_torch_neuronx from ...neuron.utils.runner import ExampleRunner from ...utils import logging from ..base import BaseOptimumCLICommand, CommandInfo if TYPE_CHECKING: from argparse import ArgumentParser logger = logging.get_logger() class CreateCustomCacheRepoCommand(BaseOptimumCLICommand): @staticmethod def parse_args(parser: "ArgumentParser"): parser.add_argument( "-n", "--name", type=str, default=CACHE_REPO_NAME, help="The name of the repo that will be used as a remote cache for the compilation files.", ) parser.add_argument( "--public", action="store_true", help="If set, the created repo will be public. By default the cache repo is private.", ) def run(self): repo_url = create_custom_cache_repo(repo_id=self.args.name, private=not self.args.public) public_or_private = "public" if self.args.public else "private" logger.info(f"Neuron cache created on the Hugging Face Hub: {repo_url.repo_id} [{public_or_private}].") logger.info(f"Neuron cache name set locally to {repo_url.repo_id} in {HF_HOME_CACHE_REPO_FILE}.") class SetCustomCacheRepoCommand(BaseOptimumCLICommand): @staticmethod def parse_args(parser: "ArgumentParser"): parser.add_argument("name", type=str, help="The name of the repo to use as remote cache.") def run(self): set_custom_cache_repo_name_in_hf_home(self.args.name) logger.info(f"Neuron cache name set locally to {self.args.name} in {HF_HOME_CACHE_REPO_FILE}.") class AddToCacheRepoCommand(BaseOptimumCLICommand): @staticmethod def parse_args(parser: "ArgumentParser"): parser.add_argument("-m", "--model", type=str, required=True, help="The name of model or path of the model.") parser.add_argument("--task", type=str, required=True, help="The task for which the model should be compiled.") # Shapes parser.add_argument( "--train_batch_size", type=int, required=True, help="The batch size to use during the model compilation for training.", ) parser.add_argument( "--eval_batch_size", type=int, default=None, help="The batch size to use during model compilation for evaluation.", ) sequence_length_group = parser.add_mutually_exclusive_group() sequence_length_group.add_argument( "--sequence_length", type=int, help="The sequence length of the model during compilation." ) seq2seq_sequence_length_group = sequence_length_group.add_argument_group() seq2seq_sequence_length_group.add_argument( "--encoder_sequence_length", type=int, help="The sequence length of the encoder part of the model during compilation.", ) seq2seq_sequence_length_group.add_argument( "--decoder_sequence_length", type=int, help="The sequence length of the decoder part of the model during compilation.", ) parser.add_argument( "--gradient_accumulation_steps", type=int, default=1, help="The number of gradient accumulation steps.." ) parser.add_argument( "--precision", choices=["fp", "bf16"], type=str, required=True, help="The precision to use during the model compilation.", ) parser.add_argument( "--num_cores", choices=list(range(1, 33)), type=int, required=True, help="The number of neuron cores to use during compilation.", ) parser.add_argument( "--example_dir", type=str, default=None, help="Path to where the example scripts are stored." ) parser.add_argument( "--max_steps", type=int, default=10, help="The maximum number of steps to run compilation for." ) def run(self): runner = ExampleRunner(self.args.model, self.args.task, example_dir=self.args.example_dir) if self.args.eval_batch_size is None: self.args.eval_batch_size = self.args.train_batch_size if self.args.sequence_length is not None: sequence_length = self.args.sequence_length elif self.args.encoder_sequence_length is None and self.args.decoder_sequence_length is None: raise ValueError( "You need to specify either sequence_length or encoder_sequence_length and decoder_sequence_length" ) elif self.args.encoder_sequence_length is None or self.args.decoder_sequence_length is None: raise ValueError("Both the encoder_sequence_length and the decoder_sequence_length must be provided.") else: sequence_length = [self.args.encoder_sequence_length, self.args.decoder_sequence_length] returncode, stdout = runner.run( self.args.num_cores, self.args.precision, self.args.train_batch_size, sequence_length, do_eval=True, eval_batch_size=self.args.eval_batch_size, gradient_accumulation_steps=self.args.gradient_accumulation_steps, num_epochs=3, max_steps=self.args.max_steps, save_steps=self.args.max_steps // 2, ) if returncode != 0: raise ValueError(f"Could not add the model to the cache. Full log:\n{stdout}.") class SynchronizeRepoCommand(BaseOptimumCLICommand): @staticmethod def parse_args(parser: "ArgumentParser"): parser.add_argument("--repo_id", type=str, default=None, help="The name of the repo to use as remote cache.") parser.add_argument( "--cache_dir", type=str, default=None, help="The cache directory that contains the compilation files." ) @requires_torch_neuronx def run(self): synchronize_hub_cache(cache_path=self.args.cache_dir, cache_repo_id=self.args.repo_id) class LookupRepoCommand(BaseOptimumCLICommand): @staticmethod def parse_args(parser: "ArgumentParser"): parser.add_argument( "model_id", type=str, help="The model_id to lookup cached versions for.", ) parser.add_argument( "--task", type=str, default=None, help="The optional task to lookup cached versions for models supporting multiple tasks.", ) parser.add_argument("--repo_id", type=str, default=None, help="The name of the repo to use as remote cache.") def _list_entries(self): entries = get_hub_cached_entries(self.args.model_id, task=self.args.task, cache_repo_id=self.args.repo_id) n_entries = len(entries) output = f"\n*** {n_entries} entrie(s) found in cache for {self.args.model_id}.***\n\n" for entry in entries: for key, value in entry.items(): output += f"\n{key}: {value}" output += "\n" print(output) def run(self): self._list_entries() class CustomCacheRepoCommand(BaseOptimumCLICommand): SUBCOMMANDS = ( CommandInfo( name="create", help="Create a model repo on the Hugging Face Hub to store Neuron X compilation files.", subcommand_class=CreateCustomCacheRepoCommand, ), CommandInfo( name="set", help="Set the name of the Neuron cache repo to use locally.", subcommand_class=SetCustomCacheRepoCommand, ), CommandInfo( name="add", help="Add a model to the cache of your choice (trainium only).", subcommand_class=AddToCacheRepoCommand, ), CommandInfo( name="synchronize", help="Synchronize the neuronx compiler cache with a hub cache repo.", subcommand_class=SynchronizeRepoCommand, ), CommandInfo( name="lookup", help="Lookup the neuronx compiler hub cache for the specified model id.", subcommand_class=LookupRepoCommand, ), )