Fix gemm tuner error mi350 (ROCm#1313)

* workaround-retry tuning when encounter invalid pointer

* workaround-retry tuning when encounter invalid pointer

* fix  lint error

* Update gemm_tuner.py

em timeout

Author: yzhou103

gradlib/gradlib/GemmTuner.py

@@ -148,6 +148,7 @@ def __init__(
         profile_file="",
         # splitK=None,
     ):
+        torch.cuda.empty_cache()
         self.m = m
         self.k = k
         self.n = n
@@ -166,19 +167,19 @@ def __init__(
         self.rocb_sols = []
         self.rtol = 1e-2
         self.atol = 1e-2
-        self.ref = self.get_gemm_ref()
+        # self.ref = self.get_gemm_ref()
         self.check_err_ratio = err_ratio
         self.splitK = None
         self.profile_file = profile_file
-        self.start = torch.cuda.Event(enable_timing=True)
-        self.end = torch.cuda.Event(enable_timing=True)
+        # self.start = torch.cuda.Event(enable_timing=True)
+        # self.end = torch.cuda.Event(enable_timing=True)
         # prefer hipblaslt unless rocblas time is less than this
         # ratio of hipblaslt time
         self.hipb_prefer_ratio = 0.995
         self.rocblas_decode = rocblas_decode
         self.mp = mp
-        self.inbpe = self.inp.element_size()
-        self.outbpe = self.ref.element_size()
+        # self.inbpe = self.inp.element_size()
+        # self.outbpe = self.ref.element_size()
         self.asm_map = {}
 
     def find_hipblas_sols(self):
@@ -379,10 +380,15 @@ def hipb_time_all_sols(self, fast_mode=0, top_sols=0):
         if fast_mode == 1:
             self.hipb_gtimedf = self.save_topn_result(ret, fast_mode, "hipblaslt")
             return []
+        print(f">>> hipblaslt top solutions, Fast Mode {fast_mode}")
         return ret
 
     def save_topn_result(self, rets, fast_mode, libtype):
         results = []
+        if not rets:
+            return pd.DataFrame(
+                columns=["solidx", "gtimems", "splitK", "err_ratio", "kernelName"]
+            )
         for info, us, err_ratio in rets:
             res_one = []
             solidx = info[1]
@@ -478,8 +484,11 @@ def rocb_time_all_sols(self, fast_mode=0, top_sols=0):
                     self.atol,
                 )
             )
-        in_data = [(len(solutions), ())]
-        ret = mp_tuner(task, in_data, self.mp, fast_mode == 1)
+        if task:
+            in_data = [(len(solutions), ())]
+            ret = mp_tuner(task, in_data, self.mp, fast_mode == 1)
+        else:
+            ret = []
         if fast_mode == 1:
             self.rocb_gtimedf = self.save_topn_result(ret, fast_mode, "rocblas")
             return []
@@ -519,6 +528,28 @@ def run_solutions(self):
         rets = self.run_best_solutions()
         return rets
 
