Merge branch 'main' into triton_env_vars

foundation-model-stack · Oct 15, 2024 · abfce93 · abfce93
2 parents 64fcd2a + d138858
commit abfce93
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/fixtures/accelerate_fsdp_defaults.yaml b/fixtures/accelerate_fsdp_defaults.yaml
@@ -14,9 +14,10 @@ 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
+  fsdp_backward_prefetch_policy: BACKWARD_PRE # for backwards compatibility
 
   # setting this to true will increase forward memory by prefetching the next FSDP all-gather, while performing
   # the current forward pass. 

diff --git a/pyproject.toml b/pyproject.toml
@@ -27,14 +27,14 @@ classifiers=[
 ]
 dependencies = [
 "numpy>=1.26.4,<2.0",
-"accelerate>=0.20.3,<0.35,!=0.34",
+"accelerate>=0.20.3,!=0.34,<1.1",
 "transformers>4.41,<4.50",
 "torch>=2.2.0,<3.0",
 "sentencepiece>=0.1.99,<0.3",
 "tokenizers>=0.13.3,<1.0",
 "tqdm>=4.66.2,<5.0",
 "trl>=0.9.3,<1.0",
-"peft>=0.8.0,<0.13",
+"peft>=0.8.0,<0.14",
 "protobuf>=5.28.0,<6.0.0",
 "datasets>=2.15.0,<3.0",
 "simpleeval>=0.9.13,<1.0",

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"