From aabed34d5cd8d06ab80f496aec9e9a912e979e15 Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Wed, 2 Sep 2020 09:38:02 -0700 Subject: [PATCH] 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 --- .../python/experimental/amp/loss_scaler.py | 17 +- .../python/experimental/checkpoint.py | 23 +-- .../experimental/model_desc_validation.py | 45 +++-- .../python/experimental/orttrainer.py | 21 ++- ...ttraining_test_orttrainer_bert_toy_onnx.py | 168 +++++++++++------- .../orttraining_test_orttrainer_frontend.py | 97 +++++++++- 6 files changed, 257 insertions(+), 114 deletions(-) diff --git a/orttraining/orttraining/python/experimental/amp/loss_scaler.py b/orttraining/orttraining/python/experimental/amp/loss_scaler.py index 711c779d66..f64db9a3fa 100644 --- a/orttraining/orttraining/python/experimental/amp/loss_scaler.py +++ b/orttraining/orttraining/python/experimental/amp/loss_scaler.py @@ -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 diff --git a/orttraining/orttraining/python/experimental/checkpoint.py b/orttraining/orttraining/python/experimental/checkpoint.py index 4bfc41067a..198e785333 100644 --- a/orttraining/orttraining/python/experimental/checkpoint.py +++ b/orttraining/orttraining/python/experimental/checkpoint.py @@ -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" diff --git a/orttraining/orttraining/python/experimental/model_desc_validation.py b/orttraining/orttraining/python/experimental/model_desc_validation.py index 05db52e5fb..f0292b0f4a 100644 --- a/orttraining/orttraining/python/experimental/model_desc_validation.py +++ b/orttraining/orttraining/python/experimental/model_desc_validation.py @@ -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: diff --git a/orttraining/orttraining/python/experimental/orttrainer.py b/orttraining/orttraining/python/experimental/orttrainer.py index 1a5267779d..985905df57 100644 --- a/orttraining/orttraining/python/experimental/orttrainer.py +++ b/orttraining/orttraining/python/experimental/orttrainer.py @@ -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 diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py index d41be439a0..23bd66a3d5 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py @@ -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, diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index 26424be89d..9813f2ea2b 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -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