address backward prefetch policy deprecation

YerevaNN · Jan 30, 2024 · b851e36 · b851e36
1 parent 32b3c9d
commit b851e36
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/config/accelerate_config.yaml b/src/config/accelerate_config.yaml
@@ -3,11 +3,11 @@ distributed_type: FSDP
 downcast_bf16: 'no'
 fsdp_config:
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_backward_prefetch: BACKWARD_PRE
   fsdp_offload_params: false
   fsdp_sharding_strategy: 1
   fsdp_state_dict_type: FULL_STATE_DICT
-  fsdp_transformer_layer_cls_to_wrap: CustomOPTForCausalLM
+  fsdp_transformer_layer_cls_to_wrap: OPTForCausalLM
 machine_rank: 0
 main_training_function: main
 mixed_precision: bf16

diff --git a/src/config/test_configs/fsdp_config.yaml b/src/config/test_configs/fsdp_config.yaml
@@ -3,7 +3,7 @@ distributed_type: FSDP
 downcast_bf16: 'no'
 fsdp_config:
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_backward_prefetch: BACKWARD_PRE
   fsdp_offload_params: false
   fsdp_sharding_strategy: 1
   fsdp_state_dict_type: FULL_STATE_DICT