add requirements

.gitmodules

@@ -4,4 +4,8 @@
 
 [submodule "lightseq"]
 	path = lightseq
-	url = https://github.com/thu-coai/lightseq-nat
+	url = https://github.com/thu-coai/lightseq-nat
+
+[submodule "cub"]
+	path = cub
+	url = https://github.com/NVIDIA/cub

README.md

@@ -27,7 +27,7 @@ This repo is modified from [``fairseq:5175fd``](https://github.com/pytorch/fairs
 ## Requirements & Installation
 
 * Python >= 3.7
-* Pytorch == 1.10.1
+* Pytorch == 1.10.1 (tested with cuda == 10.2 or 11.3)
 * gcc >= 7.0.0 (for compiling cuda operations. See FAQs if you want to use a lower gcc version)
 * ``git clone --recurse-submodules https://github.com/thu-coai/DA-Transformer.git && pip install -e .``
 * (Optional) Customized LightSeq for NAT (``cd lightseq && pip install -e .``)
@@ -43,6 +43,7 @@ Most codes of the framework are from Fairseq. We mainly add the following files.
 fs_plugins
 ├── criterions
 │   └── nat_dag_loss.py                   # DA-Transformer loss
+├── cub                                   # Requirements: Nvidia CUDA programming model
 ├── custom_ops                            # operations implementations and cuda kernels
 │   ├── dag_best_alignment.cu
 │   ├── logsoftmax_gather.cu

fs_plugins/custom_ops/dag_loss.py

@@ -16,6 +16,7 @@
 
 import os
 import math
+import sys
 
 import torch
 from torch import nn, Tensor
@@ -39,6 +40,7 @@ def get_dag_kernel():
     if dag_kernel is not None:
         return dag_kernel
     else:
+        print("Start compiling cuda operations for DA-Transformer...", file=sys.stderr, flush=True)
         dag_kernel = load(
             "dag_loss_fn",
             sources=[
@@ -47,9 +49,11 @@ def get_dag_kernel():
                 os.path.join(module_path, "dag_best_alignment.cu"),
                 os.path.join(module_path, "logsoftmax_gather.cu"),
             ],
-            extra_cflags=['-DOF_SOFTMAX_USE_FAST_MATH'],
-            extra_cuda_cflags=['-DOF_SOFTMAX_USE_FAST_MATH'],
+            extra_cflags=['-DOF_SOFTMAX_USE_FAST_MATH', '-O3'],
+            extra_cuda_cflags=['-DOF_SOFTMAX_USE_FAST_MATH', '-O3'],
+            extra_include_paths=[os.path.join(module_path, "../../cub")],
         )
+        print("Cuda operations compiling finished", file=sys.stderr, flush=True)
         return dag_kernel
 
 class DagLossFunc(Function):

setup.py

@@ -224,6 +224,7 @@ def do_setup(package_data):
             "sacrebleu[ja]",
             "tqdm",
             "bitarray",
+            "ninja"
         ],
         dependency_links=dependency_links,
         packages=find_packages(