Merge branch 'TE_develop2.8' into 'develop_v2.8'

[DCU]Fix memory overflow and test-didistributed in L1_pytorch_istributed_unittest

See merge request dcutoolkit/deeplearing/TransformerEngine!49

tests/pytorch/debug/run_distributed.py

@@ -17,6 +17,7 @@ import transformer_engine_torch as tex
 import nvdlfw_inspect.api as debug_api
 from transformer_engine.debug import set_weight_tensor_tp_group_reduce
 from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+from torch.utils.cpp_extension import IS_HIP_EXTENSION
 
 from test_numerics import (
     _emulate_linear,
@@ -47,7 +48,6 @@ TEST_NR = 0
 
 fp8_available, _ = FP8GlobalStateManager.is_fp8_available()
 
-
 def _get_tensors(parallel_mode, weight_seed=SEED, data_seed=SEED, tp_size=None, tp_rank=None):
     if tp_size is None:
         tp_size = WORLD_SIZE
@@ -72,6 +72,16 @@ def _get_tensors(parallel_mode, weight_seed=SEED, data_seed=SEED, tp_size=None,
 
 
 def _init_model(weight, parallel_mode=None, tp_group=None, name="linear"):
+    if IS_HIP_EXTENSION:
+        model = transformer_engine.pytorch.Linear(
+        IN_SIZE,
+        OUT_SIZE,
+        name=name,
+        bias=False,
+        parallel_mode=parallel_mode,
+        tp_group=(tp_group or NCCL_WORLD if parallel_mode else None),
+        )
+    else:
         model = transformer_engine.pytorch.Linear(
             IN_SIZE,
             OUT_SIZE,
@@ -363,7 +373,6 @@ def test_log_distributed(parallel_mode, gather_weight, **kwargs):
             )
     set_weight_tensor_tp_group_reduce(True)  # reset
 
-
 @run_debug_test
 def sanity_test_log_quantized_stats(parallel_mode, gather_weight, **kwargs):
     from test_log import LOG_QUANTIZED_CONFIG
@@ -580,6 +589,9 @@ def test_fake_quant_fp8(
         "dgrad_fp8": not (dgrad_weight or dgrad_grad),
         "wgrad_fp8": not (wgrad_grad or wgrad_input),
     }
+    if IS_HIP_EXTENSION:
+        if fp8_kwargs["fprop_fp8"] or fp8_kwargs["dgrad_fp8"] or fp8_kwargs["wgrad_fp8"]:
+            return      # Output type 32 (FP32) does not support int8 simulation.
     if WORLD_RANK == 0:
         fake_quant_fp8_create_config(
             fprop_inp,
@@ -667,6 +679,10 @@ if __name__ == "__main__":
     random.seed(SEED)
     _init_distributed()
 
+    if IS_HIP_EXTENSION:
+        # Output type 32 (FP32) does not support int8 simulation.
+        pass
+    else:
         test_log_expert_parallel()
         for parallel_mode in ["column", "row"]:
             for gather_weight in [True, False]:
@@ -676,6 +692,11 @@ if __name__ == "__main__":
         for parallel_mode in ["row", "column"]:
             test_disable_fp8_layer(parallel_mode)
 
+
+    if IS_HIP_EXTENSION:
+        # Output type 32 (FP32) does not support int8 simulation.
+        pass
+    else:
     # test_disable_fp8_gemms
         _run_test_with_combinations(
             test_disable_fp8_gemms, all_boolean, num_repeat=3, extra_args=["column", "row"]
@@ -690,7 +711,10 @@ if __name__ == "__main__":
         extra_args=["column", "row"],
         sample_size=20,
     )
-
+    if IS_HIP_EXTENSION:
+        # Output type 32 (FP32) does not support int8 simulation.
+        pass
+    else:
         _run_test_with_combinations(
             test_per_tensor_scaling,
             all_boolean,

tests/pytorch/distributed/run_numerics.py

@@ -509,7 +509,7 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs):
 
 def test_linear():
     """Run linear layer tests with various configurations."""
-    kwargs_list = [
+    base_kwargs_list = [
         {},
         {"bias": False},
         {"init_method": _constant},
@@ -519,6 +519,15 @@ def test_linear():
         {"delay_wgrad_compute": True},
         {"save_original_input": True},
     ]
+    #TODO:The blockwise recipe does not currently support calculations with bias set to true.
+    """
+    For AMD platforms, when the quantization recipe is fp8_block_scaling, iterate through base_kwargs_list,
+    and if the bias value is not set in kwargs or the bias value is true, set bias to false.
+    """
+    if IS_HIP_EXTENSION and QUANTIZATION == "fp8_block_scaling":
+        kwargs_list = [kwargs for kwargs in base_kwargs_list if kwargs.get("bias", True) is False]
+    else:
+        kwargs_list = base_kwargs_list
     for kwargs in kwargs_list:
         if kwargs.get("save_original_input", False) and QUANTIZATION == "fp8":
             continue
@@ -688,7 +697,7 @@ def _test_layernorm_linear(parallel_mode=None, sequence_parallel=False, **kwargs
 
 
 def test_layernorm_linear():
-    kwargs_list = [
+    base_kwargs_list = [
         {},
         {"bias": False},
         {"init_method": _constant},
@@ -699,6 +708,15 @@ def test_layernorm_linear():
         {"return_layernorm_output": True},
         {"delay_wgrad_compute": True},
     ]
+    #TODO:The blockwise recipe does not currently support calculations with bias set to true.
+    """
+    For AMD platforms, when the quantization recipe is fp8_block_scaling, iterate through base_kwargs_list,
+    and if the bias value is not set in kwargs or the bias value is true, set bias to false.
+    """
+    if IS_HIP_EXTENSION and QUANTIZATION == "fp8_block_scaling":
+        kwargs_list = [kwargs for kwargs in base_kwargs_list if kwargs.get("bias", True) is False]
+    else:
+        kwargs_list = base_kwargs_list
     for kwargs in kwargs_list:
         for parallel_mode in ["column"]:
             for sequence_parallel in [False, True]:
@@ -793,7 +811,7 @@ def _test_layernorm_mlp(set_parallel_mode=None, sequence_parallel=False, **kwarg
 
 
 def test_layernorm_mlp():
-    kwargs_list = [
+    base_kwargs_list = [
         {},
         {"init_method": _constant},
         {"output_layer_init_method": _constant},
@@ -807,7 +825,15 @@ def test_layernorm_mlp():
         {"return_layernorm_output": True},
         {"delay_wgrad_compute": True},
     ]
-
+    #TODO:The blockwise recipe does not currently support calculations with bias set to true.
+    """
+    For AMD platforms, when the quantization recipe is fp8_block_scaling, iterate through base_kwargs_list,
+    and if the bias value is not set in kwargs or the bias value is true, set bias to false.
+    """
+    if IS_HIP_EXTENSION and QUANTIZATION == "fp8_block_scaling":
+        kwargs_list = [kwargs for kwargs in base_kwargs_list if kwargs.get("bias", True) is False]
+    else:
+        kwargs_list = base_kwargs_list
     for kwargs in kwargs_list:
         for set_parallel_mode in [True]:
             for sequence_parallel in [False, True]:
@@ -882,7 +908,7 @@ def _test_transformer_layer_parallel(sequence_parallel=False, **kwargs):
 
 
 def test_transformer_layer():
-    kwargs_list = [
+    base_kwargs_list = [
         {},
         {"num_gqa_groups": 4},
         {"init_method": _constant},
@@ -902,6 +928,15 @@ def test_transformer_layer():
         {"fuse_qkv_params": True},
         {"activation": "relu"},
     ]
+    #TODO:The blockwise recipe does not currently support calculations with bias set to true.
+    """
+    For AMD platforms, when the quantization recipe is fp8_block_scaling, iterate through base_kwargs_list,
+    and if the bias value is not set in kwargs or the bias value is true, set bias to false.
+    """
+    if IS_HIP_EXTENSION and QUANTIZATION == "fp8_block_scaling":
+        kwargs_list = [kwargs for kwargs in base_kwargs_list if kwargs.get("bias", True) is False]
+    else:
+        kwargs_list = base_kwargs_list
 
     for kwargs in kwargs_list:
         for sequence_parallel in [False, True]:

tests/pytorch/distributed/test_numerics.py

@@ -9,7 +9,8 @@ from pathlib import Path
 import pytest
 import torch
 from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
-
+from torch.utils.cpp_extension import IS_HIP_EXTENSION
+import transformer_engine as te
 """
     Distributed numerics tests
 
@@ -61,4 +62,15 @@ def test_distributed(quantization):
         pytest.skip(reason_for_no_mxfp8)
     if quantization == "fp8_block_scaling" and not fp8_block_scaling_available:
         pytest.skip(reason_for_no_fp8_block_scaling)
+    if IS_HIP_EXTENSION and quantization == "fp8_block_scaling":
+        import importlib
+        ori_int8_sim_fp8 = os.environ.get("NVTE_INT8_SIM_FP8", "None")
+        os.environ["NVTE_INT8_SIM_FP8"] = "1"
+        importlib.reload(te.pytorch.fp8)
     _run_test(quantization)
+    if IS_HIP_EXTENSION and quantization == "fp8_block_scaling":
+        if ori_int8_sim_fp8 is None or ori_int8_sim_fp8 == "None":
+            os.environ["NVTE_INT8_SIM_FP8"] = "0"
+        else:
+           del os.environ["NVTE_INT8_SIM_FP8"]
+        importlib.reload(te.pytorch.fp8)

transformer_engine/common/transpose/transpose.cu

@@ -217,7 +217,13 @@ void transpose(const Tensor &input, const Tensor &noop, Tensor *output_, cudaStr
 
       // Choose between runtime-compiled or statically-compiled kernel
       const bool aligned = (row_length % THREADS_PER_WARP == 0 && num_rows % THREADS_PER_WARP == 0);
-      if (aligned && rtc::is_enabled()) {  // Runtime-compiled tuned kernel
+      //TODO:Using RTC may cause kernel crashes. Therefore, set use_rtc to true to avoid using RTC and resolve the kernel crash issue.
+#ifdef USE_ROCM
+      const bool use_rtc = false;  
+#else
+      const bool use_rtc = true;  
+#endif
+      if (aligned && rtc::is_enabled() && use_rtc) {  // Runtime-compiled tuned kernel
         // Pick kernel config
         std::vector<KernelConfig> kernel_configs;
         kernel_configs.reserve(16);