tests: Replace out of date variable with ,fixing unit tests

Signed-off-by: Will Johnson <[email protected]>
foundation-model-stack · Oct 10, 2024 · 035a3a0 · 035a3a0
1 parent 41a6e5b
commit 035a3a0
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/fixtures/accelerate_fsdp_defaults.yaml b/fixtures/accelerate_fsdp_defaults.yaml
@@ -14,7 +14,7 @@ fsdp_config:
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
 
   # this controls the FSDP pipelining
-  fsdp_backward_prefetch_policy: BACKWARD_PRE # set to BACKWARD_PRE for the most time-efficient pipeline
+  fsdp_backward_prefetch: BACKWARD_PRE # set to BACKWARD_PRE for the most time-efficient pipeline
                                               # but requires the most memory. BACKWARD_POST is the less
                                               # memory intensive option
 

diff --git a/tests/build/dummy_job_config.json b/tests/build/dummy_job_config.json
@@ -5,7 +5,7 @@
       "dynamo_use_dynamic": true,
       "num_machines": 1,
       "main_process_port": 1234,
-      "fsdp_backward_prefetch_policy": "TRANSFORMER_BASED_WRAP",
+      "fsdp_backward_prefetch": "TRANSFORMER_BASED_WRAP",
       "fsdp_sharding_strategy": 1,
       "fsdp_state_dict_type": "FULL_STATE_DICT",
       "fsdp_cpu_ram_efficient_loading": true,

diff --git a/tests/build/test_utils.py b/tests/build/test_utils.py
@@ -44,7 +44,7 @@ def test_process_accelerate_launch_args(job_config):
     args = process_accelerate_launch_args(job_config)
     # json config values used
     assert args.use_fsdp is True
-    assert args.fsdp_backward_prefetch_policy == "TRANSFORMER_BASED_WRAP"
+    assert args.fsdp_backward_prefetch == "TRANSFORMER_BASED_WRAP"
     assert args.env == ["env1", "env2"]
     assert args.training_script == "tuning.sft_trainer"
     assert args.config_file == "fixtures/accelerate_fsdp_defaults.yaml"