# Copyright (c) 2021, EleutherAI # # 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. import os import yaml import json import logging import shortuuid import copy import torch import argparse import shutil from dataclasses import dataclass from typing import List, Dict from socket import gethostname try: from typing import Literal except ImportError: from typing_extensions import Literal from deepspeed.launcher.runner import DLTS_HOSTFILE from megatron.logging import Tee from megatron.tokenizer import build_tokenizer from megatron.utils import obtain_resource_pool, expand_attention_types from .deepspeed_args import NeoXArgsDeepspeedConfig, NeoXArgsDeepspeedRunner from .neox_args import ( NeoXArgsModel, NeoXArgsTokenizer, NeoXArgsTraining, NeoXArgsParallelism, NeoXArgsLogging, NeoXArgsOther, NeoXArgsTextgen, NeoXArgsOptimizer, NeoXArgsLRScheduler, ATTENTION_TYPE_CHOICES, ) # ZERO defaults by deespeed # These values should not be changed unless defaults in deepspeed are changed # for all zero_optimization options, see https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training ZERO_DEFAULTS = { "stage": 0, "allgather_partitions": True, "reduce_scatter": True, "allgather_bucket_size": int(5e8), "overlap_comm": False, "reduce_scatter": True, "reduce_bucket_size": int(5e8), "contiguous_gradients": False, } # NeoX optimizer defaults OPT_DEFAULT = "Adam" OPT_PARAMS_DEFAULTS = { "lr": 0.001, "betas": [0.9, 0.999], "eps": 1.0e-8, "weight_decay": 0, "freeze_step": 400, "momentum": 0.0, "cuda_aware": False, } BASE_CLASSES = [ NeoXArgsDeepspeedRunner, NeoXArgsDeepspeedConfig, NeoXArgsModel, NeoXArgsLRScheduler, NeoXArgsOptimizer, NeoXArgsTokenizer, NeoXArgsTraining, NeoXArgsParallelism, NeoXArgsLogging, NeoXArgsTextgen, NeoXArgsOther, ] DEEPSPEED_ARG_CLASSES = [NeoXArgsDeepspeedRunner, NeoXArgsDeepspeedConfig] NEOX_ARG_CLASSES = [i for i in BASE_CLASSES if i not in DEEPSPEED_ARG_CLASSES] if "DLTS_HOSTFILE" in os.environ: DLTS_HOSTFILE = os.environ["DLTS_HOSTFILE"] @dataclass class NeoXArgs(*BASE_CLASSES): """ data class containing all configurations NeoXArgs inherits from a number of small configuration classes """ ############################################################################################################################ # start of instantiation def __post_init__(self): """ after initialization of default or loaded values a number of functions are performed in order to calculate values, assert consistency and do typechecking. """ if not NeoXArgs.validate_keys(): raise ValueError( self.__class__.__name__ + ".__post_init__() NeoXArgs keys cannot be validated" ) self.enable_logging() self.calculate_derived() if not self.validate_types(): raise ValueError( self.__class__.__name__ + ".__post_init__() NeoXArgs types cannot be validated" ) if not self.validate_values(): raise ValueError( self.__class__.__name__ + ".__post_init__() NeoXArgs values cannot be validated" ) def build_tokenizer(self): self.tokenizer = build_tokenizer(self) def initialize_tensorboard_writer(self): if self.tensorboard_dir and self.rank == 0: try: from torch.utils.tensorboard import SummaryWriter print("> setting tensorboard ...") self.tensorboard_writer = SummaryWriter(log_dir=self.tensorboard_dir) except (ModuleNotFoundError, ImportError): print( "WARNING: TensorBoard writing requested but is not " "available (are you using PyTorch 1.1.0 or later and do you have tensorboard installed?), " "no TensorBoard logs will be written.", flush=True, ) @classmethod def from_ymls(cls, paths_to_yml_files: List[str], overwrite_values: Dict = None): """ instantiates NeoXArgs while reading values from yml files paths_to_yml_files: list of paths to yml files overwrite_values: If provided, overwrite any values in the yamls with these values """ print(cls.__name__ + ".from_ymls() " + str(paths_to_yml_files), flush=True) # initialize an empty config dictionary to be filled by yamls config = dict() config_files = dict() # iterate of all to be loaded yaml files for conf_file_name in paths_to_yml_files: # load file with open(conf_file_name) as conf_file: conf = yaml.load(conf_file, Loader=yaml.FullLoader) # check for key duplicates and load values for conf_key, conf_value in conf.items(): if conf_key in config: raise ValueError( f"Conf file {conf_file_name} has the following duplicate keys with previously loaded file: {conf_key}" ) conf_key_converted = conf_key.replace( "-", "_" ) # TODO remove replace and update configuration files? config[conf_key_converted] = conf_value # load original config files to save unchanged with checkpoint # saving the original config retains comments filename = os.path.basename(conf_file_name) assert ( filename not in config_files ), "At least two config files have the same filename. This will result in conflicts when saving out configs with the checkpoint in one single directory. Please use unique names for configs." config_files[filename] = open(conf_file_name).read() # add config file content to neox args to make them accessible in code # this is used when saving checkpoints config["config_files"] = config_files # Configuration parameters not specified params_not_in_config = sorted( list(set(cls.__dataclass_fields__.keys()) - set(config.keys())) ) if len(params_not_in_config) > 0: logging.debug( cls.__name__ + ".from_ymls() Configuration parameters not specified (using defaults): " + ", ".join(params_not_in_config) ) if overwrite_values is not None: for k, v in overwrite_values.items(): config[k] = v # instantiate class and return # duplicate values and unrecognized keys are again checked upon instantiation return cls(**config) @classmethod def from_dict(cls, args_dict: Dict): """ instantiates NeoXArgs while reading values from input dict """ return cls(**args_dict) ############################################################################################################################ # start of command line args interface @classmethod def consume_deepy_args(cls): """ entry point for deepy.py configuring and consuming command line arguments. We can use `--wandb_group` / `--wandb_team` to overwrite those args from the command line, otherwise the value from the config is taken. """ parser = argparse.ArgumentParser( description="GPT-NeoX Configuration", allow_abbrev=False ) group = parser.add_argument_group(title="Training Configuration") group.add_argument( "user_script", type=str, help="User script to launch, followed by any required " "arguments.", ) group.add_argument( "--conf_dir", "-d", type=str, default=None, help="Directory to prefix to all configuration file paths", ) group.add_argument( "conf_file", type=str, nargs="+", help="Configuration file path. Multiple files can be provided and will be merged.", ) group = parser.add_argument_group(title="Weights and Biases monitoring args") group.add_argument( "--wandb_group", type=str, default=None, help='Weights and Biases group name - used to group together "runs".', ) group.add_argument( "--wandb_team", type=str, default=None, help="Team name for Weights and Biases.", ) group = parser.add_argument_group(title="Eval args") group.add_argument( "--eval_tasks", type=str, nargs="+", default=None, help="Optionally overwrite eval tasks to run for evaluate.py", ) group.add_argument( "--iteration", type=int, default=None, help="Iteration to load checkpoint from in evaluate.py / generate.py. If None is provided, uses the latest iteration.", ) group.add_argument( "--eval_results_prefix", type=str, default=None, help="prefix to append to eval results file", ) parser.add_argument( "-H", "--hostfile", type=str, help="Hostfile path (in MPI style) that defines the " "resource pool available to the job (e.g., " "worker-0 slots=4)", ) group = parser.add_argument_group(title="Generation args") group.add_argument( "-i", "--sample_input_file", type=str, default=None, help="Optionally overwrite `sample_input_file` for generate.py", ) group.add_argument( "-o", "--sample_output_file", type=str, default=None, help="Optionally overwrite `sample_output_file` for generate.py", ) args_parsed = parser.parse_args() # Validate user_script exists assert os.path.exists( args_parsed.user_script ), f"User script could not be found: {args_parsed.user_script}" # load config files conf_files = args_parsed.conf_file if args_parsed.conf_dir: conf_files = [os.path.join(args_parsed.conf_dir, f) for f in conf_files] # enables us to pass in `small` instead of `small.yml` conf_files = [(cf if cf.endswith(".yml") else cf + ".yml") for cf in conf_files] # determine overwrite values overwrite_values = dict() for k, v in vars(args_parsed).items(): if k not in ["conf_dir", "conf_file"] and v is not None: overwrite_values[k] = v # load args neox_args = cls.from_ymls( paths_to_yml_files=conf_files, overwrite_values=overwrite_values ) if neox_args.wandb_group is not None: # concat the wandb group name with a uid to make sure it's unique import wandb neox_args.wandb_group += "_" + wandb.util.generate_id() neox_args.print() return neox_args @classmethod def consume_neox_args(cls, overwrite_values=None): """ Deepspeed launcher needs to pass the arguments for `pretrain_gpt2.py` across to all machines. In order not to have any problems with different configs being mismatched across machines, we instead read the .yaml configuration file from the main rank, then serialize the arguments to a dictionary, which the deepspeed launcher broadcasts to all machines (`--megatron_config`). We then instantiate a new NeoXArgs from the dictionary (`.from_dict`). This should ensure args are never inconsistent across machines. """ parser = argparse.ArgumentParser( description="GPT-NeoX Configuration", allow_abbrev=False ) parser.add_argument( "--megatron_config", type=str, default=None, help="json dict dumped as string in NeoXArgs.get_deepspeed_main_args()", ) args_parsed, _ = parser.parse_known_args() megatron_config = json.loads(args_parsed.megatron_config) if overwrite_values is not None: megatron_config.update(overwrite_values) return cls.from_dict(args_dict=megatron_config) @staticmethod def convert_key_value_to_command_line_arg(k, v): if isinstance(v, bool): if v: return [f"--{k}"] else: return [] if v is None: return [] return [f"--{k}", str(v)] def get_deepspeed_main_args(self): args_list = list() # get deepspeed runner args, and only pass them in to deepspeed launcher if they differ from defaults for key, default_value in NeoXArgsDeepspeedRunner().defaults(): configured_value = getattr(self, key) if configured_value != default_value: args_list.extend( self.convert_key_value_to_command_line_arg(key, configured_value) ) if "DLTS_HOSTFILE" in os.environ: args_list.extend( self.convert_key_value_to_command_line_arg( "hostfile", os.environ["DLTS_HOSTFILE"] ) ) if "MASTER_ADDR" in os.environ: args_list.extend( self.convert_key_value_to_command_line_arg( "master_addr", os.environ["MASTER_ADDR"] ) ) if ( "--include" in args_list or "--exclude" in args_list ) and "--num_gpus" in args_list: print( "WARNING: both --include/--exclude and num_gpus were specified simultaneously - overriding num_gpus with --include/--exclude" ) # cannot specify these both simultaneously, remove num_gpus from list idx = args_list.index("--num_gpus") # pop twice, once for the arg, once for its value args_list.pop(idx) args_list.pop(idx) # add user script args_list.append(self.user_script) # get deepspeed_config args_list.append("--deepspeed_config") args_list.append(json.dumps(self.deepspeed_config)) # get all config values args_list.append("--megatron_config") neox_args = self.get_parent_class_value_dict( *self.__class__.__bases__, only_non_defaults=True ) args_list.append(json.dumps(neox_args)) return args_list ############################################################################################################################ # start of calculated properties @property def deepspeed_config(self) -> dict: """ returns a dict containing variables within deepspeed config """ return self.get_parent_class_value_dict( NeoXArgsDeepspeedConfig, only_non_defaults=True ) @property def deepspeed_runner(self) -> dict: """ returns variables within deepspeed runner """ return self.get_parent_class_value_dict(NeoXArgsDeepspeedRunner) @property def megatron_config(self) -> dict: """ returns variables within megatron args """ return self.get_parent_class_value_dict(*NEOX_ARG_CLASSES) @property def all_config(self) -> dict: """ returns variables of all args """ return self.get_parent_class_value_dict(*BASE_CLASSES) def get_parent_class_value_dict( self, *parent_classes, only_non_defaults=False ) -> dict: """ takes a sequence of parent classes and returns corresponding values (with defaults set) """ # TODO no Nones or non-defaults result = dict() for parent in parent_classes: for key, default_value in parent().defaults(): if key in ["tokenizer", "tensorboard_writer", "adlr_autoresume_object"]: continue if only_non_defaults: value = getattr(self, key) if value == default_value: continue result[key] = getattr(self, key) return result @property def params_dtype(self): """ returns the datatype on the basis of configured precision """ if self.precision == "fp16": return torch.half elif self.precision == "bfloat16": return torch.bfloat16 else: return torch.float ############################################################################################################################ # start of logging and output def enable_logging(self): """ enable Tee logs based on the configured logdir """ if self.log_dir: os.makedirs(self.log_dir, exist_ok=True) hostname = gethostname() file_prefix = os.path.join(self.log_dir, hostname) Tee(file_prefix + "_stdout.txt", err=False) Tee(file_prefix + "_stderr.txt", err=True) def print(self): """Print arguments.""" if self.rank == 0 or self.rank is None: print("-------------------- arguments --------------------", flush=True) str_list = [] for arg in vars(self): # add arg + value dots = "." * (32 - len(arg)) value = getattr(self, arg) print_str = " {} {} {}".format(arg, dots, value) # add info 'default or updated' field_def = self.__dataclass_fields__.get(arg) if field_def is not None: default_info = ( "default" if value == field_def.default else "updated" ) else: default_info = "" dots = "." * (64 - len(print_str)) print_str += dots str_list.append({"print_str": print_str, "default_info": default_info}) for arg in sorted( sorted(str_list, key=lambda x: x["print_str"].lower()), key=lambda x: x["default_info"], reverse=True, ): print(arg["print_str"] + arg["default_info"], flush=True) print("---------------- end of arguments ----------------", flush=True) ############################################################################################################################ # start of calculations and derived values def configure_distributed_args(self): """ Configures distributed training arguments from local variables set by deepspeed launcher. """ if self.deepspeed_mpi: from deepspeed.utils.distributed import mpi_discovery mpi_discovery() if self.deepspeed_slurm: os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"] os.environ["RANK"] = os.environ["SLURM_PROCID"] os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] self.update_value("local_rank", int(os.getenv("LOCAL_RANK", "0"))) self.update_value("rank", int(os.getenv("RANK", "0"))) self.update_value("world_size", int(os.getenv("WORLD_SIZE", "1"))) if self.rank == 0: print( self.__class__.__name__ + ".configure_distributed_args() using world size: {} and model-parallel size: {} ".format( self.world_size, self.model_parallel_size ), flush=True, ) @staticmethod def calculate_batch_parameters( dp_world_size, train_batch=None, micro_batch=None, grad_acc=None ): # all values are provided nothing needs to be set if train_batch is not None and micro_batch is not None and grad_acc is not None: return train_batch, micro_batch, grad_acc # gradient_accumulation_steps needs to be set elif train_batch is not None and micro_batch is not None: grad_acc = train_batch // micro_batch grad_acc //= dp_world_size # micro_batch_per_gpu needs to be set elif train_batch is not None and grad_acc is not None: micro_batch = train_batch // dp_world_size micro_batch //= grad_acc # train_batch_size needs to be set elif micro_batch is not None and grad_acc is not None: train_batch = micro_batch * grad_acc train_batch *= dp_world_size # gradient_accumulation_steps and micro_batch_per_gpus is set elif train_batch is not None: grad_acc = 1 micro_batch = train_batch // dp_world_size # train_batch_size and gradient_accumulation_step is set elif micro_batch is not None: train_batch = micro_batch * dp_world_size grad_acc = 1 # either none of the three parameters are provided or just gradient_accumulation_step is provided else: assert ( False ), "Either train_batch_size or micro_batch_per_gpu needs to be provided" return int(train_batch), int(micro_batch), int(grad_acc) @staticmethod def check_batch_parameters(dp_world_size, train_batch, micro_batch, grad_acc): assert ( train_batch > 0 ), f"Train batch size: {train_batch} has to be greater than 0" assert ( micro_batch > 0 ), f"Micro batch size per gpu: {micro_batch} has to be greater than 0" assert ( grad_acc > 0 ), f"Gradient accumulation steps: {grad_acc} has to be greater than 0" assert train_batch == micro_batch * grad_acc * dp_world_size, ( f"Check batch related parameters. train_batch_size is not equal" " to micro_batch_per_gpu * gradient_acc_step * world_size \n" f"{train_batch} != {micro_batch} * {grad_acc} * {dp_world_size}" ) def calculate_derived(self): """ Derives additional configuration values necessary for training from the current config """ # wandb # sets a unique wandb group if self.wandb_group is None: # if none is defined a uuid is set for the run self.wandb_group = shortuuid.uuid() # number of gpus # Get number of GPUs param or hostfile to determine train_batch_size global_num_gpus = getattr(self, "global_num_gpus", None) if global_num_gpus is None: if self.hostfile is not None or os.path.exists(DLTS_HOSTFILE): hostfile_path = self.hostfile or DLTS_HOSTFILE resources = obtain_resource_pool( hostfile_path, self.include or "", self.exclude or "" ) if self.num_nodes is not None and self.num_nodes > 0: resources = { k: resources[k] for k in list(resources.keys())[: self.num_nodes] } global_num_gpus = sum(map(len, resources.values())) if self.num_gpus is not None and self.num_gpus > 0: global_num_gpus = self.num_gpus * len(resources) else: global_num_gpus = torch.cuda.device_count() self.update_value("global_num_gpus", global_num_gpus) logging.info( self.__class__.__name__ + ".calculate_derived() " + f"Total number of GPUs determined to be: {global_num_gpus}" ) # get world size in the model/pipe parallel case, the actual `world size` deepspeed uses is the size of the # data-parallel group, or (num_gpus / mp_size) / pp_size pp_size = self.pipe_parallel_size pp_size = pp_size if pp_size >= 1 else 1 mp_size = self.model_parallel_size mp_size = mp_size if mp_size >= 1 else 1 self.update_value("model_parallel_size", mp_size) # pp_size and mp_size are only used here to compute dp world size and nowhere else. dp_world_size = (global_num_gpus / pp_size) / mp_size if not (dp_world_size % 1 == 0): error_message = ( self.__class__.__name__ + ".calculate_derived() " + f"(global_num_gpus / pp_size) / mp_size [({global_num_gpus} / {pp_size}) / {mp_size}] must be a whole number" ) logging.error(error_message) raise AssertionError(error_message) # Automatically derive train_batch_size = train_micro_batch_size_per_gpu*global_num_gpus*gradient_accumulation_steps ( train_batch_size, train_micro_batch_size_per_gpu, gradient_accumulation_steps, ) = self.calculate_batch_parameters( dp_world_size=dp_world_size, train_batch=self.train_batch_size, micro_batch=self.train_micro_batch_size_per_gpu, grad_acc=self.gradient_accumulation_steps, ) self.check_batch_parameters( dp_world_size=dp_world_size, train_batch=train_batch_size, micro_batch=train_micro_batch_size_per_gpu, grad_acc=gradient_accumulation_steps, ) self.update_values( { # batch size params "train_batch_size": train_batch_size, "train_micro_batch_size_per_gpu": train_micro_batch_size_per_gpu, "gradient_accumulation_steps": gradient_accumulation_steps, "batch_size": train_micro_batch_size_per_gpu, # duplicate items "gas": self.gradient_accumulation_steps, "clip_grad": self.gradient_clipping, } ) # derive steps where checkpoint should be saved if self.checkpoint_factor or self.extra_save_iters: if self.extra_save_iters: save_iters = set(self.extra_save_iters) else: save_iters = set() step = self.checkpoint_factor # don't save step 0 or 1 while step < self.train_iters: save_iters.add(step) if self.checkpoint_scale == "log": step *= self.checkpoint_factor elif self.checkpoint_scale == "linear": step += self.checkpoint_factor save_iters = list(save_iters) save_iters.sort() self.update_values( { "save_iters": save_iters, } ) # derive precision if (self.fp16 or {}).get("type", self.precision) == "bfloat16": self.update_value("precision", "bfloat16") elif (self.fp16 or {}).get("enabled", False): self.update_value("precision", "fp16") else: self.update_value("precision", "fp32") # zero optimization if self.zero_optimization is None: self.zero_optimization = copy.deepcopy( ZERO_DEFAULTS ) # a dict is overwritten and not updated key by key self.update_values( { "zero_stage": self.zero_optimization.get( "stage", ZERO_DEFAULTS["stage"] ), "zero_reduce_scatter": self.zero_optimization.get( "reduce_scatter", ZERO_DEFAULTS["reduce_scatter"] ), "zero_contiguous_gradients": self.zero_optimization.get( "contiguous_gradients", ZERO_DEFAULTS["contiguous_gradients"] ), "zero_reduce_bucket_size": self.zero_optimization.get( "reduce_bucket_size", ZERO_DEFAULTS["reduce_bucket_size"] ), "zero_allgather_bucket_size": self.zero_optimization.get( "allgather_bucket_size", ZERO_DEFAULTS["allgather_bucket_size"] ), } ) # optimizer and scheduler opt_params = self.optimizer or { "type": OPT_DEFAULT, "params": OPT_PARAMS_DEFAULTS, } self.update_values( { "optimizer_type": opt_params.get("type", OPT_DEFAULT), "lr": opt_params["params"].get("lr", OPT_PARAMS_DEFAULTS["lr"]), } ) if self.optimizer_type.lower() == "onebitadam": # onebitadam needs to instantiated by deepspeed, and so we need to pass deepspeed scheduler args # for all other optimizers, the scheduling is handled by megatron self.scheduler = { "type": "WarmupDecayLR", # for now this is the only ds scheduler offering decay "params": { "warmup_min_lr": 0, "warmup_max_lr": self.lr, "warmup_num_steps": int(self.train_iters * self.warmup), "total_num_steps": self.lr_decay_iters or self.train_iters, }, } # Fp16 loss scaling. self.update_value("dynamic_loss_scale", self.loss_scale is None) # Update 'is pipe parallel' flag # if we set pipe_parallel_size to 0 or 1, GPT2ModelPipe.to_sequential() is called, and we run training with # the sequential model without the PipelineModule wrapper to avoid the overhead it incurs self.update_value("is_pipe_parallel", self.pipe_parallel_size >= 1) # Attention config if self.attention_config is None: self.update_value("attention_config", [[["global"], self.num_layers]]) self.update_value( "attention_config", expand_attention_types(self.attention_config, self.num_layers), ) assert ( len(self.attention_config) == self.num_layers ), "Length of attention config list must equal num_layers" for item in self.attention_config: assert ( item in ATTENTION_TYPE_CHOICES ), f"Attention type {item} not recognized" if "gmlp" in self.attention_config or "amlp" in self.attention_config: assert ( not self.partition_activations ), "GMLP Blocks are not compatible with partition activations" # Sparsity config if self.sparsity_config is None: # Can't have a default value as an empty dict so need to set it here self.update_value("sparsity_config", {}) # Adding equal dataset weights if none are provided if self.train_data_paths and (self.train_data_weights is None): self.train_data_weights = [1.0] * len(self.train_data_paths) if self.valid_data_paths and (self.valid_data_weights is None): self.valid_data_weights = [1.0] * len(self.valid_data_paths) if self.test_data_paths and (self.test_data_weights is None): self.test_data_weights = [1.0] * len(self.test_data_paths) # if a sample input file is provided, default text_gen_type type to input-file if self.text_gen_type is None: if self.sample_input_file: self.update_value("text_gen_type", "input-file") else: self.update_value("text_gen_type", "unconditional") ############################################################################################################################ # start of validation functions @classmethod def validate_keys(cls): """ test that there are no duplicate arguments """ source_classes = list(cls.__bases__) defined_properties = dict() for source_class in source_classes: source_vars = list(source_class.__dataclass_fields__) for item in source_vars: if item in defined_properties.keys(): logging.error( f"({cls.__name__}) duplicate of item: {item}, in class {source_class.__name__} and {defined_properties[item]}" ) return False else: defined_properties[item] = source_class.__name__ return True def validate_values(self): # the current codebase assumes running with deepspeed only if not self.deepspeed: return False # learning rate if self.lr is None: error_message = self.__class__.__name__ + ".validate_values() lr is None" logging.error(error_message) raise ValueError(error_message) return False # required arguments required_args = [ "num_layers", "hidden_size", "num_attention_heads", "max_position_embeddings", ] for req_arg in required_args: if getattr(self, req_arg) is None: error_message = ( self.__class__.__name__ + ".validate_values() " + req_arg + " is None." ) logging.error(error_message) raise ValueError(error_message) return False # Checks. if self.hidden_size % self.num_attention_heads != 0: error_message = ( self.__class__.__name__ + ".validate_values() hidden_size must be divisible by num_attention_heads" ) logging.error(error_message) raise ValueError(error_message) return False if self.seq_length is not None: if not (self.max_position_embeddings >= self.seq_length): error_message = ( self.__class__.__name__ + ".validate_values() max_position_embeddings must be bigger or equal seq_length" ) logging.error(error_message) raise ValueError(error_message) return False if not (self.min_lr <= self.lr): error_message = ( self.__class__.__name__ + ".validate_values() min_lr must be smaller or equal lr" ) logging.error(error_message) raise ValueError(error_message) return False if ( self.save is not None and self.checkpoint_factor is None and self.extra_save_iters is None ): error_message = ( self.__class__.__name__ + ".validate_values() checkpoint_factor or extra_save_iters must be defined if save is defined" ) logging.error(error_message) raise ValueError(error_message) return False # Parameters sharing does not work with torch DDP. if (self.num_unique_layers is not None) and (self.num_layers is not None): if not (self.num_unique_layers <= self.num_layers): error_message = ( self.__class__.__name__ + ".validate_values() num-unique-layers must be smaller or equal num_layers" ) logging.error(error_message) raise ValueError(error_message) return False if not (self.num_layers % self.num_unique_layers == 0): error_message = ( self.__class__.__name__ + ".validate_values() num-layers should be divisible by num-unique-layers" ) logging.error(error_message) raise ValueError(error_message) return False if self.fp16_lm_cross_entropy and self.precision != "fp16": error_message = ( self.__class__.__name__ + ".validate_values() lm cross entropy in fp16 only support in fp16 mode." ) logging.error(error_message) raise ValueError(error_message) return False # assert that if one of train/test/valid_data_path are provided, data_path should not be has_separate_path = [ data_path is not None for data_path in [ self.train_data_paths, self.valid_data_paths, self.test_data_paths, ] ] if all(has_separate_path): assert self.data_path is None, ( "Please provide *either* `data_path` or `train/valid/test_data_path` " "in args " ) # assert that if one of train/test/valid_data_path are provided, all should be assert_error_mess = ( "One or more of train/valid/test data_path are not provided:\n\t" ) assert_error_mess += "\n\t".join( [ f"{name} data paths: {data_path}," for name, data_path in [ ["train", self.train_data_paths], ["valid", self.valid_data_paths], ["test", self.test_data_paths], ] ] ) assert any(has_separate_path) == all(has_separate_path), assert_error_mess # assert that if train / valid / test data path(s) and weights are provided, that the paths and the weights should be equal length if self.train_data_paths is not None: assert len(self.train_data_paths) == len(self.train_data_weights) if self.valid_data_paths is not None: assert len(self.valid_data_paths) == len(self.valid_data_weights) if self.test_data_paths is not None: assert len(self.test_data_paths) == len(self.test_data_weights) return True def validate_types(self): """ At runtime, checks types are actually the type specified. """ for field_name, field_def in self.__dataclass_fields__.items(): actual_value = getattr(self, field_name) if actual_value is None: continue # we allow for some values not to be configured actual_type = type(actual_value) if actual_type != field_def.type: if ( actual_type == int and field_def.type == float ): # floats should be able to be configured as ints continue # for typing.Literal (i.e a list of choices) - checks that actual value is in accepted values elif field_def.type.__origin__ == Literal: accepted_values = field_def.type.__args__ if actual_value in accepted_values: continue elif type(actual_value) == str: # case insensitive checking lowercase_accepted_values = [ i.lower() for i in accepted_values if isinstance(i, str) ] if actual_value.lower() in lowercase_accepted_values: continue logging.error( self.__class__.__name__ + ".validate_types() " + f"{field_name}: '{actual_value}' Not in accepted values: '{accepted_values}'" ) return False logging.error( self.__class__.__name__ + ".validate_types() " + f"{field_name}: '{actual_type}' instead of '{field_def.type}'" ) return False # validate deepspeed dicts for field_name in ["optimizer", "scheduler"]: value = getattr(self, field_name) if isinstance( value, dict ): # dict is checked above, only fields are checked here if "type" in value: if not isinstance(value["type"], str): logging.error( self.__class__.__name__ + ".validate_types() " + f"{field_name}: key 'type' must be a string" ) return False else: logging.error( self.__class__.__name__ + ".validate_types() " + f"{field_name}: must contain key 'type'" ) return False if "params" in value: if not isinstance(value["params"], dict): logging.error( self.__class__.__name__ + ".validate_types() " + f"{field_name}: key 'params' must be a dict" ) return False else: logging.error( self.__class__.__name__ + ".validate_types() " + f"{field_name}: must contain key 'params'" ) return False for field_name in ["fp16", "amp", "flops_profiler"]: value = getattr(self, field_name) if isinstance(value, dict): if not "enabled" in value: error_message = ( self.__class__.__name__ + ".validate_types() " + f"{field_name}: must contain key 'enabled'" ) logging.error(error_message) return False return True