add mixed data type support for GroupNorm

[mingfeima]

Stack from ghstack:

• -> add mixed data type support for GroupNorm #81852

1. If user uses amp to run bfloat16 models, torch.autocast will
   keep module paramters in acc dtype which will leave gamma andbeta
   in float while input/output will be in bfloat16.

2. If user explicitly cast the model to bfloat16,
   the input/output and gamma/beta will all be in bfloat16.

cc @VitalyFedyunin @jgong5 @XiaobingSuper @sanchitintel @ashokei @jingxu10

[facebook-github-bot]

🔗 Helpful links

• 🧪  See artifacts and rendered test results at hud.pytorch.org/pr/81852
• ✖️ Python docs build was skipped
• ✖️ C++ docs build was skipped
• ❓Need help or want to give feedback on the CI? Visit our office hours

❌ 1 New Failures

As of commit 624f804 (more details on the Dr. CI page):

Expand to see more

---

• 1/1 failures introduced in this PR

---

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages

pull / linux-focal-py3.7-gcc7 / test (default, 1, 2, linux.2xlarge) (1/1)

Step: "Test" (full log | diagnosis details)

2022-09-01T11:12:13.8946754Z AssertionError: te...than the torch result was (1.666915112918943e-05)!

2022-09-01T11:12:13.8932744Z =================================== FAILURES ===================================
2022-09-01T11:12:13.8936464Z ______ TestCommonCPU.test_python_ref__refs_native_layer_norm_cpu_float32 _______
2022-09-01T11:12:13.8941186Z [gw1] linux -- Python 3.7.13 /opt/conda/bin/python
2022-09-01T11:12:13.8942504Z Traceback (most recent call last):
2022-09-01T11:12:13.8942919Z   File "/var/lib/jenkins/workspace/test/test_ops.py", line 325, in test_python_ref
2022-09-01T11:12:13.8944395Z     self._ref_test_helper(lambda: TorchRefsMode(strict=True), device, dtype, op)
2022-09-01T11:12:13.8944806Z   File "/var/lib/jenkins/workspace/test/test_ops.py", line 308, in _ref_test_helper
2022-09-01T11:12:13.8945336Z     self.assertTrue(ref_distance <= torch_distance, msg=msg)
2022-09-01T11:12:13.8945935Z   File "/opt/conda/lib/python3.7/unittest/case.py", line 705, in assertTrue
2022-09-01T11:12:13.8946206Z     raise self.failureException(msg)
2022-09-01T11:12:13.8946754Z AssertionError: tensor(False) is not true : Reference result was farther (1.7186287103232445e-05) from the precise computation than the torch result was (1.666915112918943e-05)!
2022-09-01T11:12:13.9171705Z - generated xml file: /var/lib/jenkins/workspace/test/test-reports/python-pytest/test_ops/test_ops.xml -
2022-09-01T11:12:13.9174156Z =========================== short test summary info ============================
2022-09-01T11:12:13.9174684Z FAILED test_ops.py::TestCommonCPU::test_python_ref__refs_native_layer_norm_cpu_float32
2022-09-01T11:12:13.9178815Z !!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 1 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
2022-09-01T11:12:13.9179407Z !!!!!!!!!!!! xdist.dsession.Interrupted: stopping after 1 failures !!!!!!!!!!!!!
2022-09-01T11:12:13.9210278Z = 1 failed, 6894 passed, 4798 skipped, 105 xfailed, 121 warnings, 2 rerun in 349.85s (0:05:49) =
2022-09-01T11:12:14.0678482Z Skip info is located in the xml test reports, please either go to s3 or the hud to download them
2022-09-01T11:12:15.3915917Z Traceback (most recent call last):
2022-09-01T11:12:15.3916583Z   File "test/run_test.py", line 1065, in <module>
2022-09-01T11:12:15.3918129Z     main()

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/81852

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d7f5ccd:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/81852

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit our office hours

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mingfeima]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: Approval needed from one of the following (Rule 'superuser'):
shreyanb98, 842974287, Xirider, husthyc, aruntonic, ...

Details for Dev Infra team

Raised by workflow job

[malfet]

Small nit on the perf, we introduce unnecessary if (mixedtype) for all dtypes but BFloat16, perhaps one can refactor this code as:

template<scalar_t>
void CallGroupNorm(...)

And specialize it for BFloat16 to call mixed_type

[mingfeima]

Small nit on the perf, we introduce unnecessary if (mixedtype) for all dtypes but BFloat16, perhaps one can refactor this code as:

template<scalar_t>
void CallGroupNorm(...)

And specialize it for BFloat16 to call mixed_type

We will add minxed_type for float16 in near future for the coming new platform of Xeon (Granite Rapids).

Also I noticed that vec_scalar_t is a duplicate now, will have it replaced with opmath_t.

[mingfeima]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here