From b35992262e88dd6babeddc2c5ba5786264044a9f Mon Sep 17 00:00:00 2001
From: Dan Saunders <danjsaund@gmail.com>
Date: Mon, 31 Mar 2025 15:17:43 -0400
Subject: [PATCH] Ray train bugfix (#2458)

* fix nccl pg destroy warning

* update

* ray bugfix
---
 src/axolotl/train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/axolotl/train.py b/src/axolotl/train.py
index 56e378c7f..89f35d7eb 100644
--- a/src/axolotl/train.py
+++ b/src/axolotl/train.py
@@ -509,6 +509,7 @@ def train(
     # Save the trained model and cleanup
     save_trained_model(cfg, trainer, model, safe_serialization)
     create_model_card(cfg, trainer)
-    cleanup_distributed()
+    if not cfg.use_ray:
+        cleanup_distributed()
 
     return model, tokenizer, trainer