Raise error when using AMP on non-CUDA device#7869
Conversation
sgugger
left a comment
sgugger
left a comment
There was a problem hiding this comment.
I think the postinit is a great place for this sanity check, so the error comes fast to the user, and I don't think our current FP16 implementations (Apex and Amp) work with TPUs. I can check the latter on Monday and report back here.
If you'd rather be more forgiving, let me know. In that case we can do a logging.warning(f"AMP only supported when using CUDA. Turning off AMP because not using a CUDA device ({self.device.type})") and subsequently automatically turning off --fp16. |
|
Ok, just tested on TPUs and FP16 does not work on them indeed, though native amp fails graciously (testing on TPUs with PyTorch < 1.6 is suicide so I did not try APEX). So if you fix the styling issue, I think this is good to be merged! |
|
It has to be run by the user, I can force-push the styling on your branch if you don't have the setup to run it easily. |
Thanks but it should work fine now. Updated local dependencies and all that and reformatted. |
|
Thanks for your contribution! |
2 participants
The trainer currently implements native AMP in such a way that a GradScaler is always used. AFAIK, and as supported by this post and this report on the forums, this will only work on CUDA devices. Therefore, an error should probably be thrown when a user tries to use
--fp16alongside a non-CUDA device.This PR is currently very brief because I am uncertain about a few things:
--fp16in the trainer works with TPUs;apexand native AMP in this respect;_post_initis the right place to do an argument sanity check. Particularly because I am now callingself.devicewhich will set the device through_set_devices, even though you may (?) want to delay that as long as possible until the arguments are used in the trainer.cc @sgugger