From de4ecaa496ed841a7d1434a5c8e14914bb8b6536 Mon Sep 17 00:00:00 2001
From: jignesh2619 <136894974+jignesh2619@users.noreply.github.com>
Date: Mon, 6 Jul 2026 20:46:36 +0530
Subject: [PATCH] fix(distributed): call destroy_process_group() on exit to
 prevent NCCL resource leak

In distributed runs (accelerate launch / multi-GPU), when evaluation
hits an exception, process shutdown emits:

  WARNING: destroy_process_group() was not called before program exit,
  which can leak resources.

Root cause: Accelerator registers no atexit handler, so distributed
teardown never happens on error paths.

Fix: wrap the evaluation loop and result printing in a try/finally block.
The finally block calls accelerator.state.destroy_process_group() when
Accelerator is active, or torch.distributed.destroy_process_group()
directly when running in a pre-initialized distributed context.
wandb logging runs inside try before finally to preserve correct ordering.

Fixes #1240
---
 lmms_eval/__main__.py | 88 +++++++++++++++++++++++--------------------
 1 file changed, 47 insertions(+), 41 deletions(-)

diff --git a/lmms_eval/__main__.py b/lmms_eval/__main__.py
index d9804d98f..b8b17f1cb 100755
--- a/lmms_eval/__main__.py
+++ b/lmms_eval/__main__.py
@@ -534,47 +534,53 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None:
         else:
             is_main_process = False
 
-    for args in args_list:
-        try:
-            # if is_main_process and args.wandb_args:  # thoughtfully we should only init wandb once, instead of multiple ranks to avoid network traffics and unwanted behaviors.
-            #     wandb_logger = WandbLogger()
-
-            results, samples = cli_evaluate_single(args)
-            results_list.append(results)
-
-            if accelerator:
-                accelerator.wait_for_everyone()
-            elif torch.distributed.is_available() and torch.distributed.is_initialized():
-                torch.distributed.barrier()
-            if is_main_process and args.wandb_args:
-                try:
-                    wandb_logger.post_init(results)
-                    wandb_logger.log_eval_result()
-                    if args.wandb_log_samples and samples is not None:
-                        wandb_logger.log_eval_samples(samples)
-                except Exception as e:
-                    eval_logger.info(f"Logging to Weights and Biases failed due to {e}")
-                # wandb_logger.finish()
-
-        except Exception as e:
-            if args.verbosity == "DEBUG":
-                raise e
-            else:
-                traceback.print_exc()
-                eval_logger.error(f"Error during evaluation: {e}. Please set `--verbosity=DEBUG` to get more information.")
-                results_list.append(None)
-
-    for args, results in zip(args_list, results_list):
-        # cli_evaluate will return none if the process is not the main process (rank 0)
-        if results is not None:
-            print(f"{args.model} ({args.model_args}), gen_kwargs: ({args.gen_kwargs}), " f"limit: {args.limit}, offset: {args.offset}, num_fewshot: {args.num_fewshot}, " f"batch_size: {args.batch_size}")
-            print(get_eval_banner(branch=results.get("git_branch"), commit=results.get("git_hash")))
-            print(make_table(results))
-            if "groups" in results:
-                print(make_table(results, "groups"))
-
-    if args.wandb_args:
-        wandb_logger.run.finish()
+    try:
+        for args in args_list:
+            try:
+                # if is_main_process and args.wandb_args:  # thoughtfully we should only init wandb once, instead of multiple ranks to avoid network traffics and unwanted behaviors.
+                #     wandb_logger = WandbLogger()
+    
+                results, samples = cli_evaluate_single(args)
+                results_list.append(results)
+    
+                if accelerator:
+                    accelerator.wait_for_everyone()
+                elif torch.distributed.is_available() and torch.distributed.is_initialized():
+                    torch.distributed.barrier()
+                if is_main_process and args.wandb_args:
+                    try:
+                        wandb_logger.post_init(results)
+                        wandb_logger.log_eval_result()
+                        if args.wandb_log_samples and samples is not None:
+                            wandb_logger.log_eval_samples(samples)
+                    except Exception as e:
+                        eval_logger.info(f"Logging to Weights and Biases failed due to {e}")
+                    # wandb_logger.finish()
+    
+            except Exception as e:
+                if args.verbosity == "DEBUG":
+                    raise e
+                else:
+                    traceback.print_exc()
+                    eval_logger.error(f"Error during evaluation: {e}. Please set `--verbosity=DEBUG` to get more information.")
+                    results_list.append(None)
+    
+        for args, results in zip(args_list, results_list):
+            # cli_evaluate will return none if the process is not the main process (rank 0)
+            if results is not None:
+                print(f"{args.model} ({args.model_args}), gen_kwargs: ({args.gen_kwargs}), " f"limit: {args.limit}, offset: {args.offset}, num_fewshot: {args.num_fewshot}, " f"batch_size: {args.batch_size}")
+                print(get_eval_banner(branch=results.get("git_branch"), commit=results.get("git_hash")))
+                print(make_table(results))
+                if "groups" in results:
+                    print(make_table(results, "groups"))
+    
+        if args.wandb_args:
+            wandb_logger.run.finish()
+    finally:
+        if accelerator is not None:
+            accelerator.state.destroy_process_group()
+        elif torch.distributed.is_available() and torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
 
 
 def cli_evaluate_single(args: Union[argparse.Namespace, None] = None) -> None: