Sharded dropout.

[wujingyue]

In SDPA, dropout, right after softmax, is done head parallel. To ensure dropout works correctly, the dropout in different shards needs to be randomized in different patterns. One possible solution proposed by https://arxiv.org/pdf/1909.08053 Section B.2 is to maintain a separate random number generator for dropout within model parallel regions. This random number generator is uniquely seeded for each model parallel worker.

To work around this limitation, our tests currently set the dropout probably in SDPA to 0:

Fuser/tests/cpp/test_multidevice_transformer.cpp

Line 27 in 70263ee

constexpr double kDropoutProb = 0.1, kParamScale = 0.02, kSdpaProb = 0.0,