From 801489801d8fccd78de564d59889ea1e889ff758 Mon Sep 17 00:00:00 2001 From: Kovbo Date: Wed, 20 May 2026 21:28:26 +0000 Subject: [PATCH] Avoid Unsloth import in Megatron service --- src/art/megatron/service.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index bd9f462f5..fad4c06f6 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -1,5 +1,6 @@ import asyncio from dataclasses import dataclass, field +import gc import importlib import os from pathlib import Path @@ -18,7 +19,6 @@ from ..local.checkpoints import get_last_checkpoint_dir from ..preprocessing.pack import DiskPackedTensors from ..preprocessing.tokenize import SFTBatch -from ..unsloth.train import gc_and_empty_cuda_cache from ..utils.convert_moe_lora import convert_checkpoint_if_needed from ..utils.get_model_step import get_step_from_dir from ..utils.lifecycle import ( @@ -57,6 +57,13 @@ safe_open = safetensors.safe_open +def gc_and_empty_cuda_cache(n: int = 3) -> None: + for _ in range(n): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + class _RuntimeRequestKwargs(TypedDict, total=False): headers: dict[str, str]