Add initial support for neuron devices

[bfontain]

What does this PR do?

Add some initial support for neuron devices to composer using the pytorch xla backend.

This PR fixes a few issues:

• Renames is_tpu_installed to is_xla_installed and changes it to store state in a global (vs importing and deleting torch_xla, which could delete a needed import).
• Adds a new device type 'neuron' for training on neuron instances
• Fixes support for AMP on XLA.
• Removes use of xm.optimizer_step (which introduces a second gradient reduce given that the model gets auto wrapped in DDP/FSDP). xm.optimizer_step does 3 things: 1) gradient reduction 2) applying the optimizer 3) xm.mark_step. We already have xm.mark_step in the microbatching loop, so remove the one here improves performance slightly.

Tested using a small model on 2 neuron cores via torchrun. support for composer CLI to come.

What issue(s) does this change relate to?

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[j316chuck]

Left a few comments. Thanks for this! Overall LGTM barring nits

[j316chuck]

Thanks for adding this device @bfontain