python - 如何使用 CUDA CURAND 保存和恢复随机数生成器的状态？

Question

我正在使用我们组织内开发的大型 CUDA 矩阵库。我需要保存 CUDA RNG 的状态以获取长时间运行的模拟的快照，并能够在以后恢复它。这很简单，例如，python+numpy:

state = numpy.random.get_state()
# state is a tuple with 5 fields which can be pickled, etc.
...
numpy.random.set_state(state)

我似乎无法在 CUDA 主机 API 中找到等效功能。您可以设置种子和偏移量，但无法检索它来保存。设备 API 似乎提供了类似的东西，但是这个库使用主机 api，更改起来会很麻烦。

我正在考虑的 hack-ey 解决方案是跟踪对 RNG 的调用次数(设置种子时重置)，并简单地重复调用 RNG 函数。但是，我不确定函数参数是否必须相同，例如矩阵形状等，使其达到相同的状态。同样，如果调用次数等于初始化RNG的偏移参数，这也可以，即如果我调用RNG 200次，我可以将偏移设置为200。但是，在python中，偏移每次调用状态可以增加超过 1，所以这也可能是错误的。

感谢任何有关如何解决此问题的见解!

Answer 1

对于CURAND Host API，我相信curandSetGeneratorOffset()可能对此有用。

这是来自 curand 主机 API 文档的修改示例:

$ cat t721.cu
/*
 * This program uses the host CURAND API to generate 10
 * pseudorandom floats.  And then regenerate those same floats.
 */
#include <stdio.h>
#include <stdlib.h>
#include <cuda.h>
#include <curand.h>

#define CUDA_CALL(x) do { if((x)!=cudaSuccess) { \
    printf("Error at %s:%d\n",__FILE__,__LINE__);\
    return EXIT_FAILURE;}} while(0)
#define CURAND_CALL(x) do { if((x)!=CURAND_STATUS_SUCCESS) { \
    printf("Error at %s:%d\n",__FILE__,__LINE__);\
    return EXIT_FAILURE;}} while(0)

int main(int argc, char *argv[])
{
    size_t n = 10;
    size_t i;
    curandGenerator_t gen;
    float *devData, *hostData;

    /* Allocate n floats on host */
    hostData = (float *)calloc(n, sizeof(float));

    /* Allocate n floats on device */
    CUDA_CALL(cudaMalloc((void **)&devData, n*sizeof(float)));

    /* Create pseudo-random number generator */
    CURAND_CALL(curandCreateGenerator(&gen,
                CURAND_RNG_PSEUDO_DEFAULT));

    /* Set seed */
    CURAND_CALL(curandSetPseudoRandomGeneratorSeed(gen,
                1234ULL));
    // generator offset = 0
    /* Generate n floats on device */
    CURAND_CALL(curandGenerateUniform(gen, devData, n));
    // generator offset = n
    /* Generate n floats on device */
    CURAND_CALL(curandGenerateUniform(gen, devData, n));
    // generator offset = 2n
    /* Copy device memory to host */
    CUDA_CALL(cudaMemcpy(hostData, devData, n * sizeof(float),
        cudaMemcpyDeviceToHost));

    /* Show result */
    for(i = 0; i < n; i++) {
        printf("%1.4f ", hostData[i]);
    }
    printf("\n\n");

    CURAND_CALL(curandSetGeneratorOffset(gen, n));
    // generator offset = n
    CURAND_CALL(curandGenerateUniform(gen, devData, n));
    // generator offset = 2n
    /* Copy device memory to host */
    CUDA_CALL(cudaMemcpy(hostData, devData, n * sizeof(float),
        cudaMemcpyDeviceToHost));

    /* Show result */
    for(i = 0; i < n; i++) {
        printf("%1.4f ", hostData[i]);
    }
    printf("\n");



    /* Cleanup */
    CURAND_CALL(curandDestroyGenerator(gen));
    CUDA_CALL(cudaFree(devData));
    free(hostData);
    return EXIT_SUCCESS;
}
$ nvcc -o t721 t721.cu -lcurand
$ ./t721
0.7816 0.2338 0.6791 0.2824 0.6299 0.1212 0.4333 0.3831 0.5136 0.2987

0.7816 0.2338 0.6791 0.2824 0.6299 0.1212 0.4333 0.3831 0.5136 0.2987
$

因此，您需要跟踪生成的随机数数量(不是 RNG 函数调用的数量)直到您执行检查点时为止，并保存它。

重启时，同样初始化生成器:

    /* Create pseudo-random number generator */
    CURAND_CALL(curandCreateGenerator(&gen,
                CURAND_RNG_PSEUDO_DEFAULT));

    /* Set seed */
    CURAND_CALL(curandSetPseudoRandomGeneratorSeed(gen,
                1234ULL));

然后按先前生成的值的数量 (n) 前进:

    CURAND_CALL(curandSetGeneratorOffset(gen, n));