Fix up state with xla + performance regression

[muellerzr]

What does this PR do?

The one un-testible backend of course is the one that had issues 😓

waiting to hear back but should fix #2629, if so I'll do a patch with this PR 🤞

Fixes #2629
Fixes #2633

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@BenjaminBossan

[muellerzr]

We don't do this here (the else), it happens during set_replication later and that's not the same as setting set_device (it spawns up the TPU DDP)

[BenjaminBossan]

That's really annoying thanks for quickly fixing this, LGTM.

The check_cuda_p2p_ib_support should certainly be amended now. For the xla issue, I agree, let's wait for confirmation.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[muellerzr]

Looks to have fixed our issues here: #2629 :) Merging