Add route to modify megatron fused kernel build path (#1383)

* Add arg to configure build path

* cleanup and formatting

---------

Co-authored-by: Quentin Anthony <qganthony@yahoo.com>

Author: aflah02

configs/neox_arguments.md

@@ -684,6 +684,14 @@ Model Arguments
 
 
 
+- **fused_kernels_build_path**: typing.Optional[str]
+
+    Default = None
+
+    Optional override for the fused kernels build directory. If unset, defaults to `megatron/fused_kernels/build` relative to the package.
+
+
+
 - **fp16_lm_cross_entropy**: bool
 
     Default = False
@@ -2615,4 +2623,3 @@ Args for deepspeed runner (deepspeed.launcher.runner).
     Default = None
 
     Adds a `--account` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometimes necessary for cluster rules, or so I've heard.
-

megatron/fused_kernels/__init__.py

@@ -58,7 +58,7 @@ def load(neox_args=None):
 
     # Build path
     srcpath = pathlib.Path(__file__).parent.absolute()
-    buildpath = srcpath / "build"
+    buildpath = _get_build_path(neox_args=neox_args)
     _create_build_dir(buildpath)
 
     # Determine verbosity
@@ -156,12 +156,23 @@ def _get_cuda_bare_metal_version(cuda_dir):
     return raw_output, bare_metal_major, bare_metal_minor
 
 
+def _get_build_path(neox_args=None):
+    if neox_args is not None:
+        fused_kernels_build_path = getattr(neox_args, "fused_kernels_build_path", None)
+        if fused_kernels_build_path is not None:
+            return pathlib.Path(fused_kernels_build_path).expanduser().resolve()
+
+    srcpath = pathlib.Path(__file__).parent.absolute()
+    return srcpath / "build"
+
+
 def _create_build_dir(buildpath):
     try:
-        os.mkdir(buildpath)
-    except OSError:
-        if not os.path.isdir(buildpath):
-            print(f"Creation of the build directory {buildpath} failed")
+        pathlib.Path(buildpath).mkdir(parents=True, exist_ok=True)
+    except OSError as e:
+        raise RuntimeError(
+            f"Creation of the build directory {buildpath} failed: {e}"
+        ) from e
 
 
 def load_fused_kernels():

megatron/neox_arguments/neox_args.py

@@ -339,6 +339,11 @@ class NeoXArgsModel(NeoXArgsTemplate):
     Enable rotary embedding fusion.
     """
 
+    fused_kernels_build_path: Optional[str] = None
+    """
+    Optional override for the fused kernels build directory. If unset, defaults to `megatron/fused_kernels/build` relative to the package.
+    """
+
     fp16_lm_cross_entropy: bool = False
     """
     Move the cross entropy unreduced loss calculation for lm head to fp16.