Commit bdf3d931 authored by wenjh's avatar wenjh
Browse files

Rename tefl hygon backend


Signed-off-by: wenjh's avatarwenjh <wenjh@sugon.com>
parent 40816696
...@@ -145,10 +145,10 @@ def get_build_ext( ...@@ -145,10 +145,10 @@ def get_build_ext(
# For editable/inplace builds this is not a concern as # For editable/inplace builds this is not a concern as
# the SOs will be in a local directory anyway. # the SOs will be in a local directory anyway.
if not self.inplace: if not self.inplace:
if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))): if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
target_dir = install_dir / "transformer_engine" / lib_dir target_dir = install_dir / "transformer_engine" / lib_dir
else: else:
target_dir = install_dir / "transformer_engine_hygon" / lib_dir target_dir = install_dir / "transformer_engine_fl_hygon" / lib_dir
target_dir.mkdir(exist_ok=True, parents=True) target_dir.mkdir(exist_ok=True, parents=True)
for ext in Path(self.build_lib).glob("*.so"): for ext in Path(self.build_lib).glob("*.so"):
......
...@@ -155,7 +155,7 @@ def setup_pytorch_extension( ...@@ -155,7 +155,7 @@ def setup_pytorch_extension(
if rocm_build(): if rocm_build():
return CUDAExtension( return CUDAExtension(
name="transformer_engine_torch" if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))) else "transformer_engine_torch_hygon", name="transformer_engine_torch" if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))) else "transformer_engine_fl_torch_hygon",
sources=[str(src) for src in sources], sources=[str(src) for src in sources],
include_dirs=[str(inc) for inc in include_dirs], include_dirs=[str(inc) for inc in include_dirs],
extra_compile_args={ extra_compile_args={
......
...@@ -50,14 +50,14 @@ if rocm_build(): ...@@ -50,14 +50,14 @@ if rocm_build():
else: else:
archs = cuda_archs() archs = cuda_archs()
if bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))): if bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
common_dir = current_file_path / "transformer_engine" / "common" common_dir = current_file_path / "transformer_engine" / "common"
common_copy = current_file_path / "transformer_engine_hygon" / "common" common_copy = current_file_path / "transformer_engine_fl_hygon" / "common"
if common_copy.exists(): if common_copy.exists():
shutil.rmtree(common_copy) shutil.rmtree(common_copy)
shutil.copytree(common_dir, common_copy) shutil.copytree(common_dir, common_copy)
csrc_dir = current_file_path / "transformer_engine" / "pytorch" / "csrc" csrc_dir = current_file_path / "transformer_engine" / "pytorch" / "csrc"
csrc_copy = current_file_path / "transformer_engine_hygon" / "pytorch" / "csrc" csrc_copy = current_file_path / "transformer_engine_fl_hygon" / "pytorch" / "csrc"
if csrc_copy.exists(): if csrc_copy.exists():
shutil.rmtree(csrc_copy) shutil.rmtree(csrc_copy)
shutil.copytree(csrc_dir, csrc_copy) shutil.copytree(csrc_dir, csrc_copy)
...@@ -125,12 +125,12 @@ def setup_common_extension() -> CMakeExtension: ...@@ -125,12 +125,12 @@ def setup_common_extension() -> CMakeExtension:
if os.getenv("NVTE_USE_ROCBLAS") is not None: if os.getenv("NVTE_USE_ROCBLAS") is not None:
cmake_flags.append("-DUSE_ROCBLAS=ON") cmake_flags.append("-DUSE_ROCBLAS=ON")
if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))): if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
cmake_path = root_path / Path("transformer_engine/common") cmake_path = root_path / Path("transformer_engine/common")
else: else:
cmake_path = root_path / Path("transformer_engine_hygon/common") cmake_path = root_path / Path("transformer_engine_fl_hygon/common")
return CMakeExtension( return CMakeExtension(
name="transformer_engine" if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))) else "transformer_engine_hygon", name="transformer_engine" if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))) else "transformer_engine_fl_hygon",
cmake_path=cmake_path, cmake_path=cmake_path,
cmake_flags=cmake_flags, cmake_flags=cmake_flags,
) )
...@@ -254,7 +254,7 @@ if __name__ == "__main__": ...@@ -254,7 +254,7 @@ if __name__ == "__main__":
if "pytorch" in frameworks: if "pytorch" in frameworks:
from build_tools.pytorch import setup_pytorch_extension from build_tools.pytorch import setup_pytorch_extension
if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))): if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
ext_modules.append( ext_modules.append(
setup_pytorch_extension( setup_pytorch_extension(
"transformer_engine/pytorch/csrc", "transformer_engine/pytorch/csrc",
...@@ -265,9 +265,9 @@ if __name__ == "__main__": ...@@ -265,9 +265,9 @@ if __name__ == "__main__":
else: else:
ext_modules.append( ext_modules.append(
setup_pytorch_extension( setup_pytorch_extension(
"transformer_engine_hygon/pytorch/csrc", "transformer_engine_fl_hygon/pytorch/csrc",
current_file_path / "transformer_engine_hygon" / "pytorch" / "csrc", current_file_path / "transformer_engine_fl_hygon" / "pytorch" / "csrc",
current_file_path / "transformer_engine_hygon", current_file_path / "transformer_engine_fl_hygon",
) )
) )
if "jax" in frameworks: if "jax" in frameworks:
...@@ -281,7 +281,7 @@ if __name__ == "__main__": ...@@ -281,7 +281,7 @@ if __name__ == "__main__":
) )
) )
if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))): if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
# Configure package # Configure package
setuptools.setup( setuptools.setup(
name="transformer_engine", name="transformer_engine",
...@@ -309,12 +309,12 @@ if __name__ == "__main__": ...@@ -309,12 +309,12 @@ if __name__ == "__main__":
else: else:
# Configure package of hygon backend for TransformerEngine-FL # Configure package of hygon backend for TransformerEngine-FL
setuptools.setup( setuptools.setup(
name="transformer_engine_hygon", name="transformer_engine_fl_hygon",
version=__version__, version=__version__,
packages=setuptools.find_packages( packages=setuptools.find_packages(
include=[ include=[
"transformer_engine_hygon", "transformer_engine_fl_hygon",
"transformer_engine_hygon.*", "transformer_engine_fl_hygon.*",
], ],
), ),
extras_require=extras_require, extras_require=extras_require,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment