关于计时的一个问题

[sxysxy]

在之前的工作中，我使用fairseq.logging.meters.StopwatchMeter 计NMT预测时间、KNN检索时间，大概是这样的流程：

timerNMT.start()
x = NMTModel(input)
timerNMT.stop()
timerKNN.start()
knn_retrieve(x)
timerKNN.stop()

但是得到的结果却非常诡异，即便是非常非常小的datastore，有GPU加速的faiss进行KNN检索的总用时占了inference的大部分时间，后来我发现，https://github.com/NJUNLP/knn-box/blob/master/knnbox/retriever/utils.py 中的retrieve_k_nearest中，会临时把tensor迁移到cpu上，这个操作会先进行一个cuda同步，等待tensor的前置算子全部计算完成后，才能把它的值拷贝到cpu上。
这时我意识到，根据pytorch文档的说法，pytorch上的cuda kernel默认是异步执行的。也就是说python这边的函数返回的时候，cuda操作可能并没有完成，我把计时的流程改为

torch.cuda.synchronize()
timerNMT.start()
x = NMTModel(input)
torch.cuda.synchronize()
timerNMT.stop()
torch.cuda.synchronize()
timerKNN.start()
knn_retrieve(x)
torch.cuda.synchronize()
timerKNN.stop()

它结果就看起来靠谱多了。

另一个更简单的例子则是(运行在我的RTX8000 GPU上)：

import torch
import time
a = torch.randn(10000,100000,device='cuda')
b = torch.randn(100000,10000,device='cuda')
#torch.cuda.synchronize()
start = time.perf_counter()
c = a@b
#torch.cuda.synchronize()
print(time.perf_counter() - start)

输出 0.27494998497422785

而解除torch.cuda.synchronize()的注释后：

import torch
import time
a = torch.randn(10000,100000,device='cuda')
b = torch.randn(100000,10000,device='cuda')
torch.cuda.synchronize()
start = time.perf_counter()
c = a@b
torch.cuda.synchronize()
print(time.perf_counter() - start)

输出2.4091989540029317

也就是说如果我们想测试算法的用时，应该在计时器的开始和结束前都进行一次cuda流同步操作，以确保测得时间就是中间被测代码的实际运行到运算出结果的时间。

于是我改了一下fairseq的StopwachMeter，在start和stop中加入cuda同步操作

class StopWatchTimer:
    '''
    Timer for time measurement.
    '''
    def __init__(self, cudaSyncOnEvents=False, cudaStream : torch.cuda.Stream=None) -> None:
        '''
        Args:
            cudaSyncOnEvents: Call cuda synchronize if start、stop、reset、elapsedTime is called
        '''
        self.startTime = None 
        self.totalTime = 0
        self.itemCount = 0
        
        if cudaSyncOnEvents and (not torch.cuda.is_available()):
            raise RuntimeError("cuda is not available")
        
        self.cudaSyncOnEvents = cudaSyncOnEvents
        self.cudaStream = cudaStream
        if cudaStream:
            self.cudaSyncFunction = cudaStream.synchronize
        else:
            self.cudaSyncFunction = torch.cuda.synchronize
        
    def __cudaSync(self): 
        if self.cudaSyncOnEvents:
            self.cudaSyncFunction()
        
    def start(self):
        self.__cudaSync()
        self.startTime = time.perf_counter()
        
    def stop(self, itemCount=0):
        if self.startTime is not None:
            self.__cudaSync()
            dtime = time.perf_counter() - self.startTime
            self.totalTime += dtime
            self.itemCount += itemCount
            
    def reset(self):
        self.itemCount = 0
        self.totalTime = 0
        self.start()
        
    def elapsedTime(self) -> float:
        if self.startTime is None: 
            return 0.0
        self.__cudaSync()
        return time.perf_counter() - self.startTime

Emmm，这或许其实是fairseq的问题.....

希望或许能对一些KNN-MT的加速的工作有帮助

[sxysxy]

但另一方面，在计时时进行cuda流同步，会降低CPU（执行python代码）和GPU（执行cuda kernels）的并行程度。加入计时后的程序运行时间会有一点不可忽略的增加。

[Maxwell-Lyu]

感谢您对这个问题的深入分析，“对每步操作进行单独计时，需要在操作头尾执行sync”这个结论应该是正确的。
CUDA kernel的执行是异步的没错，事实上我们从python调用在cuda上计算的函数，几乎都不会等待计算完成。除了一些意外情况，在pytorch的repo当中也有人指出这个bug。
但是有一些情况会符合预期地，去等待执行完成，然后将结果返回，比如手动调用synchronize，或者移动结果到CPU，faiss的检索就是后者的情况。

[sxysxy]

问题在于

timerNMT.start()
x = NMTModel(input)
timerNMT.stop()
timerKNN.start()
knn_retrieve(x)
timerKNN.stop()

faiss检索的时候隐式地进行了流同步，也就是在knn_retrieve(x)的内部进行了流同步，那么x=NMTModel(input)的大部分运行时间，也会被计入timerKNN中。在最开始，我没有在timer开始前进行同步的时候，测试也“复现出”和之前很多论文中提到的，KNN检索远比NMT模型耗时的现象。但后来加入同步之后，发现KNN检索的时间远比NMT少。

PCK KNN-MT论文说他们把推理延迟降低了最多57%。但是我在KNN-BOX框架下用Vanilla KNN-MT测试IT Medical Law Koran等数据集，推理阶段KNN检索所占用的总用时都远远不到57%。

PCK KNN-MT论文公开的代码以及还有Revised-Key KNN-MT公开的代码，都是基于原始KNN-MT论文的代码实现的，它们没有单独对NMT和KNN计时的代码。不过我发现他们的代码在运行的时候发现GPU利用率极低，远不如KNN-BOX。

所以KNN-MT速度慢到底是因为facebook研究院的号称在“10亿量级的索引可以做到毫秒级检索的性能”的FAISS慢，还是原始KNN-MT的代码优化太差导致的呢？KNN-MT加速的研究工作是否还有的做呢？