Add TPUPrecisionPlugin

[carmocca]

What does this PR do?

Part of #9287

This is necessary because the optimizer_step will be moved to the PrecisionPlugin but it currently lives in the TPUAccelerator.

This PR just introduces the empty class.

Does your PR introduce any breaking changes? If yes, please list them.

None

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

[carmocca]

@SeanNaren why was this not done as

class TPUBf16PrecisionPlugin():
    precision: str = "bf16"


# accelerator connector
if self.precision == 16:
    raise Unsupported
elif self.precision == "bf16":
    return TPUBf16PrecisionPlugin

The current format can be confusing for users, perhaps making them believe AMP and TPU work together

We can easily change this, as this is unreleased

[carmocca]

Similarly, I don't get why the NativeMixedPrecisionPluginwas used for bf16 support, considering it doesn't even define the scaler. Why wasn't a Bf16PrecisionPlugin used?

[carmocca]

This comment does not block this PR.

I've opened a branch with my ideas: refactor/tpu-precision-plugin...refactor/untangle-bf16

[SeanNaren]

If the first question is "why do we make it seem precision=16 is possible for TPUs", is because this was the original default before "bf16" was introduced AFAIK. we could go ahead and deprecate precision=16 for TPUs however precision=16 doesn't even make sense for AMP, so should we deprecate this and go for precision='mixed'? I don't think so personally, and would advise against

Regarding the question of having two separate plugins for BF16 and for Native AMP, unfortunately I can't find where a discussion had place around this. I originally implemented this as two separate plugins but there was enough cross-over with the original plugin to have them in a single plugin. Maybe it makes sense to have them separate, i'm indifferent here

[kaushikb11]

LGTM!

[rohitgr7]

small comments

[codecov]

Codecov Report

Merging #10020 (742c7a6) into master (0aa220b) will increase coverage by 4%.
The diff coverage is 84%.

@@           Coverage Diff            @@
##           master   #10020    +/-   ##
========================================
+ Coverage      89%      93%    +4%     
========================================
  Files         179      180     +1     
  Lines       15810    15852    +42     
========================================
+ Hits        14021    14676   +655     
+ Misses       1789     1176   -613