From f2f64857b02d96129ba1e41fe4175d89e868f579 Mon Sep 17 00:00:00 2001 From: Sepehr Sameni Date: Mon, 15 Jun 2026 08:20:55 -0700 Subject: [PATCH] fix memory leak issue during puzzletron scoring, #1681 --- .../validate_model_defaults.yaml | 2 +- .../tools/validate_puzzle_with_multi_replacements.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml index 6b36142a3a8..ce1749d9698 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: validation +val_dataset_name: valid shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py b/modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py index a46fba52d09..3ed4b517b3e 100644 --- a/modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py +++ b/modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py @@ -20,6 +20,7 @@ # mypy: ignore-errors +import gc import json import warnings from functools import partial @@ -219,6 +220,15 @@ def validate_puzzle_solutions(args: DictConfig) -> None: torch.cuda.empty_cache() torch.cuda.synchronize() + # Stitching installs hooks that create reference cycles between the + # model and its hook closures, so the per-solution CPU model is not + # reclaimed by refcounting alone. Drop references and force a cyclic + # GC pass to avoid accumulating one full model per iteration (host + # RAM OOM / SIGKILL after ~25 solutions on an 8B model). + del stitched_model + del model + gc.collect() + dist.barrier()