an optimizer that dynamically scales the loss to prevent underflow loss scaling is a technique to prevent numeric underflow in intermediate gradients when float16 is used to prevent underflow the loss is multiplied or scaled by a certain factor called the loss scale which causes intermediate gradients to be scaled by the loss scale as well the final gradients are divided or unscaled by the loss scale to bring them back to their original value lossscaleoptimizer wraps another optimizer and applies dynamic loss scaling to it this loss scale is dynamically updated over time as follows on any train step if a nonfinite gradient is encountered the loss scale is halved and the train step is skipped if dynamicgrowthsteps have ocurred since the last time the loss scale was updated and no nonfinite gradients have occurred the loss scale is doubled args inneroptimizer the keras optimizers optimizer instance to wrap initialscale float the initial loss scale this scale will be updated during training it is recommended for this to be a very high number because a loss scale that is too high gets lowered far more quickly than a loss scale that is too low gets raised dynamicgrowthsteps int how often to update the scale upwards after every dynamicgrowthsteps steps with finite gradients the loss scale is doubled baseoptimizerkeywordargs potentially upscale loss and reset counter unscale gradients optionally build optimizer unscale gradients potentially upscale loss and reset counter if any inf or nan in grads downscale loss and reset counter tensorflow specific logic for apply which handles distribution from keras utils moduleutils import tensorflow as tf if tf distribute incrossreplicacontext raise valueerrorapply must be called in a replica context if tf internal distribute strategysupportsnomergecall self commonapplygrads trainablevariablestrainablevariables else def handlecrossreplicadistribution grads trainablevariables finiteperreplica distribution extended callforeachreplica self checkfinite argsgrads each replica computed the same finite value since grads is allreduced across replicas arbitrarily take finite from the first replica finite distribution experimentallocalresults finiteperreplica 0 def applyfn distribution extended callforeachreplica self statefulhandlefinitegrads argsgrads trainablevariables note we must call this cond in a crossreplica context distributionstrategy does not support having a cond in a replica context with a branch that calls mergecall and self optimizer applygradients calls mergecall ops cond finite applyfn self statefulhandlenonfinitegrads tf distribute getreplicacontext mergecall handlecrossreplica argsgrads trainablevariables def checkfiniteself grads tensorgrads g for g in grads if g is not none finitegrads ops allops isfiniteg for g in tensorgrads return ops allops converttotensorfinitegrads property def learningrateself return self inneroptimizer learningrate learningrate setter def learningrateself learningrate self inneroptimizer learningrate learningrate def scalelossself loss scale self dynamicscale if self built else self initialscale return loss scale def finalizevariablevaluesself varlist self inneroptimizer finalizevariablevaluesvarlist def getconfigself config super getconfig inneroptimizerconfig serializationlib serializekerasobject self inneroptimizer config update inneroptimizer inneroptimizerconfig initialscale self initialscale dynamicgrowthsteps self dynamicgrowthsteps del configlearningrate return config classmethod def fromconfigcls config customobjectsnone inneroptimizer serializationlib deserializekerasobject config popinneroptimizer customobjectscustomobjects return clsinneroptimizer config lossscaleoptimizer doc lossscaleoptimizer doc replace baseoptimizerkeywordargs optimizer baseoptimizerkeywordargs an optimizer that dynamically scales the loss to prevent underflow loss scaling is a technique to prevent numeric underflow in intermediate gradients when float16 is used to prevent underflow the loss is multiplied or scaled by a certain factor called the loss scale which causes intermediate gradients to be scaled by the loss scale as well the final gradients are divided or unscaled by the loss scale to bring them back to their original value lossscaleoptimizer wraps another optimizer and applies dynamic loss scaling to it this loss scale is dynamically updated over time as follows on any train step if a nonfinite gradient is encountered the loss scale is halved and the train step is skipped if dynamic_growth_steps have ocurred since the last time the loss scale was updated and no nonfinite gradients have occurred the loss scale is doubled args inner_optimizer the keras optimizers optimizer instance to wrap initial_scale float the initial loss scale this scale will be updated during training it is recommended for this to be a very high number because a loss scale that is too high gets lowered far more quickly than a loss scale that is too low gets raised dynamic_growth_steps int how often to update the scale upwards after every dynamic_growth_steps steps with finite gradients the loss scale is doubled base_optimizer_keyword_args potentially upscale loss and reset counter unscale gradients optionally build optimizer unscale gradients potentially upscale loss and reset counter if any inf or nan in grads downscale loss and reset counter tensorflow specific logic for apply which handles distribution each replica computed the same finite value since grads is all reduced across replicas arbitrarily take finite from the first replica note we must call this cond in a cross replica context distributionstrategy does not support having a cond in a replica context with a branch that calls merge_call and self _optimizer apply_gradients calls merge_call
from keras import backend
from keras import initializers
from keras import ops
from keras.api_export import keras_export
from keras.optimizers import optimizer
from keras.saving import serialization_lib
from keras.utils import tracking


