Aluode/PerceptionLabPortable

1# Copyright 2020 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import collections16 17from .utils import ExplicitEnum, is_torch_available, logging18 19 20if is_torch_available():21    import torch22 23 24logger = logging.get_logger(__name__)25 26 27class DebugUnderflowOverflow:28    """29    This debug class helps detect and understand where the model starts getting very large or very small, and more30    importantly `nan` or `inf` weight and activation elements.31 32    There are 2 working modes:33 34    1. Underflow/overflow detection (default)35    2. Specific batch absolute min/max tracing without detection36 37    Mode 1: Underflow/overflow detection38 39    To activate the underflow/overflow detection, initialize the object with the model :40 41    ```python42    debug_overflow = DebugUnderflowOverflow(model)43    ```44 45    then run the training as normal and if `nan` or `inf` gets detected in at least one of the weight, input or output46    elements this module will throw an exception and will print `max_frames_to_save` frames that lead to this event,47    each frame reporting48 49    1. the fully qualified module name plus the class name whose `forward` was run50    2. the absolute min and max value of all elements for each module weights, and the inputs and output51 52    For example, here is the header and the last few frames in detection report for `google/mt5-small` run in fp1653    mixed precision :54 55    ```56    Detected inf/nan during batch_number=057    Last 21 forward frames:58    abs min  abs max  metadata59    [...]60                      encoder.block.2.layer.1.DenseReluDense.wi_0 Linear61    2.17e-07 4.50e+00 weight62    1.79e-06 4.65e+00 input[0]63    2.68e-06 3.70e+01 output64                      encoder.block.2.layer.1.DenseReluDense.wi_1 Linear65    8.08e-07 2.66e+01 weight66    1.79e-06 4.65e+00 input[0]67    1.27e-04 2.37e+02 output68                      encoder.block.2.layer.1.DenseReluDense.wo Linear69    1.01e-06 6.44e+00 weight70    0.00e+00 9.74e+03 input[0]71    3.18e-04 6.27e+04 output72                      encoder.block.2.layer.1.DenseReluDense T5DenseGatedGeluDense73    1.79e-06 4.65e+00 input[0]74    3.18e-04 6.27e+04 output75                      encoder.block.2.layer.1.dropout Dropout76    3.18e-04 6.27e+04 input[0]77    0.00e+00      inf output78    ```79 80    You can see here, that `T5DenseGatedGeluDense.forward` resulted in output activations, whose absolute max value was81    around 62.7K, which is very close to fp16's top limit of 64K. In the next frame we have `Dropout` which82    renormalizes the weights, after it zeroed some of the elements, which pushes the absolute max value to more than83    64K, and we get an overflow.84 85    As you can see it's the previous frames that we need to look into when the numbers start going into very large for86    fp16 numbers.87 88    The tracking is done in a forward hook, which gets invoked immediately after `forward` has completed.89 90    By default the last 21 frames are printed. You can change the default to adjust for your needs. For example :91 92    ```python93    debug_overflow = DebugUnderflowOverflow(model, max_frames_to_save=100)94    ```95 96        To validate that you have set up this debugging feature correctly, and you intend to use it in a training that97        may take hours to complete, first run it with normal tracing enabled for one of a few batches as explained in98        the next section.99 100 101        Mode 2. Specific batch absolute min/max tracing without detection102 103        The second work mode is per-batch tracing with the underflow/overflow detection feature turned off.104 105        Let's say you want to watch the absolute min and max values for all the ingredients of each `forward` call of a106    given batch, and only do that for batches 1 and 3. Then you instantiate this class as :107 108    ```python109    debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3])110    ```111 112    And now full batches 1 and 3 will be traced using the same format as explained above. Batches are 0-indexed.113 114    This is helpful if you know that the program starts misbehaving after a certain batch number, so you can115    fast-forward right to that area.116 117 118    Early stopping:119 120    You can also specify the batch number after which to stop the training, with :121 122    ```python123    debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3], abort_after_batch_num=3)124    ```125 126    This feature is mainly useful in the tracing mode, but you can use it for any mode.127 128 129    **Performance**:130 131    As this module measures absolute `min`/``max` of each weight of the model on every forward it'll slow the training132    down. Therefore remember to turn it off once the debugging needs have been met.133 134    Args:135        model (`nn.Module`):136            The model to debug.137        max_frames_to_save (`int`, *optional*, defaults to 21):138            How many frames back to record139        trace_batch_nums(`list[int]`, *optional*, defaults to `[]`):140            Which batch numbers to trace (turns detection off)141        abort_after_batch_num  (`int``, *optional*):142            Whether to abort after a certain batch number has finished143    """144 145    def __init__(self, model, max_frames_to_save=21, trace_batch_nums=[], abort_after_batch_num=None):146        self.model = model147        self.trace_batch_nums = trace_batch_nums148        self.abort_after_batch_num = abort_after_batch_num149 150        # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence151        self.frames = collections.deque([], max_frames_to_save)152        self.frame = []153        self.batch_number = 0154        self.total_calls = 0155        self.detected_overflow = False156        self.prefix = "                 "157 158        self.analyse_model()159 160        self.register_forward_hook()161 162    def save_frame(self, frame=None):163        if frame is not None:164            self.expand_frame(frame)165        self.frames.append("\n".join(self.frame))166        self.frame = []  # start a new frame167 168    def expand_frame(self, line):169        self.frame.append(line)170 171    def trace_frames(self):172        print("\n".join(self.frames))173        self.frames = []174 175    def reset_saved_frames(self):176        self.frames = []177 178    def dump_saved_frames(self):179        print(f"\nDetected inf/nan during batch_number={self.batch_number}")180        print(f"Last {len(self.frames)} forward frames:")181        print(f"{'abs min':8} {'abs max':8} metadata")182        print("\n".join(self.frames))183        print("\n\n")184        self.frames = []185 186    def analyse_model(self):187        # extract the fully qualified module names, to be able to report at run time. e.g.:188        # encoder.block.2.layer.0.SelfAttention.o189        #190        # for shared weights only the first shared module name will be registered191        self.module_names = {m: name for name, m in self.model.named_modules()}192        # self.longest_module_name = max(len(v) for v in self.module_names.values())193 194    def analyse_variable(self, var, ctx):195        if torch.is_tensor(var):196            self.expand_frame(get_abs_min_max(var, ctx))197            if detect_overflow(var, ctx):198                self.detected_overflow = True199        elif var is None:200            self.expand_frame(f"{'None':>17} {ctx}")201        else:202            self.expand_frame(f"{'not a tensor':>17} {ctx}")203 204    def batch_start_frame(self):205        self.expand_frame(f"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***")206        self.expand_frame(f"{'abs min':8} {'abs max':8} metadata")207 208    def batch_end_frame(self):209        self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number - 1} ***\n\n")210 211    def create_frame(self, module, input, output):212        self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}")213 214        # params215        for name, p in module.named_parameters(recurse=False):216            self.analyse_variable(p, name)217 218        # inputs219        if isinstance(input, tuple):220            for i, x in enumerate(input):221                self.analyse_variable(x, f"input[{i}]")222        else:223            self.analyse_variable(input, "input")224 225        # outputs226        if isinstance(output, tuple):227            for i, x in enumerate(output):228                # possibly a tuple of tuples229                if isinstance(x, tuple):230                    for j, y in enumerate(x):231                        self.analyse_variable(y, f"output[{i}][{j}]")232                else:233                    self.analyse_variable(x, f"output[{i}]")234        else:235            self.analyse_variable(output, "output")236 237        self.save_frame()238 239    def register_forward_hook(self):240        self.model.apply(self._register_forward_hook)241 242    def _register_forward_hook(self, module):243        module.register_forward_hook(self.forward_hook)244 245    def forward_hook(self, module, input, output):246        # - input is a tuple of packed inputs (could be non-Tensors)247        # - output could be a Tensor or a tuple of Tensors and non-Tensors248 249        last_frame_of_batch = False250 251        trace_mode = self.batch_number in self.trace_batch_nums252        if trace_mode:253            self.reset_saved_frames()254 255        if self.total_calls == 0:256            self.batch_start_frame()257        self.total_calls += 1258 259        # count batch numbers - the very first forward hook of the batch will be called when the260        # batch completes - i.e. it gets called very last - we know this batch has finished261        if module == self.model:262            self.batch_number += 1263            last_frame_of_batch = True264 265        self.create_frame(module, input, output)266 267        # if last_frame_of_batch:268        #     self.batch_end_frame()269 270        if trace_mode:271            self.trace_frames()272 273        if last_frame_of_batch:274            self.batch_start_frame()275 276        if self.detected_overflow and not trace_mode:277            self.dump_saved_frames()278 279            # now we can abort, as it's pointless to continue running280            raise ValueError(281                "DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. "282                "Please scroll up above this traceback to see the activation values prior to this event."283            )284 285        # abort after certain batch if requested to do so286        if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:287            raise ValueError(288                f"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"289                f" `abort_after_batch_num={self.abort_after_batch_num}` arg"290            )291 292 293def get_abs_min_max(var, ctx):294    abs_var = var.abs()295    return f"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"296 297 298def detect_overflow(var, ctx):299    """300    Report whether the tensor contains any `nan` or `inf` entries.301 302    This is useful for detecting overflows/underflows and best to call right after the function that did some math that303    modified the tensor in question.304 305    This function contains a few other helper features that you can enable and tweak directly if you want to track306    various other things.307 308    Args:309        var: the tensor variable to check310        ctx: the message to print as a context311 312    Return:313        `True` if `inf` or `nan` was detected, `False` otherwise314    """315    detected = False316    if torch.isnan(var).any().item():317        detected = True318        print(f"{ctx} has nans")319    if torch.isinf(var).any().item():320        detected = True321        print(f"{ctx} has infs")322 323    # if needed to monitor large elements can enable the following324    if 0:  # and detected:325        n100 = var[torch.ge(var.abs(), 100)]326        if n100.numel() > 0:327            print(f"{ctx}:  n100={n100.numel()}")328        n1000 = var[torch.ge(var.abs(), 1000)]329        if n1000.numel() > 0:330            print(f"{ctx}: n1000={n1000.numel()}")331        n10000 = var[torch.ge(var.abs(), 10000)]332        if n10000.numel() > 0:333            print(f"{ctx}: n10000={n10000.numel()}")334 335    if 0:336        print(f"min={var.min():9.2e} max={var.max():9.2e}")337 338    if 0:339        print(f"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})")340 341    return detected342 343 344class DebugOption(ExplicitEnum):345    UNDERFLOW_OVERFLOW = "underflow_overflow"346    TPU_METRICS_DEBUG = "tpu_metrics_debug"347