def prepare_model(args, device): # Prepare model config = modeling.BertConfig.from_json_file(args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - (config.vocab_size % 8) model = modeling.BertForPreTraining(config) criterion = BertPretrainingCriterion(config.vocab_size, args.train_batch_size, args.max_seq_length) model.enable_apex(False) model = bert_model_with_loss(model, criterion) model = ort_supplement.create_ort_trainer(args, device, model) checkpoint = None if not args.resume_from_checkpoint: global_step = 0 else: if args.resume_step == -1 and not args.init_checkpoint: model_names = [ f for f in os.listdir(args.output_dir) if f.endswith(".pt") ] args.resume_step = max([ int(x.split('.pt')[0].split('_')[1].strip()) for x in model_names ]) global_step = args.resume_step if not args.init_checkpoint else 0 if not args.init_checkpoint: checkpoint = torch.load(os.path.join( args.output_dir, "ckpt_{}.pt".format(global_step)), map_location="cpu") else: checkpoint = torch.load(args.init_checkpoint, map_location="cpu") model.load_state_dict(checkpoint['model'], strict=False) if args.phase2 and not args.init_checkpoint: global_step -= args.phase1_end_step if is_main_process(args): print("resume step from ", args.resume_step) return model, checkpoint, global_step
def prepare_model_and_optimizer(args, device): # Prepare model config = modeling.BertConfig.from_json_file(args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - (config.vocab_size % 8) modeling.ACT2FN["bias_gelu"] = modeling.bias_gelu_training model = modeling.BertForPreTraining(config) if args.disable_weight_tying: import torch.nn as nn print ("WARNING!!!!!!! Disabling weight tying for this run") print ("BEFORE ", model.cls.predictions.decoder.weight is model.bert.embeddings.word_embeddings.weight) model.cls.predictions.decoder.weight = torch.nn.Parameter(model.cls.predictions.decoder.weight.clone().detach()) print ("AFTER ", model.cls.predictions.decoder.weight is model.bert.embeddings.word_embeddings.weight) assert (model.cls.predictions.decoder.weight is model.bert.embeddings.word_embeddings.weight) == False checkpoint = None if not args.resume_from_checkpoint: global_step = 0 else: if args.resume_step == -1 and not args.init_checkpoint: model_names = [f for f in os.listdir(args.output_dir) if f.endswith(".pt")] args.resume_step = max([int(x.split('.pt')[0].split('_')[1].strip()) for x in model_names]) global_step = args.resume_step if not args.init_checkpoint else 0 if not args.init_checkpoint: checkpoint = torch.load(os.path.join(args.output_dir, "ckpt_{}.pt".format(global_step)), map_location="cpu") else: checkpoint = torch.load(args.init_checkpoint, map_location="cpu") model.load_state_dict(checkpoint['model'], strict=False) if args.phase2 and not args.init_checkpoint: global_step -= args.phase1_end_step if is_main_process(): print("resume step from ", args.resume_step) model.to(device) param_optimizer = list(model.named_parameters()) no_decay = ['bias', 'gamma', 'beta', 'LayerNorm'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}] optimizer = FusedAdam(optimizer_grouped_parameters, lr=args.learning_rate) lr_scheduler = PolyWarmUpScheduler(optimizer, warmup=args.warmup_proportion, total_steps=args.max_steps, degree=1) if args.fp16: if args.loss_scale == 0: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale="dynamic", cast_model_outputs=torch.float16) else: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale=args.loss_scale, cast_model_outputs=torch.float16) amp._amp_state.loss_scalers[0]._loss_scale = args.init_loss_scale model.checkpoint_activations(args.checkpoint_activations) if args.resume_from_checkpoint: if args.phase2 or args.init_checkpoint: keys = list(checkpoint['optimizer']['state'].keys()) #Override hyperparameters from previous checkpoint for key in keys: checkpoint['optimizer']['state'][key]['step'] = global_step for iter, item in enumerate(checkpoint['optimizer']['param_groups']): checkpoint['optimizer']['param_groups'][iter]['step'] = global_step checkpoint['optimizer']['param_groups'][iter]['t_total'] = args.max_steps checkpoint['optimizer']['param_groups'][iter]['warmup'] = args.warmup_proportion checkpoint['optimizer']['param_groups'][iter]['lr'] = args.learning_rate optimizer.load_state_dict(checkpoint['optimizer']) # , strict=False) # Restore AMP master parameters if args.fp16: optimizer._lazy_init_maybe_master_weights() optimizer._amp_stash.lazy_init_called = True optimizer.load_state_dict(checkpoint['optimizer']) for param, saved_param in zip(amp.master_params(optimizer), checkpoint['master params']): param.data.copy_(saved_param.data) if args.local_rank != -1: if not args.allreduce_post_accumulation: model = DDP(model, message_size=250000000, gradient_predivide_factor=get_world_size()) else: flat_dist_call([param.data for param in model.parameters()], torch.distributed.broadcast, (0,) ) elif args.n_gpu > 1: model = torch.nn.DataParallel(model) criterion = BertPretrainingCriterion(config.vocab_size) if args.disable_weight_tying: # Sanity Check that new param is in optimizer print ("SANITY CHECK OPTIMIZER: ", id(model.module.cls.predictions.decoder.weight) in [id(g) for g in optimizer.param_groups[0]['params']]) assert id(model.module.cls.predictions.decoder.weight) in [id(g) for g in optimizer.param_groups[0]['params']] return model, optimizer, lr_scheduler, checkpoint, global_step, criterion
def prepare_model_and_optimizer(args, device): # Prepare model config = modeling.BertConfig.from_json_file(args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - (config.vocab_size % 8) modeling.ACT2FN["bias_gelu"] = modeling.bias_gelu_training model = modeling.BertForPreTraining(config) checkpoint = None if not args.resume_from_checkpoint: global_step = 0 else: if args.resume_step == -1 and not args.init_checkpoint: model_names = [ f for f in os.listdir(args.output_dir) if f.endswith(".pt") ] args.resume_step = max([ int(x.split('.pt')[0].split('_')[1].strip()) for x in model_names ]) global_step = args.resume_step if not args.init_checkpoint else 0 if not args.init_checkpoint: checkpoint = torch.load(os.path.join( args.output_dir, "ckpt_{}.pt".format(global_step)), map_location="cpu") else: checkpoint = torch.load(args.init_checkpoint, map_location="cpu") model.load_state_dict(checkpoint['model'], strict=False) if args.phase2 and not args.init_checkpoint: global_step -= args.phase1_end_step if is_main_process(): print("resume step from ", args.resume_step) model.to(device) # BERT modeling uses weight sharing between word embedding and prediction decoder. # So make sure the storage is pointing properly even after model is moved to device. if args.use_habana: model.cls.predictions.decoder.weight = model.bert.embeddings.word_embeddings.weight param_optimizer = list(model.named_parameters()) no_decay = ['bias', 'gamma', 'beta', 'LayerNorm'] optimizer_grouped_parameters = [{ 'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01 }, { 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] if args.use_habana: if args.use_fused_lamb: try: from hb_custom import FusedLamb except ImportError: raise ImportError("Please install hbopt.") optimizer = FusedLamb(optimizer_grouped_parameters, lr=args.learning_rate) else: optimizer = NVLAMB(optimizer_grouped_parameters, lr=args.learning_rate) else: if torch.cuda.is_available(): optimizer = FusedLAMB(optimizer_grouped_parameters, lr=args.learning_rate) else: optimizer = NVLAMB(optimizer_grouped_parameters, lr=args.learning_rate) lr_scheduler = PolyWarmUpScheduler(optimizer, warmup=args.warmup_proportion, total_steps=args.max_steps) if args.fp16: if args.loss_scale == 0: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale="dynamic", cast_model_outputs=torch.float16) else: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale=args.loss_scale, cast_model_outputs=torch.float16) amp._amp_state.loss_scalers[0]._loss_scale = args.init_loss_scale model.checkpoint_activations(args.checkpoint_activations) if args.resume_from_checkpoint: if args.phase2 or args.init_checkpoint: keys = list(checkpoint['optimizer']['state'].keys()) #Override hyperparameters from previous checkpoint for key in keys: checkpoint['optimizer']['state'][key]['step'] = global_step for iter, item in enumerate( checkpoint['optimizer']['param_groups']): checkpoint['optimizer']['param_groups'][iter][ 'step'] = global_step checkpoint['optimizer']['param_groups'][iter][ 't_total'] = args.max_steps checkpoint['optimizer']['param_groups'][iter][ 'warmup'] = args.warmup_proportion checkpoint['optimizer']['param_groups'][iter][ 'lr'] = args.learning_rate optimizer.load_state_dict(checkpoint['optimizer']) # , strict=False) # Restore AMP master parameters if args.fp16: optimizer._lazy_init_maybe_master_weights() optimizer._amp_stash.lazy_init_called = True optimizer.load_state_dict(checkpoint['optimizer']) for param, saved_param in zip(amp.master_params(optimizer), checkpoint['master params']): param.data.copy_(saved_param.data) if args.local_rank != -1: if not args.allreduce_post_accumulation: if not args.use_jit_trace: if args.use_habana: model = DDP(model) else: model = DDP(model, message_size=250000000, gradient_predivide_factor=get_world_size()) else: flat_dist_call([param.data for param in model.parameters()], torch.distributed.broadcast, (0, )) elif args.n_pu > 1: model = torch.nn.DataParallel(model) criterion = BertPretrainingCriterion(config.vocab_size) return model, optimizer, lr_scheduler, checkpoint, global_step, criterion
def main(): global timeout_sent args = parse_arguments() random.seed(args.seed + args.local_rank) np.random.seed(args.seed + args.local_rank) torch.manual_seed(args.seed + args.local_rank) torch.cuda.manual_seed(args.seed + args.local_rank) worker_init = WorkerInitObj(args.seed + args.local_rank) device, args = setup_training(args) dllogger.log(step="PARAMETER", data={"Config": [str(args)]}) # Prepare optimizer model, optimizer, lr_scheduler, checkpoint, global_step, criterion = prepare_model_and_optimizer( args, device) gradient_accumulation_steps = torch.tensor( args.gradient_accumulation_steps, dtype=torch.float32).to(device) world_size = torch.tensor(get_world_size(), dtype=torch.float32).to(device) if is_main_process(): dllogger.log(step="PARAMETER", data={"SEED": args.seed}) raw_train_start = None if args.do_train: if is_main_process(): dllogger.log(step="PARAMETER", data={"train_start": True}) dllogger.log(step="PARAMETER", data={"batch_size_per_pu": args.train_batch_size}) dllogger.log(step="PARAMETER", data={"learning_rate": args.learning_rate}) model.train() most_recent_ckpts_paths = [] average_loss = 0.0 # averaged loss every args.log_freq steps epoch = 0 training_steps = 0 model_traced = False if device.type == 'cuda': pool = ProcessPoolExecutor(1) # Note: We loop infinitely over epochs, termination is handled via iteration count while True: thread = None restored_data_loader = None if not args.resume_from_checkpoint or epoch > 0 or ( args.phase2 and global_step < 1) or args.init_checkpoint: files = [ os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if os.path.isfile(os.path.join(args.input_dir, f)) and 'training' in f ] files.sort() num_files = len(files) random.Random(args.seed + epoch).shuffle(files) f_start_id = 0 else: f_start_id = checkpoint['files'][0] files = checkpoint['files'][1:] args.resume_from_checkpoint = False num_files = len(files) # may not exist in all checkpoints epoch = checkpoint.get('epoch', 0) restored_data_loader = checkpoint.get('data_loader', None) shared_file_list = {} if torch.distributed.is_initialized( ) and get_world_size() > num_files: remainder = get_world_size() % num_files data_file = files[(f_start_id * get_world_size() + get_rank() + remainder * f_start_id) % num_files] else: data_file = files[(f_start_id * get_world_size() + get_rank()) % num_files] previous_file = data_file if restored_data_loader is None: use_pin_memory = False if args.no_cuda or args.use_habana else True num_workers = 0 if args.use_habana else 4 train_data = pretraining_dataset(data_file, args.max_predictions_per_seq) train_sampler = RandomSampler(train_data) train_dataloader = DataLoader( train_data, sampler=train_sampler, batch_size=args.train_batch_size * args.n_pu, num_workers=num_workers, worker_init_fn=worker_init, pin_memory=use_pin_memory, drop_last=True) # shared_file_list["0"] = (train_dataloader, data_file) else: train_dataloader = restored_data_loader restored_data_loader = None overflow_buf = None if args.allreduce_post_accumulation: overflow_buf = torch.cuda.IntTensor([0]) for f_id in range(f_start_id + 1, len(files)): if get_world_size() > num_files: data_file = files[(f_id * get_world_size() + get_rank() + remainder * f_id) % num_files] else: data_file = files[(f_id * get_world_size() + get_rank()) % num_files] previous_file = data_file if device.type == 'cuda': dataset_future = pool.submit(create_pretraining_dataset, data_file, args.max_predictions_per_seq, shared_file_list, args, worker_init) train_iter = tqdm(train_dataloader, desc="Iteration", disable=args.disable_progress_bar ) if is_main_process() else train_dataloader if raw_train_start is None: raw_train_start = time.time() for step, batch in enumerate(train_iter): training_steps += 1 position_ids = compute_position_ids(batch[0]) if torch.distributed.is_initialized(): torch.distributed.barrier() if args.use_habana: batch = [t.to(dtype=torch.int32) for t in batch] position_ids = position_ids.to(dtype=torch.int32) position_ids = position_ids.to(device) batch = [t.to(device) for t in batch] input_ids, segment_ids, input_mask, masked_lm_labels, next_sentence_labels = batch if args.use_jit_trace: if model_traced == False: model = torch.jit.trace(model, (input_ids, segment_ids, input_mask, position_ids), check_trace=False) model_traced = True if args.local_rank != -1 and not args.allreduce_post_accumulation: if args.use_habana: model = DDP(model) else: model = DDP(model, message_size=250000000, gradient_predivide_factor= get_world_size()) if args.local_rank != -1 and not args.allreduce_post_accumulation \ and (training_steps % args.gradient_accumulation_steps != 0): with model.no_sync(): prediction_scores, seq_relationship_score = model( input_ids, segment_ids, input_mask, position_ids) else: prediction_scores, seq_relationship_score = model( input_ids, segment_ids, input_mask, position_ids) else: if args.local_rank != -1 and not args.allreduce_post_accumulation \ and (training_steps % args.gradient_accumulation_steps != 0): with model.no_sync(): prediction_scores, seq_relationship_score = model( input_ids=input_ids, token_type_ids=segment_ids, attention_mask=input_mask, position_ids=position_ids) else: prediction_scores, seq_relationship_score = model( input_ids=input_ids, token_type_ids=segment_ids, attention_mask=input_mask, position_ids=position_ids) loss = criterion(prediction_scores, seq_relationship_score, masked_lm_labels, next_sentence_labels) if args.n_pu > 1: loss = loss.mean() # mean() to average on multi-pu. divisor = args.gradient_accumulation_steps if args.gradient_accumulation_steps > 1: if not args.allreduce_post_accumulation: # this division was merged into predivision loss = loss / gradient_accumulation_steps divisor = 1.0 if args.fp16: with amp.scale_loss( loss, optimizer, delay_overflow_check=args. allreduce_post_accumulation) as scaled_loss: scaled_loss.backward() else: loss.backward() average_loss += loss.item() if training_steps % args.gradient_accumulation_steps == 0: lr_scheduler.step() # learning rate warmup global_step = take_optimizer_step( args, optimizer, model, overflow_buf, global_step) if global_step >= args.steps_this_run or timeout_sent: train_time_raw = time.time() - raw_train_start last_num_steps = int( training_steps / args.gradient_accumulation_steps) % args.log_freq last_num_steps = args.log_freq if last_num_steps == 0 else last_num_steps average_loss = average_loss / (last_num_steps * divisor) average_loss = torch.tensor( average_loss, dtype=torch.float32).to(device) if (torch.distributed.is_initialized()): average_loss /= world_size torch.distributed.all_reduce(average_loss) final_loss = average_loss.item() if is_main_process(): dllogger.log(step=( epoch, global_step, ), data={"final_loss": final_loss}) elif training_steps % ( args.log_freq * args.gradient_accumulation_steps) == 0: if is_main_process(): dllogger.log( step=( epoch, global_step, ), data={ "average_loss": average_loss / (args.log_freq * divisor), "step_loss": loss.item() * args.gradient_accumulation_steps / divisor, "learning_rate": optimizer.param_groups[0]['lr'] }) average_loss = 0 if global_step >= args.steps_this_run or training_steps % ( args.num_steps_per_checkpoint * args. gradient_accumulation_steps) == 0 or timeout_sent: if is_main_process() and not args.skip_checkpoint: # Save a trained model dllogger.log(step="PARAMETER", data={"checkpoint_step": global_step}) model_to_save = model.module if hasattr( model, 'module' ) else model # Only save the model it-self if args.resume_step < 0 or not args.phase2: output_save_file = os.path.join( args.output_dir, "ckpt_{}.pt".format(global_step)) else: output_save_file = os.path.join( args.output_dir, "ckpt_{}.pt".format(global_step + args.phase1_end_step)) checkpoint_dict = {} if args.do_train: if args.use_habana: config = modeling.BertConfig.from_json_file( args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - ( config.vocab_size % 8) model_copy = modeling.BertForPreTraining( config) model_copy.load_state_dict( model_to_save.state_dict()) param_groups_copy = optimizer.state_dict( )['param_groups'] state_dict_copy = {} for st_key, st_val in optimizer.state_dict( )['state'].items(): st_val_copy = {} for k, v in st_val.items(): if isinstance(v, torch.Tensor): st_val_copy[k] = v.to('cpu') else: st_val_copy[k] = v state_dict_copy[ st_key] = st_val_copy optim_dict = {} optim_dict['state'] = state_dict_copy optim_dict[ 'param_groups'] = param_groups_copy checkpoint_dict = { 'model': model_copy.state_dict(), 'optimizer': optim_dict, 'files': [f_id] + files, 'epoch': epoch, 'data_loader': None if global_step >= args.max_steps else train_dataloader } elif no_cuda: checkpoint_dict = { 'model': model_to_save.state_dict(), 'optimizer': optimizer.state_dict(), 'files': [f_id] + files, 'epoch': epoch, 'data_loader': None if global_step >= args.max_steps else train_dataloader } else: checkpoint_dict = { 'model': model_to_save.state_dict(), 'optimizer': optimizer.state_dict(), 'master params': list(amp.master_params(optimizer)), 'files': [f_id] + files, 'epoch': epoch, 'data_loader': None if global_step >= args.max_steps else train_dataloader } torch.save(checkpoint_dict, output_save_file) most_recent_ckpts_paths.append( output_save_file) if len(most_recent_ckpts_paths) > 3: ckpt_to_be_removed = most_recent_ckpts_paths.pop( 0) os.remove(ckpt_to_be_removed) # Exiting the training due to hitting max steps, or being sent a # timeout from the cluster scheduler if global_step >= args.steps_this_run or timeout_sent: del train_dataloader # thread.join() return args, final_loss, train_time_raw, global_step del train_dataloader # thread.join() # Make sure pool has finished and switch train_dataloader # NOTE: Will block until complete if device.type == 'cuda': train_dataloader, data_file = dataset_future.result( timeout=None) else: train_dataloader, data_file = create_pretraining_dataset( data_file, args.max_predictions_per_seq, shared_file_list, args, worker_init) epoch += 1
MAX_PREDICTIONS_PER_SEQ = 20 MAX_SEQ_LENGTH = 128 DO_LOWER_CASE = True # LEARNING_RATE = 2e-5 # NUM_TRAIN_STEPS = 1 # NUM_WARMUP_STEPS = 10 # USE_TPU = False # BATCH_SIZE = 1 # load model bert_config = modeling.BertConfig(BERT_CONFIG_FILE) device = torch.device("cpu") model1 = modeling.BertForPreTraining(bert_config) # model2 = modeling.BertForPreTraining(bert_config) model1.load_state_dict(torch.load(INIT_CHECKPOINT_PT, map_location='cpu')) # model1.bert.from_pretrained(INIT_DIRECTORY) model1.to(device) print ('model loaded') #resolve features with open(INPUT_FILE, 'rb') as f: features = pickle.load(f) print ("%d total samples" % len(features)) all_input_ids = torch.tensor([f['input_ids'] for f in features], dtype=torch.long)
def prepare_model_and_optimizer(args, device): # Prepare model config = modeling.BertConfig.from_json_file(args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - (config.vocab_size % 8) modeling.ACT2FN["bias_gelu"] = torch.jit.script( modeling.ACT2FN["bias_gelu"]) model = modeling.BertForPreTraining(config) checkpoint = None if not args.resume_from_checkpoint: global_step = 0 else: if args.resume_step == -1 and not args.init_checkpoint: model_names = [ f for f in os.listdir(args.output_dir) if f.endswith(".pt") ] args.resume_step = max([ int(x.split('.pt')[0].split('_')[1].strip()) for x in model_names ]) global_step = args.resume_step if not args.init_checkpoint else 0 if not args.init_checkpoint: checkpoint = torch.load(os.path.join( args.output_dir, "ckpt_{}.pt".format(global_step)), map_location="cpu") else: checkpoint = torch.load(args.init_checkpoint, map_location="cpu") model.load_state_dict(checkpoint['model'], strict=False) if args.phase2 and not args.init_checkpoint: global_step -= args.phase1_end_step if is_main_process(): print("resume step from ", args.resume_step) model.to(device) param_optimizer = list(model.named_parameters()) no_decay = ['bias', 'gamma', 'beta', 'LayerNorm'] optimizer_grouped_parameters = [{ 'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01 }, { 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] optimizer = FusedLAMB(optimizer_grouped_parameters, lr=args.learning_rate) lr_scheduler = PolyWarmUpScheduler(optimizer, warmup=args.warmup_proportion, total_steps=args.max_steps) if args.fp16: if args.loss_scale == 0: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale="dynamic", cast_model_outputs=torch.float16) else: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale=args.loss_scale, cast_model_outputs=torch.float16) amp._amp_state.loss_scalers[0]._loss_scale = 2**20 if args.resume_from_checkpoint: if args.phase2 or args.init_checkpoint: keys = list(checkpoint['optimizer']['state'].keys()) #Override hyperparameters from previous checkpoint for key in keys: checkpoint['optimizer']['state'][key]['step'] = global_step for iter, item in enumerate( checkpoint['optimizer']['param_groups']): checkpoint['optimizer']['param_groups'][iter][ 'step'] = global_step checkpoint['optimizer']['param_groups'][iter][ 't_total'] = args.max_steps checkpoint['optimizer']['param_groups'][iter][ 'warmup'] = args.warmup_proportion checkpoint['optimizer']['param_groups'][iter][ 'lr'] = args.learning_rate optimizer.load_state_dict(checkpoint['optimizer']) # , strict=False) # Restore AMP master parameters if args.fp16: optimizer._lazy_init_maybe_master_weights() optimizer._amp_stash.lazy_init_called = True optimizer.load_state_dict(checkpoint['optimizer']) for param, saved_param in zip(amp.master_params(optimizer), checkpoint['master params']): param.data.copy_(saved_param.data) if args.local_rank != -1: model = DDP( model, message_size=250000000, gradient_predivide_factor=torch.distributed.get_world_size()) elif args.n_gpu > 1: model = torch.nn.DataParallel(model) criterion = BertPretrainingCriterion(config.vocab_size) return model, optimizer, lr_scheduler, checkpoint, global_step, criterion
def prepare_model_and_optimizer(args, device): # Prepare model config = modeling.BertConfig.from_json_file(args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - (config.vocab_size % 8) if args.use_sequential > 0: config.use_sequential = True else: config.use_sequential = False modeling.ACT2FN["bias_gelu"] = modeling.bias_gelu_training model = modeling.BertForPreTraining(config) model.checkpoint_activations(args.checkpoint_activations) if args.smp > 0: # SMP: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) checkpoint = None if not args.resume_from_checkpoint: global_step = 0 else: if not args.init_checkpoint: if not args.s3_checkpoint_uri: raise ValueError( "Need to set s3_checkpoint_uri, if init_checkpoint not set" ) if smp.local_rank() == 0: sync_s3_checkpoints_to_local(args.output_dir, args.s3_checkpoint_uri) smp.barrier() if args.resume_step == -1 and not args.init_checkpoint: model_names = [ f for f in os.listdir(args.output_dir) if ".pt" in f ] args.resume_step = max([ int(x.split(".pt")[0].split("_")[1].strip()) for x in model_names ]) global_step = args.resume_step if not args.init_checkpoint else 0 # SMP: Load a model that was saved with smp.save if not args.init_checkpoint: checkpoint = smp.load( os.path.join(args.output_dir, "ckpt_{}.pt".format(global_step)), partial=args.partial_checkpoint, ) else: checkpoint = smp.load(args.init_checkpoint) model.load_state_dict(checkpoint["model"], strict=False) if args.phase2 and not args.init_checkpoint: global_step -= args.phase1_end_step if is_main_process(): print("resume step from ", args.resume_step) model.to(device) param_optimizer = list(model.named_parameters()) no_decay = ["bias", "gamma", "beta", "LayerNorm"] optimizer_grouped_parameters = [ { "params": [ p for n, p in param_optimizer if not any(nd in n for nd in no_decay) ], "weight_decay": 0.01, }, { "params": [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], "weight_decay": 0.0, }, ] optimizer = FusedLAMB(optimizer_grouped_parameters, lr=args.learning_rate) if args.smp > 0: # SMP: Use Distributed Optimizer which allows the loading of optimizer state for a distributed model # Also provides APIs to obtain local optimizer state for the current mp_rank. optimizer = smp.DistributedOptimizer(optimizer) lr_scheduler = PolyWarmUpScheduler(optimizer, warmup=args.warmup_proportion, total_steps=args.max_steps) if args.fp16: if args.loss_scale == 0: model, optimizer = amp.initialize( model, optimizer, opt_level="O2", loss_scale="dynamic", cast_model_outputs=torch.float16, ) else: model, optimizer = amp.initialize( model, optimizer, opt_level="O2", loss_scale=args.loss_scale, cast_model_outputs=torch.float16, ) amp._amp_state.loss_scalers[0]._loss_scale = args.init_loss_scale if args.resume_from_checkpoint: if args.phase2 or args.init_checkpoint: keys = list(checkpoint["optimizer"]["state"].keys()) # Override hyperparameters from previous checkpoint for key in keys: checkpoint["optimizer"]["state"][key]["step"] = global_step for iter, item in enumerate( checkpoint["optimizer"]["param_groups"]): checkpoint["optimizer"]["param_groups"][iter][ "step"] = global_step checkpoint["optimizer"]["param_groups"][iter][ "t_total"] = args.max_steps checkpoint["optimizer"]["param_groups"][iter][ "warmup"] = args.warmup_proportion checkpoint["optimizer"]["param_groups"][iter][ "lr"] = args.learning_rate optimizer.load_state_dict(checkpoint["optimizer"]) # , strict=False) # Restore AMP master parameters if args.fp16: optimizer._lazy_init_maybe_master_weights() optimizer._amp_stash.lazy_init_called = True optimizer.load_state_dict(checkpoint["optimizer"]) for param, saved_param in zip(amp.master_params(optimizer), checkpoint["master params"]): param.data.copy_(saved_param.data) # if args.local_rank != -1: # if not args.allreduce_post_accumulation: # model = DDP(model, message_size=250000000, gradient_predivide_factor=get_world_size()) # else: # flat_dist_call([param.data for param in model.parameters()], torch.distributed.broadcast, (0,) ) # elif args.n_gpu > 1: # model = torch.nn.DataParallel(model) criterion = BertPretrainingCriterion(config.vocab_size) return model, optimizer, lr_scheduler, checkpoint, global_step, criterion
def prepare_model_and_optimizer(args, device, sequence_output_is_dense): # Prepare model config = modeling.BertConfig.from_json_file(args.config_file) # Padding for divisibility by 8 if config.vocab_size % 8 != 0: config.vocab_size += 8 - (config.vocab_size % 8) model = modeling.BertForPreTraining(config, sequence_output_is_dense=sequence_output_is_dense) checkpoint = None if not args.resume_from_checkpoint: global_step = 0 else: if args.resume_step == -1 and not args.init_checkpoint: model_names = [f for f in os.listdir(args.output_dir) if f.endswith(".pt")] args.resume_step = max([int(x.split('.pt')[0].split('_')[1].strip()) for x in model_names]) global_step = args.resume_step if not args.init_checkpoint else 0 if not args.init_checkpoint: checkpoint = torch.load(os.path.join(args.output_dir, "ckpt_{}.pt".format(global_step)), map_location=device) else: checkpoint = torch.load(args.init_checkpoint, map_location=device) model.load_state_dict(checkpoint['model'], strict=False) if args.phase2 and not args.init_checkpoint: global_step -= args.phase1_end_step if is_main_process(): print("resume step from ", args.resume_step) model.to(device) # If allreduce_post_accumulation_fp16 is not set, Native AMP Autocast is # used along with FP32 gradient accumulation and all-reduce if args.fp16 and args.allreduce_post_accumulation_fp16: model.half() if not args.disable_jit_fusions : model = torch.jit.script(model) param_optimizer = list(model.named_parameters()) no_decay = ['bias', 'gamma', 'beta', 'LayerNorm'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}] optimizer = FusedLAMBAMP(optimizer_grouped_parameters, lr=args.learning_rate) lr_scheduler = PolyWarmUpScheduler(optimizer, warmup=args.warmup_proportion, total_steps=args.max_steps, base_lr=args.learning_rate, device=device) grad_scaler = torch.cuda.amp.GradScaler(init_scale=args.init_loss_scale, enabled=args.fp16) model.checkpoint_activations(args.checkpoint_activations) if args.resume_from_checkpoint: # For phase2, need to reset the learning rate and step count in the checkpoint if args.phase2 or args.init_checkpoint : for group in checkpoint['optimizer']['param_groups'] : group['step'].zero_() group['lr'].fill_(args.learning_rate) else : if 'grad_scaler' in checkpoint and not args.phase2: grad_scaler.load_state_dict(checkpoint['grad_scaler']) optimizer.load_state_dict(checkpoint['optimizer']) # , strict=False) if args.local_rank != -1: # Cuda Graphs requires that DDP is captured on a side stream # It is important to synchronize the streams after the DDP initialization # so anything after sees properly initialized model weights across GPUs side_stream = torch.cuda.Stream() with torch.cuda.stream(side_stream) : model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank, bucket_cap_mb=torch.cuda.get_device_properties(device).total_memory, gradient_as_bucket_view=True) torch.cuda.current_stream().wait_stream(side_stream) from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import allreduce_hook def scale_by_grad_accum_steps_wrapper(hook: Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]) -> Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]: def scale_by_grad_accum_steps_wrapper_hook( hook_state, bucket: dist.GradBucket ) -> torch.futures.Future[torch.Tensor]: bucket.set_buffer(bucket.buffer().div_(args.gradient_accumulation_steps)) fut = hook(hook_state, bucket) return fut return scale_by_grad_accum_steps_wrapper_hook # With gradient accumulation, the DDP comm hook divides the gradients by the number # gradient accumulation steps if args.gradient_accumulation_steps > 1: model.register_comm_hook(None, scale_by_grad_accum_steps_wrapper(allreduce_hook)) optimizer.setup_fp32_params() criterion = BertPretrainingCriterion(config.vocab_size, sequence_output_is_dense=sequence_output_is_dense) if args.resume_from_checkpoint and args.init_checkpoint: start_epoch = checkpoint['epoch'] else: start_epoch = 0 return model, optimizer, grad_scaler, lr_scheduler, checkpoint, global_step, criterion, start_epoch
def main(argv): if len(argv) > 1: raise app.UsageError("Too many command-line arguments.") config = FLAGS.config input_files = sum([glob.glob(pattern) for pattern in config.input_files], []) assert input_files, "No input files!" print(f"Training with {len(input_files)} input files, including:") print(f" - {input_files[0]}") model = modeling.BertForPreTraining(config=config.model) initial_params = get_initial_params(model, init_checkpoint=config.init_checkpoint) optimizer = create_optimizer(config, initial_params) del initial_params # the optimizer takes ownership of all params output_dir = get_output_dir(config) gfile.makedirs(output_dir) # Restore from a local checkpoint, if one exists. optimizer = checkpoints.restore_checkpoint(output_dir, optimizer) if isinstance(optimizer.state, (list, tuple)): start_step = int(optimizer.state[0].step) else: start_step = int(optimizer.state.step) optimizer = optimizer.replicate() optimizer = training.harmonize_across_hosts(optimizer) data_pipeline = data.PretrainingDataPipeline( sum([glob.glob(pattern) for pattern in config.input_files], []), config.tokenizer, max_seq_length=config.max_seq_length, max_predictions_per_seq=config.max_predictions_per_seq, ) learning_rate_fn = training.create_learning_rate_scheduler( factors="constant * linear_warmup * linear_decay", base_learning_rate=config.learning_rate, warmup_steps=config.num_warmup_steps, steps_per_cycle=config.num_train_steps - config.num_warmup_steps, ) train_history = training.TrainStateHistory(learning_rate_fn) train_state = train_history.initial_state() if config.do_train: train_batch_size = config.train_batch_size if jax.host_count() > 1: assert (train_batch_size % jax.host_count() == 0 ), "train_batch_size must be divisible by number of hosts" train_batch_size = train_batch_size // jax.host_count() train_iter = data_pipeline.get_inputs(batch_size=train_batch_size, training=True) train_step_fn = training.create_train_step( model, compute_pretraining_loss_and_metrics, max_grad_norm=config.max_grad_norm, ) for step, batch in zip(range(start_step, config.num_train_steps), train_iter): optimizer, train_state = train_step_fn(optimizer, batch, train_state) if jax.host_id() == 0 and (step % config.save_checkpoints_steps == 0 or step == config.num_train_steps - 1): checkpoints.save_checkpoint(output_dir, optimizer.unreplicate(), step) config_path = os.path.join(output_dir, "config.json") if not os.path.exists(config_path): with open(config_path, "w") as f: json.dump({"model_type": "bert", **config.model}, f) tokenizer_path = os.path.join(output_dir, "sentencepiece.model") if not os.path.exists(tokenizer_path): shutil.copy(config.tokenizer, tokenizer_path) # With the current Rust data pipeline code, running more than one pipeline # at a time will lead to a hang. A simple workaround is to fully delete the # training pipeline before potentially starting another for evaluation. del train_iter if config.do_eval: eval_iter = data_pipeline.get_inputs(batch_size=config.eval_batch_size) eval_iter = itertools.islice(eval_iter, config.max_eval_steps) eval_fn = training.create_eval_fn(model, compute_pretraining_stats, sample_feature_name="input_ids") eval_stats = eval_fn(optimizer, eval_iter) eval_metrics = { "loss": jnp.mean(eval_stats["loss"]), "masked_lm_loss": jnp.mean(eval_stats["masked_lm_loss"]), "next_sentence_loss": jnp.mean(eval_stats["next_sentence_loss"]), "masked_lm_accuracy": jnp.sum(eval_stats["masked_lm_correct"]) / jnp.sum(eval_stats["masked_lm_total"]), "next_sentence_accuracy": jnp.sum(eval_stats["next_sentence_correct"]) / jnp.sum(eval_stats["next_sentence_total"]), } eval_results = [] for name, val in sorted(eval_metrics.items()): line = f"{name} = {val:.06f}" print(line, flush=True) eval_results.append(line) eval_results_path = os.path.join(output_dir, "eval_results.txt") with gfile.GFile(eval_results_path, "w") as f: for line in eval_results: f.write(line + "\n")