pytorch
59a5be3b - add mixed data type support for GroupNorm backward on CPU (#88663)

Commit
2 years ago
add mixed data type support for GroupNorm backward on CPU (#88663) ### Motivation Amp provides convenience methods for mixed precision. If users use amp to run bfloat16 models, torch.autocast will keep module parameters in acc dtype which will leave gamma and beta in float while input/output will be in bfloat16. The same goes for backward: parameters are in float, and X & dX & dY are in bfloat16. Mixed data type support for GroupNorm backward is also needed for model training with GroupNorm. ### Testing Single socket (28cores): * Contiguous: shape | forward / s | forward / s | backward / s | backward / s -- | -- | -- | -- | --   | fp32 | mixed fp32 bf16 | fp32 | mixed fp32 bf16 [10, 128, 20, 20] | 3.08E-05 | 3.50E-05 | 8.06E-05 | 7.69E-05 [10, 128, 50, 50] | 0.000121 | 0.000114 | 0.000358 | 0.000203 * Channels Last (inputs and outputs will be converted to contiguous): shape | forward / s | forward / s | backward / s | backward / s -- | -- | -- | -- | --   | fp32 | mixed fp32 bf16 | fp32 | mixed fp32 bf16 [10, 128, 20, 20] | 4.04E-05 | 4.41E-05 | 0.000226 | 0.000305 [10, 128, 50, 50] | 0.000169 | 0.000166 | 0.001628 | 0.001169 Single core: * Contiguous: shape | forward / s | forward / s | backward / s | backward / s -- | -- | -- | -- | --   | fp32 | mixed fp32 bf16 | fp32 | mixed fp32 bf16 [10, 128, 20, 20] | 2.38E-04 | 2.51E-04 | 5.94E-04 | 4.50E-04 [10, 128, 50, 50] | 0.00171 | 0.001395 | 0.0044455 | 0.00243 * Channels Last (inputs and outputs will be converted to contiguous): shape | forward / s | forward / s | backward / s | backward / s -- | -- | -- | -- | --   | fp32 | mixed fp32 bf16 | fp32 | mixed fp32 bf16 [10, 128, 20, 20] | 2.28E-04 | 3.26E-04 | 0.0016528 | 0.003165 [10, 128, 50, 50] | 0.001788 | 0.001302 | 0.0276621 | 0.019447 Pull Request resolved: https://github.com/pytorch/pytorch/pull/88663 Approved by: https://github.com/jgong5, https://github.com/mingfeima, https://github.com/malfet
Author
Committer
Parents
Loading