distributed logging

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/entrypoints/utils.py

@@ -10,7 +10,7 @@
 import os
 from pathlib import PosixPath
 
-from compressed_tensors.offload import from_accelerate
+from compressed_tensors.offload import from_accelerate, is_distributed
 from loguru import logger
 from transformers import (
     AutoConfig,
@@ -26,6 +26,7 @@
     RecipeArguments,
 )
 from llmcompressor.core import reset_session
+from llmcompressor.logger import configure_distributed_logger
 from llmcompressor.pytorch.model_load.helpers import parse_dtype
 from llmcompressor.transformers.compression.compressed_tensors_utils import (
     modify_save_pretrained,
@@ -52,6 +53,9 @@ def pre_process(
     Raises:
         FileNotFoundError: If the model or processor path is invalid.
     """
+    # Detect distributed, update logger
+    if is_distributed():
+        configure_distributed_logger()
 
     # Initialize model
     if isinstance(model_args.model, (str, PosixPath)):

src/llmcompressor/logger.py

@@ -41,9 +41,10 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
+import torch.distributed as dist
 from loguru import logger
 
-__all__ = ["LoggerConfig", "configure_logger", "logger"]
+__all__ = ["LoggerConfig", "configure_logger", "logger", "configure_distributed_logger"]
 
 
 # used by `support_log_once``
@@ -53,14 +54,20 @@
 @dataclass
 class LoggerConfig:
     disabled: bool = False
-    clear_loggers: bool = True
     console_log_level: Optional[str] = "INFO"
     log_file: Optional[str] = None
     log_file_level: Optional[str] = None
     metrics_disabled: bool = False
+    rank: Optional[int] = None
 
 
-def configure_logger(config: Optional[LoggerConfig] = None) -> None:
+# global config
+LOGGER_CONFIG = LoggerConfig()
+
+
+def configure_logger(
+    logger_config: LoggerConfig = LOGGER_CONFIG, clear_loggers: bool = False
+):
     """
     Configure the logger for LLM Compressor.
 
@@ -72,13 +79,10 @@ def configure_logger(config: Optional[LoggerConfig] = None) -> None:
     :param config: The configuration for the logger to use.
     :type config: LoggerConfig
     """
-    logger_config = config or LoggerConfig()
 
     # env vars get priority
     if (disabled := os.getenv("LLM_COMPRESSOR_LOG_DISABLED")) is not None:
         logger_config.disabled = disabled.lower() == "true"
-    if (clear_loggers := os.getenv("LLM_COMPRESSOR_CLEAR_LOGGERS")) is not None:
-        logger_config.clear_loggers = clear_loggers.lower() == "true"
     if (console_log_level := os.getenv("LLM_COMPRESSOR_LOG_LEVEL")) is not None:
         logger_config.console_log_level = console_log_level.upper()
     if (log_file := os.getenv("LLM_COMPRESSOR_LOG_FILE")) is not None:
@@ -92,15 +96,22 @@ def configure_logger(config: Optional[LoggerConfig] = None) -> None:
 
     logger.enable("llmcompressor")
 
-    if logger_config.clear_loggers:
+    # reset logger configuration
+    if clear_loggers:
         logger.remove()
 
+    # set format (optionally adding rank)
+    format = "{time:YYYY-MM-DDTHH:mm:ss.SSSS} | {function} | {level} - {message}"
+    if logger_config.rank is not None:
+        logger.configure(extra={"rank": dist.get_rank()})
+        format = "[Rank {extra[rank]}] " + format
+
     if logger_config.console_log_level:
         # log as a human readable string with the time, function, level, and message
         logger.add(
             sys.stdout,
             level=logger_config.console_log_level.upper(),
-            format="{time} | {function} | {level} - {message}",
+            format=format,
             filter=support_log_once,
         )
 
@@ -112,6 +123,7 @@ def configure_logger(config: Optional[LoggerConfig] = None) -> None:
             log_file,
             level=log_file_level.upper(),
             serialize=True,
+            format=format,
             filter=support_log_once,
         )
 
@@ -121,6 +133,10 @@ def configure_logger(config: Optional[LoggerConfig] = None) -> None:
     # initialize metric logger on loguru
     logger.level("METRIC", no=38, color="<yellow>", icon="📈")
 
+    # set global value for later calls
+    global LOGGER_CONFIG
+    LOGGER_CONFIG = logger_config
+
 
 def support_log_once(record: Dict[str, Any]) -> bool:
     """
@@ -146,14 +162,11 @@ def support_log_once(record: Dict[str, Any]) -> bool:
     return True
 
 
+def configure_distributed_logger(logger_config: LoggerConfig = LOGGER_CONFIG):
+    logger_config.rank = dist.get_rank()
+    configure_logger(logger_config, clear_loggers=True)
+
+
 # invoke logger setup on import with default values enabling console logging with INFO
 # and disabling file logging
-configure_logger(
-    config=LoggerConfig(
-        disabled=False,
-        clear_loggers=True,
-        console_log_level="INFO",
-        log_file=None,
-        log_file_level=None,
-    )
-)
+configure_logger()

src/llmcompressor/modifiers/pruning/sparsegpt/base.py

@@ -141,7 +141,7 @@ def compress_modules(self):
                     dampening_frac=self.dampening_frac,
                     preserve_sparsity_mask=self.preserve_sparsity_mask,
                 )
-                comp_logger.set_loss(loss)
+                comp_logger.set_results(name="SGPT", loss=loss)
 
             update_offload_parameter(module, "weight", sparsified_weight)
 

src/llmcompressor/modifiers/quantization/gptq/base.py

@@ -302,7 +302,7 @@ def compress_module_list(self, module_list):
             num_samples = self._num_samples[module]
             quant_args = getattr_chain(module, "quantization_scheme.weights")
 
-            logger.info(f"Quantizing {name} using {num_samples} samples")
+            logger.info(f"Quantizing {name} using {int(num_samples)} samples")
             with (
                 torch.no_grad(),
                 align_module_device(module),
@@ -316,7 +316,7 @@ def compress_module_list(self, module_list):
                     blocksize=self.block_size,
                     percdamp=self.dampening_frac,
                 )
-                comp_logger.set_loss(loss)
+                comp_logger.set_results(name="GPTQ", loss=loss)
 
             for attr, val in q_param_dict.items():
                 update_offload_parameter(module, attr, val)

src/llmcompressor/utils/metric_logging.py

@@ -6,7 +6,7 @@
 """
 
 import time
-from typing import Iterable
+from typing import Iterable, Optional
 
 import torch
 from compressed_tensors.offload import is_distributed
@@ -28,27 +28,34 @@ class CompressionLogger:
     def __init__(self, module: torch.nn.Module):
         self.module = module
         self.start_tick = None
-        self.loss = None
 
-    def set_loss(self, loss: float):
-        self.loss = loss
+        self._name = None
+        self._loss = None
+
+    def set_results(
+        self,
+        name: Optional[str] = None,
+        loss: Optional[float] = None,
+    ):
+        self._name = name
+        self._loss = loss
 
     def __enter__(self) -> "CompressionLogger":
         self.start_tick = time.time()
         return self
 
     def __exit__(self, _exc_type, _exc_val, _exc_tb):
         stop_tick = time.time()
-        patch = logger.patch(lambda r: r.update(function="compress"))
 
-        if self.start_tick is not None:
-            patch.log("METRIC", f"time {(stop_tick - self.start_tick):.2f}s")
-        if self.loss is not None:
-            patch.log("METRIC", f"error {self.loss:.2f}")
+        patch = logger.patch(lambda r: r.update(function=(self._name or "compress")))
+
+        patch.log("METRIC", f"time {(stop_tick - self.start_tick):.2f}s")
+        if self._loss is not None:
+            patch.log("METRIC", f"error {self._loss:.2f}")
 
         for device_id in _get_visible_devices():
-            max_memory = torch.cuda.max_memory_allocated(device_id)
-            used_memory = torch.cuda.get_device_properties(device_id).total_memory
+            used_memory = torch.cuda.max_memory_allocated(device_id)
+            max_memory = torch.cuda.get_device_properties(device_id).total_memory
             perc_used = 100 * used_memory / max_memory
             patch.log(
                 "METRIC",