mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-23 19:32:23 +00:00
Port legacy checkpoint API into new front-end (#4855)
* Port legacy checkpoint API into new front-end
This PR also fixes:
* Warnings on ORTTrainer for improper tensor copies
* Inaccurate LRScheduler tests using wrong LR
* Stale DeepSpeed documentation
* Minor code refactoring for Toy BERT tests
* Move experimental state_dict() and load_state_dict() into checkpoint ns
This commit is contained in:
Thiago Crepaldi
GitHub
parent
75ad7be336
commit
7cc88ef7ed
6 changed files with 594 additions and 104 deletions
280
orttraining/orttraining/python/experimental/checkpoint.py
Normal file
280
orttraining/orttraining/python/experimental/checkpoint.py
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
from collections import OrderedDict
|
||||
import numpy as np
|
||||
import onnx
|
||||
import os
|
||||
import torch
|
||||
import warnings
|
||||
|
||||
|
||||
################################################################################
|
||||
# Experimental Checkpoint APIs
|
||||
################################################################################
|
||||
|
||||
|
||||
def experimental_state_dict(ort_trainer):
|
||||
if not ort_trainer._training_session:
|
||||
warnings.warn("ONNX Runtime training session is not initialized yet. "
|
||||
"Please run train_step or eval_step at least once before calling state_dict().")
|
||||
return {}
|
||||
|
||||
# extract trained weights
|
||||
session_state = ort_trainer._training_session.get_state()
|
||||
torch_state = {}
|
||||
for name in session_state:
|
||||
torch_state[name] = torch.from_numpy(session_state[name])
|
||||
|
||||
# extract untrained weights and buffer
|
||||
for n in ort_trainer._onnx_model.graph.initializer:
|
||||
if n.name not in torch_state:
|
||||
torch_state[n.name] = torch.from_numpy(np.array(onnx.numpy_helper.to_array(n)))
|
||||
|
||||
# Need to remove redundant initializers and name suffices to map back to original torch state names
|
||||
torch_state_to_return = {key: torch_state[key] for key in ort_trainer._original_model_state_keys if key in torch_state} \
|
||||
if ort_trainer._original_model_state_keys else torch_state
|
||||
return torch_state_to_return
|
||||
|
||||
def experimental_load_state_dict(ort_trainer, state_dict, strict=False):
|
||||
# Note: It may happen ONNX model has not yet been initialized
|
||||
# In this case we cache a reference to desired state and delay the restore until after initialization
|
||||
# Unexpected behavior will result if the user changes the reference before initialization
|
||||
if not ort_trainer._training_session:
|
||||
ort_trainer.state_dict_ = state_dict
|
||||
ort_trainer.strict_ = strict
|
||||
return
|
||||
|
||||
# update onnx model from loaded state dict
|
||||
cur_initializers_names = [n.name for n in ort_trainer._onnx_model.graph.initializer]
|
||||
new_initializers = {}
|
||||
|
||||
for name in state_dict:
|
||||
if name in cur_initializers_names:
|
||||
new_initializers[name] = state_dict[name].numpy()
|
||||
elif strict:
|
||||
raise RuntimeError("Checkpoint tensor: {} is not present in the model.".format(name))
|
||||
|
||||
ort_trainer._update_onnx_model_initializers(new_initializers)
|
||||
|
||||
# create new session based on updated onnx model
|
||||
ort_trainer.state_dict_ = None
|
||||
ort_trainer._init_session()
|
||||
|
||||
# load training state
|
||||
session_state = {name:state_dict[name].numpy() for name in state_dict}
|
||||
ort_trainer._training_session.load_state(session_state, strict)
|
||||
|
||||
|
||||
def experimental_save_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix="ORT_checkpoint", checkpoint_state_dict=None):
|
||||
if checkpoint_state_dict == None:
|
||||
checkpoint_state_dict = {'model': experimental_state_dict(ort_trainer)}
|
||||
else:
|
||||
checkpoint_state_dict.update({'model': experimental_state_dict(ort_trainer)})
|
||||
|
||||
assert os.path.exists(checkpoint_dir), f"checkpoint_dir ({checkpoint_dir}) directory doesn't exist"
|
||||
|
||||
checkpoint_name = _get_checkpoint_name(checkpoint_prefix,
|
||||
ort_trainer.options.distributed.deepspeed_zero_optimization.stage,
|
||||
ort_trainer.options.distributed.world_rank,
|
||||
ort_trainer.options.distributed.world_size)
|
||||
checkpoint_file = os.path.join(checkpoint_dir, checkpoint_name)
|
||||
if os.path.exists(checkpoint_file):
|
||||
msg = f"{checkpoint_file} already exists, overwriting."
|
||||
warnings.warn(msg)
|
||||
torch.save(checkpoint_state_dict, checkpoint_file)
|
||||
|
||||
|
||||
def experimental_load_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix="ORT_checkpoint", strict=False):
|
||||
checkpoint_files = _list_checkpoint_files(
|
||||
checkpoint_dir, checkpoint_prefix)
|
||||
is_partitioned = False
|
||||
if len(checkpoint_files) > 1:
|
||||
msg = (f"Found more than one file with prefix {checkpoint_prefix} in directory {checkpoint_dir}.",
|
||||
"Attempting to load ZeRO checkpoint.")
|
||||
warnings.warn(msg)
|
||||
is_partitioned = True
|
||||
if (not ort_trainer.options.distributed.deepspeed_zero_optimization.stage) and is_partitioned:
|
||||
return _load_multi_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix, strict)
|
||||
else:
|
||||
return _load_single_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix, is_partitioned, strict)
|
||||
|
||||
|
||||
################################################################################
|
||||
# Helper functions
|
||||
################################################################################
|
||||
|
||||
|
||||
def _load_single_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix, is_partitioned, strict):
|
||||
checkpoint_name = _get_checkpoint_name(
|
||||
checkpoint_prefix, is_partitioned, ort_trainer.options.distributed.world_rank, ort_trainer.options.distributed.world_size)
|
||||
checkpoint_file = os.path.join(checkpoint_dir, checkpoint_name)
|
||||
|
||||
if is_partitioned:
|
||||
assert_msg = f"Couldn't find checkpoint file {checkpoint_file}.",\
|
||||
"Optimizer partitioning is enabled using ZeRO. Please make sure the checkpoint file exists ",\
|
||||
f"for rank {ort_trainer.options.distributed.world_rank} of {ort_trainer.options.distributed.world_size}"
|
||||
else:
|
||||
assert_msg = f"Couldn't find checkpoint file {checkpoint_file}."
|
||||
assert os.path.exists(checkpoint_file), assert_msg
|
||||
|
||||
checkpoint_state = torch.load(checkpoint_file, map_location='cpu')
|
||||
experimental_load_state_dict(ort_trainer, checkpoint_state['model'], strict=strict)
|
||||
del(checkpoint_state['model'])
|
||||
return checkpoint_state
|
||||
|
||||
|
||||
def _load_multi_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix, strict):
|
||||
checkpoint_files = _list_checkpoint_files(checkpoint_dir, checkpoint_prefix)
|
||||
|
||||
ckpt_agg = _CombineZeroCheckpoint(checkpoint_files)
|
||||
aggregate_state_dict = ckpt_agg.aggregate_checkpoints()
|
||||
|
||||
experimental_load_state_dict(ort_trainer, aggregate_state_dict, strict=strict)
|
||||
|
||||
# aggregate other keys in the state_dict.
|
||||
# Values will be overwritten for matching keys among workers
|
||||
all_checkpoint_states = dict()
|
||||
for checkpoint_file in checkpoint_files:
|
||||
checkpoint_state = torch.load(checkpoint_file, map_location='cpu')
|
||||
del(checkpoint_state['model'])
|
||||
all_checkpoint_states.update(checkpoint_state)
|
||||
return all_checkpoint_states
|
||||
|
||||
|
||||
def _list_checkpoint_files(checkpoint_dir, checkpoint_prefix, extension='.ort.pt'):
|
||||
ckpt_file_names = [f for f in os.listdir(checkpoint_dir) if f.startswith(checkpoint_prefix)]
|
||||
ckpt_file_names = [f for f in ckpt_file_names if f.endswith(extension)]
|
||||
ckpt_file_names = [os.path.join(checkpoint_dir, f) for f in ckpt_file_names]
|
||||
|
||||
assert len(ckpt_file_names) > 0, f"No checkpoint found with prefix '{checkpoint_prefix}' at '{checkpoint_dir}'"
|
||||
return ckpt_file_names
|
||||
|
||||
|
||||
def _get_checkpoint_name(prefix, is_partitioned, world_rank=None, world_size=None):
|
||||
SINGLE_CHECKPOINT_FILENAME = '{prefix}.ort.pt'
|
||||
MULTIPLE_CHECKPOINT_FILENAME = '{prefix}.ZeRO.{world_rank}.{world_size}.ort.pt'
|
||||
|
||||
if is_partitioned:
|
||||
filename = MULTIPLE_CHECKPOINT_FILENAME.format(prefix=prefix, world_rank=world_rank, world_size=(world_size-1))
|
||||
else:
|
||||
filename = SINGLE_CHECKPOINT_FILENAME.format(prefix=prefix)
|
||||
return filename
|
||||
|
||||
|
||||
class _CombineZeroCheckpoint(object):
|
||||
def __init__(self, checkpoint_files, clean_state_dict=None):
|
||||
|
||||
assert len(checkpoint_files) > 0, "No checkpoint files passed"
|
||||
self.checkpoint_files = checkpoint_files
|
||||
self.clean_state_dict = clean_state_dict
|
||||
self.world_size = int(self.checkpoint_files[0].split('ZeRO')[1].split('.')[2]) + 1
|
||||
assert len(self.checkpoint_files) == self.world_size, f"Could not find {self.world_size} files"
|
||||
self.weight_shape_map = dict()
|
||||
|
||||
def _is_sharded(self, name):
|
||||
if '_view_' in name:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _has_fp16_weights(self, state_dict):
|
||||
for k in state_dict.keys():
|
||||
if k.endswith('_fp16'):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _split_moment_name(self, name):
|
||||
name_split = name.split('_view_')
|
||||
if(len(name_split) > 1):
|
||||
view_num = int(name_split[1])
|
||||
else:
|
||||
view_num = None
|
||||
weight_name = name_split[0].split('Moment_')[1][2:]
|
||||
moment_num = int(name_split[0].split('Moment_')[1][0])
|
||||
return moment_num, weight_name, view_num
|
||||
|
||||
def _update_weight_statistics(self, name, value):
|
||||
self.weight_shape_map[name] = value.size() # original shape of tensor
|
||||
|
||||
def _reshape_tensors(self, state_dict, fp16):
|
||||
for k, v in state_dict.items():
|
||||
if k.startswith('Moment_'):
|
||||
_, weight_name, _ = self._split_moment_name(k)
|
||||
set_size = self.weight_shape_map[weight_name]
|
||||
state_dict[k] = v.reshape(set_size)
|
||||
state_dict[weight_name] = state_dict[weight_name].reshape(set_size)
|
||||
return state_dict
|
||||
|
||||
def aggregate_checkpoints(self):
|
||||
checkpoint_dir = os.path.dirname(self.checkpoint_files[0])
|
||||
checkpoint_prefix = self.checkpoint_files[0].split('.ZeRO')[0]
|
||||
self.aggregate_state_dict = dict()
|
||||
|
||||
is_fp16 = False
|
||||
weight_offset = dict()
|
||||
for i in range(self.world_size):
|
||||
checkpoint_name = _get_checkpoint_name(checkpoint_prefix, True, i, self.world_size)
|
||||
rank_state_dict = torch.load(checkpoint_name, map_location=torch.device("cpu"))
|
||||
if 'model' in rank_state_dict:
|
||||
rank_state_dict = rank_state_dict['model']
|
||||
|
||||
if self.clean_state_dict:
|
||||
rank_state_dict = self.clean_state_dict(rank_state_dict)
|
||||
|
||||
if i == 0:
|
||||
is_fp16 = self._has_fp16_weights(rank_state_dict)
|
||||
|
||||
for k, v in rank_state_dict.items():
|
||||
if k.startswith('Moment_'):
|
||||
moment_num, weight_name, view_num = self._split_moment_name(k)
|
||||
|
||||
if self._is_sharded(k):
|
||||
clean_name = 'Moment_' + str(moment_num) + '_' + weight_name
|
||||
if clean_name in self.aggregate_state_dict:
|
||||
# Found a previous shard of the moment, concatenate shards ordered by ranks
|
||||
self.aggregate_state_dict[clean_name] = torch.cat((self.aggregate_state_dict[clean_name], v), 0)
|
||||
else:
|
||||
self.aggregate_state_dict[clean_name] = v
|
||||
else:
|
||||
# Moment is not sharded, add as is
|
||||
self.aggregate_state_dict[k] = v
|
||||
|
||||
if is_fp16 and moment_num == 1:
|
||||
# FP32 weights are sharded, patch together based on moments
|
||||
if view_num == 0:
|
||||
# This FP32 weight's first shard is present on this rank,
|
||||
# flatten and add the weight's first view
|
||||
self.aggregate_state_dict[weight_name] = rank_state_dict[weight_name].view(-1)
|
||||
self._update_weight_statistics(weight_name, rank_state_dict[weight_name])
|
||||
weight_offset[weight_name] = v.numel()
|
||||
|
||||
elif view_num == 1:
|
||||
# This FP32 weight is carryforward from previous rank
|
||||
# Get start and end of weight slice to be updated from this rank
|
||||
weight_start = weight_offset[weight_name]
|
||||
weight_end = weight_start + v.numel()
|
||||
|
||||
if weight_start:
|
||||
old_value = self.aggregate_state_dict[weight_name]
|
||||
new_value = rank_state_dict[weight_name].view(-1)
|
||||
# patch the weight together
|
||||
self.aggregate_state_dict[weight_name] = torch.cat((old_value[:weight_start],
|
||||
new_value[weight_start:weight_end],
|
||||
old_value[weight_end:]), 0)
|
||||
|
||||
# update offset for next view
|
||||
weight_offset[weight_name] = weight_end
|
||||
|
||||
elif k.startswith('Update_Count'):
|
||||
clean_name = k.split('_view_')[0]
|
||||
# add a single copy of the 'Update_Count' tensor for current weight
|
||||
if clean_name not in self.aggregate_state_dict:
|
||||
self.aggregate_state_dict[clean_name] = v
|
||||
|
||||
else:
|
||||
if k not in self.aggregate_state_dict:
|
||||
self.aggregate_state_dict[k] = v
|
||||
if not (k.endswith('_fp16') or k == 'Step'):
|
||||
# FP32 Weight
|
||||
self._update_weight_statistics(k, v)
|
||||
|
||||
final_state_dict = self._reshape_tensors(
|
||||
self.aggregate_state_dict, is_fp16)
|
||||
return final_state_dict
|
||||
|
|
@ -4,6 +4,7 @@ import os
|
|||
import onnx
|
||||
import torch
|
||||
from inspect import signature
|
||||
import warnings
|
||||
|
||||
import onnxruntime as ort
|
||||
from . import _utils, amp, optim, postprocess, ORTTrainerOptions
|
||||
|
|
@ -194,7 +195,12 @@ class ORTTrainer(object):
|
|||
ort.set_cuda_mem_limit(self.options.device.mem_limit)
|
||||
ort.set_cuda_device_id(_utils.get_device_index(self.options.device.id))
|
||||
|
||||
# TODO: thiagofc: Checkpoint related for redesign
|
||||
self._original_model_state_keys = list(model.state_dict().keys()) if hasattr(model, 'state_dict') else []
|
||||
self._state_dict = None
|
||||
|
||||
self._train_step_info = TrainStepInfo(self.optim_config)
|
||||
self._training_session = None
|
||||
self._init_session()
|
||||
|
||||
def eval_step(self, *args, **kwargs):
|
||||
|
|
@ -246,6 +252,7 @@ class ORTTrainer(object):
|
|||
results = [session_run_results[o_desc.name] for o_desc in outputs_desc]
|
||||
return results[0] if len (results) == 1 else results
|
||||
|
||||
|
||||
def save_as_onnx(self, path):
|
||||
r"""Persists ONNX model into :py:attr:`path`
|
||||
|
||||
|
|
@ -395,7 +402,7 @@ class ORTTrainer(object):
|
|||
# Input can be a list or dict
|
||||
is_list_input = (match
|
||||
or len(input_names) >= len(ordered_input_list)
|
||||
or not all(x in ordered_list_kes for x in input_names))
|
||||
or not all(x in ordered_input_list for x in input_names))
|
||||
|
||||
class CombineTorchModelLossFn(torch.nn.Module):
|
||||
def __init__(self, model, loss_fn, input_names):
|
||||
|
|
@ -643,6 +650,11 @@ class ORTTrainer(object):
|
|||
self._model_desc_outputs_with_gradient_accumulation = [
|
||||
*self.model_desc.outputs, self.model_desc.gradient_accumulation]
|
||||
|
||||
# TODO: thiagofc: Checkpoint related for redesign
|
||||
if self._state_dict:
|
||||
self.load_state_dict(self._state_dict, self._load_state_dict_strict)
|
||||
self._state_dict = None
|
||||
|
||||
def _prepare_model_input(self, inputs_desc, lr, loss_scale, *inputs, **kwargs):
|
||||
# Normalize input to tuple of samples
|
||||
if type(inputs) == tuple and len(inputs) == 1 and type(inputs[0]) == list:
|
||||
|
|
@ -665,7 +677,7 @@ class ORTTrainer(object):
|
|||
# Append loss scale
|
||||
if loss_scale:
|
||||
assert self.options.mixed_precision.enabled, "Loss scale cannot be used without mixed precision"
|
||||
loss_scale = torch.tensor(loss_scale)
|
||||
loss_scale = loss_scale.clone().detach()
|
||||
input += (loss_scale, )
|
||||
extra_inputs += 1
|
||||
|
||||
|
|
|
|||
|
|
@ -76,11 +76,18 @@ class ORTTrainerOptions(object):
|
|||
'type' : 'boolean',
|
||||
'default' : False
|
||||
},
|
||||
'deepspeed_zero_stage' : {
|
||||
'type' : 'integer',
|
||||
'min' : 0,
|
||||
'max' : 1,
|
||||
'default' : 0
|
||||
'deepspeed_zero_optimization' : {
|
||||
'type' : 'dict',
|
||||
'default': {},
|
||||
'required': False,
|
||||
'schema': {
|
||||
'stage': {
|
||||
'type': 'integer',
|
||||
'min': 0,
|
||||
'max': 1,
|
||||
'default': 0
|
||||
},
|
||||
}
|
||||
},
|
||||
'enable_adasum' : {
|
||||
'type' : 'boolean',
|
||||
|
|
@ -183,8 +190,10 @@ class ORTTrainerOptions(object):
|
|||
distributed.allreduce_post_accumulation (bool, default is False):
|
||||
True enables overlap of AllReduce with computation, while False,
|
||||
postpone AllReduce until all gradients are ready
|
||||
distributed.deepspeed_zero_stage (int, default is 0):
|
||||
select which stage of DeepSpeed ZeRO technique to use. Stage 0 means disabled.
|
||||
distributed.deepspeed_zero_optimization:
|
||||
DeepSpeed ZeRO options.
|
||||
distributed.deepspeed_zero_optimization.stage (int, default is 0):
|
||||
select which stage of DeepSpeed ZeRO to use. Stage 0 means disabled.
|
||||
distributed.enable_adasum (bool, default is False):
|
||||
enable `Adasum <https://github.com/horovod/horovod/pull/1484>`_
|
||||
algorithm for AllReduce
|
||||
|
|
|
|||
|
|
@ -82,3 +82,17 @@ def _assert_state_dict_weights(state_dict_a, state_dict_b, verbose, rtol, atol):
|
|||
if verbose:
|
||||
print(f'Weight name: {a_name}: absolute difference: {np.abs(np_a_vals-np_b_vals).max()}')
|
||||
assert_allclose(a_val, b_val, rtol=rtol, atol=atol, err_msg=f"Weight mismatch for {a_name}")
|
||||
|
||||
# TODO: thiagofc: Checkpoint related for redesign
|
||||
def _get_name(name):
|
||||
if os.path.exists(name):
|
||||
return name
|
||||
rel = os.path.join("testdata", name)
|
||||
if os.path.exists(rel):
|
||||
return rel
|
||||
this = os.path.dirname(__file__)
|
||||
data = os.path.join(this, "..", "testdata")
|
||||
res = os.path.join(data, name)
|
||||
if os.path.exists(res):
|
||||
return res
|
||||
raise FileNotFoundError("Unable to find '{0}' or '{1}' or '{2}'".format(name, rel, res))
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
# generate sample input for our example
|
||||
import inspect
|
||||
import onnx
|
||||
import os
|
||||
|
|
@ -6,15 +5,16 @@ import math
|
|||
import pytest
|
||||
import copy
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from numpy.testing import assert_allclose
|
||||
|
||||
from onnxruntime import set_seed
|
||||
import onnxruntime
|
||||
from onnxruntime.capi.ort_trainer import IODescription as Legacy_IODescription,\
|
||||
ModelDescription as Legacy_ModelDescription,\
|
||||
LossScaler as Legacy_LossScaler,\
|
||||
ORTTrainer as Legacy_ORTTrainer
|
||||
from onnxruntime.experimental import _utils, amp, optim, orttrainer, TrainStepInfo,\
|
||||
from onnxruntime.experimental import _utils, amp, checkpoint, optim, orttrainer, TrainStepInfo,\
|
||||
model_desc_validation as md_val,\
|
||||
orttrainer_options as orttrainer_options
|
||||
|
||||
|
|
@ -27,9 +27,10 @@ import _test_helpers
|
|||
|
||||
|
||||
def generate_random_input_from_model_desc(desc, seed=1, device = "cuda:0"):
|
||||
'''Generates a sample input for the BERT model using the model desc.'''
|
||||
'''Generates a sample input for the BERT model using the model desc'''
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
dtype = torch.int64
|
||||
vocab_size = 30528
|
||||
num_classes = [vocab_size, 2, 2, vocab_size, 2]
|
||||
|
|
@ -49,6 +50,7 @@ def generate_random_input_from_model_desc(desc, seed=1, device = "cuda:0"):
|
|||
|
||||
def bert_model_description(dynamic_shape=True):
|
||||
'''Creates the model description dictionary with static dimensions'''
|
||||
|
||||
if dynamic_shape:
|
||||
model_desc = {'inputs': [('input_ids', ['batch_size', 'seq_len']),
|
||||
('segment_ids', ['batch_size', 'seq_len'],),
|
||||
|
|
@ -67,8 +69,10 @@ def bert_model_description(dynamic_shape=True):
|
|||
'outputs': [('loss', [], True)]}
|
||||
return model_desc
|
||||
|
||||
|
||||
def optimizer_parameters(model):
|
||||
'''A method to assign different hyper parameters for different model parameter groups'''
|
||||
|
||||
no_decay_keys = ["bias", "gamma", "beta", "LayerNorm"]
|
||||
no_decay_param_group = []
|
||||
for initializer in model.graph.initializer:
|
||||
|
|
@ -77,11 +81,13 @@ def optimizer_parameters(model):
|
|||
params = [{'params': no_decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6}]
|
||||
return params
|
||||
|
||||
|
||||
def load_bert_onnx_model():
|
||||
bert_onnx_model_path = os.path.join('..', '..', '..', 'onnxruntime', 'test', 'testdata', "bert_toy_postprocessed.onnx")
|
||||
model = onnx.load(bert_onnx_model_path)
|
||||
return model
|
||||
|
||||
|
||||
class CustomLossScaler(amp.LossScaler):
|
||||
def __init__(self, loss_scale=float(1 << 16)):
|
||||
super().__init__(loss_scale)
|
||||
|
|
@ -108,10 +114,10 @@ class LegacyCustomLossScaler():
|
|||
def update_loss_scale(self, is_all_finite):
|
||||
self.loss_scale_ *= 0.9
|
||||
|
||||
def legacy_model_params(device = torch.device("cuda", 0)):
|
||||
|
||||
def legacy_model_params(lr, device = torch.device("cuda", 0)):
|
||||
legacy_model_desc = legacy_bert_model_description()
|
||||
learning_rate_description = legacy_ort_trainer_learning_rate_description()
|
||||
lr = 0.001
|
||||
learning_rate = torch.tensor([lr]).to(device)
|
||||
return (legacy_model_desc, learning_rate_description, learning_rate)
|
||||
|
||||
|
|
@ -122,9 +128,11 @@ def legacy_model_params(device = torch.device("cuda", 0)):
|
|||
else:
|
||||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6}
|
||||
|
||||
|
||||
def legacy_ort_trainer_learning_rate_description():
|
||||
return Legacy_IODescription('Learning_Rate', [1, ], torch.float32)
|
||||
|
||||
|
||||
def legacy_bert_model_description():
|
||||
vocab_size = 30528
|
||||
input_ids_desc = Legacy_IODescription('input_ids', ['batch', 'max_seq_len_in_batch'])
|
||||
|
|
@ -137,12 +145,15 @@ def legacy_bert_model_description():
|
|||
return Legacy_ModelDescription([input_ids_desc, segment_ids_desc, input_mask_desc, masked_lm_labels_desc,
|
||||
next_sentence_labels_desc], [loss_desc])
|
||||
|
||||
|
||||
def legacy_constant_lr_scheduler_1(global_step):
|
||||
return legacy_constant_lr_scheduler(global_step, 1.0)
|
||||
|
||||
|
||||
def legacy_constant_lr_scheduler_5(global_step):
|
||||
return legacy_constant_lr_scheduler(global_step, 0.5)
|
||||
|
||||
|
||||
def legacy_constant_lr_scheduler(global_step, initial_lr):
|
||||
warmup = 0.5
|
||||
total_steps = 10
|
||||
|
|
@ -156,6 +167,7 @@ def legacy_constant_lr_scheduler(global_step, initial_lr):
|
|||
lr *= warmup_val
|
||||
return lr
|
||||
|
||||
|
||||
def legacy_cosine_lr_scheduler(global_step):
|
||||
initial_lr = 1.0
|
||||
warmup = 0.5
|
||||
|
|
@ -170,6 +182,7 @@ def legacy_cosine_lr_scheduler(global_step):
|
|||
lr *= warmup_val
|
||||
return lr
|
||||
|
||||
|
||||
def legacy_linear_lr_scheduler(global_step):
|
||||
initial_lr = 1.0
|
||||
warmup = 0.5
|
||||
|
|
@ -184,6 +197,7 @@ def legacy_linear_lr_scheduler(global_step):
|
|||
lr *= warmup_val
|
||||
return lr
|
||||
|
||||
|
||||
def legacy_poly_lr_scheduler(global_step):
|
||||
initial_lr = 1.0
|
||||
warmup = 0.5
|
||||
|
|
@ -199,21 +213,24 @@ def legacy_poly_lr_scheduler(global_step):
|
|||
lr *= warmup_val
|
||||
return lr
|
||||
|
||||
|
||||
def legacy_optim_params_a(name):
|
||||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6}
|
||||
|
||||
|
||||
def legacy_optim_params_b(name):
|
||||
params = ['bert.embeddings.LayerNorm.bias', 'bert.embeddings.LayerNorm.weight']
|
||||
if name in params:
|
||||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6}
|
||||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6}
|
||||
|
||||
|
||||
def legacy_optim_params_c(name):
|
||||
params_group = optimizer_parameters(load_bert_onnx_model())
|
||||
if name in params_group[0]['params']:
|
||||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6}
|
||||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6}
|
||||
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Testing starts here #########################################################
|
||||
|
|
@ -230,28 +247,29 @@ def testToyBERTModelSimpleTrainStep(dynamic_shape):
|
|||
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({})
|
||||
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
for i in range(10):
|
||||
# Generate random sample batch using dimensions
|
||||
sample_input = generate_random_input_from_model_desc(model_desc)
|
||||
|
||||
output = trainer.train_step(*sample_input)
|
||||
assert output.shape == torch.Size([])
|
||||
assert output.shape == torch.Size([])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expected_losses", [
|
||||
([10.988012313842773, 10.99226188659668, 11.090812683105469, 11.042860984802246, 10.988919258117676,
|
||||
11.105875015258789, 10.981894493103027, 11.081543922424316, 10.997451782226562, 11.10739517211914])
|
||||
])
|
||||
def testToyBERTDeterministicCheck(expected_losses):
|
||||
# Common setup
|
||||
train_steps = 10
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# Modeling
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
params = optimizer_parameters(model)
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
|
|
@ -262,19 +280,18 @@ def testToyBERTDeterministicCheck(expected_losses):
|
|||
'id': device,
|
||||
},
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
|
||||
# Train
|
||||
experimental_losses = []
|
||||
for i in range(train_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
|
||||
experimental_losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
|
||||
# Check output
|
||||
_test_helpers.assert_model_outputs(experimental_losses, expected_losses)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("initial_lr, lr_scheduler, expected_learning_rates, expected_losses", [
|
||||
(1.0, optim.lr_scheduler.ConstantWarmupLRScheduler, [0.18181818181818182, 0.06611570247933884, 0.03606311044327573, 0.026227716686018716, 0.02384337880547156,\
|
||||
0.02384337880547156, 0.02384337880547156, 0.02384337880547156, 0.02384337880547156, 0.02384337880547156],
|
||||
|
|
@ -298,12 +315,16 @@ def testToyBERTDeterministicCheck(expected_losses):
|
|||
10.974217414855957, 10.96664810180664, 11.193868637084961, 11.14560604095459, 11.097070693969727])
|
||||
])
|
||||
def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rates, expected_losses):
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
# Common setup
|
||||
device = 'cuda'
|
||||
total_steps = 10
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# Modeling
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
optim_config = optim.LambConfig(lr=initial_lr)
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
|
|
@ -314,22 +335,20 @@ def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rate
|
|||
},
|
||||
'lr_scheduler' : lr_scheduler(total_steps=total_steps, warmup=0.5)
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
|
||||
# Train
|
||||
losses = []
|
||||
learning_rates = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
learning_rates.append(trainer.options.lr_scheduler.get_last_lr()[0])
|
||||
|
||||
|
||||
# Check output
|
||||
_test_helpers.assert_model_outputs(learning_rates, expected_learning_rates)
|
||||
_test_helpers.assert_model_outputs(losses, expected_losses, rtol=1e-6)
|
||||
|
||||
|
||||
|
||||
# Dynamic Loss Scaler implemented implicitly
|
||||
@pytest.mark.parametrize("loss_scaler, expected_losses", [
|
||||
|
|
@ -341,13 +360,16 @@ def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rate
|
|||
11.105667114257812, 10.981982231140137, 11.081765174865723, 10.997125625610352, 11.107298851013184])
|
||||
])
|
||||
def testToyBERTModelMixedPrecisionLossScaler(loss_scaler, expected_losses):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# Modeling
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
|
|
@ -361,18 +383,17 @@ def testToyBERTModelMixedPrecisionLossScaler(loss_scaler, expected_losses):
|
|||
'loss_scaler': loss_scaler
|
||||
}
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
|
||||
# Train
|
||||
losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
|
||||
losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
|
||||
_test_helpers.assert_model_outputs(losses, expected_losses, rtol=1e-5)
|
||||
# Check output
|
||||
_test_helpers.assert_model_outputs(losses, expected_losses, rtol=1e-4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gradient_accumulation_steps, expected_losses", [
|
||||
(1, [10.988012313842773, 10.99226188659668, 11.090812683105469, 11.042860984802246, 10.988919258117676,
|
||||
|
|
@ -383,13 +404,16 @@ def testToyBERTModelMixedPrecisionLossScaler(loss_scaler, expected_losses):
|
|||
11.112862586975098, 10.996183395385742, 11.072013854980469, 11.00184154510498, 11.097928047180176])
|
||||
])
|
||||
def testToyBERTModelGradientAccumulation(gradient_accumulation_steps, expected_losses):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
device = "cuda"
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# Modeling
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
|
|
@ -402,17 +426,172 @@ def testToyBERTModelGradientAccumulation(gradient_accumulation_steps, expected_l
|
|||
'gradient_accumulation_steps' : gradient_accumulation_steps
|
||||
},
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
|
||||
# Train
|
||||
losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
|
||||
_test_helpers.assert_model_outputs(losses, expected_losses)
|
||||
|
||||
# Check output
|
||||
_test_helpers.assert_model_outputs(losses, expected_losses, rtol=1e-6)
|
||||
|
||||
|
||||
def testToyBertCheckpointBasic():
|
||||
# Common setup
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({'debug' : {'deterministic_compute': True}})
|
||||
|
||||
# Create ORTTrainer and save initial state in a dict
|
||||
model = load_bert_onnx_model()
|
||||
model_desc = bert_model_description()
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
sd = checkpoint.experimental_state_dict(trainer)
|
||||
|
||||
## All initializers must be present in the state_dict
|
||||
## when the specified model for ORTTRainer is an ONNX model
|
||||
for param in trainer._onnx_model.graph.initializer:
|
||||
assert param.name in sd
|
||||
|
||||
## Modify one of the state values and load into ORTTrainer
|
||||
sd['bert.encoder.layer.0.attention.output.LayerNorm.weight'] += 10
|
||||
checkpoint.experimental_load_state_dict(trainer, sd)
|
||||
|
||||
## Save a checkpoint
|
||||
ckpt_dir = _test_helpers._get_name("ort_ckpt")
|
||||
checkpoint.experimental_save_checkpoint(trainer, ckpt_dir, 'bert_toy_save_test')
|
||||
del trainer
|
||||
del model
|
||||
|
||||
# Create a new ORTTrainer and load the checkpoint from previous ORTTrainer
|
||||
model2 = load_bert_onnx_model()
|
||||
model_desc2 = bert_model_description()
|
||||
trainer2 = orttrainer.ORTTrainer(model2, model_desc2, optim_config, options=opts)
|
||||
checkpoint.experimental_load_checkpoint(trainer2, ckpt_dir, 'bert_toy_save_test')
|
||||
loaded_sd = checkpoint.experimental_state_dict(trainer2)
|
||||
|
||||
# Assert whether original state and the one loaded from checkpoint matches
|
||||
for k,v in loaded_sd.items():
|
||||
assert torch.all(torch.eq(v, sd[k]))
|
||||
|
||||
|
||||
def testToyBertCheckpointLoadZero():
|
||||
# Common setup
|
||||
rtol = 1e-03
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({'debug' : {'deterministic_compute': True},
|
||||
'device' : {'id' : device},
|
||||
'distributed' : {'allreduce_post_accumulation' : True}})
|
||||
|
||||
# Create ORTTrainer and save initial state in a dict
|
||||
model = load_bert_onnx_model()
|
||||
model_desc = bert_model_description()
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
ckpt_dir = _test_helpers._get_name("ort_ckpt")
|
||||
checkpoint.experimental_load_checkpoint(trainer, ckpt_dir, 'bert_toy_lamb')
|
||||
|
||||
# Expected values
|
||||
expected_eval_loss = [10.997552871]
|
||||
input_ids = torch.tensor([[26598],[21379],[19922],[ 5219],[ 5644],[20559],[23777],[25672],[22969],[16824],[16822],[635],[27399],[20647],[18519],[15546]], device=device)
|
||||
segment_ids = torch.tensor([[0],[1],[0],[1],[0],[0],[1],[0],[0],[1],[1],[0],[0],[1],[1],[1]], device=device)
|
||||
input_mask = torch.tensor([[0],[0],[0],[0],[1],[1],[1],[0],[1],[1],[0],[0],[0],[1],[0],[0]], device=device)
|
||||
masked_lm_labels = torch.tensor([[25496],[16184],[11005],[16228],[14884],[21660],[ 8678],[23083],[ 4027],[ 8397],[11921],[ 1333],[26482],[ 1666],[17925],[27978]], device=device)
|
||||
next_sentence_labels = torch.tensor([0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 1, 1, 0], device=device)
|
||||
|
||||
# Actual values
|
||||
actual_eval_loss = trainer.eval_step(input_ids, segment_ids, input_mask, masked_lm_labels, next_sentence_labels)
|
||||
actual_eval_loss = actual_eval_loss.cpu().numpy().item(0)
|
||||
|
||||
# Check results
|
||||
assert_allclose(expected_eval_loss, actual_eval_loss, rtol=rtol)
|
||||
|
||||
|
||||
def testToyBertStateDictWrapModelLossFn():
|
||||
# Common setup
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# Modeling
|
||||
class LinearModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.linear = torch.nn.Linear(2, 4)
|
||||
def forward(self, y=None, x=None):
|
||||
if y is not None:
|
||||
return self.linear(x) + y
|
||||
else:
|
||||
return self.linear(x) + torch.ones(2, 4)
|
||||
pt_model = LinearModel()
|
||||
model_desc = {'inputs' : [('x', [2, 2]),
|
||||
('label', [2, ])],
|
||||
'outputs' : [('loss', [], True),
|
||||
('output', [2, 4])]}
|
||||
optim_config = optim.SGDConfig(lr=0.02)
|
||||
def loss_fn(x, label):
|
||||
return F.nll_loss(F.log_softmax(x, dim=1), label)
|
||||
trainer = orttrainer.ORTTrainer(pt_model, model_desc, optim_config, loss_fn=loss_fn)
|
||||
|
||||
# Compare resulting state_dict keys before train
|
||||
state_dict = checkpoint.experimental_state_dict(trainer)
|
||||
assert state_dict == {}
|
||||
|
||||
# Executing train_step() once
|
||||
data = torch.randn(2, 2)
|
||||
label = torch.tensor([0, 1], dtype=torch.int64)
|
||||
trainer.train_step(x=data, label=label)
|
||||
|
||||
# Compare resulting state_dict keys after train
|
||||
state_dict = checkpoint.experimental_state_dict(trainer)
|
||||
assert state_dict.keys() == {'linear.bias', 'linear.weight'}
|
||||
|
||||
|
||||
def testToyBertCheckpointFrozenWeights():
|
||||
# Common setup
|
||||
seed = 1
|
||||
total_steps = 10
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
opts = orttrainer.ORTTrainerOptions({'debug' : {'deterministic_compute': True},
|
||||
'utils' : {'frozen_weights' : ['bert.encoder.layer.0.attention.self.value.weight']}})
|
||||
|
||||
# Create ORTTrainer and save initial state in a dict
|
||||
model = load_bert_onnx_model()
|
||||
model_desc = bert_model_description()
|
||||
optim_config = optim.LambConfig()
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
# Train for a few steps
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, seed)
|
||||
_ = trainer.train_step(*sample_input)
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, seed + total_steps + 1)
|
||||
# Evaluate once to get a base loss
|
||||
loss = trainer.eval_step(*sample_input)
|
||||
# Save checkpoint
|
||||
state_dict = checkpoint.experimental_state_dict(trainer)
|
||||
|
||||
# Load previous state into another instance of ORTTrainer
|
||||
model2 = load_bert_onnx_model()
|
||||
model_desc2 = bert_model_description()
|
||||
optim_config2 = optim.LambConfig()
|
||||
trainer2 = orttrainer.ORTTrainer(model2, model_desc2, optim_config2, options=opts)
|
||||
checkpoint.experimental_load_state_dict(trainer2, state_dict)
|
||||
# Evaluate once to get a base loss
|
||||
ckpt_loss = trainer2.eval_step(*sample_input)
|
||||
|
||||
# Must match as both trainers have the same dict state
|
||||
assert_allclose(loss.cpu(), ckpt_loss.cpu())
|
||||
loaded_state_dict = checkpoint.experimental_state_dict(trainer2)
|
||||
assert state_dict.keys() == loaded_state_dict.keys()
|
||||
|
||||
|
||||
###############################################################################
|
||||
|
|
@ -421,14 +600,16 @@ def testToyBERTModelGradientAccumulation(gradient_accumulation_steps, expected_l
|
|||
|
||||
|
||||
def testToyBERTModelLegacyExperimentalBasicTraining():
|
||||
# Common setup
|
||||
train_steps = 10
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# EXPERIMENTAL API
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
params = optimizer_parameters(model)
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
|
|
@ -439,23 +620,17 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
'id': device,
|
||||
},
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
experimental_losses = []
|
||||
for i in range(train_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
|
||||
experimental_losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
|
||||
# LEGACY IMPLEMENTATION
|
||||
device = torch.device(device)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params()
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
||||
onnxruntime.set_seed(seed)
|
||||
device = torch.device(device)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(lr=0.001)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
None,
|
||||
learning_rate_description,
|
||||
|
|
@ -463,11 +638,12 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
legacy_losses = []
|
||||
for i in range(train_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
legacy_sample_input = [*sample_input, learning_rate]
|
||||
leg_loss = legacy_trainer.train_step(*sample_input, learning_rate)
|
||||
legacy_losses.append(leg_loss.cpu().item())
|
||||
|
||||
legacy_losses.append(legacy_trainer.train_step(legacy_sample_input).cpu().item())
|
||||
# Check results
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, True, rtol=1e-5)
|
||||
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, True, rtol=1e-4)
|
||||
|
||||
@pytest.mark.parametrize("initial_lr, lr_scheduler, legacy_lr_scheduler", [
|
||||
(1.0, optim.lr_scheduler.ConstantWarmupLRScheduler, legacy_constant_lr_scheduler_1),
|
||||
|
|
@ -477,13 +653,20 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
(1.0, optim.lr_scheduler.PolyWarmupLRScheduler, legacy_poly_lr_scheduler),
|
||||
])
|
||||
def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, legacy_lr_scheduler):
|
||||
# EXPERIMENTAL API
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
############################################################################
|
||||
# These tests require hard-coded values for 'total_steps' and 'initial_lr' #
|
||||
############################################################################
|
||||
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
|
||||
# EXPERIMENTAL API
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
optim_config = optim.LambConfig(lr=initial_lr)
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
|
|
@ -494,25 +677,18 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
|
|||
},
|
||||
'lr_scheduler' : lr_scheduler(total_steps=total_steps, warmup=0.5)
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
experimental_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
experimental_losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
|
||||
assert trainer.options.lr_scheduler.get_last_lr()[0] == legacy_lr_scheduler(i)
|
||||
|
||||
# LEGACY IMPLEMENTATION
|
||||
device = torch.device(device)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params()
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
||||
onnxruntime.set_seed(seed)
|
||||
device = torch.device(device)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(initial_lr)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
None,
|
||||
learning_rate_description,
|
||||
|
|
@ -522,24 +698,29 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
|
|||
legacy_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
leg_loss = legacy_trainer.train_step(*sample_input)
|
||||
legacy_losses.append(leg_loss.cpu().item())
|
||||
|
||||
legacy_losses.append(legacy_trainer.train_step(sample_input).cpu().item())
|
||||
|
||||
# Check results
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("loss_scaler, legacy_loss_scaler", [
|
||||
(None, Legacy_LossScaler("ort_test_input_loss_scaler", True)),
|
||||
(amp.DynamicLossScaler(), Legacy_LossScaler("ort_test_input_loss_scaler", True)),
|
||||
(CustomLossScaler(), LegacyCustomLossScaler())
|
||||
])
|
||||
def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, legacy_loss_scaler):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
device = "cuda"
|
||||
seed = 1
|
||||
|
||||
# EXPERIMENTAL IMPLEMENTATION
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
|
|
@ -553,11 +734,7 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
'loss_scaler': loss_scaler
|
||||
}
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
experimental_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
|
|
@ -565,10 +742,9 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
|
||||
# LEGACY IMPLEMENTATION
|
||||
device = torch.device(device)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params()
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(optim_config.lr)
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
||||
onnxruntime.set_seed(seed)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
None,
|
||||
learning_rate_description,
|
||||
|
|
@ -579,10 +755,10 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
legacy_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
legacy_sample_input = [*sample_input, learning_rate]
|
||||
leg_loss = legacy_trainer.train_step(*sample_input, learning_rate)
|
||||
legacy_losses.append(leg_loss.cpu().item())
|
||||
|
||||
legacy_losses.append(legacy_trainer.train_step(legacy_sample_input).cpu().item())
|
||||
|
||||
# Check results
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, rtol=1e-5)
|
||||
|
||||
|
||||
|
|
@ -592,13 +768,16 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
(7)
|
||||
])
|
||||
def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation_steps):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
device = "cuda"
|
||||
seed = 1
|
||||
|
||||
# EXPERIMENTAL IMPLEMENTATION
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
|
|
@ -611,22 +790,18 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
|
|||
'gradient_accumulation_steps' : gradient_accumulation_steps
|
||||
},
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
experimental_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
experimental_losses.append(trainer.train_step(*sample_input).cpu().item())
|
||||
loss = trainer.train_step(*sample_input)
|
||||
experimental_losses.append(loss.cpu().item())
|
||||
|
||||
# LEGACY IMPLEMENTATION
|
||||
device = torch.device(device)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params()
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
||||
onnxruntime.set_seed(seed)
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(optim_config.lr)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
None,
|
||||
learning_rate_description,
|
||||
|
|
@ -636,8 +811,8 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
|
|||
legacy_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
legacy_sample_input = [*sample_input, learning_rate]
|
||||
leg_loss = legacy_trainer.train_step(*sample_input, learning_rate)
|
||||
legacy_losses.append(leg_loss.cpu().item())
|
||||
|
||||
legacy_losses.append(legacy_trainer.train_step(legacy_sample_input).cpu().item())
|
||||
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses)
|
||||
# Check results
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, rtol=1e-6)
|
||||
|
|
|
|||
|
|
@ -659,7 +659,7 @@ def testORTDeterministicCompute(seed, device):
|
|||
data, targets = batcher_fn(train_data, 0)
|
||||
_ = first_trainer.train_step(data, targets)
|
||||
assert first_trainer._onnx_model is not None
|
||||
|
||||
|
||||
# Setup for the second ORTTRainer run
|
||||
torch.manual_seed(seed)
|
||||
set_seed(seed)
|
||||
|
|
|
|||
Loading…
Reference in a new issue