Remove optimizer form deepspeed config and let engine be destroyed wi…

…thin train at end
truefoundry · May 3, 2024 · 028db8e · 028db8e
1 parent 5a5d7f1
commit 028db8e
Show file tree

Hide file tree

Showing 10 changed files with 1 addition and 84 deletions.
diff --git a/deepspeed_configs/0_ds_z0_config.json b/deepspeed_configs/0_ds_z0_config.json
@@ -11,15 +11,6 @@
     "bf16": {
       "enabled": "auto"
     },
-    "optimizer": {
-      "type": "AdamW",
-      "params": {
-        "lr": "auto",
-        "betas": "auto",
-        "eps": "auto",
-        "weight_decay": "auto"
-      }
-    },
     "activation_checkpointing": {
       "partition_activations": false,
       "cpu_checkpointing": false

diff --git a/deepspeed_configs/1_ds_z1_config.json b/deepspeed_configs/1_ds_z1_config.json
@@ -11,15 +11,6 @@
     "bf16": {
       "enabled": "auto"
     },
-    "optimizer": {
-      "type": "AdamW",
-      "params": {
-        "lr": "auto",
-        "betas": "auto",
-        "eps": "auto",
-        "weight_decay": "auto"
-      }
-    },
     "activation_checkpointing": {
       "partition_activations": false,
       "cpu_checkpointing": false

diff --git a/deepspeed_configs/2_ds_z1_offload_optimizer_config.json b/deepspeed_configs/2_ds_z1_offload_optimizer_config.json
@@ -11,15 +11,6 @@
     "bf16": {
       "enabled": "auto"
     },
-    "optimizer": {
-      "type": "AdamW",
-      "params": {
-        "lr": "auto",
-        "betas": "auto",
-        "eps": "auto",
-        "weight_decay": "auto"
-      }
-    },
     "activation_checkpointing": {
       "partition_activations": false,
       "cpu_checkpointing": false

diff --git a/deepspeed_configs/3_ds_z2_config.json b/deepspeed_configs/3_ds_z2_config.json
@@ -11,15 +11,6 @@
   "bf16": {
     "enabled": "auto"
   },
-  "optimizer": {
-    "type": "AdamW",
-    "params": {
-      "lr": "auto",
-      "betas": "auto",
-      "eps": "auto",
-      "weight_decay": "auto"
-    }
-  },
   "activation_checkpointing": {
     "partition_activations": false,
     "cpu_checkpointing": false

diff --git a/deepspeed_configs/4_ds_z2_offload_optimizer_config.json b/deepspeed_configs/4_ds_z2_offload_optimizer_config.json
@@ -11,15 +11,6 @@
   "bf16": {
     "enabled": "auto"
   },
-  "optimizer": {
-    "type": "AdamW",
-    "params": {
-      "lr": "auto",
-      "betas": "auto",
-      "eps": "auto",
-      "weight_decay": "auto"
-    }
-  },
   "activation_checkpointing": {
     "partition_activations": false,
     "cpu_checkpointing": false

diff --git a/deepspeed_configs/5_ds_z3_config.json b/deepspeed_configs/5_ds_z3_config.json
@@ -11,15 +11,6 @@
   "bf16": {
     "enabled": "auto"
   },
-  "optimizer": {
-    "type": "AdamW",
-    "params": {
-      "lr": "auto",
-      "betas": "auto",
-      "eps": "auto",
-      "weight_decay": "auto"
-    }
-  },
   "activation_checkpointing": {
     "partition_activations": false,
     "cpu_checkpointing": false

diff --git a/deepspeed_configs/6_ds_z3_offload_param_config.json b/deepspeed_configs/6_ds_z3_offload_param_config.json
@@ -11,15 +11,6 @@
   "bf16": {
     "enabled": "auto"
   },
-  "optimizer": {
-    "type": "AdamW",
-    "params": {
-      "lr": "auto",
-      "betas": "auto",
-      "eps": "auto",
-      "weight_decay": "auto"
-    }
-  },
   "activation_checkpointing": {
     "partition_activations": false,
     "cpu_checkpointing": false

diff --git a/deepspeed_configs/7_ds_z3_offload_optimizer_config.json b/deepspeed_configs/7_ds_z3_offload_optimizer_config.json
@@ -11,15 +11,6 @@
   "bf16": {
     "enabled": "auto"
   },
-  "optimizer": {
-    "type": "AdamW",
-    "params": {
-      "lr": "auto",
-      "betas": "auto",
-      "eps": "auto",
-      "weight_decay": "auto"
-    }
-  },
   "activation_checkpointing": {
     "partition_activations": false,
     "cpu_checkpointing": false

diff --git a/deepspeed_configs/8_ds_z3_offload_param_offload_optimizer_config.json b/deepspeed_configs/8_ds_z3_offload_param_offload_optimizer_config.json
@@ -11,15 +11,6 @@
   "bf16": {
     "enabled": "auto"
   },
-  "optimizer": {
-    "type": "AdamW",
-    "params": {
-      "lr": "auto",
-      "betas": "auto",
-      "eps": "auto",
-      "weight_decay": "auto"
-    }
-  },
   "activation_checkpointing": {
     "partition_activations": false,
     "cpu_checkpointing": false

diff --git a/monkey_patch.py b/monkey_patch.py
@@ -112,9 +112,7 @@ def patched_pretrain_hooks(cfg, trainer):
 
 def patched_post_train_hooks(cfg, trainer):
     # type: (DictDefault, AxolotlTrainer) -> None
-    if trainer.args.deepspeed and hasattr(trainer, "deepspeed") and hasattr(trainer.deepspeed, "destroy"):
-        trainer.deepspeed.destroy()
-    trainer.accelerator.free_memory()
+    pass
 
 
 def monkey_patch_axolotl_internals():