use env to skip PJRT initialize

[zpcore]

We skip the PJRT Megascale initialization by controlling the env.

This is a temporary fix, and is supposed to be rolled back.

Check #8609 (comment) for detailed motivation.

[tengyifei]

Suggest landing this after the libtpu change is approved.

[zpcore]

@tengyifei , shall we cherry pick this PR to 2.6 release?

[tengyifei]

@zpcore cherrypicking is fine with me.

[zpcore]

The test failed. I think it is due to pytorch/pytorch#142859.

They have reverted the PR.

[tengyifei]

Ack

[tengyifei]

@zpcore thanks. next step is to follow the process in #8455 to create a cherrypick PR

[bhavya01]

Retrospective LGTM!

[miladm]

Thanks @zpcore - can we add enough details to PR descriptions to help folks without context understand the intent of the contribution more clearly please?

[zpcore]

Thanks @zpcore - can we add enough details to PR descriptions to help folks without context understand the intent of the contribution more clearly please?

The issue in multipod run is that MegascaleXLA(MXLA) will trigger device discovery when we initialize PJRT runtime with the TPU backend. With the introduction of Pallas kernel, we did an extra MXLA trigger when call jax.jit(). Thus there are more than one device discovery been executed. Everytime we execute device discovery, all device will be assigned a ID. This will cause device ID mismatch after the second device discovery thus cause the device communication hang.

The hacky way to fix the issue is to use enviroment variable to control the device discovery won't be triggered when call jax.jit. This PR works with the fix we made internally in libtpu source code:

 const char* skip_megascale_pjrt_client = std::getenv("SKIP_MEGASCALE_PJRT_CLIENT");
  bool skip_megascale = false;
  if (skip_megascale_pjrt_client != nullptr) {
    skip_megascale = true;
  }
  if (absl::GetFlag(FLAGS_megascale_num_slices) != 1 && !skip_megascale) {
    client = xla::MegaScalePjRtClient::CreateMegaScalePjRtClient(
        std::move(tpu_client));
    ...
  }

With this fix, MegaScalePjRtClient will only be triggered in place (e.g.,)

xla/torch_xla/core/xla_model.py

Line 93 in 82d3504

devices = torch_xla._XLAC._xla_get_devices()

,
where we call runtime::GetComputationClient() and initialize the client.