Hey @Caesar, have you tried the native AMP in PyTorch? I haven’t tried that with model parallel yet, but if it does not work, that will be a bug that we need to fix.
https://pytorch.org/docs/stable/amp.html 6
https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples 1