From ef76f4628445c367405bfed948013250247be0f9 Mon Sep 17 00:00:00 2001 From: Howard Liberty Date: Fri, 29 Mar 2024 17:37:20 -0700 Subject: [PATCH] Fix is_fsdp_enabled() logic and make low_cpu_mem_usage by default --- src/transformers/modeling_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 19aab734784a4f..c05c1a061d859c 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -134,7 +134,6 @@ def is_fsdp_enabled(): torch.distributed.is_available() and torch.distributed.is_initialized() and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 - and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 ) @@ -2869,7 +2868,10 @@ def from_pretrained( adapter_name = kwargs.pop("adapter_name", "default") use_flash_attention_2 = kwargs.pop("use_flash_attention_2", False) - if is_fsdp_enabled(): + # FSDP_CPU_RAM_EFFICIENT_LOADING is set by Accelerate launcher not by TrainingArguments. + # Unless it is turned off explicitly low_cpu_mem_usage should be turned on by default + # when FSDP is on since sync_module_states is also default on. + if is_fsdp_enabled() and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "True")) == 1: low_cpu_mem_usage = True if use_auth_token is not None: