def train(): parser = ArgumentParser() parser.add_argument("--basedir", type=str) parser.add_argument("--train_file", type=str, required=True, help='File path to use for train file') parser.add_argument("--valid_file", type=str, required=True, help='File path to use for valid file') parser.add_argument("--dataset_key", default="paired", help="dataset key for basedir") parser.add_argument( "--embed_type", type=str, default='default', choices=["default", "positional", "learned-positional"], help="register label of the embeddings") parser.add_argument("--d_model", type=int, default=512, help="Model dimension (and embedding dsz)") parser.add_argument("--d_ff", type=int, default=2048, help="FFN dimension") parser.add_argument( "--d_k", type=int, default=None, help="Dimension per head. Use if num_heads=1 to reduce dims") parser.add_argument("--num_heads", type=int, default=8, help="Number of heads") parser.add_argument("--num_layers", type=int, default=8, help="Number of layers") parser.add_argument("--windowed_ra", type=str2bool, default=False, help="whether prevent attention beyond rpr_k") parser.add_argument("--num_train_workers", type=int, default=4, help="Number train workers") parser.add_argument("--nctx", type=int, default=256, help="Max input length") parser.add_argument("--pattern", default='*.json', help="Glob pattern for data") parser.add_argument("--batch_size", type=int, default=256, help="Batch Size") parser.add_argument("--subword_model_file", type=str, help="The BPE model file", required=True) parser.add_argument("--subword_vocab_file", type=str, help="The BPE subword vocab", required=True) parser.add_argument("--dropout", type=float, default=0.1, help="Dropout") parser.add_argument("--lr_scheduler", type=str, default='cosine', help="The type of learning rate decay scheduler") parser.add_argument("--lr_decay_steps", type=int, help="decay steps of lr scheduler") parser.add_argument("--lr_decay_rate", type=float, help="decay rate of lr scheduler") parser.add_argument("--lr_alpha", type=float, help="parameter alpha for cosine decay scheduler") parser.add_argument("--optim", default="adamw", type=str, help="Optimizer to use (defaults to adamw)") parser.add_argument("--lr", type=float, default=4.0e-4, help="Learning rate") parser.add_argument("--clip", type=float, default=1.0, help="Clipping gradient norm") parser.add_argument("--weight_decay", type=float, default=1.0e-2, help="Weight decay") parser.add_argument("--epochs", type=int, default=32, help="Num training epochs") parser.add_argument( "--restart_from", type=str, help="Option allows you to restart from a previous checkpoint") parser.add_argument( "--restart_tt", type=str, help="Optional param for legacy checkpoints (step|epoch)") parser.add_argument("--warmup_steps", type=int, default=10000, help="Num warmup steps") parser.add_argument("--saves_per_epoch", type=int, default=10, help="The number of checkpoints to save per epoch") parser.add_argument("--reduction_d_k", type=int, default=64, help="Dimensions of Key and Query in the single headed" "reduction layers") parser.add_argument( "--stacking_layers", type=int, nargs='+', default=[1024, 1024, 1024], help="Hidden sizes of the dense stack (ff2 from the convert paper)") parser.add_argument("--ff_pdrop", type=float, default=0.1, help="Dropout in the dense stack") parser.add_argument("--reader_type", type=str, default='preprocessed', choices=['ntp', 'nsp', 'preprocessed']) parser.add_argument("--model_type", default="dual-encoder", choices=["dual-encoder", "encoder-decoder"]) parser.add_argument("--loss", type=str, default='all', choices=['triplet', 'all']) parser.add_argument( '--rpr_k', help= 'Relative attention positional sizes pass 0 if you dont want relative attention', type=int, default=[8], nargs='+') parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)") parser.add_argument("--distributed", type=str2bool, default=False, help="Are we doing distributed training?") parser.add_argument( "--local_rank", type=int, default=-1, help= "Local rank for distributed training (-1 means use the environment variables to find)" ) args = parser.parse_args() if args.basedir is None: args.basedir = '{}-{}-paired-{}-bpe-{}'.format(args.model_type, args.reader_type, args.dataset_key, os.getpid()) logging.basicConfig( level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) num_gpus = get_num_gpus_multiworker() args.distributed = args.distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") if args.distributed: args.device, updated_local_rank = init_distributed(args.local_rank) args.local_rank = updated_local_rank reader = MultiFileDatasetReader(args.nctx, args.subword_model_file, args.subword_vocab_file, args.pattern, reader_type=args.reader_type) vocab = reader.build_vocab() # If we are not using chars, then use 'x' for both input and output preproc_data = baseline.embeddings.load_embeddings( 'x', dsz=args.d_model, known_vocab=vocab['x'], preserve_vocab_indices=True, embed_type=args.embed_type) vocabs = preproc_data['vocab'] os.makedirs(args.basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs, os.path.join(args.basedir, 'vocabs.json')) embeddings = preproc_data['embeddings'] logger.info("Loaded embeddings") train_set = reader.load(args.train_file, vocabs) valid_set = reader.load(args.valid_file, vocabs, distribute=False, shuffle=False) train_loader = DataLoader(train_set, batch_size=args.batch_size, num_workers=args.num_train_workers) valid_loader = DataLoader(valid_set, batch_size=args.batch_size) logger.info("Loaded datasets") logger.info("Using embedding type [%s]", args.embed_type) if len(args.rpr_k) == 0 or args.rpr_k[0] < 1: rpr_k = None elif len(args.rpr_k) == 1: rpr_k = args.rpr_k[0] else: rpr_k = args.rpr_k model = create_model(embeddings, d_model=args.d_model, d_ff=args.d_ff, dropout=args.dropout, num_heads=args.num_heads, num_layers=args.num_layers, model_type=args.model_type, rpr_k=rpr_k, d_k=args.d_k, reduction_d_k=args.reduction_d_k, stacking_layers=args.stacking_layers, ff_pdrop=args.ff_pdrop, windowed_ra=args.windowed_ra, logger=logger) model.to(args.device) loss_function = model.create_loss(args.loss) loss_function.to(args.device) logger.info("Loaded model and loss") # according to pytorch, len(train_loader) will return len(train_set) when train_set is IterableDataset, so manually # correct it here steps_per_epoch = len(train_loader) // (args.batch_size * num_gpus) valid_steps = len(valid_loader) // args.batch_size update_on = steps_per_epoch // args.saves_per_epoch report_on = max(10, update_on) // 10 logger.info( f"Steps per epoch per GPU: {steps_per_epoch}. Saving checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(args.lr_scheduler, args.lr, steps_per_epoch, args.epochs, logger, decay_steps=args.lr_decay_steps, decay_rate=args.lr_decay_rate, alpha=args.lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(args.warmup_steps, lr=args.lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=args.lr) global_step = 0 start_epoch = 0 if args.restart_from: model.load_state_dict(torch.load(args.restart_from)) vec = args.restart_from.split("-") if args.restart_tt: tick_type = args.restart_tt else: tick_type = vec[-2] step_num = int(vec[-1].split(".")[0]) if tick_type == 'epoch': start_epoch = step_num global_step = start_epoch * steps_per_epoch else: start_epoch = step_num // steps_per_epoch global_step = step_num logger.info( "Restarting from a previous checkpoint %s.\n\tStarting at global_step=%d, epoch=%d", args.restart_from, global_step, start_epoch + 1) optimizer = OptimizerManager(model, global_step, optim=args.optim, lr=args.lr, lr_function=lr_sched, weight_decay=args.weight_decay) logger.info("Model has {:,} parameters".format( sum(p.numel() for p in model.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if args.distributed: model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank) logger.info("Model located on %d", args.local_rank) model_base = os.path.join(args.basedir, 'checkpoint') steps = global_step for epoch in range(start_epoch, args.epochs): avg_loss = Average('average_train_loss') metrics = {} optimizer.zero_grad() start = time.time() model.train() train_itr = iter(train_loader) for i in range(steps_per_epoch): batch = next(train_itr) steps += 1 x, y = batch inputs = x.to(args.device) labels = y.to(args.device) loss = loss_function(inputs, labels) loss.backward() avg_loss.update(loss.item()) torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() optimizer.zero_grad() if (i + 1) % report_on == 0: logging.info(avg_loss) if (i + 1) % update_on == 0 and args.local_rank < 1: elapsed = (time.time() - start) / 60 logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) save_checkpoint(model, model_base, steps, tick_type='step') # How much time elapsed in minutes elapsed = (time.time() - start) / 60 train_avg_loss = avg_loss.avg # This is the average training token-level loss across all machines # This is the token-level training perplexity metrics['train_elapsed_min'] = elapsed metrics['average_train_loss'] = train_avg_loss if args.local_rank < 1: avg_valid_loss = Average('average_valid_loss') start = time.time() model.eval() valid_itr = iter(valid_loader) for j in range(valid_steps): with torch.no_grad(): x, y = batch inputs = x.to(args.device) labels = y.to(args.device) loss = loss_function(inputs, labels) avg_valid_loss.update(loss.item()) valid_avg_loss = avg_valid_loss.avg elapsed = (time.time() - start) / 60 metrics['valid_elapsed_min'] = elapsed metrics['average_valid_loss'] = valid_avg_loss logger.info(metrics) save_checkpoint(model, model_base, epoch, tick_type='epoch')
def train(): parser = ArgumentParser() parser.add_argument("--basedir", type=str) parser.add_argument("--train_file", type=str, required=True, help='File path to use for train file') parser.add_argument("--valid_file", type=str, required=True, help='File path to use for valid file') parser.add_argument("--dataset_key", default="tlm", help="dataset key for basedir") parser.add_argument( "--embed_type", type=str, default='default', choices=["default", "positional", "learned-positional"], help="register label of the embeddings") parser.add_argument("--d_model", type=int, default=512, help="Model dimension (and embedding dsz)") parser.add_argument("--d_ff", type=int, default=2048, help="FFN dimension") parser.add_argument( "--d_k", type=int, default=None, help="Dimension per head. Use if num_heads=1 to reduce dims") parser.add_argument("--num_heads", type=int, default=8, help="Number of heads") parser.add_argument("--num_layers", type=int, default=8, help="Number of layers") parser.add_argument("--num_train_workers", type=int, default=4, help="Number train workers") parser.add_argument("--nctx", type=int, default=256, help="Max input length") parser.add_argument("--file_type", default='json', help="Glob pattern for data") parser.add_argument("--batch_size", type=int, default=256, help="Batch Size") parser.add_argument("--subword_model_file", type=str, help="The BPE model file", required=True) parser.add_argument("--subword_vocab_file", type=str, help="The BPE subword vocab", required=True) parser.add_argument("--dropout", type=float, default=0.1, help="Dropout") parser.add_argument("--ffn_pdrop", type=float, default=0.0, help="Dropout in the dense stack") parser.add_argument("--layer_drop", type=float, default=0.0, help="LayerDrop to apply") parser.add_argument("--lr_scheduler", type=str, default='cosine', help="The type of learning rate decay scheduler") parser.add_argument("--lr_decay_steps", type=int, help="decay steps of lr scheduler") parser.add_argument("--lr_decay_rate", type=float, help="decay rate of lr scheduler") parser.add_argument("--lr_alpha", type=float, help="parameter alpha for cosine decay scheduler") parser.add_argument("--optim", default="adamw", type=str, help="Optimizer to use (defaults to adamw)") parser.add_argument("--lr", type=float, default=4.0e-4, help="Learning rate") parser.add_argument("--clip", type=float, default=1.0, help="Clipping gradient norm") parser.add_argument("--weight_decay", type=float, default=1.0e-2, help="Weight decay") parser.add_argument("--epochs", type=int, default=32, help="Num training epochs") parser.add_argument( "--restart_from", type=str, help="Option allows you to restart from a previous checkpoint") parser.add_argument("--restart_tt", type=str, help="Optional param for legacy checkpoints", choices=['step', 'epoch', 'ignore']) parser.add_argument("--warmup_steps", type=int, default=10000, help="Num warmup steps") parser.add_argument("--saves_per_epoch", type=int, default=10, help="The number of checkpoints to save per epoch") parser.add_argument("--mlm", type=str2bool, default=True, help="Use Masked Language Model (MLM) objective") parser.add_argument("--preprocessed", type=str2bool, default=True, help="Has the data already been preprocessed?") parser.add_argument( '--rpr_k', help= 'Relative attention positional sizes pass 0 if you dont want relative attention', type=int, default=[8], nargs='+') parser.add_argument( '--rpr_value_on', type=str2bool, default=True, help= "In relative attention, whether add positional correction to values in addition to the " "correction to attention matrix") parser.add_argument("--windowed_ra", type=str2bool, default=False, help="whether prevent attention beyond rpr_k") parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)") parser.add_argument("--distributed", type=str2bool, default=False, help="Are we doing distributed training?") parser.add_argument( "--local_rank", type=int, default=-1, help= "Local rank for distributed training (-1 means use the environment variables to find)" ) args = parser.parse_args() if args.basedir is None: args.basedir = 'lm-{}-bpe-{}'.format(args.dataset_key, os.getpid()) logging.basicConfig( level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) num_gpus = get_num_gpus_multiworker() args.distributed = args.distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") do_on_demand_masking = args.mlm and not args.preprocessed if do_on_demand_masking: logger.info(f"On-demand masking is turned on") if args.distributed: args.device, updated_local_rank = init_distributed(args.local_rank) args.local_rank = updated_local_rank if args.file_type == 'tfrecord': reader_type = 'tfrecord' elif args.preprocessed: reader_type = 'preprocessed' else: reader_type = 'lang' reader = MultiFileDatasetReader( src_nctx=args.nctx, model_file=args.subword_model_file, vocab_file=args.subword_vocab_file, file_type=args.file_type, reader_type=reader_type, record_keys=['x', 'y'] if args.mlm else ['x']) # This looks a bit funny but the streaming reader ignores our vocab and gives us the one from the subword_model # However, we do need to get counts from our dataset for validation so we can calculate the perplexity vocab = reader.build_vocab([args.valid_file]) # If we are not using chars, then use 'x' for both input and output preproc_data = baseline.embeddings.load_embeddings( 'x', dsz=args.d_model, known_vocab=vocab['x'], preserve_vocab_indices=True, embed_type=args.embed_type) vocabs = preproc_data['vocab'] os.makedirs(args.basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs, os.path.join(args.basedir, 'vocabs.json')) embeddings = {'x': preproc_data['embeddings']} logger.info("Loaded embeddings") train_set = reader.load(args.train_file, vocabs) valid_set = reader.load(args.valid_file, vocabs, distribute=False, shuffle=False) train_loader = DataLoader(train_set, batch_size=args.batch_size, num_workers=args.num_train_workers) valid_loader = DataLoader(valid_set, batch_size=args.batch_size) logger.info("Loaded datasets") logger.info("Using embedding type [%s]", args.embed_type) if args.mlm: mask_from = vocabs vocab_size = len(mask_from) mask_value = mask_from.get("[MASK]") if mask_value == -1: logger.error( "We could not find a suitable masking token in the vocab") return if len(args.rpr_k) == 0 or args.rpr_k[0] < 1: rpr_k = None elif len(args.rpr_k) == 1: rpr_k = args.rpr_k[0] else: rpr_k = args.rpr_k TLM = TransformerMaskedLanguageModel if args.mlm else TransformerLanguageModel model = TLM.create(embeddings, hsz=args.d_model, d_ff=args.d_ff, tie_weights=True, dropout=args.dropout, gpu=False, num_heads=args.num_heads, layers=args.num_layers, rpr_k=rpr_k, d_k=args.d_k, ffn_pdrop=args.ffn_pdrop, windowed_ra=args.windowed_ra, rpr_value_on=args.rpr_value_on, layer_drop=args.layer_drop, src_keys=['x'], tgt_key='x') model.to(args.device) loss_function = model.create_loss() loss_function.to(args.device) logger.info("Loaded model and loss") steps_per_epoch = len(train_loader) // num_gpus valid_steps = len(valid_loader) update_on = steps_per_epoch // args.saves_per_epoch report_on = max(10, update_on) // 10 logger.info( f"Steps per epoch per GPU: {steps_per_epoch}. Saving checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(args.lr_scheduler, args.lr, steps_per_epoch, args.epochs, logger, decay_steps=args.lr_decay_steps, decay_rate=args.lr_decay_rate, alpha=args.lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(args.warmup_steps, lr=args.lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=args.lr) global_step = 0 start_epoch = 0 if args.restart_from: if args.restart_from.endswith('npz'): load_tlm_npz(model, args.restart_from) else: model.load_state_dict(torch.load(args.restart_from)) vec = args.restart_from.split("-") if args.restart_tt: tick_type = args.restart_tt else: tick_type = vec[-2] step_num = int(vec[-1].split(".")[0]) if tick_type == 'epoch': start_epoch = step_num global_step = start_epoch * steps_per_epoch elif tick_type == 'step': start_epoch = step_num // steps_per_epoch global_step = step_num else: logger.warning( f"The previous tick was {step_num} but command-line specifies to ignore, setting to 0" ) logger.info( "Restarting from a previous checkpoint %s.\n\tStarting at global_step=%d, epoch=%d", args.restart_from, global_step, start_epoch + 1) optimizer = OptimizerManager(model, global_step, optim=args.optim, lr=args.lr, lr_function=lr_sched, weight_decay=args.weight_decay) logger.info("Model has {:,} parameters".format( sum(p.numel() for p in model.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if args.distributed: # This program assume pure data parallelism, each model is on a single gpu # If we wanted to support model and data parallelism we would need to update # the selection of gpus based on rank, it would need to select multiple ids # based on rank, here we select only a single gpu and use it for input and # output. model = DistributedDataParallel(model, device_ids=[args.device], output_device=args.device) logger.info("Model located on %s", args.device) model_base = os.path.join(args.basedir, 'checkpoint') steps = global_step timer = Timer() for epoch in range(start_epoch, args.epochs): avg_loss = Average('average_train_loss') metrics = {} optimizer.zero_grad() timer.start() model.train() train_itr = iter(train_loader) for i in range(steps_per_epoch): batch = next(train_itr) steps += 1 x, y = batch inputs = x.to(args.device) labels = y.to(args.device) if do_on_demand_masking: inputs, labels, _ = on_demand_mlm_masking( inputs, labels, mask_value, vocab_size) inputs = {'x': inputs} labels = labels.transpose(0, 1).contiguous() logits = model(inputs, None)[0].transpose(0, 1).contiguous() if args.mlm: loss = loss_function(logits, labels) else: shift_logits = logits[:-1] shift_labels = labels[1:] loss = loss_function(shift_logits, shift_labels) loss.backward() avg_loss.update(loss.item()) torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() optimizer.zero_grad() if (i + 1) % report_on == 0: logging.info(avg_loss) if (i + 1) % update_on == 0 and args.local_rank < 1: elapsed = timer.elapsed(True) logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) logging.info('LR: %f', optimizer.current_lr) save_checkpoint(model, model_base, steps, tick_type='step') # How much time elapsed in minutes elapsed = timer.elapsed(True) train_token_loss = avg_loss.avg # This is the average training token-level loss across all machines # This is the token-level training perplexity train_token_ppl = math.exp(train_token_loss) metrics['train_elapsed_min'] = elapsed metrics['average_train_loss'] = train_token_loss metrics['train_ppl'] = train_token_ppl if args.local_rank < 1: avg_valid_loss = Average('average_valid_loss') timer.start() model.eval() valid_itr = iter(valid_loader) for j in range(valid_steps): batch = next(valid_itr) with torch.no_grad(): x, y = batch inputs = x.to(args.device) labels = y.to(args.device) if do_on_demand_masking: inputs, labels, _ = on_demand_mlm_masking( inputs, labels, mask_value, vocab_size) inputs = {'x': inputs} labels = labels.transpose(0, 1).contiguous() logits = model(inputs, None)[0].transpose(0, 1).contiguous() if args.mlm: loss = loss_function(logits, labels) else: shift_logits = logits[:-1] shift_labels = labels[1:] loss = loss_function(shift_logits, shift_labels) avg_valid_loss.update(loss.item()) valid_token_loss = avg_valid_loss.avg valid_token_ppl = math.exp(valid_token_loss) metrics['valid_elapsed_min'] = timer.elapsed(True) metrics['average_valid_loss'] = valid_token_loss metrics['average_valid_word_ppl'] = valid_token_ppl logger.info(metrics) save_checkpoint(model, model_base, epoch, save_npz=True)
def run(basedir=None, train_file=None, valid_file=None, dataset_key='paired', embed_type='default', d_model=512, d_ff=2048, d_k=None, num_heads=8, num_layers=8, num_train_workers=4, nctx=256, tgt_nctx=None, file_type='json', record_keys=['x', 'y'], batch_size=256, subword_model_file=None, subword_vocab_file=None, dropout=0.1, lr_scheduler='cosine', lr_decay_steps=None, lr_decay_rate=None, lr_alpha=None, optim='adamw', lr=4.0e-4, clip=1.0, weight_decay=1.0e-2, epochs=32, restart_from=None, restart_tt=None, warmup_steps=10000, saves_per_epoch=10, layer_drop=0.0, reader_type='preprocessed', src_begin_tok=[], src_end_tok=['<EOS>'], tgt_begin_tok=['<GO>'], tgt_end_tok=['<EOS>'], lower=False, rpr_k=[8], device='cuda', distributed=False, local_rank=-1, save_npz=False, extra_tokens=["[CLS]", "[MASK]"], subword_type='bpe', label_smoothing=None, ra_type=None, transformer_type=None, **kwargs): if basedir is None: basedir = f's2s-{reader_type}-paired-{dataset_key}-bpe-{os.getpid()}' logging.basicConfig( level=logging.INFO if local_rank in [-1, 0] else logging.WARN) num_gpus = get_num_gpus_multiworker() distributed = distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") if distributed: device, updated_local_rank = init_distributed(local_rank) local_rank = updated_local_rank if not tgt_nctx: tgt_nctx = nctx reader = MultiFileDatasetReader(nctx, tgt_nctx, src_begin_tok, src_end_tok, tgt_begin_tok, tgt_end_tok, subword_model_file, subword_vocab_file, file_type, reader_type=reader_type, record_keys=record_keys, lower=lower, extra_tokens=extra_tokens, subword_type=subword_type) vocab = reader.build_vocab() # If we are not using chars, then use 'x' for both input and output preproc_data = baseline.embeddings.load_embeddings( 'x', dsz=d_model, known_vocab=vocab['x'], preserve_vocab_indices=True, embed_type=embed_type) vocabs = preproc_data['vocab'] os.makedirs(basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs, os.path.join(basedir, 'vocabs.json')) embeddings = preproc_data['embeddings'] logger.info("Loaded embeddings") train_set = reader.load(train_file, vocabs) valid_set = reader.load(valid_file, vocabs, distribute=False, shuffle=False) train_loader = DataLoader(train_set, batch_size=batch_size, num_workers=num_train_workers) valid_loader = DataLoader(valid_set, batch_size=batch_size) logger.info("Loaded datasets") logger.info("Using embedding type [%s]", embed_type) if len(rpr_k) == 0 or rpr_k[0] < 1: rpr_k = None elif len(rpr_k) == 1: rpr_k = rpr_k[0] else: rpr_k = rpr_k hps = { "dsz": d_model, "hsz": d_model, "d_ff": d_ff, "dropout": dropout, "num_heads": num_heads, "layers": num_layers, "encoder_type": "transformer", "decoder_type": "transformer", "src_lengths_key": "x_lengths", "d_k": d_k, "layer_drop": layer_drop, "rpr_k": rpr_k, "ra_type": ra_type, "transformer_type": transformer_type } model = TiedEmbeddingsSeq2SeqModel({'x': embeddings}, None, **hps) model.to(device) loss_function = model.create_loss(label_smoothing=label_smoothing) loss_function.to(device) logger.info("Created model and loss") steps_per_epoch = len(train_loader) // num_gpus valid_steps = len(valid_loader) update_on = steps_per_epoch // saves_per_epoch report_on = max(10, update_on) // 10 logger.info( f"Steps per epoch per GPU: {steps_per_epoch}. Saving checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(lr_scheduler, lr, steps_per_epoch, epochs, logger, decay_steps=lr_decay_steps, decay_rate=lr_decay_rate, alpha=lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(warmup_steps, lr=lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=lr) global_step = 0 start_epoch = 0 if restart_from: global_step, start_epoch = reload_from_checkpoint( restart_from, restart_tt, model, steps_per_epoch) logger.info( "Restarting from a previous checkpoint %s.\n\tStarting at global_step=%d, epoch=%d", restart_from, global_step, start_epoch + 1) optimizer = OptimizerManager(model, global_step, optim=optim, lr=lr, lr_function=lr_sched, weight_decay=weight_decay) logger.info("Model has {:,} parameters".format( sum(p.numel() for p in model.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if distributed: model = DistributedDataParallel(model, device_ids=[device], output_device=device) logger.info("Model located on %d", local_rank) model_base = os.path.join(basedir, 'checkpoint') steps = global_step timer = Timer() for epoch in range(start_epoch, epochs): avg_loss = Average('average_train_loss') metrics = {} optimizer.zero_grad() timer.start() model.train() train_itr = iter(train_loader) for i in range(steps_per_epoch): batch = next(train_itr) steps += 1 x, y = batch loss = run_step(x, y, model, loss_function, distributed) loss.backward() avg_loss.update(loss.item()) torch.nn.utils.clip_grad_norm_(model.parameters(), clip) optimizer.step() optimizer.zero_grad() if (i + 1) % report_on == 0: logging.info(avg_loss) if (i + 1) % update_on == 0 and local_rank < 1: elapsed = timer.elapsed(True) logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) logging.info('LR: %f', optimizer.current_lr) save_checkpoint(model, model_base, steps, tick_type='step', save_npz=save_npz) # How much time elapsed in minutes elapsed = timer.elapsed(True) train_avg_loss = avg_loss.avg # This is the average training token-level loss across all machines # This is the token-level training perplexity metrics['train_elapsed_min'] = elapsed metrics['average_train_loss'] = train_avg_loss if local_rank < 1: avg_valid_loss = Average('average_valid_loss') timer.start() model.eval() valid_itr = iter(valid_loader) for j in range(valid_steps): with torch.no_grad(): batch = next(valid_itr) x, y = batch loss = run_step(x, y, model, loss_function, distributed) avg_valid_loss.update(loss.item()) valid_avg_loss = avg_valid_loss.avg elapsed = timer.elapsed(True) metrics['valid_elapsed_min'] = elapsed metrics['average_valid_loss'] = valid_avg_loss logger.info(metrics) save_checkpoint(model, model_base, epoch, tick_type='epoch', save_npz=save_npz)
def run(basedir=None, train_file=None, valid_file=None, dataset_key='tlm', embed_type='default', d_model=512, d_ff=2048, d_k=None, num_heads=8, num_layers=8, num_train_workers=4, nctx=256, file_type='json', batch_size=256, subword_model_file=None, subword_vocab_file=None, dropout=0.1, ffn_pdrop=0.0, layer_drop=0.0, lr_scheduler='cosine', lr_decay_steps=None, lr_decay_rate=None, lr_alpha=0.0, optim='adamw', lr=4.0e-4, clip=1.0, weight_decay=1.0e-2, epochs=32, restart_from=None, restart_tt=None, warmup_steps=10000, saves_per_epoch=10, mlm=True, preprocessed=True, rpr_k=[8], rpr_value_on=False, windowed_ra=False, device="cuda", distributed=False, local_rank=-1, extra_tokens=["[CLS]", "[MASK]"], do_early_stopping=False, model_type='transformer-mlm', modules=[], ra_type=None, transformer_type=None, **kwargs): if basedir is None: basedir = 'lm-{}-bpe-{}'.format(dataset_key, os.getpid()) logging.basicConfig( level=logging.INFO if local_rank in [-1, 0] else logging.WARN) for module in modules: import_user_module(module) num_gpus = get_num_gpus_multiworker() distributed = distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") do_on_demand_masking = mlm and not preprocessed if do_on_demand_masking: logger.info(f"On-demand masking is turned on") if distributed: device, updated_local_rank = init_distributed(local_rank) local_rank = updated_local_rank if file_type == 'tfrecord': reader_type = 'tfrecord' elif preprocessed: reader_type = 'preprocessed' else: reader_type = 'lang' reader = MultiFileDatasetReader(src_nctx=nctx, model_file=subword_model_file, vocab_file=subword_vocab_file, file_type=file_type, reader_type=reader_type, record_keys=['x', 'y'] if mlm else ['x'], extra_tokens=extra_tokens) # This looks a bit funny but the streaming reader ignores our vocab and gives us the one from the subword_model # However, we do need to get counts from our dataset for validation so we can calculate the perplexity vocab = reader.build_vocab([valid_file]) # If we are not using chars, then use 'x' for both input and output preproc_data = baseline.embeddings.load_embeddings( 'x', dsz=d_model, known_vocab=vocab['x'], preserve_vocab_indices=True, embed_type=embed_type) vocabs = preproc_data['vocab'] os.makedirs(basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs, os.path.join(basedir, 'vocabs.json')) embeddings = {'x': preproc_data['embeddings']} logger.info("Loaded embeddings") train_set = reader.load(train_file, vocabs) valid_set = reader.load(valid_file, vocabs, distribute=False, shuffle=False) train_loader = DataLoader(train_set, batch_size=batch_size, num_workers=num_train_workers) valid_loader = DataLoader(valid_set, batch_size=batch_size) logger.info("Loaded datasets") logger.info("Using embedding type [%s]", embed_type) if 'mlm' in model_type: mask_from = vocabs vocab_size = len(mask_from) mask_value = mask_from.get("[MASK]") if mask_value == -1: logger.error( "We could not find a suitable masking token in the vocab") return if len(rpr_k) == 0 or rpr_k[0] < 1: rpr_k = None elif len(rpr_k) == 1: rpr_k = None if rpr_k[0] == 0 else rpr_k[0] if ra_type != None and ra_type != 'shaw' and rpr_k is not None: print( f"Relative attention mismatch. You requested {ra_type} with rpr set. Setting it to 0" ) rpr_k = None model = create_lang_model( embeddings, hsz=d_model, nctx=nctx, # Only for gMLP d_ff=d_ff, tie_weights=True, dropout=dropout, gpu=False, num_heads=num_heads, layers=num_layers, rpr_k=rpr_k, d_k=d_k, ffn_pdrop=ffn_pdrop, windowed_ra=windowed_ra, rpr_value_on=rpr_value_on, layer_drop=layer_drop, model_type=model_type, ra_type=ra_type, transformer_type=transformer_type, src_keys=['x'], tgt_key='x') model.to(device) loss_function = model.create_loss() loss_function.to(device) logger.info("Loaded model and loss") steps_per_epoch = len(train_loader) // num_gpus update_on = steps_per_epoch // saves_per_epoch report_on = max(10, update_on) // 10 logger.info( f"Steps per epoch per GPU: {steps_per_epoch}. Saving checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(lr_scheduler, lr, steps_per_epoch, epochs, logger, decay_steps=lr_decay_steps, decay_rate=lr_decay_rate, alpha=lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(warmup_steps, lr=lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=lr) global_step = 0 start_epoch = 0 if restart_from: if restart_from.endswith('npz'): load_tlm_npz(model, restart_from) else: model.load_state_dict(torch.load(restart_from)) vec = restart_from.split("-") if restart_tt: tick_type = restart_tt else: tick_type = vec[-2] step_num = int(vec[-1].split(".")[0]) if tick_type == 'epoch': start_epoch = step_num global_step = start_epoch * steps_per_epoch elif tick_type == 'step': start_epoch = step_num // steps_per_epoch global_step = step_num else: logger.warning( f"The previous tick was {step_num} but command-line specifies to ignore, setting to 0" ) logger.info( "Restarting from a previous checkpoint %s.\n\tStarting at global_step=%d, epoch=%d", restart_from, global_step, start_epoch + 1) optimizer = OptimizerManager(model, global_step, optim=optim, lr=lr, lr_function=lr_sched, weight_decay=weight_decay) logger.info("Model has {:,} parameters".format( sum(p.numel() for p in model.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if distributed: # This program assume pure data parallelism, each model is on a single gpu # If we wanted to support model and data parallelism we would need to update # the selection of gpus based on rank, it would need to select multiple ids # based on rank, here we select only a single gpu and use it for input and # output. model = DistributedDataParallel(model, device_ids=[device], output_device=device, find_unused_parameters=True) logger.info("Model located on %s", device) model_base = os.path.join(basedir, 'checkpoint') steps = global_step best_valid_loss = np.inf timer = Timer() for epoch in range(start_epoch, epochs): avg_loss = Average('average_train_loss') metrics = {} optimizer.zero_grad() timer.start() model.train() train_itr = iter(train_loader) for i in range(steps_per_epoch): batch = next(train_itr) steps += 1 x, y = batch inputs = x.to(device) labels = y.to(device) if do_on_demand_masking: inputs, labels, _ = on_demand_mlm_masking( inputs, labels, mask_value, vocab_size) inputs = {'x': inputs} labels = labels.contiguous() logits = model(inputs, None)[0].contiguous() if mlm: loss = loss_function(logits, labels) else: shift_logits = logits[:, -1] shift_labels = labels[:, 1:] loss = loss_function(shift_logits, shift_labels) loss.backward() avg_loss.update(loss.item()) torch.nn.utils.clip_grad_norm_(model.parameters(), clip) optimizer.step() optimizer.zero_grad() if (i + 1) % report_on == 0: logging.info(avg_loss) if (i + 1) % update_on == 0 and local_rank < 1: elapsed = timer.elapsed(True) logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) logging.info('LR: %f', optimizer.current_lr) if not do_early_stopping: save_checkpoint(model, model_base, steps, tick_type='step') else: valid_token_loss = validate(model, loss_function, valid_loader, avg_loss, timer, metrics, do_on_demand_masking, mlm, mask_value, vocab_size, device) if valid_token_loss < best_valid_loss: best_valid_loss = valid_token_loss logger.info( f"New best valid loss: {best_valid_loss}. Saving checkpoint..." ) save_checkpoint(model, model_base, steps, tick_type='step') model.train() if not do_early_stopping: _ = validate(model, loss_function, valid_loader, avg_loss, timer, metrics, do_on_demand_masking, mlm, mask_value, vocab_size, device) save_checkpoint(model, model_base, epoch, tick_type='epoch')
def train(): parser = ArgumentParser() parser.add_argument("--basedir", type=str) parser.add_argument("--dataset_key", type=str, default='wikitext-2', help="key from DATASETS global") parser.add_argument("--train_file", type=str, help='Optional file path to use for train file') parser.add_argument("--valid_file", type=str, help='Optional file path to use for valid file') parser.add_argument("--dataset_cache", type=str, default=os.path.expanduser('~/.bl-data'), help="Path or url of the dataset cache") parser.add_argument("--cache_features", type=str2bool, default=True) parser.add_argument( "--embed_type", type=str, default='default', help= "register label of the embeddings, so far support positional or learned-positional" ) parser.add_argument("--d_model", type=int, default=410, help="Model dimension (and embedding dsz)") parser.add_argument("--d_ff", type=int, default=2100, help="FFN dimension") parser.add_argument( "--d_k", type=int, default=None, help="Dimension per head. Use if num_heads=1 to reduce dims") parser.add_argument("--num_heads", type=int, default=10, help="Number of heads") parser.add_argument("--num_layers", type=int, default=16, help="Number of layers") parser.add_argument("--nctx", type=int, default=256, help="Max input length") parser.add_argument("--batch_size", type=int, default=8, help="Batch Size") parser.add_argument("--tokens", choices=["words", "chars", "bpe", "wordpiece"], default="wordpiece", help="What tokens to use") parser.add_argument("--subword_model_file", type=str, help="If using subwords, pass this", default='bert-base-uncased') parser.add_argument( "--subword_vocab_file", type=str, help="If using subwords with separate vocab file, pass here") parser.add_argument( "--subword_special_tokens", type=str, nargs='*', help= "When using wordpiece vectorizer, this list provide special tokens to the never_split " "argument of BertTokenizer. These special tokens should also be in the customized vocab " "file so that they have their indices.") parser.add_argument("--dropout", type=float, default=0.1, help="Dropout") parser.add_argument("--lr_scheduler", type=str, help="The type of learning rate decay scheduler", default='cosine') parser.add_argument("--lr_decay_steps", type=int, help="decay steps of lr scheduler") parser.add_argument("--lr_decay_rate", type=float, help="decay rate of lr scheduler") parser.add_argument("--lr_alpha", type=float, help="parameter alpha for cosine decay scheduler") parser.add_argument("--optim", default="adam", type=str, help="Optimizer to use (defaults to adam)") parser.add_argument("--lr", type=float, default=4.0e-4, help="Learning rate") parser.add_argument("--clip", type=float, default=0.25, help="Clipping gradient norm") parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay") parser.add_argument("--epochs", type=int, default=20, help="Num training epochs") parser.add_argument( "--restart_from", type=str, help="Option allows you to restart from a previous checkpoint") parser.add_argument( "--restart_tt", type=str, help="Optional param for legacy checkpoints (step|epoch)") parser.add_argument("--warmup_steps", type=int, default=1000, help="Num warmup steps") parser.add_argument( "--saves_per_epoch", type=int, default=5, help="The number of checkpoints to save within an epoch") parser.add_argument("--mlm", type=str2bool, default=False, help="Use Masked Language Model (MLM) objective") parser.add_argument( '--rpr_k', help= 'Relative attention positional sizes pass 0 if you dont want relative attention', type=int, default=[8], nargs='+') parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)") parser.add_argument("--distributed", type=str2bool, default=False, help="Are we doing distributed training?") parser.add_argument( "--local_rank", type=int, default=-1, help= "Local rank for distributed training (-1 means use the environment variables to find)" ) parser.add_argument("--chars_per_word", type=int, default=40, help="How many max characters per word") args = parser.parse_args() if args.train_file and not args.valid_file: logger.error( "If you provide a train_file, you must provide a valid_file") return if not args.train_file and args.valid_file: logger.error( "If you provide a valid_file, you must also provide a train_file") return if args.tokens == "chars" and args.mlm: logger.error( "Character composition cannot currently be used with the MLM objective" ) if args.basedir is None: args.basedir = 'transformer-{}-{}-{}'.format(args.dataset_key, args.tokens, os.getpid()) logging.basicConfig( level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) logger.info("Cache directory [%s]", args.dataset_cache) num_gpus = get_num_gpus_multiworker() args.distributed = args.distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") if args.distributed: args.device = init_distributed(args.local_rank) if args.train_file: dataset = { 'train_file': args.train_file, 'valid_file': args.valid_file } else: dataset = DataDownloader(DATASETS[args.dataset_key], args.dataset_cache).download() if args.subword_special_tokens is None: special_tokens = () else: special_tokens = tuple(args.subword_special_tokens) reader = create_reader(args.tokens, args.nctx, args.chars_per_word, args.subword_model_file, args.subword_vocab_file, special_tokens) preproc_data = load_embed_and_vocab(args.tokens, reader, dataset, args.dataset_key, args.embed_type, args.d_model, args.cache_features) vocabs = preproc_data['vocabs'] if args.mlm: mask_from = vocabs['x'] vocab_size = len(mask_from) mask_value = mask_from.get("[MASK]", mask_from.get("<MASK>", -1)) if mask_value == -1: logger.error( "We could not find a suitable masking token in the vocab") return os.makedirs(args.basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs['x'], os.path.join(args.basedir, 'vocabs.json')) embeddings = preproc_data['embeddings'] valid_num_words = preproc_data['valid_num_words'] tgt_key = preproc_data['tgt_key'] logger.info("Loaded embeddings") train_set = load_data_caching(args.tokens, reader, dataset, 'train_file', vocabs, args.cache_features, logger) valid_set = load_data_caching(args.tokens, reader, dataset, 'valid_file', vocabs, args.cache_features, logger) logger.info("valid. tokens [%s], valid. words [%s]", valid_set.tensors[-1].numel(), valid_num_words) train_sampler = torch.utils.data.distributed.DistributedSampler( train_set) if args.distributed else None train_loader = DataLoader(train_set, sampler=train_sampler, batch_size=args.batch_size, shuffle=(not args.distributed)) valid_loader = DataLoader(valid_set, batch_size=args.batch_size, shuffle=False) logger.info("Loaded datasets") if len(args.rpr_k) == 0 or args.rpr_k[0] < 1: rpr_k = None elif len(args.rpr_k) == 1: rpr_k = args.rpr_k[0] else: rpr_k = args.rpr_k TLM = TransformerMaskedLanguageModel if args.mlm else TransformerLanguageModel model = TLM.create(embeddings, hsz=args.d_model, d_ff=args.d_ff, tie_weights=(args.tokens != 'chars'), dropout=args.dropout, gpu=False, num_heads=args.num_heads, layers=args.num_layers, rpr_k=rpr_k, d_k=args.d_k, src_keys=['x'], tgt_key=tgt_key) model.to(args.device) loss_function = model.create_loss() loss_function.to(args.device) logger.info("Loaded model and loss") # in this case (train_loader is not iterator) the division by number of gpus is automatically taken care of by # torch.DataLoader steps_per_epoch = len(train_loader) update_on = steps_per_epoch // args.saves_per_epoch report_on = update_on // 10 logger.info( f"Steps per epoch per GPU: {steps_per_epoch}. Saving a checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(args.lr_scheduler, args.lr, steps_per_epoch, args.epochs, logger, decay_steps=args.lr_decay_steps, decay_rate=args.lr_decay_rate, alpha=args.lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(args.warmup_steps, lr=args.lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=args.lr) global_step = 0 start_epoch = 0 if args.restart_from: model.load_state_dict(torch.load(args.restart_from)) vec = args.restart_from.split("-") if args.restart_tt: tick_type = args.restart_tt else: tick_type = vec[-2] step_num = int(vec[-1].split(".")[0]) if tick_type == 'epoch': start_epoch = step_num global_step = start_epoch * steps_per_epoch else: start_epoch = step_num // steps_per_epoch global_step = step_num logger.info( "Restarting from a previous checkpoint %s.\n\tStarting at global_step=%d, epoch=%d", args.restart_from, global_step, start_epoch + 1) optimizer = OptimizerManager(model, global_step, optim=args.optim, lr=args.lr, lr_function=lr_sched, weight_decay=args.weight_decay) logger.info("Model has {:,} parameters".format( sum(p.numel() for p in model.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if args.distributed: # This program assume pure data parallelism, each model is on a single gpu # If we wanted to support model and data parallelism we would need to update # the selection of gpus based on rank, it would need to select multiple ids # based on rank, here we select only a single gpu and use it for input and # output. model = DistributedDataParallel(model, device_ids=[args.device], output_device=args.device) logger.info("Model located on %s", args.device) # This is the training loop steps = global_step model_base = os.path.join(args.basedir, 'checkpoint') for epoch in range(start_epoch, args.epochs): avg_loss = Average('average_train_loss') metrics = {} optimizer.zero_grad() if args.distributed: train_sampler.set_epoch(epoch) start = time.time() model.train() for i, batch in enumerate(train_loader): steps += 1 x, y = batch inputs = x.to(args.device) labels = y.to(args.device) if args.mlm: inputs, labels, _ = on_demand_mlm_masking( inputs, labels, mask_value, vocab_size) inputs = {'x': inputs} labels = labels.transpose(0, 1).contiguous() logits = model(inputs, None)[0].transpose(0, 1).contiguous() if args.mlm: loss = loss_function(logits, labels) else: shift_logits = logits[:-1] shift_labels = labels[1:] loss = loss_function(shift_logits, shift_labels) loss.backward() avg_loss.update(loss.item()) torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() optimizer.zero_grad() if (i + 1) % report_on == 0: logging.info(avg_loss) if (i + 1) % update_on == 0 and args.local_rank < 1: elapsed = (time.time() - start) / 60 logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) save_checkpoint(model, model_base, steps, tick_type='step') # How much time elapsed in minutes elapsed = (time.time() - start) / 60 train_token_loss = avg_loss.avg # This is the average training token-level loss across all machines # This is the token-level training perplexity train_token_ppl = math.exp(train_token_loss) metrics['train_elapsed_min'] = elapsed metrics['average_train_loss'] = train_token_loss metrics['train_ppl'] = train_token_ppl model_base = os.path.join(args.basedir, 'checkpoint') avg_valid_loss = Average('average_valid_loss') start = time.time() model.eval() for batch in valid_loader: with torch.no_grad(): x, y = batch inputs = x.to(args.device) labels = y.to(args.device) if args.mlm: inputs, labels, _ = on_demand_mlm_masking( inputs, labels, mask_value, vocab_size) inputs = {'x': inputs} labels = labels.transpose(0, 1).contiguous() logits = model(inputs, None)[0].transpose(0, 1).contiguous() if args.mlm: loss = loss_function(logits, labels) else: shift_logits = logits[:-1] shift_labels = labels[1:] loss = loss_function(shift_logits, shift_labels) avg_valid_loss.update(loss.item()) valid_token_loss = avg_valid_loss.avg valid_token_ppl = math.exp(valid_token_loss) elapsed = (time.time() - start) / 60 metrics['valid_elapsed_min'] = elapsed metrics['average_valid_loss'] = valid_token_loss if args.tokens in ['bpe', 'wordpiece']: metrics['valid_token_ppl'] = valid_token_ppl metrics['average_valid_word_ppl'] = math.exp( valid_token_loss * valid_set.tensors[-1].numel() / valid_num_words) else: metrics['average_valid_word_ppl'] = valid_token_ppl logger.info(metrics) if args.local_rank < 1: save_checkpoint(model, model_base, epoch, save_npz=True)
def train(): parser = ArgumentParser() parser.add_argument("--basedir", type=str) parser.add_argument("--train_file", type=str, help='Optional file path to use for train file') parser.add_argument("--valid_file", type=str, help='Optional file path to use for valid file') parser.add_argument("--preprocessed", type=str2bool, default=True, help="Has the data already been preprocessed?") parser.add_argument("--gen_d_model", type=int, default=256, help="Model dimension (and embedding dsz)") parser.add_argument("--gen_d_ff", type=int, default=1024, help="FFN dimension") parser.add_argument( "--gen_d_k", type=int, default=None, help="Dimension per head. Use if num_heads=1 to reduce dims") parser.add_argument("--gen_num_heads", type=int, default=8, help="Number of heads") parser.add_argument("--gen_num_layers", type=int, default=8, help="Number of layers") parser.add_argument("--gen_dropout", type=float, default=0.1, help="Dropout") parser.add_argument( '--gen_rpr_k', help= 'Relative attention positional sizes pass 0 if you dont want relative attention', type=int, default=[8], nargs='+') parser.add_argument("--discrim_d_model", type=int, default=512, help="Model dimension (and embedding dsz)") parser.add_argument("--discrim_d_ff", type=int, default=2048, help="FFN dimension") parser.add_argument( "--discrim_d_k", type=int, default=None, help="Dimension per head. Use if num_heads=1 to reduce dims") parser.add_argument("--discrim_num_heads", type=int, default=8, help="Number of heads") parser.add_argument("--discrim_num_layers", type=int, default=8, help="Number of layers") parser.add_argument("--discrim_dropout", type=float, default=0.1, help="Dropout") parser.add_argument( '--discrim_rpr_k', help= 'Relative attention positional sizes pass 0 if you dont want relative attention', type=int, default=[8], nargs='+') parser.add_argument("--num_train_workers", type=int, default=4, help="Number train workers") parser.add_argument( "--nctx", type=int, default=256, help="Max context length (for both encoder and decoder)") parser.add_argument( "--embed_type", type=str, default='default', choices=["default", "positional", "learned-positional"], help="register label of the embeddings") parser.add_argument( "--pattern", default='*.json', help= "Glob pattern for files, defaults to *.json if preprocessed, *.txt otherwise" ) parser.add_argument("--batch_size", type=int, default=256, help="Batch Size") parser.add_argument("--dataset_key", default="reddit", help="dataset key for basedir") parser.add_argument("--subword_model_file", type=str, required=True) parser.add_argument("--subword_vocab_file", type=str, required=True) parser.add_argument("--lr_scheduler", type=str, default='cosine', help="The type of learning rate decay scheduler") parser.add_argument("--lr_decay_steps", type=int, help="decay steps of lr scheduler") parser.add_argument("--lr_decay_rate", type=float, help="decay rate of lr scheduler") parser.add_argument("--lr_alpha", type=float, help="parameter alpha for cosine decay scheduler") parser.add_argument("--optim", default="adam", type=str, help="Optimizer to use (defaults to adam)") parser.add_argument("--lr", type=float, default=4.0e-4, help="Learning rate") parser.add_argument("--clip", type=float, default=1.0, help="Clipping gradient norm") parser.add_argument("--gen_loss_scale", type=float, default=50.0, help="Scaling for loss function") parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay") parser.add_argument("--epochs", type=int, default=32, help="Num training epochs") parser.add_argument( "--restart_from", type=str, help= "Option allows you to restart from the latest checkpoint in a directory" ) parser.add_argument( "--restart_tt", type=str, choices=['step', 'epoch'], default='step', help="Optional param for legacy checkpoints (step|epoch)") parser.add_argument("--warmup_steps", type=int, default=10000, help="Num warmup steps") parser.add_argument("--saves_per_epoch", type=int, default=100, help="The number of checkpoints to save per epoch") parser.add_argument("--print", type=str2bool, default=True, help="Print some output") parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)") parser.add_argument("--distributed", type=str2bool, default=False, help="Are we doing distributed training?") parser.add_argument( "--local_rank", type=int, default=-1, help= "Local rank for distributed training (-1 means use the environment variables to find)" ) args = parser.parse_args() if args.train_file and not args.valid_file: logger.error( "If you provide a train_file, you must provide a valid_file") return if not args.train_file and args.valid_file: logger.error( "If you provide a valid_file, you must also provide a train_file") return if args.basedir is None: args.basedir = 'gd-{}-bpe-{}'.format(args.dataset_key, os.getpid()) logging.basicConfig( format="%(name)s: %(levelname)s: %(message)s", level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) num_gpus = get_num_gpus_multiworker() args.distributed = args.distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") if args.distributed: args.device, args.local_rank = init_distributed(args.local_rank) if not args.preprocessed: reader_type = "lang" args.pattern = "*.txt" else: reader_type = "preprocessed" reader = MultiFileDatasetReader(args.nctx, args.subword_model_file, args.subword_vocab_file, args.pattern, reader_type=reader_type) # just return the vocab from the BPE vectorizer vocab = reader.build_vocab([]) gen_embed = baseline.embeddings.load_embeddings('x', dsz=args.gen_d_model, known_vocab=vocab['x'], embed_type=args.embed_type) vocabs = gen_embed['vocab'] index2word = revlut(vocabs) discrim_embed = baseline.embeddings.load_embeddings( 'x', dsz=args.discrim_d_model, known_vocab=vocab['x'], embed_type=args.embed_type) os.makedirs(args.basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs, os.path.join(args.basedir, 'vocabs.json')) gen_embeddings = {'x': gen_embed['embeddings']} discrim_embeddings = {'x': discrim_embed['embeddings']} logger.info("Loaded embeddings") train_set = reader.load(args.train_file, vocabs) valid_set = reader.load(args.valid_file, vocabs) train_loader = DataLoader(train_set, batch_size=args.batch_size, num_workers=args.num_train_workers) valid_loader = DataLoader(valid_set, batch_size=args.batch_size) train_steps_per_epoch = len(train_loader) // (args.batch_size * num_gpus) valid_steps_per_epoch = len(valid_loader) // args.batch_size logger.info("Loaded datasets") logger.info("Using embedding type [%s]", args.embed_type) mask_value = vocabs.get("[MASK]", vocabs.get("<MASK>", -1)) if mask_value == -1: logger.error("We could not find a suitable masking token in the vocab") return os.makedirs(args.basedir, exist_ok=True) vocab_size = len(vocabs) if len(args.gen_rpr_k) == 0 or args.gen_rpr_k[0] < 1: gen_rpr_k = None elif len(args.gen_rpr_k) == 1: gen_rpr_k = args.gen_rpr_k[0] else: gen_rpr_k = args.gen_rpr_k if len(args.gen_rpr_k) == 0 or args.discrim_rpr_k[0] < 1: discrim_rpr_k = None elif len(args.discrim_rpr_k) == 1: discrim_rpr_k = args.discrim_rpr_k[0] else: discrim_rpr_k = args.discrim_rpr_k gen_model = TransformerMaskedLanguageModel.create( gen_embeddings, hsz=args.gen_d_model, d_ff=args.gen_d_ff, tie_weights=True, dropout=args.gen_dropout, num_heads=args.gen_num_heads, layers=args.gen_num_layers, rpr_k=gen_rpr_k, d_k=args.gen_d_k, src_keys=['x'], tgt_key='x') discrim_model = TransformerDiscriminator(discrim_embeddings, d_model=args.discrim_d_model, d_ff=args.discrim_d_ff, dropout=args.discrim_dropout, num_heads=args.discrim_num_heads, layers=args.discrim_num_layers, activation='gelu', layer_norm_eps=1.0e-12, rpr_k=discrim_rpr_k, d_k=args.discrim_d_k) gen_model.to(args.device) gen_loss_fn = gen_model.create_loss() discrim_model.to(args.device) discrim_loss_fn = discrim_model.create_loss() logger.info("Loaded model and loss") update_on = train_steps_per_epoch // args.saves_per_epoch report_on = update_on // 10 logger.info( f"Steps per epoch per GPU: {train_steps_per_epoch}. Saving checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(args.lr_scheduler, args.lr, train_steps_per_epoch, args.epochs, logger, decay_steps=args.lr_decay_steps, decay_rate=args.lr_decay_rate, alpha=args.lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(args.warmup_steps, lr=args.lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=args.lr) global_step = 0 start_epoch = 0 if args.restart_from: if not os.path.isdir(args.restart_from): raise Exception( f"Cannot restart from {args.restart_from}, directory not found" ) tick_type = args.restart_tt discrim_latest, step_num = find_latest_checkpoint( args.restart_from, wildcard=f'checkpoint-discrim-{tick_type}') gen_latest, _ = find_latest_checkpoint( args.restart_from, wildcard=f'checkpoint-gen-{tick_type}') discrim_model.load_state_dict(torch.load(discrim_latest)) gen_model.load_state_dict(torch.load(gen_latest)) if tick_type == 'step': start_epoch = step_num // train_steps_per_epoch global_step = step_num else: start_epoch = step_num global_step = train_steps_per_epoch * start_epoch parameters = list(discrim_model.parameters()) + list( gen_model.parameters()) optz = OptimizerManager(parameters, global_step, optim=args.optim, lr=args.lr, lr_function=lr_sched, weight_decay=args.weight_decay) logger.info("Generator has {:,} parameters".format( sum(p.numel() for p in gen_model.parameters() if p.requires_grad))) logger.info("Discriminator has {:,} parameters".format( sum(p.numel() for p in discrim_model.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if args.distributed: # This program assume pure data parallelism, each model is on a single gpu # If we wanted to support model and data parallelism we would need to update # the selection of gpus based on rank, it would need to select multiple ids # based on rank, here we select only a single gpu and use it for input and # output. gen_model = DistributedDataParallel(gen_model, device_ids=[args.device], output_device=args.device) discrim_model = DistributedDataParallel(discrim_model, device_ids=[args.device], output_device=args.device) logger.info("Model located on %s", args.device) # This is the training loop steps = global_step model_base = os.path.join(args.basedir, 'checkpoint') discrim_base = f'{model_base}-discrim' gen_base = f'{model_base}-gen' do_on_demand_masking = not args.preprocessed if do_on_demand_masking: logger.info(f"On-demand masking is turned on") timer = Timer() for epoch in range(start_epoch, args.epochs): gen_model.train() discrim_model.train() avg_gen_loss = Average('average_train_gen_loss') avg_discrim_loss = Average('average_train_discrim_loss') avg_discrim_acc = Average('average_train_discrim_acc') avg_train_loss = Average('average5_train_loss') metrics = {} optz.zero_grad() timer.start() print(f'Starting epoch {epoch + 1}') train_iter = iter(train_loader) valid_iter = iter(valid_loader) for i in range(train_steps_per_epoch): steps += 1 x, y = next(train_iter) do_report = (i + 1) % report_on == 0 and args.print gen_loss_step, discrim_loss_step, acc = gen_vs_discrim( x, y, args.device, gen_model, gen_loss_fn, discrim_model, discrim_loss_fn, mask_value, vocab_size, index2word, do_report, do_on_demand_masking) avg_gen_loss.update(gen_loss_step.item()) total_loss_step = gen_loss_step + args.gen_loss_scale * discrim_loss_step total_loss_step.backward() avg_discrim_loss.update(discrim_loss_step.item()) avg_train_loss.update(total_loss_step.item()) avg_discrim_acc.update(acc) torch.nn.utils.clip_grad_norm_(parameters, args.clip) optz.step() optz.zero_grad() if (i + 1) % report_on == 0: logging.info('Loss g=%f, d=%f total=%f, Per token acc=%f', avg_gen_loss.avg, avg_discrim_loss.avg, avg_train_loss.avg, avg_discrim_acc.avg) if (i + 1) % update_on == 0 and args.local_rank < 1: elapsed = timer.elapsed(True) logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) logging.info('LR: %f', optz.current_lr) save_checkpoint(gen_model, gen_base, steps, tick_type='step') save_checkpoint(discrim_model, discrim_base, steps, tick_type='step') # How much time elapsed in minutes elapsed = timer.elapsed(True) # This is the average training token-level loss across all machines # This is the token-level training perplexity metrics['train_elapsed_min'] = elapsed metrics['average_train_gen_loss'] = avg_gen_loss.avg metrics['average_train_discrim_loss'] = avg_discrim_loss.avg metrics[ 'average_train_discrim_per_token_accuracy'] = avg_discrim_acc.avg metrics['average_train_loss'] = avg_train_loss.avg if args.local_rank < 1: avg_valid_gen_loss = Average('average_valid_gen_loss') avg_valid_discrim_loss = Average('average_valid_discrim_loss') avg_valid_discrim_acc = Average('average_valid_discrim_acc') avg_valid_loss = Average('average_valid_loss') timer.start() gen_model.eval() discrim_model.eval() for i in range(valid_steps_per_epoch): with torch.no_grad(): x, y = next(valid_iter) do_report = (i + 1) % report_on == 0 and args.print gen_loss_step, discrim_loss_step, acc = gen_vs_discrim( x, y, args.device, gen_model, gen_loss_fn, discrim_model, discrim_loss_fn, mask_value, vocab_size, index2word, do_report, do_on_demand_masking) avg_valid_gen_loss.update(gen_loss_step.item()) avg_valid_discrim_acc.update(acc) avg_valid_discrim_loss.update(discrim_loss_step.item()) total_loss_step = gen_loss_step + args.gen_loss_scale * discrim_loss_step avg_valid_loss.update(total_loss_step.item()) elapsed = timer.elapsed(True) metrics['valid_elapsed_min'] = elapsed metrics['average_valid_gen_loss'] = avg_valid_gen_loss.avg metrics['average_valid_discrim_loss'] = avg_valid_discrim_loss.avg metrics[ 'average_valid_discrim_per_token_accuracy'] = avg_valid_discrim_acc.avg metrics['average_valid_loss'] = avg_valid_loss.avg logger.info(metrics) save_checkpoint(discrim_model, discrim_base, epoch, tick_type='epoch', save_npz=True) save_checkpoint(gen_model, gen_base, epoch, tick_type='epoch', save_npz=True)
def train(): parser = ArgumentParser() parser.add_argument("--basedir", type=str) parser.add_argument("--train_file", type=str, required=True, help='File path to use for train file') parser.add_argument("--valid_file", type=str, required=True, help='File path to use for valid file') parser.add_argument("--dataset_key", default="paired", help="dataset key for basedir") parser.add_argument( "--embed_type", type=str, default='default', choices=["default", "positional", "learned-positional"], help="register label of the embeddings") parser.add_argument("--d_model", type=int, default=512, help="Model dimension (and embedding dsz)") parser.add_argument("--d_ff", type=int, default=2048, help="FFN dimension") parser.add_argument( "--d_k", type=int, default=None, help="Dimension per head. Use if num_heads=1 to reduce dims") parser.add_argument("--num_heads", type=int, default=8, help="Number of heads") parser.add_argument("--num_layers", type=int, default=8, help="Number of layers") parser.add_argument("--windowed_ra", type=str2bool, default=False, help="whether prevent attention beyond rpr_k") parser.add_argument("--num_train_workers", type=int, default=4, help="Number train workers") parser.add_argument("--nctx", type=int, default=256, help="Max input length") parser.add_argument("--tgt_nctx", type=int, help="Max output length, default to args.nctx") parser.add_argument("--file_type", default='json', help="Suffix for data") parser.add_argument("--record_keys", default=['x', 'y'], nargs='+') parser.add_argument("--batch_size", type=int, default=256, help="Batch Size") parser.add_argument("--subword_model_file", type=str, help="The BPE model file", required=True) parser.add_argument("--subword_vocab_file", type=str, help="The BPE subword vocab", required=True) parser.add_argument("--dropout", type=float, default=0.1, help="Dropout") parser.add_argument("--lr_scheduler", type=str, default='cosine', help="The type of learning rate decay scheduler") parser.add_argument("--lr_decay_steps", type=int, help="decay steps of lr scheduler") parser.add_argument("--lr_decay_rate", type=float, help="decay rate of lr scheduler") parser.add_argument("--lr_alpha", type=float, help="parameter alpha for cosine decay scheduler") parser.add_argument("--optim", default="adamw", type=str, help="Optimizer to use (defaults to adamw)") parser.add_argument("--lr", type=float, default=4.0e-4, help="Learning rate") parser.add_argument("--clip", type=float, default=1.0, help="Clipping gradient norm") parser.add_argument("--weight_decay", type=float, default=1.0e-2, help="Weight decay") parser.add_argument("--epochs", type=int, default=32, help="Num training epochs") parser.add_argument( "--restart_from", type=str, help="Option allows you to restart from a previous checkpoint") parser.add_argument( "--restart_tt", type=str, help="Optional param for legacy checkpoints (step|epoch)") parser.add_argument("--warmup_steps", type=int, default=10000, help="Num warmup steps") parser.add_argument("--saves_per_epoch", type=int, default=10, help="The number of checkpoints to save per epoch") parser.add_argument("--reduction_d_k", type=int, default=64, help="Dimensions of Key and Query in the single headed" "reduction layers") parser.add_argument( "--reduction_type", type=str, default="2ha", help="If using a dual encoder, specifies the reduction type") parser.add_argument( "--unfreeze_after_step", default=0, type=int, help= "Unfreeze encoders after step, ignored if we dont have a checkpoint") parser.add_argument( "--stacking_layers", type=int, nargs='+', default=[], help="Hidden sizes of the dense stack (ff2 from the convert paper)") parser.add_argument("--layer_drop", type=float, default=0.0, help="LayerDrop to apply") parser.add_argument("--ff_pdrop", type=float, default=0.1, help="Dropout in the dense stack") parser.add_argument("--reader_type", type=str, default='preprocessed', choices=['ntp', 'nsp', 'preprocessed', 'tfrecord']) parser.add_argument( "--model_type", default="dual-encoder", choices=["dual-encoder", "encoder-decoder", "transformer-bow"]) parser.add_argument("--src_begin_tok", type=str, nargs='+', default=[]) parser.add_argument("--src_end_tok", type=str, nargs='+', default=['<EOS>']) parser.add_argument("--tgt_begin_tok", type=str, nargs='+', default=['<GO>']) parser.add_argument("--tgt_end_tok", type=str, nargs='+', default=['<EOS>']) parser.add_argument('--lower', type=baseline.str2bool, default=False) parser.add_argument( "--loss", type=str, default='symmetric', choices=['triplet', 'all', 'all_mean', 'contrastive', 'symmetric']) parser.add_argument( "--learn_temp", type=str2bool, default=True, help= "If 'constrastive' or 'symmetric' loss, should we learn the temperature scaling" ) parser.add_argument( "--init_temp", type=float, help="Initialize the temperature for 'contrastive' or 'symmetric' loss" ) parser.add_argument( '--rpr_k', help= 'Relative attention positional sizes pass 0 if you dont want relative attention', type=int, default=[8], nargs='+') parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device (cuda or cpu)") parser.add_argument("--distributed", type=str2bool, default=False, help="Are we doing distributed training?") parser.add_argument( "--local_rank", type=int, default=-1, help= "Local rank for distributed training (-1 means use the environment variables to find)" ) parser.add_argument("--save_npz", type=str2bool, default=False, help="Whether save npz checkpoint") args = parser.parse_args() if args.basedir is None: args.basedir = '{}-{}-paired-{}-bpe-{}'.format(args.model_type, args.reader_type, args.dataset_key, os.getpid()) logging.basicConfig( level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) num_gpus = get_num_gpus_multiworker() args.distributed = args.distributed or num_gpus > 1 logger.info(f"Using {num_gpus} GPUs in this job.") if args.distributed: args.device, updated_local_rank = init_distributed(args.local_rank) args.local_rank = updated_local_rank if not args.tgt_nctx: args.tgt_nctx = args.nctx reader = MultiFileDatasetReader(args.nctx, args.tgt_nctx, args.src_begin_tok, args.src_end_tok, args.tgt_begin_tok, args.tgt_end_tok, args.subword_model_file, args.subword_vocab_file, args.file_type, reader_type=args.reader_type, record_keys=args.record_keys, lower=args.lower) vocab = reader.build_vocab() # If we are not using chars, then use 'x' for both input and output preproc_data = baseline.embeddings.load_embeddings( 'x', dsz=args.d_model, known_vocab=vocab['x'], preserve_vocab_indices=True, embed_type=args.embed_type) vocabs = preproc_data['vocab'] os.makedirs(args.basedir, exist_ok=True) # We want to make sure to save our input vocab into the basedir for reuse later write_json(vocabs, os.path.join(args.basedir, 'vocabs.json')) embeddings = preproc_data['embeddings'] logger.info("Loaded embeddings") train_set = reader.load(args.train_file, vocabs) valid_set = reader.load(args.valid_file, vocabs, distribute=False, shuffle=False) train_loader = DataLoader(train_set, batch_size=args.batch_size, num_workers=args.num_train_workers) valid_loader = DataLoader(valid_set, batch_size=args.batch_size) logger.info("Loaded datasets") logger.info("Using embedding type [%s]", args.embed_type) if len(args.rpr_k) == 0 or args.rpr_k[0] < 1: rpr_k = None elif len(args.rpr_k) == 1: rpr_k = args.rpr_k[0] else: rpr_k = args.rpr_k model = create_model(embeddings, d_model=args.d_model, d_ff=args.d_ff, dropout=args.dropout, num_heads=args.num_heads, num_layers=args.num_layers, model_type=args.model_type, rpr_k=rpr_k, d_k=args.d_k, reduction_d_k=args.reduction_d_k, stacking_layers=args.stacking_layers, ff_pdrop=args.ff_pdrop, windowed_ra=args.windowed_ra, reduction_type=args.reduction_type, layer_drop=args.layer_drop, logger=logger) model.to(args.device) if args.model_type == 'encoder-decoder': run_step = run_step_s2s else: run_step = run_step_dual logger.info( f"Creating {args.loss}, init temperature: {args.init_temp}, learnable: {args.learn_temp}" ) loss_function = model.create_loss(loss_type=args.loss, init_temp=args.init_temp, learn_temp=args.learn_temp) loss_function.to(args.device) logger.info("Created model and loss") steps_per_epoch = len(train_loader) // num_gpus valid_steps = len(valid_loader) update_on = steps_per_epoch // args.saves_per_epoch report_on = max(10, update_on) // 10 logger.info( f"Steps per epoch per GPU: {steps_per_epoch}. Saving checkpoint every {update_on} steps." ) lr_decay = get_lr_decay(args.lr_scheduler, args.lr, steps_per_epoch, args.epochs, logger, decay_steps=args.lr_decay_steps, decay_rate=args.lr_decay_rate, alpha=args.lr_alpha) linear_warmup = WarmupLinearSchedulerPyTorch(args.warmup_steps, lr=args.lr) lr_sched = CompositeLRScheduler(linear_warmup, lr_decay, lr=args.lr) global_step = 0 start_epoch = 0 if args.restart_from: if args.unfreeze_after_step > 0 and args.model_type == "dual-encoder": logger.info(f"Encoders will be frozen until step %d", args.unfreeze_after_step) global_step, start_epoch = reload_from_checkpoint( args.model_type, args.restart_from, args.restart_tt, model, steps_per_epoch) logger.info( "Restarting from a previous checkpoint %s.\n\tStarting at global_step=%d, epoch=%d", args.restart_from, global_step, start_epoch + 1) target = model if args.model_type == 'encoder-decoder' else loss_function optimizer = OptimizerManager(target, global_step, optim=args.optim, lr=args.lr, lr_function=lr_sched, weight_decay=args.weight_decay) logger.info("Model has {:,} parameters".format( sum(p.numel() for p in target.parameters() if p.requires_grad))) # Prepare model for distributed training if needed if args.distributed: model = DistributedDataParallel(model, device_ids=[args.device], output_device=args.device) logger.info("Model located on %d", args.local_rank) model_base = os.path.join(args.basedir, 'checkpoint') steps = global_step timer = Timer() for epoch in range(start_epoch, args.epochs): avg_loss = Average('average_train_loss') metrics = {} optimizer.zero_grad() timer.start() model.train() train_itr = iter(train_loader) for i in range(steps_per_epoch): batch = next(train_itr) if steps > args.unfreeze_after_step and hasattr( model, 'freeze') and model.freeze: logging.info("Unfreezing encoders at step %d", steps) model.freeze = False steps += 1 x, y = batch loss = run_step(x, y, model, loss_function, args.device, args.distributed) loss.backward() avg_loss.update(loss.item()) torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() optimizer.zero_grad() if (i + 1) % report_on == 0: logging.info(avg_loss) if (i + 1) % update_on == 0 and args.local_rank < 1: elapsed = timer.elapsed(True) logging.info('elapsed time this epoch %d min', elapsed) logging.info('elapsed step time %f steps/min', i / elapsed) logging.info('LR: %f', optimizer.current_lr) save_checkpoint(model, model_base, steps, tick_type='step', save_npz=args.save_npz) # How much time elapsed in minutes elapsed = timer.elapsed(True) train_avg_loss = avg_loss.avg # This is the average training token-level loss across all machines # This is the token-level training perplexity metrics['train_elapsed_min'] = elapsed metrics['average_train_loss'] = train_avg_loss if args.local_rank < 1: avg_valid_loss = Average('average_valid_loss') timer.start() model.eval() valid_itr = iter(valid_loader) for j in range(valid_steps): with torch.no_grad(): batch = next(valid_itr) x, y = batch loss = run_step(x, y, model, loss_function, args.device, args.distributed) avg_valid_loss.update(loss.item()) valid_avg_loss = avg_valid_loss.avg elapsed = timer.elapsed(True) metrics['valid_elapsed_min'] = elapsed metrics['average_valid_loss'] = valid_avg_loss logger.info(metrics) save_checkpoint(model, model_base, epoch, tick_type='epoch', save_npz=args.save_npz)