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,13 +243,22 @@ if __name__ == "__main__":
             if "pytorch" in frameworks:
                 from build_tools.pytorch import setup_pytorch_extension
 
-                ext_modules.append(
-                    setup_pytorch_extension(
-                        "transformer_engine/pytorch/csrc",
-                        current_file_path / "transformer_engine" / "pytorch" / "csrc",
-                        current_file_path / "transformer_engine",
+                if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))):
+                    ext_modules.append(
+                        setup_pytorch_extension(
+                            "transformer_engine/pytorch/csrc",
+                            current_file_path / "transformer_engine" / "pytorch" / "csrc",
+                            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,27 +270,62 @@ if __name__ == "__main__":
                     )
                 )
 
-    # Configure package
-    setuptools.setup(
-        name="transformer_engine",
-        version=__version__,
-        packages=setuptools.find_packages(
-            include=[
-                "transformer_engine",
-                "transformer_engine.*",
-                "transformer_engine/build_tools",
-            ],
-        ),
-        extras_require=extras_require,
-        description="Transformer acceleration library",
-        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,
-    )
+    if not bool(int(os.getenv("NVTE_HYGON_BACKEND", "0"))):
+        # Configure package
+        setuptools.setup(
+            name="transformer_engine",
+            version=__version__,
+            packages=setuptools.find_packages(
+                include=[
+                    "transformer_engine",
+                    "transformer_engine.*",
+                    "transformer_engine/build_tools",
+                ],
+            ),
+            extras_require=extras_require,
+            description="Transformer acceleration library",
+            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,
+        )
+    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,
+        )