[FEATURE] fp16 support for FSDP ZeRO-1 #53
Conversation
Merging Upstream Ratex
…(), changed shard parameter generator method to sharded_paramaters()
…instance variable to constructor, renamed _get_shard to _shard_tensor
…test for single GPU
…s excluded when padding is needed due to errors
|
Some tests using fp16 parameters are failing. The loss difference is close between the use of FSDP fp32 buffer and the RAF buffer with no ZeRO-1 and optimizer ZeRO-1 is about 0.002 for these tests. However, the tolerance is currently set to 1e-10. Either more exploration is needed to find the reason for the difference, or the tolerance could be made more lenient. |
|
The e-3 difference does not seems to be ok as we just discussed. Let's spend more time to figure out the issue. |
|
@shawnakdeb any updates on this? |
No updates yet, I did not have stable internet access for the past week, but I will continue work on the issue now. |
I did just realize that with my amazon id closed, I no longer have access to my multi-GPU environment to run the script. Is there a way for us to get around this? |
I think, you can still access the instance if you have still have the key. Let's communicate this offline. Will send you an email. |
2 participants
Description
Support for fp16 parameter type. Shards are kept in fp32 for the optimizer while output and gradient calculations are done in original fp16 Dtype.
cc @awslabs/raf-reviewer, @zachzzc, @zhen-jia, @XinweiFu