Support FSDP

[tocean]

Description
Support FSDP with FP8.

Major Revision

• Add fsdp package
• Add mnist example
• Add FSDPAdam and FSDPAdamW optimizer
• Add document
• Add UT

[tocean]

Log for mnist_fsdp.py:
Train Epoch: 1 Loss: 0.536632
Test set: Average loss: 0.1431, Accuracy: 9581/10000 (95.81%)

Train Epoch: 2 Loss: 0.176441
Test set: Average loss: 0.0877, Accuracy: 9729/10000 (97.29%)

Train Epoch: 3 Loss: 0.127415
Test set: Average loss: 0.0687, Accuracy: 9793/10000 (97.93%)

Train Epoch: 4 Loss: 0.105661
Test set: Average loss: 0.0608, Accuracy: 9813/10000 (98.13%)

Train Epoch: 5 Loss: 0.096828
Test set: Average loss: 0.0551, Accuracy: 9826/10000 (98.26%)

Train Epoch: 6 Loss: 0.090231
Test set: Average loss: 0.0527, Accuracy: 9829/10000 (98.29%)

Train Epoch: 7 Loss: 0.083397
Test set: Average loss: 0.0506, Accuracy: 9833/10000 (98.33%)

Train Epoch: 8 Loss: 0.081701
Test set: Average loss: 0.0497, Accuracy: 9833/10000 (98.33%)

Train Epoch: 9 Loss: 0.081912
Test set: Average loss: 0.0488, Accuracy: 9839/10000 (98.39%)

Train Epoch: 10 Loss: 0.079299
Test set: Average loss: 0.0487, Accuracy: 9841/10000 (98.41%)

Train Epoch: 11 Loss: 0.078325
Test set: Average loss: 0.0482, Accuracy: 9837/10000 (98.37%)

Train Epoch: 12 Loss: 0.077337
Test set: Average loss: 0.0481, Accuracy: 9838/10000 (98.38%)

Train Epoch: 13 Loss: 0.077516
Test set: Average loss: 0.0479, Accuracy: 9836/10000 (98.36%)

Train Epoch: 14 Loss: 0.076482
Test set: Average loss: 0.0479, Accuracy: 9837/10000 (98.37%)

[wkcn]

Good job!

I have the following questions:

1. Which line of the code synchroinize the scaling factor (meta.scale or meta.scale_inv) across GPUs?
2. Gradient accumulation seems to be not supported yet. We can raise a NotImplemented exception when the steps of gradient accumulation is larger than 1.
3. Is there any comparision on memory footprint between BF16-FSDP and FP8-FSDP?

[tocean]

1. Which line of the code synchroinize the scaling factor (meta.scale or meta.scale_inv) across GPUs?
2. Gradient accumulation seems to be not supported yet. We can raise a NotImplemented exception when the steps of gradient accumulation is larger than 1.
3. Is there any comparision on memory footprint between BF16-FSDP and FP8-FSDP?

For these 3 questions:

1. See FSDPAdamW::step. In this function, all_reduce of amax is called.
2. Add checking loigc in _fp8_post_backward_hook.
3. I checked the memory saving use T5 example here. The model I use is t5-3b. The memory footprint for bf16, FP32, MS-AMP is 28GB, 40GB and 34GB respectively. It is a bit strange that MS-AMP uses more memory than BF16.

[wkcn]

1. Which line of the code synchroinize the scaling factor (meta.scale or meta.scale_inv) across GPUs?
2. Gradient accumulation seems to be not supported yet. We can raise a NotImplemented exception when the steps of gradient accumulation is larger than 1.
3. Is there any comparision on memory footprint between BF16-FSDP and FP8-FSDP?

For these 3 questions:

1. See FSDPAdamW::step. In this function, all_reduce of amax is called.
2. Add checking loigc in _fp8_post_backward_hook.
3. I checked the memory saving use T5 example here. The model I use is t5-3b. The memory footprint for bf16, FP32, MS-AMP is 28GB, 40GB and 34GB respectively. It is a bit strange that MS-AMP uses more memory than BF16.

Regarding to the answer 3, it may be related to the argument mixed_precision of FSDP.

When mixed_precision is not None, FSDP will create a FP32 master weight for FP8Linear. It leads to duplicated master weights in FSDP and FP8Optimizer.

[wkcn]

LGTM. Thanks!

[penghouwen]

an initial version of FSDP in MSAMP

[penghouwen]

In the future, we will do:

• gradient accumulation
• speed optimization
• acc calibration

[tocean]

In the future, we will do:

• gradient accumulation
• speed optimization
• acc calibration

Sure. Let's improve it in next iteration.