From 537f3b92203890b4750b09cf36ba4cdf003fd8b8 Mon Sep 17 00:00:00 2001 From: Kelvin Cheng Date: Mon, 19 May 2025 02:57:07 -0400 Subject: [PATCH 1/2] Automatic log saving inside oneshot. Removed redundant logs from stage.py Signed-off-by: Kelvin Cheng --- src/llmcompressor/entrypoints/oneshot.py | 12 ++++++++++++ src/llmcompressor/modifiers/stage.py | 2 -- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/llmcompressor/entrypoints/oneshot.py b/src/llmcompressor/entrypoints/oneshot.py index 67c1e2ba2..677447680 100644 --- a/src/llmcompressor/entrypoints/oneshot.py +++ b/src/llmcompressor/entrypoints/oneshot.py @@ -1,5 +1,8 @@ +import os +from datetime import datetime from typing import Optional +from loguru import logger from torch.utils.data import DataLoader from transformers import PreTrainedModel @@ -100,6 +103,15 @@ def __init__( :param output_dir: Path to save the output model after carrying out oneshot """ + base_log_path = "sparse_logs" # TODO: maybe make this an input argument? + if base_log_path: + os.makedirs(base_log_path, exist_ok=True) + date_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + logger.add( + f"{base_log_path}/oneshot_{date_str}.log", + level="DEBUG", + ) + model_args, dataset_args, recipe_args, _, output_dir = parse_args(**kwargs) self.model_args = model_args diff --git a/src/llmcompressor/modifiers/stage.py b/src/llmcompressor/modifiers/stage.py index fb160955b..def52ad8b 100644 --- a/src/llmcompressor/modifiers/stage.py +++ b/src/llmcompressor/modifiers/stage.py @@ -67,7 +67,6 @@ def initialize(self, state: "State", **kwargs): modifier.initialize(state, **kwargs) if accelerator: accelerator.wait_for_everyone() - state.loggers.system.info(tag="stage", string="Modifiers initialized") def finalize(self, state: "State", **kwargs): """ @@ -88,7 +87,6 @@ def finalize(self, state: "State", **kwargs): accelerator.wait_for_everyone() self.applied = True - state.loggers.system.info(tag="stage", string="Modifiers finalized") def update_event(self, state: "State", event: "Event", **kwargs): """ From 29f572d70202d9cc58961ac17cda2e81cead8395 Mon Sep 17 00:00:00 2001 From: Kelvin Cheng Date: Mon, 19 May 2025 15:14:12 -0400 Subject: [PATCH 2/2] Make log dir configurable; allow disabling by passing None Signed-off-by: Kelvin Cheng --- src/llmcompressor/entrypoints/oneshot.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/llmcompressor/entrypoints/oneshot.py b/src/llmcompressor/entrypoints/oneshot.py index 677447680..bedca7392 100644 --- a/src/llmcompressor/entrypoints/oneshot.py +++ b/src/llmcompressor/entrypoints/oneshot.py @@ -85,6 +85,7 @@ class Oneshot: def __init__( self, + log_dir: Optional[str] = "sparse_logs", **kwargs, ): """ @@ -103,12 +104,12 @@ def __init__( :param output_dir: Path to save the output model after carrying out oneshot """ - base_log_path = "sparse_logs" # TODO: maybe make this an input argument? - if base_log_path: - os.makedirs(base_log_path, exist_ok=True) + # Set up logging + if log_dir: + os.makedirs(log_dir, exist_ok=True) date_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") logger.add( - f"{base_log_path}/oneshot_{date_str}.log", + f"{log_dir}/oneshot_{date_str}.log", level="DEBUG", )