def prepare_model(args, device): config = BertConfig.from_pretrained(args.bert_model, cache_dir=args.cache_dir) # config.num_hidden_layers = 12 if args.force_num_hidden_layers: logger.info("Modifying model config with num_hidden_layers to %d", args.force_num_hidden_layers) config.num_hidden_layers = args.force_num_hidden_layers model = BertForPreTraining(config) if args.init_state_dict is not None: model.load_state_dict(args.init_state_dict) model_desc = bert_model_description(config) lr_scheduler = LinearWarmupLRScheduler(total_steps=int(args.max_steps), warmup=args.warmup_proportion) loss_scaler = amp.DynamicLossScaler() if args.fp16 else None options = orttrainer.ORTTrainerOptions({'batch': { 'gradient_accumulation_steps': args.gradient_accumulation_steps}, 'device': {'id': str(device)}, 'mixed_precision': { 'enabled': args.fp16, 'loss_scaler': loss_scaler}, 'graph_transformer': { 'attn_dropout_recompute': args.attn_dropout_recompute, 'gelu_recompute': args.gelu_recompute, 'transformer_layer_recompute': args.transformer_layer_recompute, }, 'debug': {'deterministic_compute': True, }, 'utils': { 'grad_norm_clip': True}, 'distributed': { 'world_rank': max(0, args.local_rank), 'world_size': args.world_size, 'local_rank': max(0, args.local_rank), 'allreduce_post_accumulation': args.allreduce_post_accumulation, 'deepspeed_zero_optimization': {'stage': args.deepspeed_zero_stage}, 'enable_adasum': False}, 'lr_scheduler': lr_scheduler }) param_optimizer = list(model.named_parameters()) no_decay_keys = ["bias", "gamma", "beta", "LayerNorm"] params = [{ 'params': [n for n, p in param_optimizer if any(no_decay_key in n for no_decay_key in no_decay_keys)], "alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6}, { 'params': [n for n, p in param_optimizer if not any(no_decay_key in n for no_decay_key in no_decay_keys)], "alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6}] optim_config = optim.AdamConfig(params=params, lr=2e-5, do_bias_correction=True) model = orttrainer.ORTTrainer(model, model_desc, optim_config, options=options) return model
def convert_ckpt_compatible(ckpt_path, config_path): ckpt = torch.load(ckpt_path, map_location='cpu') keys = list(ckpt.keys()) for key in keys: if 'LayerNorm' in key: if 'gamma' in key: ckpt[key.replace('gamma', 'weight')] = ckpt.pop(key) else: ckpt[key.replace('beta', 'bias')] = ckpt.pop(key) model_config = BertConfig.from_json_file(config_path) model = BertForPreTraining(model_config) model.load_state_dict(ckpt) new_ckpt = model.bert.state_dict() return new_ckpt
def convert_multibert_checkpoint_to_pytorch(tf_checkpoint_path, config_path, save_path): tf_path = os.path.abspath(tf_checkpoint_path) logger.info(f"Converting TensorFlow checkpoint from {tf_path}") # Load weights from TF model init_vars = tf.train.list_variables(tf_path) names = [] arrays = [] config = BertConfig.from_pretrained(config_path) model = BertForPreTraining(config) layer_nums = [] for full_name, shape in init_vars: array = tf.train.load_variable(tf_path, full_name) names.append(full_name) split_names = full_name.split("/") for name in split_names: if name.startswith("layer_"): layer_nums.append(int(name.split("_")[-1])) arrays.append(array) logger.info(f"Read a total of {len(arrays):,} layers") name_to_array = dict(zip(names, arrays)) # Check that number of layers match assert config.num_hidden_layers == len(list(set(layer_nums))) state_dict = model.state_dict() # Need to do this explicitly as it is a buffer position_ids = state_dict["bert.embeddings.position_ids"] new_state_dict = {"bert.embeddings.position_ids": position_ids} # Encoder Layers for weight_name in names: pt_weight_name = weight_name.replace("kernel", "weight").replace("gamma", "weight").replace("beta", "bias") name_split = pt_weight_name.split("/") for name_idx, name in enumerate(name_split): if name.startswith("layer_"): name_split[name_idx] = name.replace("_", ".") if name_split[-1].endswith("embeddings"): name_split.append("weight") if name_split[0] == "cls": if name_split[-1] == "output_bias": name_split[-1] = "bias" if name_split[-1] == "output_weights": name_split[-1] = "weight" if name_split[-1] == "weight" and name_split[-2] == "dense": name_to_array[weight_name] = name_to_array[weight_name].T pt_weight_name = ".".join(name_split) new_state_dict[pt_weight_name] = torch.from_numpy(name_to_array[weight_name]) new_state_dict["cls.predictions.decoder.weight"] = new_state_dict["bert.embeddings.word_embeddings.weight"].clone() new_state_dict["cls.predictions.decoder.bias"] = new_state_dict["cls.predictions.bias"].clone().T # Load State Dict model.load_state_dict(new_state_dict) # Save PreTrained logger.info(f"Saving pretrained model to {save_path}") model.save_pretrained(save_path) return model
def train(args): if not os.path.exists(args.save_dir): os.mkdir(args.save_dir) if args.gpu != '-1' and torch.cuda.is_available(): device = torch.device('cuda') torch.cuda.set_rng_state(torch.cuda.get_rng_state()) torch.backends.cudnn.deterministic = True else: device = torch.device('cpu') config = { 'train': { 'unchanged_variable_weight': 0.1, 'buffer_size': 5000 }, 'encoder': { 'type': 'SequentialEncoder' }, 'data': { 'vocab_file': 'data/vocab.bpe10000/vocab' } } train_set = Dataset('data/preprocessed_data/train-shard-*.tar') dev_set = Dataset('data/preprocessed_data/dev.tar') vocab = Vocab.load('data/vocab.bpe10000/vocab') if args.decoder: vocab_size = len(vocab.all_subtokens) + 1 else: vocab_size = len(vocab.source_tokens) + 1 max_iters = args.max_iters lr = args.lr warm_up = args.warm_up batch_size = 4096 effective_batch_size = args.batch_size max_embeds = 1000 if args.decoder else 512 bert_config = BertConfig(vocab_size=vocab_size, max_position_embeddings=max_embeds, num_hidden_layers=6, hidden_size=256, num_attention_heads=4) model = BertForPreTraining(bert_config) if args.restore: state_dict = torch.load(os.path.join(args.save_dir, args.res_name)) model.load_state_dict(state_dict['model']) batch_count = state_dict['step'] epoch = state_dict['epoch'] model.train() model.to(device) if len(args.gpu) > 1 and device == torch.device('cuda'): model = nn.DataParallel(model) def lr_func(step): if step > warm_up: return (max_iters - step) / (max_iters - warm_up) else: return (step / warm_up) optimizer = torch.optim.Adam(model.parameters(), lr=lr, eps=1e-6, weight_decay=0.01) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_func, last_epoch=-1) loss_fn = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none') if args.restore: optimizer.load_state_dict(state_dict['optim']) scheduler.load_state_dict(state_dict['scheduler']) batch_count = 0 epoch = 0 cum_loss = 0.0 while True: # load training dataset, which is a collection of ASTs and maps of gold-standard renamings train_set_iter = train_set.batch_iterator( batch_size=batch_size, return_examples=False, config=config, progress=True, train=True, max_seq_len=512, num_readers=args.num_readers, num_batchers=args.num_batchers) epoch += 1 print("Epoch {}".format(epoch)) loss = 0 num_seq = 0 optimizer.zero_grad() for batch in train_set_iter: if args.decoder: input_ids = batch.tensor_dict['prediction_target'][ 'src_with_true_var_names'] else: input_ids = batch.tensor_dict['src_code_tokens'] attention_mask = torch.ones_like(input_ids) attention_mask[input_ids == 0] = 0.0 assert torch.max(input_ids) < vocab_size assert torch.min(input_ids) >= 0 if input_ids.shape[0] > max_embeds: print( "Warning - length {} is greater than max length {}. Skipping." .format(input_ids.shape[0], max_embeds)) continue input_ids, labels = mask_tokens(inputs=input_ids, mask_token_id=vocab_size - 1, vocab_size=vocab_size, mlm_probability=0.15) input_ids[attention_mask == 0] = 0 labels[attention_mask == 0] = -100 if torch.cuda.is_available(): input_ids = input_ids.cuda() labels = labels.cuda() attention_mask = attention_mask.cuda() outputs = model(input_ids=input_ids, attention_mask=attention_mask, masked_lm_labels=labels) unreduced_loss = loss_fn( outputs[0].view(-1, bert_config.vocab_size), labels.view(-1)).reshape(labels.shape) / ( torch.sum(labels != -100, axis=1).unsqueeze(1) + 1e-7) loss += unreduced_loss.sum() num_seq += input_ids.shape[0] if num_seq > effective_batch_size: batch_count += 1 loss /= num_seq cum_loss += loss.item() if batch_count % 20 == 0: print("{} batches, Loss : {:.4}, LR : {:.6}".format( batch_count, cum_loss / 20, scheduler.get_lr()[0])) cum_loss = 0.0 if batch_count % 10000 == 0: fname1 = os.path.join( args.save_dir, 'bert_{}_step_{}.pth'.format( ('decoder' if args.decoder else 'encoder'), batch_count)) fname2 = os.path.join( args.save_dir, 'bert_{}.pth'.format( ('decoder' if args.decoder else 'encoder'), batch_count)) state = { 'epoch': epoch, 'step': batch_count, 'model': model.module.state_dict(), 'optim': optimizer.state_dict(), 'scheduler': scheduler.state_dict() } torch.save(state, fname1) torch.save(state, fname2) print("Saved file to path {}".format(fname1)) print("Saved file to path {}".format(fname2)) loss.backward() optimizer.step() scheduler.step() optimizer.zero_grad() loss = 0 num_seq = 0 if batch_count == max_iters: print(f'[Learner] Reached max iters', file=sys.stderr) exit() print("Max_len = {}".format(max_len)) break