Add end_training/destroy_pg to everything and unpin numpy (#3030)

* Add end_training/destroy_pg to everything

* Carry over to AcceleratorState

* If forked, ignore

* More numpy fun

* Skip only init

Author: muellerzr

examples/by_feature/automatic_gradient_accumulation.py

@@ -217,6 +217,7 @@ def inner_training_loop(batch_size):
     # And call it at the end with no arguments
     # Note: You could also refactor this outside of your training loop function
     inner_training_loop()
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/checkpointing.py

@@ -276,6 +276,7 @@ def training_function(config, args):
             if args.output_dir is not None:
                 output_dir = os.path.join(args.output_dir, output_dir)
             accelerator.save_state(output_dir)
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/cross_validation.py

@@ -255,6 +255,7 @@ def training_function(config, args):
     preds = torch.stack(test_predictions, dim=0).sum(dim=0).div(int(args.num_folds)).argmax(dim=-1)
     test_metric = metric.compute(predictions=preds, references=test_references)
     accelerator.print("Average test metrics from all folds:", test_metric)
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/ddp_comm_hook.py

@@ -192,6 +192,7 @@ def training_function(config, args):
         eval_metric = metric.compute()
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}:", eval_metric)
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/deepspeed_with_config_support.py

@@ -716,6 +716,7 @@ def group_texts(examples):
 
         with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
             json.dump({"perplexity": perplexity, "eval_loss": eval_loss.item()}, f)
+    accelerator.end_training()
 
 
 if __name__ == "__main__":

examples/by_feature/early_stopping.py

@@ -222,6 +222,7 @@ def training_function(config, args):
 
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}:", eval_metric)
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/fsdp_with_peak_mem_tracking.py

@@ -399,8 +399,7 @@ def collate_fn(examples):
                 step=epoch,
             )
 
-    if args.with_tracking:
-        accelerator.end_training()
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/gradient_accumulation.py

@@ -197,6 +197,7 @@ def training_function(config, args):
         eval_metric = metric.compute()
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}:", eval_metric)
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/local_sgd.py

@@ -202,6 +202,7 @@ def training_function(config, args):
         eval_metric = metric.compute()
         # Use accelerator.print to print only on the main process.
         accelerator.print(f"epoch {epoch}:", eval_metric)
+    accelerator.end_training()
 
 
 def main():

examples/by_feature/megatron_lm_gpt_pretraining.py

@@ -703,6 +703,7 @@ def group_texts(examples):
 
         with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
             json.dump({"perplexity": perplexity}, f)
+    accelerator.end_training()
 
 
 if __name__ == "__main__":