Merge branch 'develop_v2.10' into release_v2.10

build_tools/build_ext.py

@@ -145,7 +145,10 @@ def get_build_ext(
             # For editable/inplace builds this is not a concern as
             # the SOs will be in a local directory anyway.
             if not self.inplace:
+                if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
                     target_dir = install_dir / "transformer_engine" / lib_dir
+                else:
+                    target_dir = install_dir / "transformer_engine_fl_hygon" / lib_dir
                 target_dir.mkdir(exist_ok=True, parents=True)
 
                 for ext in Path(self.build_lib).glob("*.so"):

build_tools/pytorch.py

@@ -155,7 +155,7 @@ def setup_pytorch_extension(
 
     if rocm_build():
         return CUDAExtension(
-            name="transformer_engine_torch",
+            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],
             include_dirs=[str(inc) for inc in include_dirs],
             extra_compile_args={

setup.py

@@ -50,6 +50,17 @@ if rocm_build():
 else:
     archs = cuda_archs()
 
+if bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
+    common_dir = current_file_path / "transformer_engine" / "common"
+    common_copy = current_file_path / "transformer_engine_fl_hygon" / "common"
+    if common_copy.exists():
+        shutil.rmtree(common_copy)
+    shutil.copytree(common_dir, common_copy)
+    csrc_dir = current_file_path / "transformer_engine" / "pytorch" / "csrc"
+    csrc_copy = current_file_path / "transformer_engine_fl_hygon" / "pytorch" / "csrc"
+    if csrc_copy.exists():
+        shutil.rmtree(csrc_copy)
+    shutil.copytree(csrc_dir, csrc_copy)
 
 class TimedBdist(bdist_wheel):
     """Helper class to measure build time"""
@@ -114,9 +125,13 @@ def setup_common_extension() -> CMakeExtension:
         if os.getenv("NVTE_USE_ROCBLAS") is not None:
             cmake_flags.append("-DUSE_ROCBLAS=ON")
 
+    if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
+        cmake_path = root_path / Path("transformer_engine/common")
+    else:
+        cmake_path = root_path / Path("transformer_engine_fl_hygon/common")
     return CMakeExtension(
-        name="transformer_engine",
-        cmake_path=root_path / Path("transformer_engine/common"),
+        name="transformer_engine" if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))) else "transformer_engine_fl_hygon",
+        cmake_path=cmake_path,
         cmake_flags=cmake_flags,
     )
 
@@ -239,6 +254,7 @@ if __name__ == "__main__":
             if "pytorch" in frameworks:
                 from build_tools.pytorch import setup_pytorch_extension
 
+                if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
                     ext_modules.append(
                         setup_pytorch_extension(
                             "transformer_engine/pytorch/csrc",
@@ -246,6 +262,14 @@ if __name__ == "__main__":
                             current_file_path / "transformer_engine",
                         )
                     )
+                else:
+                    ext_modules.append(
+                        setup_pytorch_extension(
+                            "transformer_engine_fl_hygon/pytorch/csrc",
+                            current_file_path / "transformer_engine_fl_hygon" / "pytorch" / "csrc",
+                            current_file_path / "transformer_engine_fl_hygon",
+                        )
+                    )
             if "jax" in frameworks:
                 from build_tools.jax import setup_jax_extension
 
@@ -257,6 +281,7 @@ if __name__ == "__main__":
                     )
                 )
 
+    if not bool(int(os.getenv("TEFL_HYGON_BACKEND", "0"))):
         # Configure package
         setuptools.setup(
             name="transformer_engine",
@@ -281,3 +306,27 @@ if __name__ == "__main__":
             include_package_data=include_package_data,
             package_data=package_data,
         )
+    else:
+        # Configure package of hygon backend for TransformerEngine-FL
+        setuptools.setup(
+            name="transformer_engine_fl_hygon",
+            version=__version__,
+            packages=setuptools.find_packages(
+                include=[
+                    "transformer_engine_fl_hygon",
+                    "transformer_engine_fl_hygon.*",
+                ],
+            ),
+            extras_require=extras_require,
+            description="Transformer acceleration library for TransformerEngine-FL",
+            long_description=long_description,
+            long_description_content_type="text/x-rst",
+            ext_modules=ext_modules,
+            cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist},
+            python_requires=f">={min_python_version_str()}",
+            classifiers=["Programming Language :: Python :: 3"],
+            install_requires=install_requires,
+            license_files=("LICENSE",),
+            include_package_data=include_package_data,
+            package_data=package_data,
+        )