torch.use_deterministic_algorithms(True) の動作確認

torch.use_deterministic_algorithms(mode) で PyTorch を決定論的モードにすると何が変わるのか確認します。ドキュメントには決定論的モードで決定論的になる関数が色々と載っていますが、ここでは torch.Tensor.index_add_ で確認します。

以下のスクリプトは、index_add_ によって、64 個の正規乱数を、長さ 2 のゼロベクトル dst の第 1 成分と第 2 成分のどちらかにランダムに足します。長さ 64 の正規乱数列 src と長さ 64 のターゲットインデクス列 index_ は torch.manual_seed(0) によっていつも同じになります。であれば、正規乱数を足した後の dst もいつも同じになってほしいですが、CUDA テンソルへの index_add_ が並列に実行されるために、足される順序が前後することがあります。浮動小数を足す順序が前後すると、丸め誤差によって結果が異なってきます。これを避けたいときは、決定論的モードにします。決定論的モードにすると、同じハード＆ソフト環境であれば結果が同じになります。

script.py

import torch


def test(mode=True):
    print(f'\n{mode=}')
    torch.use_deterministic_algorithms(mode)
    torch.manual_seed(0)
    n = 64
    m = 2
    src = torch.randn(n, device='cuda')
    index_ = torch.randint(0, m, (n,), device='cuda')
    for i in range(5):
        dst = torch.zeros(m, device='cuda')
        dst.index_add_(0, index_, src)
        print(f'{i}.  {dst[0].item():.9f}  {dst[1].item():.9f}')


if __name__ == '__main__':
    test(False)
    test(True)

実際、スクリプト実行結果は以下となり、torch.use_deterministic_algorithms(False) としたときは 5 回の dst 作成結果が僅かに揃わず、torch.use_deterministic_algorithms(True) としたときは 5 回の dst 作成結果が揃っています。

mode=False
0.  -0.030139029  -1.037073374
1.  -0.030139349  -1.037073135
2.  -0.030139446  -1.037073374
3.  -0.030139349  -1.037073135
4.  -0.030139327  -1.037073612

mode=True
0.  -0.030139051  -1.037073493
1.  -0.030139051  -1.037073493
2.  -0.030139051  -1.037073493
3.  -0.030139051  -1.037073493
4.  -0.030139051  -1.037073493