From 2d7e92f24252c0c33692edaaf126f09e706865df Mon Sep 17 00:00:00 2001 From: Wonjoo Lee Date: Tue, 5 Sep 2023 12:42:48 -0700 Subject: [PATCH] Enable megacore_dense by default (#5520) (#5531) Summary: This change enables megacore_dense by default to allow asynchorous cc ops especailly for GSPMD. Test Plan: CI Co-authored-by: Jiewen Tan --- torch_xla/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py index de99bd0d0bdd..cec3b267227f 100644 --- a/torch_xla/__init__.py +++ b/torch_xla/__init__.py @@ -65,6 +65,7 @@ def _setup_default_env(): _set_missing_env('TPU_ML_PLATFORM', 'PyTorch/XLA') if server_is_alive(): _set_missing_env('XRT_START_LOCAL_SERVER', '0') + _set_missing_env('TPU_MEGACORE', 'megacore_dense') _fd, _tmp_fname = -1, '' @@ -170,4 +171,4 @@ def _init_xla_lazy_backend(): # TODO @wonjoo come up with a long term fix in Dynamo. torch._dynamo.config.automatic_dynamic_shapes = False -from .stablehlo import save_as_stablehlo, save_torch_model_as_stablehlo \ No newline at end of file +from .stablehlo import save_as_stablehlo, save_torch_model_as_stablehlo