def main(args):
import_user_module(args)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
oom_batch = task.dataset('train').get_dummy_batch(1, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(args, task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if args.fp16:
trainer = FP16Trainer(args, task, model, criterion)
else:
if torch.cuda.get_device_capability(0)[0] >= 7:
print(
'| NOTICE: your device may support faster training with --fp16'
)
trainer = Trainer(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
load_checkpoint(args, trainer, epoch_itr)
# Send a dummy batch to warm the caching allocator
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
trainer.dummy_train_step(dummy_batch)
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch <= max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
## "Prune" heads (actually mask but shh...)
#if len(args.transformer_mask_heads) > 0:
# # Determine which head to prune
# to_prune = parse_head_pruning_descriptors(
# args.transformer_mask_heads,
# reverse_descriptors=args.transformer_mask_all_but_one_head,
# n_heads=model.encoder.layers[0].self_attn.num_heads
# )
# print(to_prune)
# # Apply pruning
# mask_heads(model, to_prune, args.transformer_mask_rescale)
# Save initial model
initial = os.path.join(args.save_dir, "checkpoint_initial.pt")
trainer.save_checkpoint(initial, {})
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
#if epoch_itr.epoch % args.save_interval == 0:
# save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
# ***********************************Below Changed******************************
# save checkpoint
#if epoch_itr.epoch % args.save_interval == 0:
save_interval = 5 # prune and save checkpoint for every five epoch
if epoch_itr.epoch % save_interval == 0: #****** changed
# do prunning before saving
prune2(args, task, model, trainer, epoch_itr) #****** changed2
# save checkpoint
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
prune2(args, task, model, trainer, epoch_itr
) #****** changed2 do last prunning on the last chekcpoint saved
# ***********************************Above Changed******************************
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
import_user_module(args)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
#args.distributed_world_size = 1
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(args, task)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset(args.train_subset).get_dummy_batch(
args.max_tokens, max_positions)
oom_batch = task.dataset(args.train_subset).get_dummy_batch(
1, max_positions)
# Build trainer
print("Building trainer...")
trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
print("Initialize dataloader...")
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=args.required_batch_size_multiple,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# Initialize distributed training (after data loading)
print("Initialize distributed training (after data loading)...")
if init_distributed:
import socket
args.distributed_rank = distributed_utils.distributed_init(args)
print('| initialized host {} as rank {}'.format(
socket.gethostname(), args.distributed_rank))
# Load the latest checkpoint if one is available
print("Load the latest checkpoint if one is available...")
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
if args.reset_target_embedding:
trainer.init_meters(args)
print("reset trainer.meters")
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
if args.distributed_rank == 0:
if os.path.basename(args.save_dir) != "":
log_file = os.path.join(
args.save_dir,
"({0})-params.log".format(os.path.basename(args.save_dir)))
else:
log_file = os.path.join(
args.save_dir,
"({0})-params.log".format(args.save_dir.split('/')[-2]))
# create log file
args.log_file = log_file
if os.path.exists(log_file):
w = open(log_file, "a+", encoding="utf-8")
else:
w = open(log_file, "w", encoding="utf-8")
w.write(str(args).replace(", ", ",\n") + "\n")
w.write(str(model) + "\n")
w.flush()
w.close()
print("saving params file into{}...".format(log_file))
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
summary_writer = SummaryWriter(log_dir=args.save_dir,
enable=args.distributed_rank == 0)
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
first_train = True
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
first_train = False
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
if not hasattr(save_checkpoint, 'not_best'):
save_checkpoint.not_best = 0
if not args.no_first_valid and first_train:
valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets,
True, summary_writer)
if args.finetune_params != '':
print("| train parameters.")
for name, param in trainer.model.named_parameters():
if trainer.should_train(name):
print(name)
print("| fixed parameters.")
for name, param in trainer.model.named_parameters():
if not trainer.should_train(name):
print(name)
if args.start_ckpt != '':
save_checkpoint.not_best = 0
save_checkpoint.best = 9999
print("| train begin.")
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr, summary_writer)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(
args, trainer, task, epoch_itr, valid_subsets,
epoch_itr.epoch % args.test_bleu_interval == 0, summary_writer)
if args.early_stop > 0:
if hasattr(save_checkpoint,
'best') and valid_losses[0] > save_checkpoint.best:
save_checkpoint.not_best += 1
print("| Not the best ckpt... not best:",
save_checkpoint.not_best)
if save_checkpoint.not_best > args.early_stop:
print("| Early stop...")
break
else:
save_checkpoint.not_best = 0
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
os.system("ps aux | grep redis-server | awk '{print $2}' | xargs kill")
if args.save_output:
save_expert_outputs(args, task, trainer)
def prune(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', "valid"])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {},'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
print('| Optimizer {}'.format(trainer.optimizer.__class__.__name__))
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
prune_meter = StopwatchMeter()
prune_meter.start()
# Estimate head importance scores
head_importance, head_stats = estimate_head_importance(
args, trainer, task, epoch_itr)
prune_meter.stop()
print('| done estimating head importance in {:.1f} seconds'.format(
prune_meter.sum))
torch.save(head_stats,
f"{os.path.dirname(args.restore_file)}/heads_stats.bin")
# Print
print("Head importances")
print("Encoder self attention")
for layer in range(head_importance["encoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_self"][layer]))
print("Encoder decoder attention")
for layer in range(head_importance["encoder_decoder"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_decoder"][layer]))
print("Decoder self attention")
for layer in range(head_importance["decoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["decoder_self"][layer]))
# Print sorted pruning profile
encoder_self_profile = get_profile(head_importance["encoder_self"],
prefix="E")
encoder_decoder_profile = get_profile(head_importance["encoder_decoder"],
prefix="A")
decoder_self_profile = get_profile(head_importance["decoder_self"],
prefix="D")
# Join all
all_profiles = {}
if not (args.decoder_self_only or args.encoder_decoder_only):
all_profiles.update(encoder_self_profile)
if not (args.encoder_self_only or args.decoder_self_only):
all_profiles.update(encoder_decoder_profile)
if not (args.encoder_self_only or args.encoder_decoder_only):
all_profiles.update(decoder_self_profile)
sorted_profiles = sorted(all_profiles.items(),
key=lambda x: x[1],
reverse=args.one_minus)
print("Heads sorted by importance:")
print(" ".join(p for p, _ in sorted_profiles))
print("Sorted head importance scores:")
print(" ".join(f"{v.data:.5f}" for _, v in sorted_profiles))
if args.only_importance:
return
tot_n_heads = len(sorted_profiles)
# Eval pruning
if args.one_head:
kept_layers = set()
to_prune_profile = []
for p, _ in reversed(sorted_profiles):
layer_name = ":".join(p.split(":")[:-1])
if layer_name not in kept_layers:
kept_layers.add(layer_name)
continue
else:
to_prune_profile.insert(0, p)
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
return
for i in range(0, 10):
n_to_prune = int(ceil(tot_n_heads * i / 10))
to_prune_profile = [p for p, _ in sorted_profiles[:n_to_prune]]
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
print('| num. model params: {}'.format(sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates() < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
mlperf_compliance.mlperf_log.LOGGER.propagate = False
# framework = f'Pytorch NGC {os.environ["NVIDIA_PYTORCH_VERSION"]}'
# mlperf_submission_log(
# benchmark=mlperf_compliance.constants.TRANSFORMER,
# framework=framework)
mlperf_compliance.mlperf_log.setdefault(
root_dir=os.path.dirname(os.path.abspath(__file__)),
benchmark=mlperf_compliance.constants.TRANSFORMER,
stack_offset=1,
extra_print=False)
mlperf_print(key=mlperf_compliance.constants.INIT_START,
log_all_ranks=True)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# preinit and warmup streams/groups for allreduce communicators
allreduce_communicators = None
if args.distributed_world_size > 1 and args.enable_parallel_backward_allred_opt:
allreduce_groups = [
torch.distributed.new_group()
for _ in range(args.parallel_backward_allred_cuda_nstreams)
]
allreduce_streams = [
torch.cuda.Stream()
for _ in range(args.parallel_backward_allred_cuda_nstreams)
]
for group, stream in zip(allreduce_groups, allreduce_streams):
with torch.cuda.stream(stream):
torch.distributed.all_reduce(torch.cuda.FloatTensor(1),
group=group)
allreduce_communicators = (allreduce_groups, allreduce_streams)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
mlperf_print(key=mlperf_compliance.constants.GLOBAL_BATCH_SIZE,
value=args.max_tokens * args.distributed_world_size)
mlperf_print(key=mlperf_compliance.constants.OPT_NAME,
value=args.optimizer)
assert (len(args.lr) == 1)
mlperf_print(key=mlperf_compliance.constants.OPT_BASE_LR,
value=args.lr[0] if len(args.lr) == 1 else args.lr)
mlperf_print(key=mlperf_compliance.constants.OPT_LR_WARMUP_STEPS,
value=args.warmup_updates)
assert (args.max_source_positions == args.max_target_positions)
mlperf_print(key=mlperf_compliance.constants.MAX_SEQUENCE_LENGTH,
value=args.max_target_positions)
mlperf_print(key=mlperf_compliance.constants.OPT_ADAM_BETA_1,
value=eval(args.adam_betas)[0])
mlperf_print(key=mlperf_compliance.constants.OPT_ADAM_BETA_2,
value=eval(args.adam_betas)[1])
mlperf_print(key=mlperf_compliance.constants.OPT_ADAM_EPSILON,
value=args.adam_eps)
pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
result = torch.cuda.cudart().cudaDeviceSetLimit(ctypes.c_int(0x05),
ctypes.c_int(128))
result = torch.cuda.cudart().cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
# torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if args.fp16:
trainer = FP16Trainer(args,
task,
model,
criterion,
allreduce_communicators=allreduce_communicators)
else:
if torch.cuda.get_device_capability(0)[0] >= 7:
print(
'| NOTICE: your device may support faster training with --fp16'
)
trainer = Trainer(args,
task,
model,
criterion,
allreduce_communicators=None)
#if (args.online_eval or args.target_bleu) and not args.remove_bpe:
# args.remove_bpe='@@ '
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
# Send a dummy batch to warm the caching allocator
dummy_batch = language_pair_dataset.get_dummy_batch_isolated(
args.max_tokens, max_positions, 8)
trainer.dummy_train_step(dummy_batch)
# Train until the learning rate gets too small or model reaches target score
max_epoch = args.max_epoch if args.max_epoch >= 0 else math.inf
max_update = args.max_update or math.inf
tgt_bleu = args.target_bleu or math.inf
current_bleu = 0.0
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
# mlperf compliance synchronization
if args.distributed_world_size > 1:
assert (torch.distributed.is_initialized())
torch.distributed.all_reduce(torch.cuda.FloatTensor(1))
torch.cuda.synchronize()
mlperf_print(key=mlperf_compliance.constants.INIT_STOP, sync=True)
mlperf_print(key=mlperf_compliance.constants.RUN_START, sync=True)
# second sync after RUN_START tag is printed.
# this ensures no rank touches data until after RUN_START tag is printed.
barrier()
# Load dataset splits
load_dataset_splits(task, ['train', 'test'])
ctr = 0
class DummyEpochBatchIterator:
def __init__(self, epoch=0):
self.epoch = epoch
epoch_itr = DummyEpochBatchIterator(0)
# Main training loop
while lr >= args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update and current_bleu < tgt_bleu:
first_epoch = epoch_itr.epoch + 1
mlperf_print(key=mlperf_compliance.constants.BLOCK_START,
metadata={
'first_epoch_num': first_epoch,
'epoch_count': 1
},
sync=True)
mlperf_print(key=mlperf_compliance.constants.EPOCH_START,
metadata={'epoch_num': first_epoch},
sync=True)
start = time.time()
gc.disable()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset(args.train_subset),
dataloader_num_workers=args.dataloader_num_workers,
dataloader_pin_memory=args.enable_dataloader_pin_memory,
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
epoch=epoch_itr.epoch if ctr is not 0 else 0,
bucket_growth_factor=args.bucket_growth_factor,
seq_len_multiple=args.seq_len_multiple,
batching_scheme=args.batching_scheme,
batch_multiple_strategy=args.batch_multiple_strategy,
)
print("got epoch iterator", time.time() - start)
# Load the latest checkpoint if one is available
if ctr is 0:
load_checkpoint(args, trainer, epoch_itr)
# train for one epoch
start = time.time()
#exit(1)
train(args, trainer, task, epoch_itr)
print("epoch time ", time.time() - start)
start = time.time()
mlperf_print(key=mlperf_compliance.constants.EPOCH_STOP,
metadata={'epoch_num': first_epoch},
sync=True)
#if epoch_itr.epoch % args.validate_interval == 0:
# valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
# Eval BLEU score
if args.online_eval or (not tgt_bleu is math.inf):
current_bleu = score(args, trainer, task, epoch_itr,
args.gen_subset)
mlperf_print(key=mlperf_compliance.tags.EVAL_ACCURACY,
value=str(current_bleu),
metadata={'epoch_num': first_epoch})
gc.enable()
# Only use first validation loss to update the learning rate
#lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# Save checkpoint
#if epoch_itr.epoch % args.save_interval == 0:
# save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
ctr = ctr + 1
print("validation and scoring ", time.time() - start)
mlperf_print(key=mlperf_compliance.constants.BLOCK_STOP,
metadata={'first_epoch_num': first_epoch},
sync=True)
train_meter.stop()
status = 'success' if current_bleu >= tgt_bleu else 'aborted'
mlperf_print(key=mlperf_compliance.constants.RUN_STOP,
metadata={'status': status})
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
if args.distributed_world_size > 1:
assert (torch.distributed.is_initialized())
torch.distributed.broadcast(torch.tensor([1], device="cuda"), 0)
torch.cuda.synchronize()
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
result = torch.cuda.cudart().cudaDeviceSetLimit(ctypes.c_int(0x05),
ctypes.c_int(128))
result = torch.cuda.cudart().cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if args.fp16:
trainer = FP16Trainer(args, task, model, criterion)
else:
if torch.cuda.get_device_capability(0)[0] >= 7:
print(
'| NOTICE: your device may support faster training with --fp16'
)
trainer = Trainer(args, task, model, criterion)
if (args.online_eval or args.target_bleu) and not args.remove_bpe:
args.remove_bpe = '@@ '
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
max_positions = trainer.get_model().max_positions()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
load_checkpoint(args, trainer, epoch_itr)
# Send a dummy batch to warm the caching allocator
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
trainer.dummy_train_step(dummy_batch)
# Train until the learning rate gets too small or model reaches target score
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
tgt_bleu = args.target_bleu or math.inf
current_bleu = 0.0
best_bleu = 0.0
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr >= args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update and current_bleu < tgt_bleu:
# train for one epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# Eval BLEU score
if args.online_eval or (not tgt_bleu is math.inf):
current_bleu, current_sc_bleu = score(args, trainer, task,
epoch_itr, args.gen_subset)
if current_bleu > best_bleu:
best_bleu = current_bleu
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
# Only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# Save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
from mlperf_compliance.mlperf_log import transformer_print
transformer_print(
key=mlperf_log.RUN_CLEAR_CACHES
) #before this tag we should run clearing caches on the host
# mlperf compliance synchronization
if args.distributed_world_size > 1:
assert (torch.distributed.is_initialized())
torch.distributed.broadcast(torch.tensor([1], device="cuda"), 0)
torch.cuda.synchronize()
transformer_print(key=mlperf_log.RUN_START)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
transformer_print(key=mlperf_log.OPT_NAME, value=args.optimizer)
transformer_print(key=mlperf_log.OPT_LR, value=args.lr)
transformer_print(key=mlperf_log.OPT_HP_ADAM_BETA1,
value=eval(args.adam_betas)[0])
transformer_print(key=mlperf_log.OPT_HP_ADAM_BETA2,
value=eval(args.adam_betas)[1])
transformer_print(key=mlperf_log.OPT_HP_ADAM_EPSILON, value=args.adam_eps)
pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
result = torch.cuda.cudart().cudaDeviceSetLimit(ctypes.c_int(0x05),
ctypes.c_int(128))
result = torch.cuda.cudart().cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
torch.manual_seed(args.seed)
transformer_print(key=mlperf_log.RUN_SET_RANDOM_SEED, value=args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
transformer_print(key=mlperf_log.MODEL_HP_SEQ_BEAM_SEARCH,
value={
'alpha': args.lenpen,
'beam_size': args.beam,
'extra_decode_length': args.max_len_b,
'vocab_size': task.target_dictionary.__len__()
})
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if args.fp16:
trainer = FP16Trainer(args, task, model, criterion)
else:
if torch.cuda.get_device_capability(0)[0] >= 7:
print(
'| NOTICE: your device may support faster training with --fp16'
)
trainer = Trainer(args, task, model, criterion)
if (args.online_eval or args.target_bleu) and not args.remove_bpe:
args.remove_bpe = '@@ '
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
transformer_print(key=mlperf_log.INPUT_BATCH_SIZE, value=args.max_tokens)
transformer_print(key=mlperf_log.INPUT_ORDER)
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
# Send a dummy batch to warm the caching allocator
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
trainer.dummy_train_step(dummy_batch)
# Train until the learning rate gets too small or model reaches target score
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
tgt_bleu = args.target_bleu or math.inf
current_bleu = 0.0
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
ctr = 0
class DummyEpochBatchIterator:
def __init__(self, epoch=0):
self.epoch = epoch
epoch_itr = DummyEpochBatchIterator(0)
transformer_print(key=mlperf_log.TRAIN_LOOP)
while lr >= args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update and current_bleu < tgt_bleu:
transformer_print(key=mlperf_log.TRAIN_EPOCH, value=epoch_itr.epoch)
import time
start = time.time()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
epoch=epoch_itr.epoch if ctr is not 0 else 0)
print("got epoch iterator", time.time() - start)
# Load the latest checkpoint if one is available
if ctr is 0:
load_checkpoint(args, trainer, epoch_itr)
# train for one epoch
start = time.time()
train(args, trainer, task, epoch_itr)
print("epoch time ", time.time() - start)
start = time.time()
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# Eval BLEU score
transformer_print(key=mlperf_log.EVAL_START, value=epoch_itr.epoch)
if args.online_eval or (not tgt_bleu is math.inf):
current_bleu = score(args, trainer, task, epoch_itr,
args.gen_subset)
transformer_print(key=mlperf_log.EVAL_ACCURACY,
value={
'epoch': epoch_itr.epoch,
'value': current_bleu
})
transformer_print(key=mlperf_log.EVAL_TARGET, value=tgt_bleu)
transformer_print(key=mlperf_log.EVAL_STOP, value=epoch_itr.epoch)
# Only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# Save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
ctr = ctr + 1
print("validation and scoring ", time.time() - start)
train_meter.stop()
transformer_print(key=mlperf_log.RUN_STOP)
transformer_print(key=mlperf_log.RUN_FINAL)
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', 'valid'])
# print("<AFTER>load_dataset_splits")
# Build model and criterion
model = task.build_model(args)
print('| num. model params: {}'.format(sum(p.numel() for p in model.parameters())))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# print("<AFTER>build_model")
# Validation
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
val_criterion = task.build_criterion(args, 'label_smoothed_cross_entropy')
val_trainer = Trainer(args, task, model, val_criterion)
class_pretrain_flag = False
mt_pretrain_flag = False
# Pre-training on CNN discriminator and Seq2Seq recontruction
if args.task == 'style_transfer':
# classification pretrain
criterion = task.build_criterion(args, 'classification')
trainer = Trainer(args, task, model, criterion)
print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
max_positions = trainer.get_model().max_positions()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset('train'),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
load_checkpoint(args, trainer, epoch_itr, load_optim=True, find_best=args.restore_best)
max_epoch = args.pre_train_max_epoch
while epoch_itr.epoch < max_epoch:
class_pretrain_flag = True
# train for one epoch
train(args, trainer, task, epoch_itr)
valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
# save to checkpoint
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
print("Done classification pretrain")
# MT pretrain
criterion = task.build_criterion(args, 'style_transfer_pretrain')
trainer = Trainer(args, task, model, criterion)
print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
# Load the latest checkpoint if one is available
if epoch_itr.epoch <= args.pre_train_max_epoch:
load_checkpoint(args, trainer, epoch_itr, load_optim=False, find_best=True)
epoch_itr.epoch = args.pre_train_max_epoch
save_checkpoint.best = float("inf")
else:
load_checkpoint(args, trainer, epoch_itr, load_optim=True, find_best=args.restore_best)
# Send a dummy batch to warm the caching allocator
dummy_batch = task.dataset('train').get_dummy_batch(args.max_tokens, max_positions)
trainer.dummy_train_step(dummy_batch)
max_epoch = 2 * args.pre_train_max_epoch
while epoch_itr.epoch < max_epoch:
mt_pretrain_flag = True
# train for one epoch
train(args, trainer, task, epoch_itr)
valid_losses = validate(args, val_trainer, task, epoch_itr, valid_subsets)
# save to checkpoint
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
print("Done MT pretrain")
# Training
if args.task == 'style_transfer':
criterion_name = "style_transfer_train"
print("Loading plain data")
load_dataset_splits(task, ['plain'])
else:
criterion_name = None
criterion = task.build_criterion(args, criterion_name)
print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
trainer = Trainer(args, task, model, criterion)
# Load the latest checkpoint if one is available
if epoch_itr.epoch <= 2*args.pre_train_max_epoch:
load_checkpoint(args, trainer, epoch_itr, load_optim=False,
fix_discriminator=True, find_best=True)
else:
load_checkpoint(args, trainer, epoch_itr, load_optim=True,
fix_discriminator=True, find_best=args.restore_best)
print("# WARNING: Loading checkpoint with optimizer")
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
if args.task == 'style_transfer':
src_plain_epoch_iter = data.EpochBatchIterator(
dataset=task.dataset('plain')[0],
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
trg_plain_epoch_iter = data.EpochBatchIterator(
dataset=task.dataset('plain')[1],
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
pre_train_max_epoch = 2 * args.pre_train_max_epoch
# Train until the learning rate gets too small
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
while lr > args.min_lr and epoch_itr.epoch < (max_epoch + pre_train_max_epoch) and trainer.get_num_updates() < max_update:
# train for one epoch
train(args, trainer, task, epoch_itr,
use_plain=(args.task=='style_transfer'),
src_plain_epoch_iter=src_plain_epoch_iter,
trg_plain_epoch_iter=trg_plain_epoch_iter,
)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(
args, val_trainer, task, epoch_itr, valid_subsets)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args):
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
mllog.config(filename=os.path.join(
os.path.dirname(os.path.abspath(__file__)), 'transformer.log'))
mllogger = mllog.get_mllogger()
mllogger.logger.propagate = False
log_start(key=constants.INIT_START, log_all_ranks=True)
# preinit and warmup streams/groups for allreduce communicators
allreduce_communicators = None
if args.distributed_world_size > 1 and args.enable_parallel_backward_allred_opt:
allreduce_groups = [
torch.distributed.new_group()
for _ in range(args.parallel_backward_allred_cuda_nstreams)
]
allreduce_streams = [
torch.cuda.Stream()
for _ in range(args.parallel_backward_allred_cuda_nstreams)
]
for group, stream in zip(allreduce_groups, allreduce_streams):
with torch.cuda.stream(stream):
torch.distributed.all_reduce(torch.cuda.FloatTensor(1),
group=group)
allreduce_communicators = (allreduce_groups, allreduce_streams)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
log_event(key=constants.GLOBAL_BATCH_SIZE,
value=args.max_tokens * args.distributed_world_size)
log_event(key=constants.OPT_NAME, value=args.optimizer)
assert (len(args.lr) == 1)
log_event(key=constants.OPT_BASE_LR,
value=args.lr[0] if len(args.lr) == 1 else args.lr)
log_event(key=constants.OPT_LR_WARMUP_STEPS, value=args.warmup_updates)
assert (args.max_source_positions == args.max_target_positions)
log_event(key=constants.MAX_SEQUENCE_LENGTH,
value=args.max_target_positions,
metadata={'method': 'discard'})
log_event(key=constants.OPT_ADAM_BETA_1, value=eval(args.adam_betas)[0])
log_event(key=constants.OPT_ADAM_BETA_2, value=eval(args.adam_betas)[1])
log_event(key=constants.OPT_ADAM_EPSILON, value=args.adam_eps)
log_event(key=constants.SEED, value=args.seed)
# L2 Sector Promotion
pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
result = ctypes.CDLL('libcudart.so').cudaDeviceSetLimit(
ctypes.c_int(0x05), ctypes.c_int(128))
result = ctypes.CDLL('libcudart.so').cudaDeviceGetLimit(
pValue, ctypes.c_int(0x05))
worker_seeds, shuffling_seeds = setup_seeds(
args.seed,
args.max_epoch + 1,
torch.device('cuda'),
args.distributed_rank,
args.distributed_world_size,
)
worker_seed = worker_seeds[args.distributed_rank]
print(
f'Worker {args.distributed_rank} is using worker seed: {worker_seed}')
torch.manual_seed(worker_seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Build trainer
if args.fp16:
if args.distributed_weight_update != 0:
from fairseq.fp16_trainer import DistributedFP16Trainer
trainer = DistributedFP16Trainer(
args,
task,
model,
criterion,
allreduce_communicators=allreduce_communicators)
else:
from fairseq.fp16_trainer import FP16Trainer
trainer = FP16Trainer(
args,
task,
model,
criterion,
allreduce_communicators=allreduce_communicators)
else:
if torch.cuda.get_device_capability(0)[0] >= 7:
print(
'| NOTICE: your device may support faster training with --fp16'
)
trainer = Trainer(args,
task,
model,
criterion,
allreduce_communicators=None)
#if (args.online_eval or args.target_bleu) and not args.remove_bpe:
# args.remove_bpe='@@ '
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Initialize dataloader
max_positions = trainer.get_model().max_positions()
# Send a dummy batch to warm the caching allocator
dummy_batch = language_pair_dataset.get_dummy_batch_isolated(
args.max_tokens, max_positions, 8)
trainer.dummy_train_step(dummy_batch)
# Train until the learning rate gets too small or model reaches target score
max_epoch = args.max_epoch if args.max_epoch >= 0 else math.inf
max_update = args.max_update or math.inf
tgt_bleu = args.target_bleu or math.inf
current_bleu = 0.0
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
# mlperf compliance synchronization
if args.distributed_world_size > 1:
assert (torch.distributed.is_initialized())
torch.distributed.all_reduce(torch.cuda.FloatTensor(1))
torch.cuda.synchronize()
log_end(key=constants.INIT_STOP, sync=False)
log_start(key=constants.RUN_START, sync=True)
# second sync after RUN_START tag is printed.
# this ensures no rank touches data until after RUN_START tag is printed.
barrier()
# Load dataset splits
load_dataset_splits(task, ['train', 'test'])
log_event(key=constants.TRAIN_SAMPLES,
value=len(task.dataset(args.train_subset)),
sync=False)
log_event(key=constants.EVAL_SAMPLES,
value=len(task.dataset(args.gen_subset)),
sync=False)
ctr = 0
start = time.time()
epoch_itr = data.EpochBatchIterator(
dataset=task.dataset(args.train_subset),
dataloader_num_workers=args.dataloader_num_workers,
dataloader_pin_memory=args.enable_dataloader_pin_memory,
max_tokens=args.max_tokens,
max_sentences=args.max_sentences_valid,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seeds=shuffling_seeds,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
epoch=epoch_itr.epoch if ctr is not 0 else 0,
bucket_growth_factor=args.bucket_growth_factor,
seq_len_multiple=args.seq_len_multiple,
batching_scheme=args.batching_scheme,
batch_multiple_strategy=args.batch_multiple_strategy,
)
print("got epoch iterator", time.time() - start)
# Main training loop
while lr >= args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update and current_bleu < tgt_bleu:
first_epoch = epoch_itr.epoch + 1
log_start(key=constants.BLOCK_START,
metadata={
'first_epoch_num': first_epoch,
'epoch_count': 1
},
sync=False)
log_start(key=constants.EPOCH_START,
metadata={'epoch_num': first_epoch},
sync=False)
gc.disable()
# Load the latest checkpoint if one is available
if ctr is 0:
load_checkpoint(args, trainer, epoch_itr)
# train for one epoch
start = time.time()
#exit(1)
train(args, trainer, task, epoch_itr, shuffling_seeds)
print("epoch time ", time.time() - start)
start = time.time()
log_end(key=constants.EPOCH_STOP,
metadata={'epoch_num': first_epoch},
sync=False)
# Eval BLEU score
if args.online_eval or (not tgt_bleu is math.inf):
current_bleu = score(args, trainer, task, epoch_itr,
args.gen_subset)
log_event(key=constants.EVAL_ACCURACY,
value=float(current_bleu) / 100.0,
metadata={'epoch_num': first_epoch})
gc.enable()
# Only use first validation loss to update the learning rate
#lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# Save checkpoint
#if epoch_itr.epoch % args.save_interval == 0:
# save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
ctr = ctr + 1
print("validation and scoring ", time.time() - start)
log_end(key=constants.BLOCK_STOP,
metadata={'first_epoch_num': first_epoch},
sync=False)
train_meter.stop()
status = 'success' if current_bleu >= tgt_bleu else 'aborted'
log_end(key=constants.RUN_STOP, metadata={'status': status})
print('| done training in {:.1f} seconds'.format(train_meter.sum))
def main(args, init_distributed=False):
import_user_module(args)
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# 加载数据集
load_dataset_splits(task, ['train', 'valid'])
# 初始化分布式训练
if init_distributed:
import socket
args.distributed_rank = distributed_utils.distributed_init(args)
print('| initialized host {} as rank {}'.format(
socket.gethostname(), args.distributed_rank))
# build模型和损失
model = task.build_model(args)
#损失函数, eg: CrossEntropyCriterion()
criterion = task.build_criterion(args)
# 打印模型架构
print(model)
print('| model {}, criterion {}'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {} (num. trained: {})'.format(
sum(p.numel() for p in model.parameters()),
sum(p.numel() for p in model.parameters() if p.requires_grad),
))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
oom_batch = task.dataset('train').get_dummy_batch(1, max_positions)
#是否从预训练模型加载参数
model.copy_pretrained_params(args)
# 构建一个builder
trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# 初始化 dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
)
# 加载最新的checkpoint(如果有), 继续训练模型
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# 训练直到学习率变得太小
max_epoch = args.max_epoch or math.inf
max_update = args.max_update or math.inf
lr = trainer.get_lr()
train_meter = StopwatchMeter()
train_meter.start()
valid_losses = [None]
valid_subsets = args.valid_subset.split(',')
while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates(
) < max_update:
# 训练一个epoch
train(args, trainer, task, epoch_itr)
if epoch_itr.epoch % args.validate_interval == 0:
valid_losses = validate(args, trainer, task, epoch_itr,
valid_subsets)
# ema process
if not args.no_ema:
old_data = ema_restore(trainer.ema, trainer.model)
valid_losses_ema = validate(args, trainer, task, epoch_itr,
valid_subsets)
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args,
trainer,
epoch_itr,
valid_losses_ema[0],
suffix='ema')
ema_reverse(trainer.ema, trainer.model, old_data)
# only use first validation loss to update the learning rate
lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
# save checkpoint
if epoch_itr.epoch % args.save_interval == 0:
save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
train_meter.stop()
print('| done training in {:.1f} seconds'.format(train_meter.sum))