# Copyright (c) 2021 Graphcore Ltd. All rights reserved. import os import sys import yaml import argparse import torch import popdist import popdist.poptorch import horovod.torch as hvd config_file = os.path.join(os.path.dirname(__file__), "configs.yml") def str_to_bool(value): if isinstance(value, bool) or value is None: return value if value.lower() in {"false", "f", "0", "no", "n"}: return False elif value.lower() in {"true", "t", "1", "yes", "y"}: return True raise argparse.ArgumentTypeError(f"{value} is not a valid boolean value") def init_popdist(args): hvd.init() args.use_popdist = True if popdist.getNumTotalReplicas() != args.replication_factor: print(f"The number of replicas is overridden by PopRun. " f"The new value is {popdist.getNumTotalReplicas()}.") args.replication_factor = int(popdist.getNumLocalReplicas()) args.popdist_rank = popdist.getInstanceIndex() args.popdist_size = popdist.getNumInstances() def parse_args(args=None): pparser = argparse.ArgumentParser("DALL-E Configuration name", add_help=False) pparser.add_argument("--config", type=str, help="Configuration Name", default="unit_test") pargs, remaining_args = pparser.parse_known_args(args=args) config_name = pargs.config parser = argparse.ArgumentParser( "PopTorch mini DALL-E", add_help=True, formatter_class=argparse.ArgumentDefaultsHelpFormatter ) # Execution parser.add_argument("--batch-size", type=int, help="Set the micro batch-size") parser.add_argument("--epochs", type=int, help="Number of training epochs") parser.add_argument("--device-iterations", type=int, help="Number of batches per training step") parser.add_argument("--replication-factor", type=int, default=1, help="Number of replicas") parser.add_argument( "--gradient-accumulation", type=int, help="Number of gradients to accumulate before updating the weights." ) parser.add_argument( "--stochastic-rounding", type=str_to_bool, nargs="?", const=True, default=False, help="Enable stochastic rounding", ) parser.add_argument( "--embedding-serialization-factor", type=int, default=1, help="Matmul serialization factor the embedding layers" ) parser.add_argument( "--enable-half-partials", type=str_to_bool, nargs="?", const=True, default=False, help="Enable half partials for matmuls and convolutions globally", ) parser.add_argument("--enable-rts", type=str_to_bool, nargs="?", const=True, default=False, help="Enabling RTS") parser.add_argument( "--optimizer-state-offchip", type=str_to_bool, nargs="?", const=True, default=False, help="Set the tensor storage location for optimizer state to be offchip.", ) parser.add_argument("--ipus-per-replica", type=int, default=1, help="Number of IPUs required by each replica") parser.add_argument( "--layers-per-ipu", type=int, nargs="+", default=[0, 0, 8, 8], help="number of layers placed on each IPU" ) parser.add_argument("--cls-ipu-id", type=int, default=None, help="IPU id of classification layer") parser.add_argument("--embedding-ipu-id", type=int, default=0, help="IPU id of image embedding and text embedding") parser.add_argument( "--matmul-proportion", type=float, nargs="+", default=[0.6], help="Relative IPU memory proportion size allocated for matmul", ) parser.add_argument( "--async-dataloader", type=str_to_bool, nargs="?", const=True, default=True, help="Enable asynchronous mode in the DataLoader", ) parser.add_argument("--random-seed", type=int, default=42, help="Seed for RNG") parser.add_argument("--fp16", action="store_true", help="Use fp16, otherwise use fp32") # Optimizer parser.add_argument( "--optimizer", type=str, choices=["Adam", "AdamW"], default="Adam", help="optimizer to use for the training" ) parser.add_argument( "--learning-rate", type=float, help="Learning rate value for constant schedule, maximum for linear schedule." ) parser.add_argument( "--lr-scheduler", type=str, choices=["ReduceLROnPlateau", "multi_step", "constant"], help="Learning rate scheduler", ) parser.add_argument("--loss-scaling", type=float, help="Loss scaling factor (recommend using powers of 2)") parser.add_argument("--enable-als", action="store_true", help="Enable automatic loss scaling") parser.add_argument("--weight-decay", type=float, default=0, help="Set the weight decay") parser.add_argument( "--enable-half-first-order-momentum", type=str_to_bool, nargs="?", const=True, default=False, help="Use float16 for the first order momentum in the optimizer.", ) # Model parser.add_argument("--hidden-size", type=int, help="The size of the hidden state of the transformer layers") parser.add_argument("--text-seq-len", type=int, help="The max text sequence length") parser.add_argument("--num-hidden-layers", type=int, help="The number of transformer layers") parser.add_argument("--num-attention-heads", type=int, help="Set the number of heads in self attention") parser.add_argument("--dim-head", type=int, help="The number of head dimension in self attention") parser.add_argument("--ff-dropout", type=float, nargs="?", const=True, help="Attention dropout probability") parser.add_argument("--attn-dropout", type=float, nargs="?", const=True, help="Feed forward dropout probability") parser.add_argument( "--sandwich-norm", type=str_to_bool, nargs="?", const=True, default=False, help="Use Sandwich LayerNorm" ) parser.add_argument("--loss-img-weight", default=7, type=int, help="Image loss weight") parser.add_argument( "--attn-types", default="full", type=str, help="comma separated list of attention types(full, axial_row, axial_col, conv_like).", ) parser.add_argument("--bpe-path", type=str, help="Path to BPE json file") parser.add_argument( "--truncate-captions", action="store_true", help="Captions passed in which exceed the max token length will be truncated.", ) # Dataset parser.add_argument( "--byteio", type=str_to_bool, nargs="?", const=True, default=False, help="Use byte data format to transfer image data from host to IPU", ) parser.add_argument("--input-folder", type=str, default=None, help="Path to folder of images and text for training") parser.add_argument( "--generated-data", action="store_true", help="Random data created on CPU ram, dataloader feeds the same data to IPU every step", ) parser.add_argument( "--synthetic-data", action="store_true", help="Random data created on the IPU, no host to IPU IO" ) # Checkpointing parser.add_argument( "--checkpoint-output-dir", type=str, default="", help="Directory where checkpoints will be saved to.\ This can be either an absolute or relative path.", ) parser.add_argument( "--checkpoint-save-steps", default=1000, type=int, help="Option to checkpoint model after n steps." ) parser.add_argument( "--pretrained-checkpoint", type=str, help="Checkpoint to be retrieved for further training. This can" "be either an absolute or relative path to the checkpoint file.", ) parser.add_argument("--vae_path", type=str, help="path to trained discrete VAE") parser.add_argument( "--vqgan_model_path", type=str, default=None, help="path to trained VQGAN weights. This should be a .ckpt file." ) parser.add_argument( "--vqgan_config_path", type=str, default=None, help="path to trained VQGAN config. This should be a .yaml file." ) # Misc parser.add_argument("--dataloader-workers", type=int, help="The number of dataloader workers") parser.add_argument( "--wandb", type=str_to_bool, nargs="?", const=True, default=False, help="Enabling logging to Weights and Biases" ) parser.add_argument("--wandb-project-name", default="miniDALL-E", help="Wandb project name") parser.add_argument("--wandb-run-name", default=None, help="Weights & Biases run name.") parser.add_argument( "--compile-only", action="store_true", help="Create an offline IPU target that can only be used for offline compilation.", ) parser.add_argument( "--executable-cache-dir", type=str, default="", help="Directory where Poplar executables are cached. If set, recompilation of identical graphs can be avoided. " "Required for both saving and loading executables.", ) parser.add_argument("--profile-dir", type=str, help="Directory for profiling results") # Load the yaml yaml_args = dict() if config_name is not None: with open(config_file, "r") as f: try: yaml_args.update(**yaml.safe_load(f)[config_name]) except yaml.YAMLError as exc: print(exc) sys.exit(1) # Check the yaml args are valid known_args = set(vars(parser.parse_args(""))) unknown_args = set(yaml_args) - known_args if unknown_args: print(f" Warning: Unknown arg(s) in config file: {unknown_args}") parser.set_defaults(**yaml_args) args = parser.parse_args(remaining_args) # Initialise PopDist if popdist.isPopdistEnvSet(): init_popdist(args) hvd.broadcast(torch.Tensor([args.random_seed]), root_rank=0) else: args.use_popdist = False # Expand matmul_proportion input into list representation if isinstance(args.matmul_proportion, float): args.matmul_proportion = [args.matmul_proportion] * args.ipus_per_replica if len(args.matmul_proportion) != args.ipus_per_replica: if len(args.matmul_proportion) == 1: args.matmul_proportion = args.matmul_proportion * args.ipus_per_replica else: raise ValueError( f"Length of matmul_proportion doesn't match ipus_per_replica: " f"{args.matmul_proportion} vs {args.ipus_per_replica}" ) if args.synthetic_data or args.generated_data: print( "Warning: One of --synthetic-data or --generated-data have been " "provided. This will override the value passed to --input-folder, " "and real data will not be used." ) args.input_folder = None if args.synthetic_data and args.generated_data: print( "Warning: Both of --synthetic-data and --generated-data have been " "provided. Only one can be provided. Exiting." ) sys.exit(1) return args def sync_metrics(outputs, factor=1, average=True): if popdist.isPopdistEnvSet(): if isinstance(outputs, float): return float(hvd.allreduce(torch.Tensor([outputs]), average=average).item()) else: return [hvd.allreduce(output.div(factor), average=average).mean().item() for output in outputs] else: if isinstance(outputs, float): return outputs else: return [output.div(factor).mean().item() for output in outputs]