mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-11 17:48:34 +00:00
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:
parent
a3c95374c3
commit
f3b0c93a45
3 changed files with 20 additions and 20 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue