[Random] Random state management

[comaniac]

Description

• Add random state management to deal with the requirements of using the same or different random seeds in a TP group. The implementation is based on the one in Megatron-LM.
• Add an activation checkpointing with the consideration of random states.
• Add an op DropoutWithTensorParallel that can be replaced by users when writing a schedule.
• Add unit tests.
• Disable randomly plugin in pytest. This plugin makes our random seed setup in the test fixture useless.

Notes:

1. We now offer an API set_random_seed for users to call in the training script. Users have to manually call it and specify the rank of 3D parallelism.
2. All changes in this PR will have no impact if set_random_seed is not called in advance.
3. Fidelity testing shows the updated GPT schedule with 3D parallelism could align the loss to ZeRO-3 (with and without activation checkpointing), but flash attention has to be disabled.
4. I'll update the flash attention to the latest one and see if the problem will gone.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @szhengac @chhzh123

[chhzh123]

LGTM. Thanks @comaniac.

[szhengac]

if i understand correctly. all the DP ranks also use the same seed so the loss wouldn't be right, but we only use this script for benchmarking

[szhengac]

LGTM