def testORTTransformerModelExport(seed, device):
# Common setup
optim_config = optim.LambConfig()
opts = orttrainer.ORTTrainerOptions({
'debug': {
'check_model_export': True,
},
'device': {
'id': device,
}
})
# Setup for the first ORTTRainer run
torch.manual_seed(seed)
set_seed(seed)
model, model_desc, my_loss, batcher_fn, train_data, val_data, _ = _load_pytorch_transformer_model(
device)
first_trainer = orttrainer.ORTTrainer(model,
model_desc,
optim_config,
loss_fn=my_loss,
options=opts)
data, targets = batcher_fn(train_data, 0)
_ = first_trainer.train_step(data, targets)
assert first_trainer._onnx_model is not None
def create_orttrainer_and_save_checkpoint(device,
trainer_opts,
checkpoint_dir,
state_dict_key_name='state_dict',
use_lamb=True):
learning_rate = 0.1
seed = 1
torch.manual_seed(seed)
set_seed(seed)
optim_config = optim.LambConfig(
lr=learning_rate) if use_lamb else optim.AdamConfig(lr=learning_rate)
model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(
device)
trainer = orttrainer.ORTTrainer(
model,
model_desc,
optim_config,
loss_fn=loss_fn,
options=orttrainer.ORTTrainerOptions(trainer_opts))
if 'distributed' in trainer_opts:
train_data = next(
islice(
_chunkify(train_data,
trainer_opts['distributed']['world_size']),
trainer_opts['distributed']['world_rank'], None))
# run train steps
_train(trainer, train_data, batcher_fn)
# save current model parameters as a checkpoint
if checkpoint_dir:
_save(trainer, checkpoint_dir, state_dict_key_name)
def create_orttrainer_and_load_checkpoint(device, trainer_opts, checkpoint_dir):
"""Instantiate and load checkpoint into trainer
- Instantiates the ORTTrainer with given input trainer_opts configuration for a simple transformer model
- Loads the checkpoint from directory checkpoint_dir into the trainer
- Runs eval_step on the trainer so the trainer onnx graph is initialized
- Returns the trainer state_dict and the pytorch model
"""
seed = 1
torch.manual_seed(seed)
set_seed(seed)
# PyTorch transformer model setup
learning_rate = 0.1
optim_config = optim.LambConfig(lr=learning_rate)
model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts))
# load checkpoint into trainer
checkpoint.experimental_load_checkpoint(trainer, checkpoint_dir)
# run an eval step to innitialize the graph
torch.manual_seed(seed)
set_seed(seed)
data, targets = batcher_fn(train_data, 0)
trainer.eval_step(data, targets)
return checkpoint.experimental_state_dict(trainer), model
def load_model_optim_state_and_eval(device, trainer_opts, use_lamb=True):
learning_rate = 0.1
seed = 1
torch.manual_seed(seed)
set_seed(seed)
optim_config = optim.LambConfig(
lr=learning_rate) if use_lamb else optim.AdamConfig(lr=learning_rate)
model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(
device)
trainer = orttrainer.ORTTrainer(
model,
model_desc,
optim_config,
loss_fn=loss_fn,
options=orttrainer.ORTTrainerOptions(trainer_opts))
# load dummy state
dummy_init_state = generate_dummy_optim_state(model, optim_config)
checkpoint._experimental_load_optimizer_state(trainer, dummy_init_state)
# run an eval step to innitialize the graph
data, targets = batcher_fn(train_data, 0)
trainer.eval_step(data, targets)
return dummy_init_state, checkpoint.experimental_state_dict(trainer)
def _create_trainer(zero_enabled=False):
"""Cerates a simple ORTTrainer for ORTTrainer functional tests"""
device = 'cuda'
optim_config = optim.LambConfig(lr=0.1)
opts = {'device': {'id': device}, 'debug': {'deterministic_compute': True}}
if zero_enabled:
opts['distributed'] = {
'world_rank': 0,
'world_size': 1,
'allreduce_post_accumulation': True,
'deepspeed_zero_optimization': {
'stage': 1
}
}
model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(
device)
trainer = orttrainer.ORTTrainer(model,
model_desc,
optim_config,
loss_fn=loss_fn,
options=orttrainer.ORTTrainerOptions(opts))
return trainer