Enable bfloat16 for hardtanh_backward_cuda

[cchan]

I'm not sure why this was left out in the first place as all adjacent operations have both Half and BFloat16. Things seem to work as expected and this enables relu6 to be used in bfloat16 training. Hardtanh backward is super simple and precision is not relevant.

import torch
x_fp32 = torch.tensor([-1,2,4,7], requires_grad=True, dtype=torch.float32, device="cuda")
x_bf16 = torch.tensor([-1,2,4,7], requires_grad=True, dtype=torch.bfloat16, device="cuda")
torch.nn.functional.relu6(x_fp32).sum().backward()
torch.nn.functional.relu6(x_bf16).sum().backward()
assert (x_fp32.grad == x_bf16.grad).all()

Previously would fail with:

Traceback (most recent call last):
  File "test_hardtanh_patch.py", line 5, in <module>
    torch.nn.functional.relu6(x_bf16).sum().backward()
  File ".../lib/python3.8/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File ".../lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: "hardtanh_backward_cuda" not implemented for 'BFloat16'

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 b5ab130:
💚 Looks good so far! There are no failures yet. 💚

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cchan / name: Clive Chan (b4444c8)

[cchan]

cc @rohithkrn - just checking, do you recall if there was a reason to not add hardtanh_backward, like some test not passing? or just an oversight in this: #32065

[cchan]

@pytorchbot rebase

[pytorch-bot]

You don't have permissions to rebase this PR, only people with write permissions may rebase PRs.

[rohithkrn]

cc @rohithkrn - just checking, do you recall if there was a reason to not add hardtanh_backward, like some test not passing? or just an oversight in this: #32065

@cchan honestly don't remember. It's been a while.

[Chillee]

@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

[pytorchmergebot]

Merge failed

Reason: This PR is too stale; the last push date was more than 3 days ago. Please rebase and try again. You can rebase by leaving the following comment on this PR:
@pytorchbot rebase

Details for Dev Infra team

Raised by workflow job

[cchan]

@pytorchbot rebase

[pytorch-bot]

You don't have permissions to rebase this PR, only people with write permissions may rebase PRs.

[ngimel]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot successfully started a rebase job. Check the current status here

[pytorchmergebot]

Successfully rebased patch-5 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout patch-5 && git pull --rebase)

[ngimel]

@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