[Enhance][Temp] add environ var that serves as switch for attn_lse hook

Author: nil0x9

xtuner/v1/utils/internal_metrics.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+import os
 import torch
 import torch.distributed as dist
 from torch import nn
@@ -72,6 +73,8 @@ def calculate_module_weight_rms(self, module: nn.Module, layer_name: str, dtype:
 
     def register_attn_extra_info_hook(self, module: nn.Module):
         """Register attention extra info hook as a forward hook"""
+        if os.getenv("DISABLE_ATTN_MONITOR_HOOK", "0") == "1":
+            return
         def hook(module, input, output):
             extra_info = output[1]
             if extra_info.get("softmax_lse", None) is not None: