From 7cc88ef7ed6ba0cb48688c2d17592e627ff456fa Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Wed, 19 Aug 2020 14:27:28 -0700 Subject: [PATCH] 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 --- .../python/experimental/checkpoint.py | 280 ++++++++++++++ .../python/experimental/orttrainer.py | 16 +- .../python/experimental/orttrainer_options.py | 23 +- .../orttraining/test/python/_test_helpers.py | 14 + ...ttraining_test_orttrainer_bert_toy_onnx.py | 363 +++++++++++++----- .../orttraining_test_orttrainer_frontend.py | 2 +- 6 files changed, 594 insertions(+), 104 deletions(-) create mode 100644 orttraining/orttraining/python/experimental/checkpoint.py diff --git a/orttraining/orttraining/python/experimental/checkpoint.py b/orttraining/orttraining/python/experimental/checkpoint.py new file mode 100644 index 0000000000..026ef7e8af --- /dev/null +++ b/orttraining/orttraining/python/experimental/checkpoint.py @@ -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 diff --git a/orttraining/orttraining/python/experimental/orttrainer.py b/orttraining/orttraining/python/experimental/orttrainer.py index 55e116eadf..f6d9fb79bf 100644 --- a/orttraining/orttraining/python/experimental/orttrainer.py +++ b/orttraining/orttraining/python/experimental/orttrainer.py @@ -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 diff --git a/orttraining/orttraining/python/experimental/orttrainer_options.py b/orttraining/orttraining/python/experimental/orttrainer_options.py index 2e86719388..6e160cc375 100644 --- a/orttraining/orttraining/python/experimental/orttrainer_options.py +++ b/orttraining/orttraining/python/experimental/orttrainer_options.py @@ -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 `_ algorithm for AllReduce diff --git a/orttraining/orttraining/test/python/_test_helpers.py b/orttraining/orttraining/test/python/_test_helpers.py index d59c12a1f8..1b7fee0ec1 100644 --- a/orttraining/orttraining/test/python/_test_helpers.py +++ b/orttraining/orttraining/test/python/_test_helpers.py @@ -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)) 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 bbf68037ff..09c81d8826 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 @@ -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) diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index bbad604428..24723a41b2 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -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)