+    def cleanup(self):
+        if hasattr(self, "inp"):
+            del self.inp
+        if hasattr(self, "weights"):
+            del self.weights
+        if hasattr(self, "bias") and self.bias is not None:
+            del self.bias
+        if hasattr(self, "blob"):
+            cpu_blob = self.blob.cpu()
+            del cpu_blob
+
+    def cleanup(self):
+        if hasattr(self, "inp"):
+            del self.inp
+        if hasattr(self, "weights"):
+            del self.weights
+        if hasattr(self, "bias") and self.bias is not None:
+            del self.bias
+        if hasattr(self, "blob"):
+            cpu_blob = self.blob.cpu()
+            del cpu_blob
+
 
 class GemmTuner(GemmCommonTuner):
     ARG_DEFAULTS = {
@@ -597,6 +628,7 @@ def __init__(
 
         self.hipb_prefer_ratio = 0.995
         self.cu_num = self.get_cu_num()
+        self.gemmobj = None
 
     def calculate_perf(
         self,
@@ -708,7 +740,6 @@ def tune(self, untunedf, tunedf, args):
             ds = df.loc[i, :]
             indtype = ds["dtype"]
             outdtype = ds["outdtype"]
-
             gemmobj = Gemm(
                 ds["M"],
                 ds["N"],
@@ -722,9 +753,11 @@ def tune(self, untunedf, tunedf, args):
                 err_ratio=args.errRatio,
                 profile_file=args.profile_file,
             )
+
             ret.extend(gemmobj.run_solutions())
+            gemmobj.cleanup()
             del gemmobj
-            torch.cuda.empty_cache()
+
         return ret
 
     def processResult(self, rets, fast_mode):
@@ -819,7 +852,6 @@ def post_process(self, rets, args, topk=-1, fast_mode=False):
             if best_gtimedfs.empty:
                 best_gtimedfs = resultdf1
             else:
-                print("concat ", resultdf1)
                 best_gtimedfs = pd.concat([best_gtimedfs, resultdf1], ignore_index=True)
 
             print(f"{key} >>> Fastest Solution is \n {resultdf1}", flush=True)

gradlib/gradlib/gemm_tuner.py

@@ -28,6 +28,8 @@
 from GemmTuner import GemmTuner
 
 import time
+import multiprocessing as mp
+import gc
 
 aiter.rocb_create_extension()
 aiter.hipb_create_extension()
@@ -89,7 +91,7 @@ def load_input_gemms(input_file):
         return
 
 
-if __name__ == "__main__":
+def runGemmTuner():
     gtuner = GemmTuner()
     ext_group = gtuner.parser.add_argument_group("extra parameters")
     ext_group.add_argument(
@@ -117,7 +119,6 @@ def load_input_gemms(input_file):
         help="Tensor parallelism to be used.",
     )
     args = gtuner.parse_args()
-
     if args.outdtype is None:
         args.outdtype = args.indtype
     indtype = get_dtype(args.indtype)
@@ -130,9 +131,7 @@ def load_input_gemms(input_file):
             print(">>> Warning! NO MODEL SPECIFIED. Tuning for LL2 13B TP1")
             # LL2 13B sizes
             mksets = [(15360, 5120), (5120, 5120), (27648, 5120), (5120, 13824)]
-
             gtuner.add_gemm(m=32000, n=1, k=5120, indtype=indtype)  # logits gemm
-
         else:
             mksets, hidden_size, dtype = generate_mk_sets(args.model_dir, args.tp)
             gtuner.add_gemm(
@@ -141,11 +140,62 @@ def load_input_gemms(input_file):
                 k=hidden_size,
                 indtype=dtype,
             )  # TODO: Handle cases where vocab_size is not divisible by tp
-
             for n in sorted(nsets):
                 for m, k in mksets:
                     gtuner.add_gemm(m, n, k, indtype=dtype)
         gtuner.untunedf.to_csv("./tmp_untuned.csv", index=False)
         args.untune_file = "./tmp_untuned.csv"
-
     gtuner.run(args)
+
+
+def clean():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if hasattr(torch.cuda, "memory_allocated"):
+        torch.cuda.synchronize()
+    try:
+        if hasattr(mp, "resource_tracker"):
+            mp.resource_tracker.ensure_running()
+            # clean  leaked semaphore objects
+            if hasattr(mp.resource_tracker, "_CLEANUP_FUNCS"):
+                # be careful
+                for name in list(mp.resource_tracker._CLEANUP_FUNCS.keys()):
+                    try:
+                        mp.resource_tracker._CLEANUP_FUNCS.pop(name)()
+                    except:
+                        pass
+    except Exception as e:
+        print(f"Resource cleanup warning: {e}")
+
+
+if __name__ == "__main__":
+    retries = 0
+    MAX_TRY = 30
+    mp.set_start_method("spawn", force=True)
+    while retries <= MAX_TRY:
+        try:
+            process = mp.Process(target=runGemmTuner, args=(), daemon=False)
+            process.start()
+            process.join()
+            if process.exitcode != 0:
+                time.sleep(0.5 * retries)
+                print(
+                    "!Error when run GemmTuner process exitcode is ", process.exitcode
+                )
+                clean()
+                retries += 1
+            else:
+                break
+        except Exception as e:
+            print(f"Process creation failed: {e}")
+            retries += 1
+            clean()
+            time.sleep(1)
+        finally:
+            if process and process.is_alive():
+                process.terminate()
+                process.join(timeout=5)
+
+    clean()
+    print(f"retried num is {retries}")