diff --git a/orttraining/orttraining/python/experimental/model_desc_validation.py b/orttraining/orttraining/python/experimental/model_desc_validation.py index 71af68449d..693d27524e 100644 --- a/orttraining/orttraining/python/experimental/model_desc_validation.py +++ b/orttraining/orttraining/python/experimental/model_desc_validation.py @@ -5,7 +5,7 @@ from ._utils import static_vars LEARNING_RATE_IO_DESCRIPTION_NAME = "__learning_rate" -IS_FINITE_IO_DESCRIPTION_NAME = "__is_finite" +ALL_FINITE_IO_DESCRIPTION_NAME = "__all_finite" LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME = "__loss_scale_input_name" GRADIENT_ACCUMULATION_IO_DESCRIPTION_NAME = "__gradient_accumulation_name" @@ -49,7 +49,7 @@ class _ORTTrainerModelDesc(object): else: self._validated['outputs'][idx] = self._OutputDescription(*output) - # Hard-code learning rate, is_finite descriptors + # Hard-code learning rate, all_finite descriptors self.learning_rate = self._InputDescriptionTyped(LEARNING_RATE_IO_DESCRIPTION_NAME, [1], torch.float32) # Convert dict in object @@ -94,11 +94,11 @@ class _ORTTrainerModelDesc(object): pretty_msg += f'(name={self.learning_rate.name}, shape={self.learning_rate.shape}, dtype={self.learning_rate.dtype})' # Mixed precision - if getattr(self, IS_FINITE_IO_DESCRIPTION_NAME, None) or getattr(self, LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME, None): + if getattr(self, ALL_FINITE_IO_DESCRIPTION_NAME, None) or getattr(self, LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME, None): pretty_msg += '\nMixed Precision:' - if getattr(self, IS_FINITE_IO_DESCRIPTION_NAME, None): + if getattr(self, ALL_FINITE_IO_DESCRIPTION_NAME, None): pretty_msg += '\n\tis gradients finite: ' - pretty_msg += f'(name={self.is_finite.name}, shape={self.is_finite.shape}, dtype={self.is_finite.dtype})' + pretty_msg += f'(name={self.all_finite.name}, shape={self.all_finite.shape}, dtype={self.all_finite.dtype})' if getattr(self, LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME, None): pretty_msg += '\n\tloss scale input name: ' pretty_msg += f'(name={self.loss_scale_input.name}, shape={self.loss_scale_input.shape}, dtype={self.loss_scale_input.dtype})' @@ -156,12 +156,12 @@ class _ORTTrainerModelDesc(object): self._add_output_description(self, name, [1], False, torch.bool, None, GRADIENT_ACCUMULATION_IO_DESCRIPTION_NAME) @property - def is_finite(self): - return getattr(self, IS_FINITE_IO_DESCRIPTION_NAME, None) + def all_finite(self): + return getattr(self, ALL_FINITE_IO_DESCRIPTION_NAME, None) - @is_finite.setter - def is_finite(self, name): - self._add_output_description(self, name, [1], False, torch.bool, None, IS_FINITE_IO_DESCRIPTION_NAME) + @all_finite.setter + def all_finite(self, name): + self._add_output_description(self, name, [1], False, torch.bool, None, ALL_FINITE_IO_DESCRIPTION_NAME) @property def loss_scale_input(self): diff --git a/orttraining/orttraining/python/experimental/orttrainer.py b/orttraining/orttraining/python/experimental/orttrainer.py index f44ec3dca4..55e116eadf 100644 --- a/orttraining/orttraining/python/experimental/orttrainer.py +++ b/orttraining/orttraining/python/experimental/orttrainer.py @@ -316,7 +316,7 @@ class ORTTrainer(object): outputs_desc = self._model_desc_outputs_with_gradient_accumulation elif self.options.mixed_precision.enabled: mixed_precision_without_fetches = True - outputs_desc = self._model_desc_outputs_with_is_finite + outputs_desc = self._model_desc_outputs_with_all_finite # Update Learning Rate if Necessary if self.options.lr_scheduler: @@ -346,8 +346,8 @@ class ORTTrainer(object): # because all_fp32_gradients_finite is still in the feed. self._train_io_binding.clear_binding_outputs() - is_all_finite = session_run_results[self.model_desc.is_finite.name] - self._train_step_info.is_finite = is_all_finite + is_all_finite = session_run_results[self.model_desc.all_finite.name] + self._train_step_info.all_finite = is_all_finite if loss_scaler: loss_scaler.update(self._train_step_info) if is_all_finite: @@ -626,8 +626,8 @@ class ORTTrainer(object): self.model_desc.loss_scale_input = self._training_session.loss_scale_input_name self._model_desc_inputs_with_lr_and_loss_scale = [ *self._model_desc_inputs_with_lr, self.model_desc.loss_scale_input] - self.model_desc.is_finite = _utils.get_all_gradients_finite_name_from_session(self._training_session) - self._model_desc_outputs_with_is_finite = [*self.model_desc.outputs, self.model_desc.is_finite] + self.model_desc.all_finite = _utils.get_all_gradients_finite_name_from_session(self._training_session) + self._model_desc_outputs_with_all_finite = [*self.model_desc.outputs, self.model_desc.all_finite] elif self.options.mixed_precision.loss_scaler: raise ValueError("Loss Scaler cannot be specified when Mixed Precision is not enabled") diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index 9dbd1301bd..bbad604428 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -156,11 +156,11 @@ def testORTTrainerModelDescValidSchemas(input_dict, input_dtype, output_dtype): is_loss = input_dict['outputs'][idx][2] if len(input_dict['outputs'][idx]) == 3 else False assert is_loss == o_desc.is_loss - # Set is_finite name and check its description - model_description.is_finite = md_val.IS_FINITE_IO_DESCRIPTION_NAME - assert model_description.is_finite.name == md_val.IS_FINITE_IO_DESCRIPTION_NAME - assert model_description.is_finite.shape == [1] - assert model_description.is_finite.dtype == torch.bool + # Set all_finite name and check its description + model_description.all_finite = md_val.ALL_FINITE_IO_DESCRIPTION_NAME + assert model_description.all_finite.name == md_val.ALL_FINITE_IO_DESCRIPTION_NAME + assert model_description.all_finite.shape == [1] + assert model_description.all_finite.dtype == torch.bool # Set loss_scale_input and check its description model_description.loss_scale_input = md_val.LOSS_SCALE_INPUT_IO_DESCRIPTION_NAME