@keras_export(
    [
        "keras.optimizers.LossScaleOptimizer",
        "keras.mixed_precision.LossScaleOptimizer",
    ]
)
class LossScaleOptimizer(optimizer.Optimizer):
    

    def __init__(
        self,
        inner_optimizer,
        initial_scale=2.0**15,
        dynamic_growth_steps=2000,
        **kwargs,
    ):
        if not kwargs.pop("dynamic", True):
            raise ValueError(
                "LossScaleOptimizer no longer suports `dynamic=False`. "
                "Instead, simply set `loss_scale_factor` directly on the "
                "`inner_optimizer`."
            )
        super().__init__(learning_rate=0.0, **kwargs)
        self.inner_optimizer = inner_optimizer
        self.initial_scale = initial_scale
        self.dynamic_growth_steps = dynamic_growth_steps

    @tracking.no_automatic_dependency_tracking
    def build(self, var_list):
        self.step_counter = self.add_variable(
            shape=(),
            dtype="int",
            initializer=initializers.Zeros(),
            name="step_counter",
        )
        self.dynamic_scale = self.add_variable(
            shape=(),
            dtype="float32",
            initializer=initializers.Constant(self.initial_scale),
            name="dynamic_scale",
        )
        self.inner_optimizer.build(var_list)
        self.built = True

    @property
    def variables(self):
        return self._variables + self.inner_optimizer.variables

    def stateless_apply(self, optimizer_variables, grads, trainable_variables):
        if not self.built:
            raise ValueError(
                f"To call `stateless_apply`, {self.__class__.__name__} "
                "must be built (i.e. its variables must have been created). "
                "You can build it via `optimizer.build(trainable_variables)`."
            )
        finite = self.check_finite(grads)
        return ops.cond(
            finite,
            lambda: self._stateless_handle_finite_grads(
                optimizer_variables, grads, trainable_variables
            ),
            lambda: self._stateless_handle_non_finite_grads(
                optimizer_variables, trainable_variables
            ),
        )

    def _stateless_handle_finite_grads(
        self, optimizer_variables, grads, trainable_variables
    ):
        def upscale():
            mapping = list(zip(self.variables, optimizer_variables))
            with backend.StatelessScope(state_mapping=mapping) as scope:
                self.step_counter.assign(0)
                self.dynamic_scale.assign(self.dynamic_scale * 2.0)
            return [scope.get_current_value(v) for v in self._variables]

        def increment():
            mapping = list(zip(self.variables, optimizer_variables))
            with backend.StatelessScope(state_mapping=mapping) as scope:
                self.step_counter.assign_add(1)
            return [scope.get_current_value(v) for v in self._variables]

        mapping = list(zip(self.variables, optimizer_variables))
        with backend.StatelessScope(state_mapping=mapping):
            
            own_variables = ops.cond(
                ops.equal(self.step_counter, self.dynamic_growth_steps - 1),
                upscale,
                increment,
            )

            
            scale = self.dynamic_scale
            unscaled_grads = [
                g if g is None else ops.divide(g, scale) for g in grads
            ]
            (
                new_trainable_variables,
                new_inner_variables,
            ) = self.inner_optimizer.stateless_apply(
                self.inner_optimizer.variables,
                unscaled_grads,
                trainable_variables,
            )

        new_optimizer_variables = own_variables + new_inner_variables
        return new_trainable_variables, new_optimizer_variables

    def _stateless_handle_non_finite_grads(
        self, optimizer_variables, trainable_variables
    ):
        mapping = list(zip(self.variables, optimizer_variables))
        with backend.StatelessScope(state_mapping=mapping) as scope:
            self.step_counter.assign(0)
            self.dynamic_scale.assign(self.dynamic_scale / 2.0)
        new_optimizer_variables = []
        for v in self.variables:
            new_optimizer_variables.append(scope.get_current_value(v))
        return trainable_variables, new_optimizer_variables

    def apply(self, grads, trainable_variables=None):
        
        if not self.built:
            with backend.name_scope(self.name, caller=self):
                self.build(trainable_variables)
            self.built = True

        if backend.backend() == "tensorflow":
            self._tf_apply(grads, trainable_variables)
        else:
            self._common_apply(grads, trainable_variables)

    def _stateful_handle_finite_grads(self, grads, trainable_variables):
        scale = self.dynamic_scale
        
        unscaled_grads = [
            g if g is None else ops.divide(g, scale) for g in grads
        ]
        self.inner_optimizer.apply(
            unscaled_grads, trainable_variables=trainable_variables
        )

        def upscale():
            self.step_counter.assign(0)
            self.dynamic_scale.assign(self.dynamic_scale * 2.0)

        def increment():
            self.step_counter.assign_add(1)

        
        ops.cond(
            ops.equal(self.step_counter, self.dynamic_growth_steps - 1),
            upscale,
            increment,
        )

    def _stateful_handle_non_finite_grads(self):
        
        self.step_counter.assign(0)
        self.dynamic_scale.assign(self.dynamic_scale / 2.0)

    def _common_apply(self, grads, trainable_variables=None):
        finite = self.check_finite(grads)
        ops.cond(
            finite,
            lambda: self._stateful_handle_finite_grads(
                grads, trainable_variables
            ),
            self._stateful_handle_non_finite_grads,
        )

    def _tf_apply(self, grads, trainable_variables=None):
        
        from keras.utils.module_utils import tensorflow as tf

        if tf.distribute.in_cross_replica_context():
            raise ValueError("apply() must be called in a replica context.")

        if tf.__internal__.distribute.strategy_supports_no_merge_call():
            self._common_apply(grads, trainable_variables=trainable_variables)
        else:

            def _handle_cross_replica(distribution, grads, trainable_variables):
                finite_per_replica = (
                    distribution.extended.call_for_each_replica(
                        self.check_finite, args=(grads,)
                    )
                )
                
                
                
                finite = distribution.experimental_local_results(
                    finite_per_replica
                )[0]

                def apply_fn():
                    distribution.extended.call_for_each_replica(
                        self._stateful_handle_finite_grads,
                        args=(grads, trainable_variables),
                    )

                
                
                
                
                ops.cond(
                    finite, apply_fn, self._stateful_handle_non_finite_grads
                )

            tf.distribute.get_replica_context().merge_call(
                _handle_cross_replica, args=(grads, trainable_variables)
            )

    def check_finite(self, grads):
        tensor_grads = [g for g in grads if g is not None]
        finite_grads = [ops.all(ops.isfinite(g)) for g in tensor_grads]
        return ops.all(ops.convert_to_tensor(finite_grads))

    @property
    def learning_rate(self):
        return self.inner_optimizer.learning_rate

    @learning_rate.setter
    def learning_rate(self, learning_rate):
        self.inner_optimizer.learning_rate = learning_rate

    def scale_loss(self, loss):
        scale = self.dynamic_scale if self.built else self.initial_scale
        return loss * scale

    def finalize_variable_values(self, var_list):
        self.inner_optimizer.finalize_variable_values(var_list)

    def get_config(self):
        config = super().get_config()
        inner_optimizer_config = serialization_lib.serialize_keras_object(
            self.inner_optimizer
        )
        config.update(
            {
                "inner_optimizer": inner_optimizer_config,
                "initial_scale": self.initial_scale,
                "dynamic_growth_steps": self.dynamic_growth_steps,
            }
        )
        del config["learning_rate"]
        return config

    @classmethod
    def from_config(cls, config, custom_objects=None):
        inner_optimizer = serialization_lib.deserialize_keras_object(
            config.pop("inner_optimizer"),
            custom_objects=custom_objects,
        )
        return cls(inner_optimizer, **config)


LossScaleOptimizer.__doc__ = LossScaleOptimizer.__doc__.replace(
    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
)