




'''simple docstring'''







import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
						import torch


_lowerCamelCase				:		Optional[int]				      =    logging.get_logger(__name__)





class 	lowerCamelCase_				:




					'''simple docstring'''






					def __init__(       self,    __UpperCamelCase,    __UpperCamelCase=21,    __UpperCamelCase=[],    __UpperCamelCase=None					):
							_snake_case      			=			model
							_snake_case      			=			trace_batch_nums
							_snake_case      			=			abort_after_batch_num

							# keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
							_snake_case      			=			collections.deque([],    _lowerCamelCase					)
							_snake_case      			=			[]
							_snake_case      			=			0
							_snake_case      			=			0
							_snake_case      			=			False
							_snake_case      			=			"                 "

							self.analyse_model()

							self.register_forward_hook()






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase=None					):
							if frame is not None:
									self.expand_frame(_lowerCamelCase					)
							self.frames.append("\n".join(self.frame					)					)
							_snake_case      			=			[]  # start a new frame






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							self.frame.append(_lowerCamelCase					)






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






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			[]






					def 							__SCREAMING_SNAKE_CASE      (       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"					)
							_snake_case      			=			[]






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			{m: name for name, m in self.model.named_modules()}
							# self.longest_module_name = max(len(v) for v in self.module_names.values())






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






					def 							__SCREAMING_SNAKE_CASE      (       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 							__SCREAMING_SNAKE_CASE      (       self					):
							self.expand_frame(f'''{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n'''					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase					):
							self.expand_frame(f'''{self.prefix} {self.module_names[module]} {module.__class__.__name__}'''					)

							# params
							for name, p in module.named_parameters(recurse=_lowerCamelCase					):
									self.analyse_variable(_lowerCamelCase,    _lowerCamelCase					)

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

							# outputs
							if isinstance(_lowerCamelCase,    _lowerCamelCase					):
									for i, x in enumerate(_lowerCamelCase					):
											# possibly a tuple of tuples
											if isinstance(_lowerCamelCase,    _lowerCamelCase					):
													for j, y in enumerate(_lowerCamelCase					):
															self.analyse_variable(_lowerCamelCase,    f'''output[{i}][{j}]'''					)
											else:
													self.analyse_variable(_lowerCamelCase,    f'''output[{i}]'''					)
							else:
									self.analyse_variable(_lowerCamelCase,    "output"					)

							self.save_frame()






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






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							module.register_forward_hook(self.forward_hook					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase					):
							_snake_case      			=			False

							_snake_case      			=			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

							# 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 module == self.model:
									self.batch_number += 1
									_snake_case      			=			True

							self.create_frame(_lowerCamelCase,    _lowerCamelCase,    _lowerCamelCase					)

							# if last_frame_of_batch:
							#     self.batch_end_frame()

							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()

									# now we can abort, as it's pointless to continue running
									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."					)

							# abort after certain batch if requested to do so
							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 						_lowercase						(      lowercase_		,				lowercase_    ):



		'''simple docstring'''







		_snake_case      			=			var.abs()
		return f'''{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}'''






def 						_lowercase						(      lowercase_		,				lowercase_    ):



		'''simple docstring'''







		_snake_case      			=			False
		if torch.isnan(__UpperCamelCase    ).any().item():
				_snake_case      			=			True
				print(f'''{ctx} has nans'''    )
		if torch.isinf(__UpperCamelCase    ).any().item():
				_snake_case      			=			True
				print(f'''{ctx} has infs'''    )

		# if needed to monitor large elements can enable the following
		if 0:  # and detected:
				_snake_case      			=			var[torch.ge(var.abs()		,				100    )]
				if naaa.numel() > 0:
						print(f'''{ctx}:  n100={naaa.numel()}'''    )
				_snake_case      			=			var[torch.ge(var.abs()		,				1000    )]
				if naaaa.numel() > 0:
						print(f'''{ctx}: n1000={naaaa.numel()}'''    )
				_snake_case      			=			var[torch.ge(var.abs()		,				10000    )]
				if naaaaa.numel() > 0:
						print(f'''{ctx}: n10000={naaaaa.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 	lowerCamelCase_				(			_lowerCAmelCase				):




					'''simple docstring'''






					lowercase__       =					"underflow_overflow"
					lowercase__       =					"tpu_metrics_debug"













"""simple docstring"""


import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
       import torch

       from transformers import GPTaLMHeadModel


@require_torch
@require_sentencepiece
@require_tokenizers
class    snake_case						(				unittest.TestCase		):


   '''simple docstring'''

   @slow
   @unittest.skip('''Model is not available.'''    )
   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     __A		     =					'''nvidia/megatron-gpt2-345m/'''
     if "MYDIR" in os.environ:
       __A		     =					os.path.join(os.environ['''MYDIR'''],			_lowerCamelCase    )
     __A		     =					GPTaLMHeadModel.from_pretrained(_lowerCamelCase    )
     model.to(_lowerCamelCase    )
     model.half()

     __A		     =					torch.tensor(
         [[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]],			device=_lowerCamelCase,			dtype=torch.long,			)

     with torch.no_grad():
       __A		     =					model(_lowerCamelCase    ).logits

     __A		     =					torch.Size((1, 9, 5_02_57)    )
     self.assertEqual(output.shape,			_lowerCamelCase    )

     __A		     =					torch.tensor(
         [
             4.94_14,
             -0.29_20,
             -1.21_48,
             -4.02_73,
             -0.51_61,
             -5.21_09,
             -1.24_12,
             -1.83_01,
             -1.77_34,
             -4.71_48,
             -0.23_17,
             -1.08_11,
             -2.17_77,
             0.41_41,
             -3.79_69,
             -4.05_86,
             -2.53_32,
             -3.38_09,
             4.38_67,
         ],			device=_lowerCamelCase,			dtype=torch.half,			)

     for i in range(19    ):
       __A     ,       __A		     =					8 * i // 17, 27_92 * i  # along the diagonal
       __A     ,       __A		     =					output[0, r, c], expected_diag[i]
       __A		     =					f'row={r} col={c} computed={computed} expected={expected}'
       self.assertAlmostEqual(_lowerCamelCase,			_lowerCamelCase,			delta=1e-4,			msg=_lowerCamelCase    )




