Trainer config fixes

Author: matthayes

config/ds_z3_bf16_config.json

@@ -29,7 +29,7 @@
       "stage3_param_persistence_threshold": "auto",
       "stage3_max_live_parameters": 1e9,
       "stage3_max_reuse_distance": 1e9,
-      "stage3_gather_16bit_weights_on_model_save": false
+      "stage3_gather_16bit_weights_on_model_save": true
     },
     "gradient_accumulation_steps": "auto",
     "gradient_clipping": "auto",

training/trainer.py

@@ -181,8 +181,8 @@ def train(
         evaluation_strategy="steps",
         eval_steps=10,
         save_strategy="steps",
-        save_steps=1000,
-        save_total_limit=None,
+        save_steps=200,
+        save_total_limit=1,
         load_best_model_at_end=True,
         report_to="tensorboard",
         disable_tqdm=True,
@@ -204,6 +204,9 @@ def train(
     logger.info("Training")
     trainer.train()
 
+    logger.info(f"Saving Model to {local_output_dir}")
+    trainer.save_model(output_dir=local_output_dir)
+
     if dbfs_output_dir:
         logger.info(f"Saving Model to {dbfs_output_dir}")
         trainer.save_model(output_dir=dbfs_output_dir)