From 557d9f34e04044222d6afbe4a0a50acb040b65e9 Mon Sep 17 00:00:00 2001 From: Pei Zhang Date: Wed, 22 Jan 2025 21:41:23 -0800 Subject: [PATCH] use env to skip PJRT initialize (#8609) --- torch_xla/experimental/custom_kernel.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py index 4067b464cdee..10036aeb9ca7 100644 --- a/torch_xla/experimental/custom_kernel.py +++ b/torch_xla/experimental/custom_kernel.py @@ -138,10 +138,12 @@ def trace_pallas(kernel: Callable, return trace_pallas_arg_to_payload[hash_key], tensor_args # Here we ignore the kwargs for execution as most of the time, the kwargs is only used in traced code. + os.environ['SKIP_MEGASCALE_PJRT_CLIENT'] = 'true' ir = jax.jit( kernel, static_argnums=static_argnums, static_argnames=static_argnames).lower(*jax_args, **kwargs).compiler_ir() payload = _extract_backend_config(ir) + os.environ.pop('SKIP_MEGASCALE_PJRT_CLIENT', None) if use_cache: # if we reach here it means we have a cache miss.