mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-13 18:08:13 +00:00
Fix checkpoint API and improve loss scaler handling (#4950)
This PR also includes:
* More LossScaler tests
* Minor LossScaler improvement
* Check model after extra post processing
* Improve basic training tests to include all optimizers
* Set rtol=1e-7 tolerance for Legacy vs Experimental frontend API tests
* Increase number of training tests for Legacy vs Experimental tests
* Minor refactoring on existing tests
* Fix Checkpoint API for Gradient Accumulation / fp16 scenarios
This commit is contained in:
Thiago Crepaldi
GitHub
parent
eebc2cccce
commit
aabed34d5c
6 changed files with 257 additions and 114 deletions
|
|
@ -8,8 +8,10 @@ class LossScaler(object):
|
|||
"""
|
||||
|
||||
def __init__(self, loss_scale):
|
||||
assert isinstance(loss_scale, (int, float)) and loss_scale > 0, "'loss_scale' must be a positive float"
|
||||
self._input_name = None
|
||||
self._loss_scale = loss_scale
|
||||
self._loss_scale = float(loss_scale)
|
||||
self._initial_loss_scale = float(loss_scale)
|
||||
|
||||
@property
|
||||
def input_name(self):
|
||||
|
|
@ -27,12 +29,12 @@ class LossScaler(object):
|
|||
|
||||
@loss_scale.setter
|
||||
def loss_scale(self, loss_scale):
|
||||
assert isinstance(loss_scale, float) and loss_scale > 0, "'loss_scale' must be a positive float"
|
||||
self._loss_scale = loss_scale
|
||||
assert isinstance(loss_scale, (int, float)) and loss_scale > 0, "'loss_scale' must be a positive float"
|
||||
self._loss_scale = float(loss_scale)
|
||||
|
||||
def reset(self):
|
||||
r"""Resets loss scaler internal state"""
|
||||
raise NotImplementedError
|
||||
self._loss_scale = self._initial_loss_scale
|
||||
|
||||
def update(self, train_step_info):
|
||||
r"""Updates loss based on user input and training session info
|
||||
|
|
@ -93,15 +95,16 @@ class DynamicLossScaler(LossScaler):
|
|||
self.up_scale_window = up_scale_window
|
||||
self.min_loss_scale = min_loss_scale
|
||||
self.max_loss_scale = max_loss_scale
|
||||
|
||||
self._initial_loss_scale = loss_scale
|
||||
self._stable_steps_count = 0
|
||||
|
||||
def reset(self):
|
||||
self.loss_scale = self._initial_loss_scale
|
||||
super().reset()
|
||||
self._stable_steps_count = 0
|
||||
|
||||
def update(self, train_step_info):
|
||||
if not self.automatic_update:
|
||||
return self.loss_scale
|
||||
|
||||
if train_step_info.all_finite:
|
||||
self._stable_steps_count += 1
|
||||
|
||||
|
|
|
|||
|
|
@ -11,11 +11,11 @@ import warnings
|
|||
################################################################################
|
||||
|
||||
|
||||
def experimental_state_dict(ort_trainer):
|
||||
def experimental_state_dict(ort_trainer, include_optimizer_state=True):
|
||||
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 {}
|
||||
return ort_trainer._state_dict
|
||||
|
||||
# extract trained weights
|
||||
session_state = ort_trainer._training_session.get_state()
|
||||
|
|
@ -28,10 +28,11 @@ def experimental_state_dict(ort_trainer):
|
|||
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
|
||||
# Need to remove redundant (optimizer) initializers to map back to original torch state names
|
||||
if not include_optimizer_state and ort_trainer._torch_state_dict_keys:
|
||||
return {key: torch_state[key] for key in ort_trainer._torch_state_dict_keys if key in torch_state}
|
||||
return torch_state
|
||||
|
||||
|
||||
def experimental_load_state_dict(ort_trainer, state_dict, strict=False):
|
||||
# Note: It may happen ONNX model has not yet been initialized
|
||||
|
|
@ -42,7 +43,7 @@ def experimental_load_state_dict(ort_trainer, state_dict, strict=False):
|
|||
ort_trainer._load_state_dict_strict = strict
|
||||
return
|
||||
|
||||
# update onnx model from loaded state dict
|
||||
# Update onnx model from loaded state dict
|
||||
cur_initializers_names = [n.name for n in ort_trainer._onnx_model.graph.initializer]
|
||||
new_initializers = {}
|
||||
|
||||
|
|
@ -63,11 +64,11 @@ def experimental_load_state_dict(ort_trainer, state_dict, strict=False):
|
|||
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)}
|
||||
def experimental_save_checkpoint(ort_trainer, checkpoint_dir, checkpoint_prefix="ORT_checkpoint", checkpoint_state_dict=None, include_optimizer_state=True):
|
||||
if checkpoint_state_dict is None:
|
||||
checkpoint_state_dict = {'model': experimental_state_dict(ort_trainer, include_optimizer_state)}
|
||||
else:
|
||||
checkpoint_state_dict.update({'model': experimental_state_dict(ort_trainer)})
|
||||
checkpoint_state_dict.update({'model': experimental_state_dict(ort_trainer, include_optimizer_state)})
|
||||
|
||||
assert os.path.exists(checkpoint_dir), f"checkpoint_dir ({checkpoint_dir}) directory doesn't exist"
|
||||
|
||||
|
|
|
|||
|
|
@ -152,7 +152,7 @@ class _ORTTrainerModelDesc(object):
|
|||
|
||||
@gradient_accumulation.setter
|
||||
def gradient_accumulation(self, name):
|
||||
self._add_output_description(self, name, [1], False, torch.bool, None, GRADIENT_ACCUMULATION_IO_DESCRIPTION_NAME)
|
||||
self._add_output_description(self, name, [1], False, torch.bool, None, GRADIENT_ACCUMULATION_IO_DESCRIPTION_NAME, ignore_duplicate=True)
|
||||
|
||||
@property
|
||||
def all_finite(self):
|
||||
|
|
@ -160,7 +160,7 @@ class _ORTTrainerModelDesc(object):
|
|||
|
||||
@all_finite.setter
|
||||
def all_finite(self, name):
|
||||
self._add_output_description(self, name, [1], False, torch.bool, None, ALL_FINITE_IO_DESCRIPTION_NAME)
|
||||
self._add_output_description(self, name, [1], False, torch.bool, None, ALL_FINITE_IO_DESCRIPTION_NAME, ignore_duplicate=True)
|
||||
|
||||
@property
|
||||
def loss_scale_input(self):
|
||||
|
|
@ -168,9 +168,9 @@ class _ORTTrainerModelDesc(object):
|
|||
|
||||
@loss_scale_input.setter
|
||||
def loss_scale_input(self, name):
|
||||
self._add_input_description(self, name, [], torch.float32, LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME)
|
||||
self._add_input_description(self, name, [], torch.float32, LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME, ignore_duplicate=True)
|
||||
|
||||
def _add_input_description(self, node, name, shape, dtype=None, attr_name=None):
|
||||
def _add_input_description(self, node, name, shape, dtype=None, attr_name=None, ignore_duplicate=False):
|
||||
'''Add a new input description into the node object
|
||||
|
||||
If 'dtype' is specified, a typed input description namedtuple(name, shape, dtype) is created.
|
||||
|
|
@ -184,15 +184,22 @@ class _ORTTrainerModelDesc(object):
|
|||
shape (list): shape of input description
|
||||
dtype (torch.dtype): input data type
|
||||
attr_name (str, default is None): friendly name to allow direct access to the output description
|
||||
ignore_duplicate (bool, default is False): silently skips addition of duplicate inputs
|
||||
'''
|
||||
|
||||
assert isinstance(name, str) and len(name) > 0, "'name' is an invalid input name"
|
||||
not_found = True
|
||||
if not ignore_duplicate:
|
||||
if id(node) == id(self.inputs):
|
||||
not_found = all([name not in i_desc.name for i_desc in node])
|
||||
assert not_found, f"'name' {name} already exists in the inputs description"
|
||||
else:
|
||||
not_found = attr_name not in dir(self)
|
||||
assert not_found, f"'attr_name' {attr_name} already exists in the 'node'"
|
||||
elif not not_found:
|
||||
return
|
||||
assert isinstance(shape, list) and all([(isinstance(dim, int) or (isinstance(dim, str) and len(dim) > 0))\
|
||||
for dim in shape]), "'shape' must be a list of int or str with length at least 1"
|
||||
if id(node) == id(self.inputs):
|
||||
assert all([name not in i_desc.name for i_desc in node]), f"'name' {name} already exists in the inputs description"
|
||||
else:
|
||||
assert attr_name not in dir(self), f"'attr_name' {attr_name} already exists in the 'node'"
|
||||
assert dtype is None or isinstance(dtype, torch.dtype), "'dtype' must be either None or a torch.dtype type"
|
||||
if dtype:
|
||||
new_input_desc = self._InputDescriptionTyped(name, shape, dtype)
|
||||
|
|
@ -205,7 +212,7 @@ class _ORTTrainerModelDesc(object):
|
|||
assert isinstance(attr_name, str) and len(attr_name) > 0, "Invalid 'attr_name'"
|
||||
setattr(node, attr_name, new_input_desc)
|
||||
|
||||
def _add_output_description(self, node, name, shape, is_loss, dtype=None, dtype_amp=None, attr_name=None):
|
||||
def _add_output_description(self, node, name, shape, is_loss, dtype=None, dtype_amp=None, attr_name=None, ignore_duplicate=False):
|
||||
'''Add a new output description into the node object as a tuple
|
||||
|
||||
When (name, shape, is_loss, dtype) is specified, a typed output description is created
|
||||
|
|
@ -221,6 +228,7 @@ class _ORTTrainerModelDesc(object):
|
|||
dtype (torch.dtype): input data type
|
||||
dtype_amp (torch.dtype, default is None): input data type for evaluation with mixed precision.
|
||||
attr_name (str, default is None): friendly name to allow direct access to the output description
|
||||
ignore_duplicate (bool, default is False): silently skips addition of duplicate outputs
|
||||
'''
|
||||
|
||||
assert isinstance(name, str) and len(name) > 0, "'name' is an invalid output name"
|
||||
|
|
@ -228,13 +236,18 @@ class _ORTTrainerModelDesc(object):
|
|||
for dim in shape]), "'shape' must be a list of int or str with length at least 1"
|
||||
assert isinstance(is_loss, bool), "'is_loss' must be a bool"
|
||||
|
||||
if id(node) == id(self.outputs):
|
||||
assert all([name not in o_desc.name for o_desc in node]), f"'name' {name} already exists in the outputs description"
|
||||
is_loss_count = 1 if is_loss else 0
|
||||
assert all([o_desc.is_loss is False for o_desc in node]) if is_loss else True,\
|
||||
"Only one 'is_loss' is supported at outputs description"
|
||||
else:
|
||||
assert attr_name not in dir(self), f"'attr_name' {attr_name} already exists in the 'node'"
|
||||
not_found = True
|
||||
if not ignore_duplicate:
|
||||
if id(node) == id(self.outputs):
|
||||
not_found = all([name not in o_desc.name for o_desc in node])
|
||||
assert not_found, f"'name' {name} already exists in the outputs description"
|
||||
assert all([not o_desc.is_loss for o_desc in node]) if is_loss else True,\
|
||||
"Only one 'is_loss' is supported at outputs description"
|
||||
else:
|
||||
not_found = attr_name not in dir(self)
|
||||
assert not_found, f"'attr_name' {attr_name} already exists in the 'node'"
|
||||
elif not not_found:
|
||||
return
|
||||
|
||||
assert dtype is None or isinstance(dtype, torch.dtype), "'dtype' must be either None or a torch.dtype type"
|
||||
if dtype:
|
||||
|
|
|
|||
|
|
@ -144,6 +144,8 @@ class ORTTrainer(object):
|
|||
"'loss_fn' must be either 'None' or 'torch.nn.Module'"
|
||||
self._torch_model = model
|
||||
self.loss_fn = loss_fn
|
||||
# TODO: Subject to change after checkpoint redesign
|
||||
self._torch_state_dict_keys = list(model.state_dict().keys())
|
||||
elif isinstance(model, onnx.ModelProto):
|
||||
assert loss_fn is None, "'loss_fn' must not be specified when 'model' is an ONNX model"
|
||||
self._onnx_model = model
|
||||
|
|
@ -168,6 +170,7 @@ class ORTTrainer(object):
|
|||
self._onnx_model = postprocess.run_postprocess(self._onnx_model)
|
||||
if self.options._internal_use.extra_postprocess:
|
||||
self._onnx_model = self.options._internal_use.extra_postprocess(self._onnx_model)
|
||||
assert isinstance(self._onnx_model, onnx.ModelProto), "'extra_postprocess' must return a ONNX model"
|
||||
|
||||
# When input model is already ONNX (and not exported from Pytorch within ORTTrainer),
|
||||
# append 'dtype' from ONNX into model description's
|
||||
|
|
@ -195,9 +198,8 @@ 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
|
||||
# TODO: Subject to change after checkpoint redesign
|
||||
self._state_dict = {}
|
||||
|
||||
self._train_step_info = TrainStepInfo(self.optim_config)
|
||||
self._training_session = None
|
||||
|
|
@ -337,7 +339,7 @@ class ORTTrainer(object):
|
|||
if self.options.mixed_precision.enabled:
|
||||
loss_scaler = self.options.mixed_precision.loss_scaler
|
||||
assert loss_scaler, "Loss scaler is required when mixed precision is enabled"
|
||||
loss_scale = torch.tensor([loss_scaler.loss_scale])
|
||||
loss_scale = loss_scaler.loss_scale
|
||||
inputs_desc = self._model_desc_inputs_with_lr_and_loss_scale
|
||||
|
||||
# Get data. CombineTorchModelLossFn takes label as last input and outputs loss first
|
||||
|
|
@ -647,10 +649,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
|
||||
# TODO: Subject to change after checkpoint redesign
|
||||
if self._state_dict:
|
||||
checkpoint.load_state_dict(self, self._state_dict, self._load_state_dict_strict)
|
||||
self._state_dict = None
|
||||
checkpoint.experimental_load_state_dict(self, self._state_dict, self._load_state_dict_strict)
|
||||
self._state_dict_debug = self._state_dict
|
||||
self._state_dict = {}
|
||||
|
||||
def _prepare_model_input(self, inputs_desc, lr, loss_scale, *inputs, **kwargs):
|
||||
# Normalize input to tuple of samples
|
||||
|
|
@ -674,8 +677,8 @@ class ORTTrainer(object):
|
|||
# Append loss scale
|
||||
if loss_scale is not None:
|
||||
assert self.options.mixed_precision.enabled, "Loss scale cannot be used without mixed precision"
|
||||
loss_scale = loss_scale.clone().detach()
|
||||
input += (loss_scale, )
|
||||
loss_scale = torch.tensor([loss_scale])
|
||||
input += (loss_scale,)
|
||||
extra_inputs += 1
|
||||
|
||||
# Only assert length of input when fetches is not used
|
||||
|
|
|
|||
|
|
@ -70,22 +70,6 @@ def bert_model_description(dynamic_shape=True):
|
|||
return model_desc
|
||||
|
||||
|
||||
def optimizer_parameters_mutiple_groups(model):
|
||||
'''A method to assign different hyper parameters for different model parameter groups'''
|
||||
|
||||
no_decay_keys = ["bias", "gamma", "beta", "LayerNorm"]
|
||||
no_decay_param_group = []
|
||||
decay_param_group = []
|
||||
for initializer in model.graph.initializer:
|
||||
if any(key in initializer.name for key in no_decay_keys):
|
||||
no_decay_param_group.append(initializer.name)
|
||||
else:
|
||||
decay_param_group.append(initializer.name)
|
||||
params = [{'params': no_decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6},
|
||||
{'params': decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.01, "epsilon": 1e-6}]
|
||||
return params
|
||||
|
||||
|
||||
def optimizer_parameters(model):
|
||||
'''A method to assign different hyper parameters for different model parameter groups'''
|
||||
|
||||
|
|
@ -181,7 +165,7 @@ def legacy_optim_params_c(name):
|
|||
(True),
|
||||
(False)
|
||||
])
|
||||
def testToyBERTModelSimpleTrainStep(dynamic_shape):
|
||||
def testToyBERTModelBasicTraining(dynamic_shape):
|
||||
model_desc = bert_model_description(dynamic_shape)
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
|
|
@ -307,7 +291,6 @@ def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rate
|
|||
_test_helpers.assert_model_outputs(losses, expected_losses, rtol=1e-6)
|
||||
|
||||
|
||||
# Dynamic Loss Scaler implemented implicitly
|
||||
@pytest.mark.parametrize("loss_scaler, expected_losses", [
|
||||
(None, [10.98803424835205, 10.99240493774414, 11.090575218200684, 11.042827606201172, 10.988829612731934,\
|
||||
11.105679512023926, 10.981968879699707, 11.081787109375, 10.997162818908691, 11.107288360595703]),
|
||||
|
|
@ -471,13 +454,17 @@ def testToyBertCheckpointLoadZero():
|
|||
assert_allclose(expected_eval_loss, actual_eval_loss, rtol=rtol)
|
||||
|
||||
|
||||
def testToyBertStateDictWrapModelLossFn():
|
||||
@pytest.mark.parametrize("loss_scaler, optimizer_config, gradient_accumulation_steps", [
|
||||
(None, optim.AdamConfig(), 1),
|
||||
(None, optim.LambConfig(), 1),
|
||||
(None, optim.SGDConfig(), 1),
|
||||
(amp.DynamicLossScaler(), optim.AdamConfig(), 1),
|
||||
(amp.DynamicLossScaler(), optim.LambConfig(), 5),
|
||||
#(amp.DynamicLossScaler(), optim.SGDConfig(), 1), # SGD doesnt support fp16
|
||||
])
|
||||
def testToyBertStateDictWrapModelLossFn(loss_scaler, optimizer_config, gradient_accumulation_steps):
|
||||
# Common setup
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
# Modeling
|
||||
class LinearModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
|
@ -487,28 +474,58 @@ def testToyBertStateDictWrapModelLossFn():
|
|||
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)
|
||||
|
||||
# Dummy data
|
||||
data1 = torch.randn(2, 2)
|
||||
label1 = torch.tensor([0, 1], dtype=torch.int64)
|
||||
data2 = torch.randn(2, 2)
|
||||
label2 = torch.tensor([0, 1], dtype=torch.int64)
|
||||
|
||||
# Setup training based on test parameters
|
||||
opts = {'debug' : {'deterministic_compute': True},
|
||||
'batch' : { 'gradient_accumulation_steps' : gradient_accumulation_steps}}
|
||||
if loss_scaler:
|
||||
opts['mixed_precision'] = { 'enabled': True, 'loss_scaler': loss_scaler}
|
||||
opts = orttrainer.ORTTrainerOptions(opts)
|
||||
|
||||
# Training session 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
pt_model = LinearModel()
|
||||
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)
|
||||
trainer = orttrainer.ORTTrainer(pt_model, model_desc, optimizer_config, loss_fn=loss_fn, options=opts)
|
||||
|
||||
# Compare resulting state_dict keys before train
|
||||
# Check state_dict keys before train. Must be empty
|
||||
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
|
||||
# Train once and check initial state
|
||||
trainer.train_step(x=data1, label=label1)
|
||||
state_dict = checkpoint.experimental_state_dict(trainer)
|
||||
assert state_dict.keys() == {'linear.bias', 'linear.weight'}
|
||||
assert all([weight in state_dict.keys() for weight in ['linear.bias', 'linear.weight']])
|
||||
|
||||
# Initialize training session 2 from state of Training 1
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
trainer2 = orttrainer.ORTTrainer(pt_model, model_desc, optimizer_config, loss_fn=loss_fn, options=opts)
|
||||
checkpoint.experimental_load_state_dict(trainer2, state_dict)
|
||||
|
||||
# Verify state was loaded properly
|
||||
for k,v in state_dict.items():
|
||||
assert_allclose(v, trainer2._state_dict[k])
|
||||
|
||||
# Perform a second step in both training session 1 and 2 and verify they match
|
||||
trainer.train_step(x=data2, label=label2)
|
||||
state_dict = checkpoint.experimental_state_dict(trainer)
|
||||
trainer2.train_step(x=data2, label=label2)
|
||||
state_dict2 = checkpoint.experimental_state_dict(trainer2)
|
||||
for k,v in state_dict.items():
|
||||
assert_allclose(v, state_dict2[k])
|
||||
|
||||
|
||||
def testToyBertCheckpointFrozenWeights():
|
||||
|
|
@ -648,11 +665,15 @@ def testToyBERTSaveAsONNX():
|
|||
###############################################################################
|
||||
# Temporary tests comparing Legacy vs Experimental ORTTrainer APIs ############
|
||||
###############################################################################
|
||||
|
||||
|
||||
def testToyBERTModelLegacyExperimentalBasicTraining():
|
||||
@pytest.mark.parametrize("optimizer_config", [
|
||||
(optim.AdamConfig),
|
||||
(optim.LambConfig),
|
||||
(optim.SGDConfig)
|
||||
])
|
||||
def testToyBERTModelLegacyExperimentalBasicTraining(optimizer_config):
|
||||
# Common setup
|
||||
train_steps = 10
|
||||
train_steps = 512
|
||||
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
torch.manual_seed(seed)
|
||||
|
|
@ -661,8 +682,6 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
# EXPERIMENTAL API
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
params = optimizer_parameters(model)
|
||||
optim_config = optim.LambConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
'deterministic_compute': True
|
||||
|
|
@ -671,6 +690,7 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
'id': device,
|
||||
},
|
||||
})
|
||||
optim_config = optimizer_config(lr=0.01)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
experimental_losses = []
|
||||
for i in range(train_steps):
|
||||
|
|
@ -680,12 +700,24 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
# LEGACY IMPLEMENTATION
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
|
||||
if optimizer_config == optim.AdamConfig:
|
||||
legacy_optimizer = 'AdamOptimizer'
|
||||
elif optimizer_config == optim.LambConfig:
|
||||
legacy_optimizer = 'LambOptimizer'
|
||||
elif optimizer_config == optim.SGDConfig:
|
||||
legacy_optimizer = 'SGDOptimizer'
|
||||
else:
|
||||
raise RuntimeError("Invalid optimizer_config")
|
||||
|
||||
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",
|
||||
model = load_bert_onnx_model()
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(lr=optim_config.lr)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, legacy_optimizer,
|
||||
None,
|
||||
learning_rate_description,
|
||||
device)
|
||||
device,
|
||||
_use_deterministic_compute=True)
|
||||
legacy_losses = []
|
||||
for i in range(train_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
|
|
@ -693,7 +725,7 @@ def testToyBERTModelLegacyExperimentalBasicTraining():
|
|||
legacy_losses.append(leg_loss.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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("initial_lr, lr_scheduler, legacy_lr_scheduler", [
|
||||
|
|
@ -709,7 +741,7 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
|
|||
############################################################################
|
||||
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
total_steps = 128
|
||||
device = 'cuda'
|
||||
seed = 1
|
||||
warmup = 0.05
|
||||
|
|
@ -762,7 +794,7 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
|
|||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
device = torch.device(device)
|
||||
|
||||
model = load_bert_onnx_model()
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(initial_lr)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "AdamOptimizer",
|
||||
None,
|
||||
|
|
@ -787,7 +819,7 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
|
|||
])
|
||||
def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, legacy_loss_scaler):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
total_steps = 128
|
||||
device = "cuda"
|
||||
seed = 1
|
||||
|
||||
|
|
@ -796,7 +828,7 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
onnxruntime.set_seed(seed)
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
optim_config = optim.LambConfig()
|
||||
optim_config = optim.AdamConfig(lr=0.001)
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
'deterministic_compute': True
|
||||
|
|
@ -816,11 +848,12 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
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(optim_config.lr)
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
device = torch.device(device)
|
||||
model = load_bert_onnx_model()
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(optim_config.lr)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "AdamOptimizer",
|
||||
None,
|
||||
learning_rate_description,
|
||||
device,
|
||||
|
|
@ -834,7 +867,7 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
legacy_losses.append(leg_loss.cpu().item())
|
||||
|
||||
# Check results
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, rtol=1e-5)
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gradient_accumulation_steps", [
|
||||
|
|
@ -844,7 +877,7 @@ def testToyBERTModelMixedPrecisionLossScalerLegacyExperimental(loss_scaler, lega
|
|||
])
|
||||
def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation_steps):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
total_steps = 128
|
||||
device = "cuda"
|
||||
seed = 1
|
||||
|
||||
|
|
@ -853,7 +886,7 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
|
|||
onnxruntime.set_seed(seed)
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
optim_config = optim.LambConfig()
|
||||
optim_config = optim.AdamConfig()
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
'deterministic_compute': True
|
||||
|
|
@ -873,16 +906,17 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
|
|||
experimental_losses.append(loss.cpu().item())
|
||||
|
||||
# LEGACY IMPLEMENTATION
|
||||
device = torch.device(device)
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
device = torch.device(device)
|
||||
model = load_bert_onnx_model()
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(optim_config.lr)
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "AdamOptimizer",
|
||||
None,
|
||||
learning_rate_description,
|
||||
device,
|
||||
_use_deterministic_compute=True,
|
||||
gradient_accumulation_steps=gradient_accumulation_steps)
|
||||
_use_deterministic_compute = True,
|
||||
gradient_accumulation_steps = gradient_accumulation_steps)
|
||||
legacy_losses = []
|
||||
for i in range(total_steps):
|
||||
sample_input = generate_random_input_from_model_desc(model_desc, i)
|
||||
|
|
@ -890,7 +924,7 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
|
|||
legacy_losses.append(leg_loss.cpu().item())
|
||||
|
||||
# Check results
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, rtol=1e-6)
|
||||
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses)
|
||||
|
||||
@pytest.mark.parametrize("params, legacy_optim_map", [
|
||||
# Change the hyper parameters for all parameters
|
||||
|
|
@ -903,15 +937,17 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
|
|||
])
|
||||
def testToyBERTModelLegacyExperimentalCustomOptimParameters(params, legacy_optim_map):
|
||||
# Common setup
|
||||
total_steps = 10
|
||||
total_steps = 128
|
||||
device = "cuda"
|
||||
seed = 1
|
||||
|
||||
# EXPERIMENTAL API
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
model_desc = bert_model_description()
|
||||
model = load_bert_onnx_model()
|
||||
|
||||
optim_config = optim.LambConfig(params, alpha= 0.9, beta= 0.999, lambda_coef= 0.01, epsilon= 1e-6, do_bias_correction=False)
|
||||
optim_config = optim.AdamConfig(params, alpha= 0.9, beta= 0.999, lambda_coef= 0.01, epsilon= 1e-6, do_bias_correction=False)
|
||||
opts = orttrainer.ORTTrainerOptions({
|
||||
'debug' : {
|
||||
'deterministic_compute': True
|
||||
|
|
@ -920,9 +956,6 @@ def testToyBERTModelLegacyExperimentalCustomOptimParameters(params, legacy_optim
|
|||
'id': device,
|
||||
},
|
||||
})
|
||||
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
|
||||
|
||||
experimental_losses = []
|
||||
|
|
@ -931,12 +964,13 @@ def testToyBERTModelLegacyExperimentalCustomOptimParameters(params, legacy_optim
|
|||
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(trainer.optim_config.lr)
|
||||
torch.manual_seed(seed)
|
||||
onnxruntime.set_seed(seed)
|
||||
device = torch.device(device)
|
||||
model = load_bert_onnx_model()
|
||||
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(trainer.optim_config.lr)
|
||||
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
|
||||
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "AdamOptimizer",
|
||||
legacy_optim_map,
|
||||
learning_rate_description,
|
||||
device,
|
||||
|
|
|
|||
|
|
@ -876,7 +876,7 @@ def testORTTrainerLegacyAndExperimentalWeightsCheck(seed, device):
|
|||
])
|
||||
def testORTTrainerLegacyAndExperimentalPrecisionLossScaler(seed, device):
|
||||
# Common data
|
||||
total_steps = 5
|
||||
total_steps = 128
|
||||
|
||||
# Setup experimental API
|
||||
torch.manual_seed(seed)
|
||||
|
|
@ -921,7 +921,7 @@ def testORTTrainerLegacyAndExperimentalPrecisionLossScaler(seed, device):
|
|||
# Compare legacy vs experimental APIs
|
||||
assert experimental_preds_dtype == legacy_preds_dtype
|
||||
_test_helpers.assert_legacy_onnx_weights(trainer, legacy_trainer, rtol=1e-4, atol=1e-2)
|
||||
_test_helpers.assert_model_outputs(legacy_loss, experimental_loss, rtol=1e-4)
|
||||
_test_helpers.assert_model_outputs(legacy_loss, experimental_loss)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed,device,gradient_accumulation_steps,total_steps", [
|
||||
|
|
@ -966,7 +966,7 @@ def testORTTrainerLegacyAndExperimentalGradientAccumulation(seed, device, gradie
|
|||
legacy_loss.append(leg_loss.cpu())
|
||||
|
||||
# Compare legacy vs experimental APIs
|
||||
_test_helpers.assert_model_outputs(legacy_loss, experimental_loss, rtol=1e-6)
|
||||
_test_helpers.assert_model_outputs(legacy_loss, experimental_loss)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed,device,optimizer_config,lr_scheduler, get_lr_this_step", [
|
||||
|
|
@ -1008,7 +1008,7 @@ def testORTTrainerLegacyAndExperimentalLRScheduler(seed, device, optimizer_confi
|
|||
options = orttrainer.ORTTrainerOptions({'device' : {'id' : device},
|
||||
'debug' : {'deterministic_compute' : True},
|
||||
'lr_scheduler' : lr_scheduler})
|
||||
model, model_desc, my_loss, batcher_fn, train_data, val_data, _ = _load_pytorch_transformer_model(device)
|
||||
model, model_desc, my_loss, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device)
|
||||
optim_config = optimizer_config(lr=lr)
|
||||
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=my_loss, options=options)
|
||||
# Training loop
|
||||
|
|
@ -1054,3 +1054,92 @@ def testORTTrainerLegacyAndExperimentalLRScheduler(seed, device, optimizer_confi
|
|||
|
||||
# Compare legacy vs experimental APIs
|
||||
_test_helpers.assert_model_outputs(legacy_loss, experimental_loss)
|
||||
|
||||
|
||||
def testLossScalerLegacyAndExperimentalFullCycle():
|
||||
info = orttrainer.TrainStepInfo(optimizer_config=optim.LambConfig(lr=0.001), all_finite=True, fetches=[], optimization_step=0, step=0)
|
||||
new_ls = amp.DynamicLossScaler()
|
||||
old_ls = Legacy_LossScaler("ort_test_input_loss_scaler", True)
|
||||
|
||||
# Initial state
|
||||
train_step_info = orttrainer.TrainStepInfo(optim.LambConfig())
|
||||
assert_allclose(new_ls.loss_scale, old_ls.loss_scale_)
|
||||
assert new_ls.up_scale_window == old_ls.up_scale_window_
|
||||
assert_allclose(new_ls.min_loss_scale, old_ls.min_loss_scale_)
|
||||
assert_allclose(new_ls.max_loss_scale, old_ls.max_loss_scale_)
|
||||
|
||||
# Performing 9*2000 updates to cover all branches of LossScaler.update(train_step_info.all_finite=True)
|
||||
for cycles in range(1, 10):
|
||||
|
||||
# 1999 updates without overflow produces 1999 stable steps
|
||||
for i in range(1, 2000):
|
||||
new_loss_scale = new_ls.update(train_step_info)
|
||||
old_ls.update_loss_scale(train_step_info.all_finite)
|
||||
old_loss_scale = old_ls.loss_scale_
|
||||
assert new_ls._stable_steps_count == old_ls.stable_steps_
|
||||
# import pdb; pdb.set_trace()
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
# 2000th update without overflow doubles the loss and zero stable steps until max_loss_scale is reached
|
||||
new_loss_scale = new_ls.update(train_step_info)
|
||||
old_ls.update_loss_scale(train_step_info.all_finite)
|
||||
old_loss_scale = old_ls.loss_scale_
|
||||
assert new_ls._stable_steps_count == old_ls.stable_steps_
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
# After 8 cycles, loss scale should be float(1 << 16)*(2**8)
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
# After 9 cycles, loss scale reaches max_loss_scale and it is not doubled from that point on
|
||||
for count in range(1, 2050):
|
||||
new_loss_scale = new_ls.update(train_step_info)
|
||||
old_ls.update_loss_scale(train_step_info.all_finite)
|
||||
old_loss_scale = old_ls.loss_scale_
|
||||
assert new_ls._stable_steps_count == old_ls.stable_steps_
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
# Setting train_step_info.all_finite = False to test down scaling
|
||||
train_step_info.all_finite = False
|
||||
|
||||
# Performing 24 updates to half the loss scale each time
|
||||
for count in range(1, 25):
|
||||
new_loss_scale = new_ls.update(train_step_info)
|
||||
old_ls.update_loss_scale(train_step_info.all_finite)
|
||||
old_loss_scale = old_ls.loss_scale_
|
||||
assert new_ls._stable_steps_count == old_ls.stable_steps_
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
# After 24 updates with gradient overflow, loss scale is 1.0
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
# After 25 updates, min_loss_scale is reached and loss scale is not halfed from that point on
|
||||
for count in range(1, 5):
|
||||
new_loss_scale = new_ls.update(train_step_info)
|
||||
old_ls.update_loss_scale(train_step_info.all_finite)
|
||||
old_loss_scale = old_ls.loss_scale_
|
||||
assert new_ls._stable_steps_count == old_ls.stable_steps_
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
|
||||
|
||||
def testLossScalerLegacyAndExperimentalRandomAllFinite():
|
||||
new_ls = amp.DynamicLossScaler()
|
||||
old_ls = Legacy_LossScaler("ort_test_input_loss_scaler", True)
|
||||
|
||||
# Initial state
|
||||
train_step_info = orttrainer.TrainStepInfo(optim.LambConfig())
|
||||
assert_allclose(new_ls.loss_scale, old_ls.loss_scale_)
|
||||
assert new_ls.up_scale_window == old_ls.up_scale_window_
|
||||
assert_allclose(new_ls.min_loss_scale, old_ls.min_loss_scale_)
|
||||
assert_allclose(new_ls.max_loss_scale, old_ls.max_loss_scale_)
|
||||
|
||||
import random
|
||||
out = []
|
||||
for _ in range(1, 64):
|
||||
train_step_info.all_finite = bool(random.getrandbits(1))
|
||||
new_loss_scale = new_ls.update(train_step_info)
|
||||
old_ls.update_loss_scale(train_step_info.all_finite)
|
||||
old_loss_scale = old_ls.loss_scale_
|
||||
assert new_ls._stable_steps_count == old_ls.stable_steps_
|
||||
assert_allclose(new_loss_scale, old_loss_scale)
|
||||
out.append(new_loss_scale)
|
||||
assert new_loss_scale > 1e-7
|
||||
|
|
|
|||
Loading…
Reference in a new issue