Fix issue preventing loss scaler to run due (#4833)

`LossScaler.update()` was not being properly called due to the incorrect TrainStepInfo.all_finite assignment.

Additionally to this fix, _ORTTrainerModelDesc.is_finite was renamed to _ORTTrainerModelDesc.all_finite to make it more uniform with TrainStepInfo
This commit is contained in:
Thiago Crepaldi 2020-08-18 10:03:02 -07:00 committed by GitHub
parent a3c95374c3
commit f3b0c93a45
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 20 additions and 20 deletions

View file

@ -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):

View file

@ -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")

View file

@ -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