Add hygon backend for TE-FL

Signed-off-by: wenjh <wenjh@sugon.com>

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("NVTE_HYGON_BACKEND", "0"))) else "transformer_engine_torch_hygon",
             sources=[str(src) for src in sources],
             include_dirs=[str(inc) for inc in include_dirs],
             extra_compile_args={

setup.py

@@ -114,9 +114,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("NVTE_HYGON_BACKEND", "0"))):
+        cmake_path = root_path / Path("transformer_engine/common")
+    else:
+        cmake_path = root_path / Path("transformer_engine_hygon/common")
     return CMakeExtension(
-        name="transformer_engine",
-        cmake_path=root_path / Path("transformer_engine/common"),
+        name="transformer_engine" if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))) else "transformer_engine_hygon",
+        cmake_path=cmake_path,
         cmake_flags=cmake_flags,
     )
 
@@ -239,6 +243,7 @@ if __name__ == "__main__":
             if "pytorch" in frameworks:
                 from build_tools.pytorch import setup_pytorch_extension
 
+                if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))):
                     ext_modules.append(
                         setup_pytorch_extension(
                             "transformer_engine/pytorch/csrc",
@@ -246,6 +251,14 @@ if __name__ == "__main__":
                             current_file_path / "transformer_engine",
                         )
                     )
+                else:
+                    ext_modules.append(
+                        setup_pytorch_extension(
+                            "transformer_engine_hygon/pytorch/csrc",
+                            current_file_path / "transformer_engine_hygon" / "pytorch" / "csrc",
+                            current_file_path / "transformer_engine_hygon",
+                        )
+                    )
             if "jax" in frameworks:
                 from build_tools.jax import setup_jax_extension
 
@@ -257,6 +270,7 @@ if __name__ == "__main__":
                     )
                 )
 
+    if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))):
         # Configure package
         setuptools.setup(
             name="transformer_engine",
@@ -281,3 +295,37 @@ if __name__ == "__main__":
             include_package_data=include_package_data,
             package_data=package_data,
         )
+    else:
+        # Configure package of hygon backend for TransformerEngine-FL
+        common_dir = current_file_path / "transformer_engine" / "common"
+        common_copy = current_file_path / "transformer_engine_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_hygon" / "pytorch" / "csrc"
+        if csrc_copy.exists():
+            shutil.rmtree(csrc_copy)
+        shutil.copytree(csrc_dir, csrc_copy)
+        setuptools.setup(
+            name="transformer_engine_hygon",
+            version=__version__,
+            packages=setuptools.find_packages(
+                include=[
+                    "transformer_engine_hygon",
+                    "transformer_engine_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,
+        )