def train(args, model, tokenizer, query_cache, passage_cache):
""" Train the model """
logger.info("Training/evaluation parameters %s", args)
tb_writer = None
if is_first_worker():
tb_writer = SummaryWriter(log_dir=args.log_dir)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
real_batch_size = args.train_batch_size * args.gradient_accumulation_steps * \
(torch.distributed.get_world_size() if args.local_rank != -1 else 1)
optimizer_grouped_parameters = []
layer_optim_params = set()
for layer_name in [
"roberta.embeddings", "score_out", "downsample1", "downsample2",
"downsample3"
]:
layer = getattr_recursive(model, layer_name)
if layer is not None:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
if getattr_recursive(model, "roberta.encoder.layer") is not None:
for layer in model.roberta.encoder.layer:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
optimizer_grouped_parameters.append({
"params":
[p for p in model.parameters() if p not in layer_optim_params]
})
if args.optimizer.lower() == "lamb":
optimizer = Lamb(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
elif args.optimizer.lower() == "adamw":
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
else:
raise Exception(
"optimizer {0} not recognized! Can only be lamb or adamW".format(
args.optimizer))
def optimizer_to(optim, device):
for param in optim.state.values():
# Not sure there are any global tensors in the state dict
if isinstance(param, torch.Tensor):
param.data = param.data.to(device)
if param._grad is not None:
param._grad.data = param._grad.data.to(device)
elif isinstance(param, dict):
for subparam in param.values():
if isinstance(subparam, torch.Tensor):
subparam.data = subparam.data.to(device)
if subparam._grad is not None:
subparam._grad.data = subparam._grad.data.to(
device)
torch.cuda.empty_cache()
# Check if saved optimizer or scheduler states exist
if os.path.isfile(
os.path.join(args.model_name_or_path,
"optimizer.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
optimizer.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "optimizer.pt"),
map_location='cpu'))
optimizer_to(optimizer, args.device)
model.to(args.device)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True,
)
# Train
logger.info("***** Running training *****")
logger.info(" Max steps = %d", args.max_steps)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
global_step = 0
# Check if continuing training from a checkpoint
if os.path.exists(args.model_name_or_path):
# set global_step to gobal_step of last saved checkpoint from model
# path
if "-" in args.model_name_or_path:
global_step = int(
args.model_name_or_path.split("-")[-1].split("/")[0])
else:
global_step = 0
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from global step %d", global_step)
is_hypersphere_training = (args.hyper_align_weight > 0
or args.hyper_unif_weight > 0)
if is_hypersphere_training:
logger.info(
f"training with hypersphere property regularization, align weight {args.hyper_align_weight}, unif weight {args.hyper_unif_weight}"
)
if not args.dual_training:
args.dual_loss_weight = 0.0
tr_loss_dict = {}
model.zero_grad()
model.train()
set_seed(args) # Added here for reproductibility
last_ann_no = -1
train_dataloader = None
train_dataloader_iter = None
# dev_ndcg = 0
step = 0
if args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=args.max_steps)
if os.path.isfile(os.path.join(
args.model_name_or_path,
"scheduler.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
scheduler.load_state_dict(
torch.load(
os.path.join(args.model_name_or_path, "scheduler.pt")))
while global_step < args.max_steps:
if step % args.gradient_accumulation_steps == 0 and global_step % args.logging_steps == 0:
# check if new ann training data is availabe
ann_no, ann_path, ndcg_json = get_latest_ann_data(
args.ann_dir, is_grouped=(args.grouping_ann_data > 0))
if ann_path is not None and ann_no != last_ann_no:
logger.info("Training on new add data at %s", ann_path)
time.sleep(30) # wait until transmission finished
with open(ann_path, 'r') as f:
ann_training_data = f.readlines()
# marcodev_ndcg = ndcg_json['marcodev_ndcg']
logging.info(f"loading:\n{ndcg_json}")
ann_checkpoint_path = ndcg_json['checkpoint']
ann_checkpoint_no = get_checkpoint_no(ann_checkpoint_path)
aligned_size = (len(ann_training_data) //
args.world_size) * args.world_size
ann_training_data = ann_training_data[:aligned_size]
logger.info(
"Total ann queries: %d",
len(ann_training_data) if args.grouping_ann_data < 0 else
len(ann_training_data) * args.grouping_ann_data)
if args.grouping_ann_data > 0:
if args.polling_loaded_data_batch_from_group:
train_dataset = StreamingDataset(
ann_training_data,
GetGroupedTrainingDataProcessingFn_polling(
args, query_cache, passage_cache))
else:
train_dataset = StreamingDataset(
ann_training_data,
GetGroupedTrainingDataProcessingFn_origin(
args, query_cache, passage_cache))
else:
if not args.dual_training:
if args.triplet:
train_dataset = StreamingDataset(
ann_training_data,
GetTripletTrainingDataProcessingFn(
args, query_cache, passage_cache))
else:
train_dataset = StreamingDataset(
ann_training_data,
GetTrainingDataProcessingFn(
args, query_cache, passage_cache))
else:
# return quadruplet
train_dataset = StreamingDataset(
ann_training_data,
GetQuadrapuletTrainingDataProcessingFn(
args, query_cache, passage_cache))
train_dataloader = DataLoader(train_dataset,
batch_size=args.train_batch_size)
train_dataloader_iter = iter(train_dataloader)
# re-warmup
if not args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=len(ann_training_data)
if args.grouping_ann_data < 0 else
len(ann_training_data) * args.grouping_ann_data)
if args.local_rank != -1:
dist.barrier()
if is_first_worker():
# add ndcg at checkpoint step used instead of current step
for key in ndcg_json:
if "marcodev" in key:
tb_writer.add_scalar(key, ndcg_json[key],
ann_checkpoint_no)
if 'trec2019_ndcg' in ndcg_json:
tb_writer.add_scalar("trec2019_ndcg",
ndcg_json['trec2019_ndcg'],
ann_checkpoint_no)
if last_ann_no != -1:
tb_writer.add_scalar("epoch", last_ann_no,
global_step - 1)
tb_writer.add_scalar("epoch", ann_no, global_step)
last_ann_no = ann_no
try:
batch = next(train_dataloader_iter)
except StopIteration:
logger.info("Finished iterating current dataset, begin reiterate")
train_dataloader_iter = iter(train_dataloader)
batch = next(train_dataloader_iter)
# original way
if args.grouping_ann_data <= 0:
batch = tuple(t.to(args.device) for t in batch)
if args.triplet:
inputs = {
"query_ids": batch[0].long(),
"attention_mask_q": batch[1].long(),
"input_ids_a": batch[3].long(),
"attention_mask_a": batch[4].long(),
"input_ids_b": batch[6].long(),
"attention_mask_b": batch[7].long()
}
if args.dual_training:
inputs["neg_query_ids"] = batch[9].long()
inputs["attention_mask_neg_query"] = batch[10].long()
inputs["prime_loss_weight"] = args.prime_loss_weight
inputs["dual_loss_weight"] = args.dual_loss_weight
else:
inputs = {
"input_ids_a": batch[0].long(),
"attention_mask_a": batch[1].long(),
"input_ids_b": batch[3].long(),
"attention_mask_b": batch[4].long(),
"labels": batch[6]
}
else:
# the default collate_fn will convert item["q_pos"] into batch format ... I guess
inputs = {
"query_ids": batch["q_pos"][0].to(args.device).long(),
"attention_mask_q": batch["q_pos"][1].to(args.device).long(),
"input_ids_a": batch["d_pos"][0].to(args.device).long(),
"attention_mask_a": batch["d_pos"][1].to(args.device).long(),
"input_ids_b": batch["d_neg"][0].to(args.device).long(),
"attention_mask_b": batch["d_neg"][1].to(args.device).long(),
}
if args.dual_training:
inputs["neg_query_ids"] = batch["q_neg"][0].to(
args.device).long()
inputs["attention_mask_neg_query"] = batch["q_neg"][1].to(
args.device).long()
inputs["prime_loss_weight"] = args.prime_loss_weight
inputs["dual_loss_weight"] = args.dual_loss_weight
inputs["temperature"] = args.temperature
inputs["loss_objective"] = args.loss_objective_function
if is_hypersphere_training:
inputs["alignment_weight"] = args.hyper_align_weight
inputs["uniformity_weight"] = args.hyper_unif_weight
step += 1
if args.local_rank != -1:
# sync gradients only at gradient accumulation step
if step % args.gradient_accumulation_steps == 0:
outputs = model(**inputs)
else:
with model.no_sync():
outputs = model(**inputs)
else:
outputs = model(**inputs)
# model outputs are always tuple in transformers (see doc)
loss = outputs[0]
loss_item_dict = outputs[1]
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
for k in loss_item_dict:
loss_item_dict[k] = loss_item_dict[k].mean()
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
for k in loss_item_dict:
loss_item_dict[
k] = loss_item_dict[k] / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
if args.local_rank != -1:
if step % args.gradient_accumulation_steps == 0:
loss.backward()
else:
with model.no_sync():
loss.backward()
else:
loss.backward()
def incremental_tr_loss(tr_loss_dict, loss_item_dict, total_loss):
for k in loss_item_dict:
if k not in tr_loss_dict:
tr_loss_dict[k] = loss_item_dict[k].item()
else:
tr_loss_dict[k] += loss_item_dict[k].item()
if "loss_total" not in tr_loss_dict:
tr_loss_dict["loss_total"] = total_loss.item()
else:
tr_loss_dict["loss_total"] += total_loss.item()
return tr_loss_dict
tr_loss_dict = incremental_tr_loss(tr_loss_dict,
loss_item_dict,
total_loss=loss)
if step % args.gradient_accumulation_steps == 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
args.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
logs = {}
learning_rate_scalar = scheduler.get_lr()[0]
logs["learning_rate"] = learning_rate_scalar
for k in tr_loss_dict:
logs[k] = tr_loss_dict[k] / args.logging_steps
tr_loss_dict = {}
if is_first_worker():
for key, value in logs.items():
tb_writer.add_scalar(key, value, global_step)
logger.info(json.dumps({**logs, **{"step": global_step}}))
if is_first_worker(
) and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(args.output_dir,
"checkpoint-{}".format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
torch.save(args, os.path.join(output_dir, "training_args.bin"))
logger.info("Saving model checkpoint to %s", output_dir)
torch.save(optimizer.state_dict(),
os.path.join(output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(output_dir, "scheduler.pt"))
logger.info("Saving optimizer and scheduler states to %s",
output_dir)
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
tb_writer.close()
return global_step
def train(args, model, tokenizer, f, train_fn):
""" Train the model """
tb_writer = None
if is_first_worker():
tb_writer = SummaryWriter(log_dir=args.log_dir)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
real_batch_size = args.train_batch_size * args.gradient_accumulation_steps * \
(torch.distributed.get_world_size() if args.local_rank != -1 else 1)
if args.max_steps > 0:
t_total = args.max_steps
else:
t_total = args.expected_train_size // real_batch_size * args.num_train_epochs
print('????t_total', t_total)
# layerwise optimization for lamb
optimizer_grouped_parameters = []
layer_optim_params = set()
for layer_name in [
"roberta.embeddings", "score_out", "downsample1", "downsample2",
"downsample3", "embeddingHead"
]:
layer = getattr_recursive(model, layer_name)
if layer is not None:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
if getattr_recursive(model, "roberta.encoder.layer") is not None:
for layer in model.roberta.encoder.layer:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
# if getattr_recursive(model, "roberta.encoder.layer") is not None:
# for layer in model.roberta.encoder.layer:
# optimizer_grouped_parameters.append({"params": layer.parameters()})
# for p in layer.parameters():
# layer_optim_params.add(p)
optimizer_grouped_parameters.append({
"params":
[p for p in model.parameters() if p not in layer_optim_params]
})
if args.optimizer.lower() == "lamb":
optimizer = Lamb(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
elif args.optimizer.lower() == "adamw":
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
else:
raise Exception(
"optimizer {0} not recognized! Can only be lamb or adamW".format(
args.optimizer))
if args.scheduler.lower() == "linear":
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=t_total)
elif args.scheduler.lower() == "cosine":
scheduler = CosineAnnealingLR(optimizer, t_total, 1e-8)
else:
raise Exception(
"Scheduler {0} not recognized! Can only be linear or cosine".
format(args.scheduler))
# Check if saved optimizer or scheduler states exist
if os.path.isfile(os.path.join(
args.model_name_or_path, "optimizer.pt")) and os.path.isfile(
os.path.join(
args.model_name_or_path,
"scheduler.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
optimizer.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
scheduler.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True,
)
# Train!
logger.info("***** Running training *****")
logger.info(" Num Epochs = %d", args.num_train_epochs)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
epochs_trained = 0
steps_trained_in_current_epoch = 0
# Check if continuing training from a checkpoint
if os.path.exists(args.model_name_or_path):
# set global_step to gobal_step of last saved checkpoint from model path
try:
global_step = int(
args.model_name_or_path.split("-")[-1].split("/")[0])
epochs_trained = global_step // (args.expected_train_size //
args.gradient_accumulation_steps)
steps_trained_in_current_epoch = global_step % (
args.expected_train_size // args.gradient_accumulation_steps)
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from epoch %d", epochs_trained)
logger.info(" Continuing training from global step %d",
global_step)
logger.info(" Will skip the first %d steps in the first epoch",
steps_trained_in_current_epoch)
except:
logger.info(" Start training from a pretrained model")
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
train_iterator = trange(
epochs_trained,
int(args.num_train_epochs),
desc="Epoch",
disable=args.local_rank not in [-1, 0],
)
set_seed(args) # Added here for reproductibility
#print('???',args.local_rank)
#assert 1==0, "?????"
for m_epoch in train_iterator:
f.seek(0)
sds = StreamingDataset(f, train_fn)
epoch_iterator = DataLoader(sds,
batch_size=args.per_gpu_train_batch_size,
num_workers=1)
for step, batch in tqdm(enumerate(epoch_iterator),
desc="Iteration",
disable=args.local_rank not in [-1, 0]):
#assert 1==0, "?????"
# Skip past any already trained steps if resuming training
#assert 1==0, steps_trained_in_current_epoch
if not args.reset_iter:
if steps_trained_in_current_epoch > 0:
steps_trained_in_current_epoch -= 1
continue
model.train()
batch = tuple(t.to(args.device).long() for t in batch)
# print('???',*batch)
# assert 1==0, "!!!!!"
if (step + 1) % args.gradient_accumulation_steps == 0:
outputs = model(*batch)
else:
with model.no_sync():
# print('???',*batch)
# assert 1==0
outputs = model(*batch)
# model outputs are always tuple in transformers (see doc)
loss = outputs[0]
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
if (step + 1) % args.gradient_accumulation_steps == 0:
loss.backward()
else:
with model.no_sync():
loss.backward()
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if is_first_worker(
) and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(
args.output_dir, "checkpoint-{}".format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if 'fairseq' not in args.train_model_type:
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
else:
torch.save(model.state_dict(),
os.path.join(output_dir, 'model.pt'))
torch.save(args,
os.path.join(output_dir, "training_args.bin"))
logger.info("Saving model checkpoint to %s", output_dir)
torch.save(optimizer.state_dict(),
os.path.join(output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(output_dir, "scheduler.pt"))
logger.info("Saving optimizer and scheduler states to %s",
output_dir)
dist.barrier()
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
logs = {}
if args.evaluate_during_training and global_step % (
args.logging_steps_per_eval *
args.logging_steps) == 0:
model.eval()
reranking_mrr, full_ranking_mrr = passage_dist_eval(
args, model, tokenizer)
if is_first_worker():
print("Reranking/Full ranking mrr: {0}/{1}".format(
str(reranking_mrr), str(full_ranking_mrr)))
mrr_dict = {
"reranking": float(reranking_mrr),
"full_raking": float(full_ranking_mrr)
}
tb_writer.add_scalars("mrr", mrr_dict, global_step)
print(args.output_dir)
loss_scalar = (tr_loss - logging_loss) / args.logging_steps
learning_rate_scalar = scheduler.get_lr()[0]
logs["learning_rate"] = learning_rate_scalar
logs["loss"] = loss_scalar
logging_loss = tr_loss
if is_first_worker():
for key, value in logs.items():
print(key, type(value))
tb_writer.add_scalar(key, value, global_step)
tb_writer.add_scalar("epoch", m_epoch, global_step)
print(json.dumps({**logs, **{"step": global_step}}))
dist.barrier()
if args.max_steps > 0 and global_step > args.max_steps:
train_iterator.close()
break
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
tb_writer.close()
return global_step, tr_loss / global_step
def train(args, model, tokenizer, query_cache, passage_cache):
""" Train the model """
logger.info("Training/evaluation parameters %s", args)
tb_writer = None
if is_first_worker():
tb_writer = SummaryWriter(log_dir=args.log_dir)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
real_batch_size = args.train_batch_size * args.gradient_accumulation_steps * \
(torch.distributed.get_world_size() if args.local_rank != -1 else 1)
optimizer_grouped_parameters = []
layer_optim_params = set()
for layer_name in [
"roberta.embeddings", "score_out", "downsample1", "downsample2",
"downsample3"
]:
layer = getattr_recursive(model, layer_name)
if layer is not None:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
if getattr_recursive(model, "roberta.encoder.layer") is not None:
for layer in model.roberta.encoder.layer:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
optimizer_grouped_parameters.append({
"params":
[p for p in model.parameters() if p not in layer_optim_params]
})
if args.optimizer.lower() == "lamb":
optimizer = Lamb(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
elif args.optimizer.lower() == "adamw":
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
else:
raise Exception(
"optimizer {0} not recognized! Can only be lamb or adamW".format(
args.optimizer))
# Check if saved optimizer or scheduler states exist
if os.path.isfile(
os.path.join(args.model_name_or_path,
"optimizer.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
optimizer.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True,
)
# Train
logger.info("***** Running training *****")
logger.info(" Max steps = %d", args.max_steps)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
global_step = 0
# Check if continuing training from a checkpoint
if os.path.exists(args.model_name_or_path):
# set global_step to gobal_step of last saved checkpoint from model
# path
if "-" in args.model_name_or_path:
global_step = int(
args.model_name_or_path.split("-")[-1].split("/")[0])
else:
global_step = 0
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from global step %d", global_step)
tr_loss = 0.0
model.zero_grad()
model.train()
set_seed(args) # Added here for reproductibility
last_ann_no = -2
train_dataloader = None
train_dataloader_iter = None
dev_ndcg = 0
step = 0
if args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=args.max_steps)
while global_step < args.max_steps:
if step % args.gradient_accumulation_steps == 0 and global_step % args.logging_steps == 0:
# check if new ann training data is availabe
# if not os.path.exists(args.ann_dir):
# print('???',args.blob_ann_dir)
# else:
ann_no, ann_path, ndcg_json = get_latest_ann_data(args.ann_dir)
if ann_path is None:
ann_no, ann_path, ndcg_json = get_latest_ann_data(
args.blob_ann_dir)
ann_no = -1
if ann_path is not None and ann_no != last_ann_no:
logger.info("Training on new add data at %s", ann_path)
with open(ann_path, 'r') as f:
ann_training_data = f.readlines()
dev_ndcg = ndcg_json['ndcg']
ann_checkpoint_path = ndcg_json['checkpoint']
ann_checkpoint_no = get_checkpoint_no(ann_checkpoint_path)
aligned_size = (len(ann_training_data) //
args.world_size) * args.world_size
ann_training_data = ann_training_data[:aligned_size]
logger.info("Total ann queries: %d", len(ann_training_data))
if args.triplet:
train_dataset = StreamingDataset(
ann_training_data,
GetTripletTrainingDataProcessingFn(
args, query_cache, passage_cache))
else:
train_dataset = StreamingDataset(
ann_training_data,
GetTrainingDataProcessingFn(args, query_cache,
passage_cache))
train_dataloader = DataLoader(train_dataset,
batch_size=args.train_batch_size)
train_dataloader_iter = iter(train_dataloader)
# re-warmup
if not args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=len(ann_training_data))
if args.local_rank != -1:
dist.barrier()
if is_first_worker():
# add ndcg at checkpoint step used instead of current step
tb_writer.add_scalar("dev_ndcg", dev_ndcg,
ann_checkpoint_no)
if last_ann_no != -1:
tb_writer.add_scalar("epoch", last_ann_no,
global_step - 1)
tb_writer.add_scalar("epoch", ann_no, global_step)
last_ann_no = ann_no
try:
batch = next(train_dataloader_iter)
except StopIteration:
logger.info("Finished iterating current dataset, begin reiterate")
train_dataloader_iter = iter(train_dataloader)
batch = next(train_dataloader_iter)
batch = tuple(t.to(args.device) for t in batch)
step += 1
if args.triplet:
inputs = {
"query_ids": batch[0].long(),
"attention_mask_q": batch[1].long(),
"input_ids_a": batch[3].long(),
"attention_mask_a": batch[4].long(),
"input_ids_b": batch[6].long(),
"attention_mask_b": batch[7].long()
}
else:
inputs = {
"input_ids_a": batch[0].long(),
"attention_mask_a": batch[1].long(),
"input_ids_b": batch[3].long(),
"attention_mask_b": batch[4].long(),
"labels": batch[6]
}
# sync gradients only at gradient accumulation step
if step % args.gradient_accumulation_steps == 0:
outputs = model(**inputs)
else:
with model.no_sync():
outputs = model(**inputs)
# model outputs are always tuple in transformers (see doc)
loss = outputs[0]
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
if step % args.gradient_accumulation_steps == 0:
loss.backward()
else:
with model.no_sync():
loss.backward()
tr_loss += loss.item()
if step % args.gradient_accumulation_steps == 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
args.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
logs = {}
loss_scalar = tr_loss / args.logging_steps
learning_rate_scalar = scheduler.get_lr()[0]
logs["learning_rate"] = learning_rate_scalar
logs["loss"] = loss_scalar
tr_loss = 0
if is_first_worker():
for key, value in logs.items():
tb_writer.add_scalar(key, value, global_step)
logger.info(json.dumps({**logs, **{"step": global_step}}))
if is_first_worker(
) and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(args.output_dir,
"checkpoint-{}".format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob_output_dir = os.path.join(
args.blob_output_dir, "checkpoint-{}".format(global_step))
if not os.path.exists(blob_output_dir):
os.makedirs(blob_output_dir)
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
#model_to_save.save_pretrained(output_dir)
#tokenizer.save_pretrained(output_dir)
if 'fairseq' not in args.train_model_type:
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
model_to_save.save_pretrained(blob_output_dir)
tokenizer.save_pretrained(blob_output_dir)
else:
torch.save(model.state_dict(),
os.path.join(output_dir, 'model.pt'))
torch.save(model.state_dict(),
os.path.join(blob_output_dir, 'model.pt'))
torch.save(args, os.path.join(output_dir, "training_args.bin"))
torch.save(args,
os.path.join(blob_output_dir, "training_args.bin"))
logger.info("Saving model checkpoint to %s", blob_output_dir)
torch.save(optimizer.state_dict(),
os.path.join(output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(output_dir, "scheduler.pt"))
torch.save(optimizer.state_dict(),
os.path.join(blob_output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(blob_output_dir, "scheduler.pt"))
logger.info("Saving optimizer and scheduler states to %s",
output_dir)
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
tb_writer.close()
return global_step