fixes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Author: pggPL

tests/pytorch/debug/test_api_features.py

@@ -12,7 +12,7 @@
 # limitations under the License.
 
 import torch
-from transformer_engine.pytorch.tensor import Float8Tensor, Float8Quantizer
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer
 
 import nvdlfw_inspect.api as nvinspect_api
 

tests/pytorch/debug/test_numerics.py

@@ -19,7 +19,7 @@
 from transformer_engine.common.recipe import DelayedScaling, Format
 from transformer_engine.pytorch.constants import TE_DType
 from transformer_engine.pytorch.fp8 import _default_sf_compute
-from transformer_engine.pytorch.tensor import Float8Quantizer
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
 from transformer_engine.pytorch.module.base import (
     _2X_ACC_DGRAD,
     _2X_ACC_FPROP,

tests/pytorch/debug/test_sanity.py

@@ -19,7 +19,7 @@
 from transformer_engine.common.recipe import DelayedScaling, Format
 from transformer_engine.pytorch.constants import TE_DType
 from transformer_engine.pytorch.fp8 import _default_sf_compute
-from transformer_engine.pytorch.tensor import Float8Quantizer
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
 
 from test_numerics import create_config_file
 

tests/pytorch/distributed/test_fusible_ops.py

@@ -21,12 +21,8 @@
 import transformer_engine.common.recipe
 import transformer_engine.pytorch as te
 from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
-<<<<<<< HEAD
-from transformer_engine.pytorch.tensor import Float8Quantizer
-=======
 from transformer_engine.pytorch.tensor import QuantizedTensor
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
->>>>>>> origin/release_v2.0
 import transformer_engine.pytorch.ops as te_ops
 from transformer_engine.pytorch.ops._common import is_float8_tensor
 from transformer_engine.pytorch.utils import is_bf16_compatible

transformer_engine/debug/debug_quantization.py

@@ -9,6 +9,7 @@
 
 from ..pytorch.tensor.quantized_tensor import QuantizedTensor, Quantizer, _IdentityFunc, prepare_for_saving, restore_from_saved
 from transformer_engine.debug.debug_state import TEDebugState
+import transformer_engine_torch as tex
 
 """
     This file contains DebugQuantizer and DebugQuantizedTensor objects, which are wrapper along Quantizer and QuantizedTensor
@@ -235,33 +236,40 @@ def update_quantized(
         noop_flag: Optional[torch.Tensor] = None,
     ) -> QuantizedTensor:
         assert noop_flag is None, "CUDA Graphs are not supported with debug=True!"
+        iteration = nvinspect_api.DEBUG_MANAGER._trainer_iteration_count
         updated_second_gemm = False
         updated_first_gemm = False
         if self.parent_quantizer is not None:
-            if self.first_gemm_usage and self.fp8_quantize_first_gemm:
-                dst.first_gemm.quantize_(src) 
+            if dst.first_gemm_tensor is not None and self.fp8_quantize_first_gemm:
+                if hasattr(dst.first_gemm_tensor, "quantize_"):
+                    dst.first_gemm_tensor.quantize_(src, noop_flag=None)
+                else:
+                    tex.quantize(src, self.parent_quantizer, dst.first_gemm_tensor, None)
                 updated_first_gemm = True
-            elif self.second_gemm_usage and self.fp8_quantize_first_gemm:
-                dst.second_gemm.quantize_(src) 
+            if dst.second_gemm_tensor is not None and self.fp8_quantize_second_gemm:
+                if hasattr(dst.second_gemm_tensor, "quantize_"):
+                    dst.second_gemm_tensor.quantize_(src, noop_flag=None)
+                else:
+                    tex.quantize(src, self.parent_quantizer, dst.second_gemm_tensor, None)
                 updated_second_gemm = True
 
         if self.process_tensor_second_gemm:
             out = nvinspect_api.transformer_engine.process_tensor(
                 layer_name=self.layer_name, tensor_name=self.tensor_name, 
-                gemm=self.second_gemm_gemm_name, tensor=src, 
-                default_quantizer=self.parent_quantizer, out=dst.second_gemm)
+                gemm=self.second_gemm_name, tensor=src, 
+                default_quantizer=self.parent_quantizer, out=dst.second_gemm_tensor, iteration=iteration)
             assert out is None, "API call nvinspect_api.transformer_engine.process_tensor with out != None should return None"
             updated_second_gemm = True
         if self.process_tensor_first_gemm:
             nvinspect_api.transformer_engine.process_tensor(
                 layer_name=self.layer_name, tensor_name=self.tensor_name, 
-                gemm=self.process_tensor_first_gemm, default_quantizer=self.parent_quantizer, 
-                tensor=src, out=dst.first_gemm)
+                gemm=self.first_gemm_name, tensor=src,
+                default_quantizer=self.parent_quantizer, out=dst.first_gemm_tensor, iteration=iteration)
             updated_first_gemm = True
         if not updated_second_gemm:
-            dst.second_gemm.copy_(src)
+            dst.second_gemm_tensor.copy_(src)
         if updated_second_gemm and not updated_first_gemm:
-            dst.first_gemm.copy_(src)
+            dst.first_gemm_tensor.copy_(src)
             # if updated_first_gemm and updated_second_gemm, then
             # dst.second_gemm and dst.first_gemm. is the same tensor,
             # and it is already updated.
@@ -313,11 +321,12 @@ def restore_from_saved(self, tensors):
             restore_from_saved([self.first_gemm_tensor, self.second_gemm_tensor], tensors, return_saved_tensors=True)
         return saved_tensors
 
-    def _quantize(self, tensor):
+    def quantize_(self, tensor, *, noop_flag = None):
+        assert noop_flag is None, "CUDA Graphs are not supported with debug=True!"
         self.quantizer.update_quantized(tensor, self)
 
     def dequantize(self, *, dtype = torch.float32):
-        return self.first_gemm.dequantize().to(dtype)
+        return self.first_gemm_tensor.dequantize().to(dtype)
 
     def get_tensor(self, transpose:bool):
         # Is used in the python gemm() to get tensor or transpose of the tensor.

transformer_engine/debug/debug_state.py

@@ -30,6 +30,7 @@ def initialize(cls):
     def reset(cls):
         from .features.utils.stats_buffer import STATS_BUFFERS, StatsBuffers
         STATS_BUFFERS.reset()
+        cls.debug_enabled = None
         cls.layers_initialized.clear()
 
     @classmethod

transformer_engine/debug/features/api.py

@@ -100,7 +100,7 @@ def use_process_tensor(self, *args, **kwargs):
         return False
 
     def process_tensor(self, *args, **kwargs):
-        return kwargs["tensor"]
+        raise RuntimeError("use_process_tensor() returned True, process_tensor() was invoked, but it is not handled by any API.")
 
     def look_at_tensor_before_process(self, *args, **kwargs):
         pass
@@ -167,9 +167,10 @@ def routing_condition(self, api_name, config, layer_name, feature_obj, **kwargs)
         return status, modified_config
 
     def output_assertions_hook(self, api_name, ret, **kwargs):
-        if api_name in {"process_tensor"}:
-            assert type(ret) in [torch.Tensor, Float8Tensor, Float8TensorBase, MXFP8Tensor, MXFP8TensorBase], \
-                f"This API {api_name} must return a tensor."
+        pass
+        #if api_name in {"process_tensor"}:
+        #    assert type(ret) in [torch.Tensor, Float8Tensor, Float8TensorBase, MXFP8Tensor, MXFP8TensorBase], \
+        #        f"This API {api_name} must return a tensor."
 
 
 

transformer_engine/debug/features/disable_fp8_layer.py

@@ -33,7 +33,7 @@ class DisableFp8Layer:
     """
 
     @api_method
-    def fp8_gemm(self, config, layer_name, gemm):
+    def fp8_gemm(self, config, layer_name, *args, **kwargs):
         for key in config:
             if key not in ["enabled", "gemm"]:
               raise ValueError(f"[NVTORCH INSPECT ERROR] Unexpected key in config: \"{key}\".")

transformer_engine/debug/features/log_tensor_stats.py

@@ -66,9 +66,10 @@ def look_at_tensor_before_process(self, config, layer_name,
         options = (config.get('start_step', None), config.get('end_step', None), config.get('start_end_list', None),)
         skip_reduction = False
         reduction_group = nvinspect_api.get_tensor_reduction_group()
-        if self.tensor_name == "weight":
+        if tensor_name == "weight":
             if TEDebugState.weight_tensor_tp_group_reduce:
-                reduction_group = self.tp_group
+                pass
+                #reduction_group = self.tp_group
             else:
                 skip_reduction = True
 

transformer_engine/debug/features/utils/stats_buffer.py

@@ -102,7 +102,7 @@ def log(self):
             combiner = STATS[stat_name][1]
             stat_value = combiner(gathered_helper_stats)
 
-            MetricLogger.log_scalar(f"{self.layer_name}_{self.tensor_name}_{stat_name}", stat_value.float(), self.iteration)
+            MetricLogger.log_scalar(f"{self.layer_name}_{self.tensor_name}_{stat_name}", stat_value, self.iteration)
             output[(self.layer_name, self.tensor_name, stat_name, self.iteration)] = stat_value # for debuggin purpouses
         self._reset_before_next_step()
         return output