2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license this debug class helps detect and understand where the model starts getting very large or very small and more importantly nan or inf weight and activation elements there are 2 working modes 1 underflowoverflow detection default 2 specific batch absolute minmax tracing without detection mode 1 underflowoverflow detection to activate the underflowoverflow detection initialize the object with the model python debugoverflow debugunderflowoverflowmodel then run the training as normal and if nan or inf gets detected in at least one of the weight input or output elements this module will throw an exception and will print maxframestosave frames that lead to this event each frame reporting 1 the fully qualified module name plus the class name whose forward was run 2 the absolute min and max value of all elements for each module weights and the inputs and output for example here is the header and the last few frames in detection report for googlemt5small run in fp16 mixed precision detected infnan during batchnumber0 last 21 forward frames abs min abs max metadata encoder block 2 layer 1 densereludense wi0 linear 2 17e07 4 50e00 weight 1 79e06 4 65e00 input0 2 68e06 3 70e01 output encoder block 2 layer 1 densereludense wi1 linear 8 08e07 2 66e01 weight 1 79e06 4 65e00 input0 1 27e04 2 37e02 output encoder block 2 layer 1 densereludense wo linear 1 01e06 6 44e00 weight 0 00e00 9 74e03 input0 3 18e04 6 27e04 output encoder block 2 layer 1 densereludense t5densegatedgeludense 1 79e06 4 65e00 input0 3 18e04 6 27e04 output encoder block 2 layer 1 dropout dropout 3 18e04 6 27e04 input0 0 00e00 inf output you can see here that t5densegatedgeludense forward resulted in output activations whose absolute max value was around 62 7k which is very close to fp16 s top limit of 64k in the next frame we have dropout which renormalizes the weights after it zeroed some of the elements which pushes the absolute max value to more than 64k and we get an overlow as you can see it s the previous frames that we need to look into when the numbers start going into very large for fp16 numbers the tracking is done in a forward hook which gets invoked immediately after forward has completed by default the last 21 frames are printed you can change the default to adjust for your needs for example python debugoverflow debugunderflowoverflowmodel maxframestosave100 to validate that you have set up this debugging feature correctly and you intend to use it in a training that may take hours to complete first run it with normal tracing enabled for one of a few batches as explained in the next section mode 2 specific batch absolute minmax tracing without detection the second work mode is perbatch tracing with the underflowoverflow detection feature turned off let s say you want to watch the absolute min and max values for all the ingredients of each forward call of a given batch and only do that for batches 1 and 3 then you instantiate this class as python debugoverflow debugunderflowoverflowmodel tracebatchnums1 3 and now full batches 1 and 3 will be traced using the same format as explained above batches are 0indexed this is helpful if you know that the program starts misbehaving after a certain batch number so you can fastforward right to that area early stopping you can also specify the batch number after which to stop the training with python debugoverflow debugunderflowoverflowmodel tracebatchnums1 3 abortafterbatchnum3 this feature is mainly useful in the tracing mode but you can use it for any mode performance as this module measures absolute minmax of each weight of the model on every forward it ll slow the training down therefore remember to turn it off once the debugging needs have been met args model nn module the model to debug maxframestosave int optional defaults to 21 how many frames back to record tracebatchnumslistint optional defaults to which batch numbers to trace turns detection off abortafterbatchnum int optional whether to abort after a certain batch number has finished keep a lifo buffer of frames to dump as soon as infnan is encountered to give context to the problem emergence extract the fully qualified module names to be able to report at run time e g encoder block 2 layer 0 selfattention o for shared weights only the first shared module name will be registered self longestmodulename maxlenv for v in self modulenames values params inputs outputs possibly a tuple of tuples input is a tuple of packed inputs could be nontensors output could be a tensor or a tuple of tensors and nontensors count batch numbers the very first forward hook of the batch will be called when the batch completes i e it gets called very last we know this batch has finished if lastframeofbatch self batchendframe now we can abort as it s pointless to continue running abort after certain batch if requested to do so report whether the tensor contains any nan or inf entries this is useful for detecting overflowsunderflows and best to call right after the function that did some math that modified the tensor in question this function contains a few other helper features that you can enable and tweak directly if you want to track various other things args var the tensor variable to check ctx the message to print as a context return true if inf or nan was detected false otherwise if needed to monitor large elements can enable the following 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license this debug class helps detect and understand where the model starts getting very large or very small and more importantly nan or inf weight and activation elements there are 2 working modes 1 underflow overflow detection default 2 specific batch absolute min max tracing without detection mode 1 underflow overflow detection to activate the underflow overflow detection initialize the object with the model python debug_overflow debugunderflowoverflow model then run the training as normal and if nan or inf gets detected in at least one of the weight input or output elements this module will throw an exception and will print max_frames_to_save frames that lead to this event each frame reporting 1 the fully qualified module name plus the class name whose forward was run 2 the absolute min and max value of all elements for each module weights and the inputs and output for example here is the header and the last few frames in detection report for google mt5 small run in fp16 mixed precision detected inf nan during batch_number 0 last 21 forward frames abs min abs max metadata encoder block 2 layer 1 densereludense wi_0 linear 2 17e 07 4 50e 00 weight 1 79e 06 4 65e 00 input 0 2 68e 06 3 70e 01 output encoder block 2 layer 1 densereludense wi_1 linear 8 08e 07 2 66e 01 weight 1 79e 06 4 65e 00 input 0 1 27e 04 2 37e 02 output encoder block 2 layer 1 densereludense wo linear 1 01e 06 6 44e 00 weight 0 00e 00 9 74e 03 input 0 3 18e 04 6 27e 04 output encoder block 2 layer 1 densereludense t5densegatedgeludense 1 79e 06 4 65e 00 input 0 3 18e 04 6 27e 04 output encoder block 2 layer 1 dropout dropout 3 18e 04 6 27e 04 input 0 0 00e 00 inf output you can see here that t5densegatedgeludense forward resulted in output activations whose absolute max value was around 62 7k which is very close to fp16 s top limit of 64k in the next frame we have dropout which renormalizes the weights after it zeroed some of the elements which pushes the absolute max value to more than 64k and we get an overlow as you can see it s the previous frames that we need to look into when the numbers start going into very large for fp16 numbers the tracking is done in a forward hook which gets invoked immediately after forward has completed by default the last 21 frames are printed you can change the default to adjust for your needs for example python debug_overflow debugunderflowoverflow model max_frames_to_save 100 to validate that you have set up this debugging feature correctly and you intend to use it in a training that may take hours to complete first run it with normal tracing enabled for one of a few batches as explained in the next section mode 2 specific batch absolute min max tracing without detection the second work mode is per batch tracing with the underflow overflow detection feature turned off let s say you want to watch the absolute min and max values for all the ingredients of each forward call of a given batch and only do that for batches 1 and 3 then you instantiate this class as python debug_overflow debugunderflowoverflow model trace_batch_nums 1 3 and now full batches 1 and 3 will be traced using the same format as explained above batches are 0 indexed this is helpful if you know that the program starts misbehaving after a certain batch number so you can fast forward right to that area early stopping you can also specify the batch number after which to stop the training with python debug_overflow debugunderflowoverflow model trace_batch_nums 1 3 abort_after_batch_num 3 this feature is mainly useful in the tracing mode but you can use it for any mode performance as this module measures absolute min max of each weight of the model on every forward it ll slow the training down therefore remember to turn it off once the debugging needs have been met args model nn module the model to debug max_frames_to_save int optional defaults to 21 how many frames back to record trace_batch_nums list int optional defaults to which batch numbers to trace turns detection off abort_after_batch_num int optional whether to abort after a certain batch number has finished keep a lifo buffer of frames to dump as soon as inf nan is encountered to give context to the problem emergence start a new frame extract the fully qualified module names to be able to report at run time e g encoder block 2 layer 0 selfattention o for shared weights only the first shared module name will be registered self longest_module_name max len v for v in self module_names values params inputs outputs possibly a tuple of tuples input is a tuple of packed inputs could be non tensors output could be a tensor or a tuple of tensors and non tensors count batch numbers the very first forward hook of the batch will be called when the batch completes i e it gets called very last we know this batch has finished if last_frame_of_batch self batch_end_frame now we can abort as it s pointless to continue running abort after certain batch if requested to do so report whether the tensor contains any nan or inf entries this is useful for detecting overflows underflows and best to call right after the function that did some math that modified the tensor in question this function contains a few other helper features that you can enable and tweak directly if you want to track various other things args var the tensor variable to check ctx the message to print as a context return true if inf or nan was detected false otherwise if needed to monitor large elements can enable the following and detected
import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
    import torch


logger = logging.get_logger(__name__)


class DebugUnderflowOverflow:
    

    def __init__(self, model, max_frames_to_save=21, trace_batch_nums=[], abort_after_batch_num=None):
        self.model = model
        self.trace_batch_nums = trace_batch_nums
        self.abort_after_batch_num = abort_after_batch_num

        
        self.frames = collections.deque([], max_frames_to_save)
        self.frame = []
        self.batch_number = 0
        self.total_calls = 0
        self.detected_overflow = False
        self.prefix = "                 "

        self.analyse_model()

        self.register_forward_hook()

    def save_frame(self, frame=None):
        if frame is not None:
            self.expand_frame(frame)
        self.frames.append("\n".join(self.frame))
        self.frame = []  

    def expand_frame(self, line):
        self.frame.append(line)

    def trace_frames(self):
        print("\n".join(self.frames))
        self.frames = []

    def reset_saved_frames(self):
        self.frames = []

    def dump_saved_frames(self):
        print(f"\nDetected inf/nan during batch_number={self.batch_number}")
        print(f"Last {len(self.frames)} forward frames:")
        print(f"{'abs min':8} {'abs max':8} metadata")
        print("\n".join(self.frames))
        print("\n\n")
        self.frames = []

    def analyse_model(self):
        
        
        
        
        self.module_names = {m: name for name, m in self.model.named_modules()}
        

    def analyse_variable(self, var, ctx):
        if torch.is_tensor(var):
            self.expand_frame(get_abs_min_max(var, ctx))
            if detect_overflow(var, ctx):
                self.detected_overflow = True
        elif var is None:
            self.expand_frame(f"{'None':>17} {ctx}")
        else:
            self.expand_frame(f"{'not a tensor':>17} {ctx}")

    def batch_start_frame(self):
        self.expand_frame(f"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***")
        self.expand_frame(f"{'abs min':8} {'abs max':8} metadata")

    def batch_end_frame(self):
        self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n")

    def create_frame(self, module, input, output):
        self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}")

        
        for name, p in module.named_parameters(recurse=False):
            self.analyse_variable(p, name)

        
        if isinstance(input, tuple):
            for i, x in enumerate(input):
                self.analyse_variable(x, f"input[{i}]")
        else:
            self.analyse_variable(input, "input")

        
        if isinstance(output, tuple):
            for i, x in enumerate(output):
                
                if isinstance(x, tuple):
                    for j, y in enumerate(x):
                        self.analyse_variable(y, f"output[{i}][{j}]")
                else:
                    self.analyse_variable(x, f"output[{i}]")
        else:
            self.analyse_variable(output, "output")

        self.save_frame()

    def register_forward_hook(self):
        self.model.apply(self._register_forward_hook)

    def _register_forward_hook(self, module):
        module.register_forward_hook(self.forward_hook)

    def forward_hook(self, module, input, output):
        
        

        last_frame_of_batch = False

        trace_mode = True if self.batch_number in self.trace_batch_nums else False
        if trace_mode:
            self.reset_saved_frames()

        if self.total_calls == 0:
            self.batch_start_frame()
        self.total_calls += 1

        
        
        if module == self.model:
            self.batch_number += 1
            last_frame_of_batch = True

        self.create_frame(module, input, output)

        
        

        if trace_mode:
            self.trace_frames()

        if last_frame_of_batch:
            self.batch_start_frame()

        if self.detected_overflow and not trace_mode:
            self.dump_saved_frames()

            
            raise ValueError(
                "DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. "
                "Please scroll up above this traceback to see the activation values prior to this event."
            )

        
        if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
            raise ValueError(
                f"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"
                f" `abort_after_batch_num={self.abort_after_batch_num}` arg"
            )


def get_abs_min_max(var, ctx):
    abs_var = var.abs()
    return f"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"


def detect_overflow(var, ctx):
    
    detected = False
    if torch.isnan(var).any().item():
        detected = True
        print(f"{ctx} has nans")
    if torch.isinf(var).any().item():
        detected = True
        print(f"{ctx} has infs")

    
    if 0:  
        n100 = var[torch.ge(var.abs(), 100)]
        if n100.numel() > 0:
            print(f"{ctx}:  n100={n100.numel()}")
        n1000 = var[torch.ge(var.abs(), 1000)]
        if n1000.numel() > 0:
            print(f"{ctx}: n1000={n1000.numel()}")
        n10000 = var[torch.ge(var.abs(), 10000)]
        if n10000.numel() > 0:
            print(f"{ctx}: n10000={n10000.numel()}")

    if 0:
        print(f"min={var.min():9.2e} max={var.max():9.2e}")

    if 0:
        print(f"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})")

    return detected


class DebugOption(ExplicitEnum):
    UNDERFLOW_OVERFLOW = "underflow_overflow"
    TPU_METRICS_DEBUG = "tpu_metrics_debug"