modelscope / modelscope/ms-swift

如何将自定义指标保存至tensorboard和swanlab

Open
#7,111 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

Describe the bug
预训练LLM时,在callback.py中自定义了显存/设备资源监控回调类,统计至state: TrainerState中,在logging.jsonl和控制台均打印出相关指标信息,但无法保存至tensorboard和swanlab中,如何解决?
另外,对于自定义的meric(如ppl计算),可以保存至tensorboard和swanlab中

Your hardware and system info
ms-swift:3.12.0dev0
npu:910B

Additional context
自定义类代码如下

class DeviceResourceMonitorCallback(TrainerCallback):
    """显存/设备资源监控回调类
    
    核心功能:实时监控显存/内存占用、设备温度/功耗(针对GPU/NPU),避免OOM或硬件过载。
    触发时机:每步/每轮结束后。
    """

    def __init__(
        self,
        memory_warning_threshold: float = 0.9,
        temperature_warning_threshold: float = 80.0
    ):
        """
        Args:
            memory_warning_threshold: 显存使用率警告阈值 (0-1)
            temperature_warning_threshold: GPU温度警告阈值 (摄氏度)
        """
        self.memory_warning_threshold = memory_warning_threshold
        self.temperature_warning_threshold = temperature_warning_threshold
        self.resource_stats = {}
        self.device_type = None
        self.pynvml_available = False
        self.psutil_available = False
        self._init_monitoring_libs()

    def _init_monitoring_libs(self):
        """初始化监控库"""
        # 检测pynvml可用性 (用于GPU温度/功耗监控)
        try:
            import pynvml
            pynvml.nvmlInit()
            self.pynvml_available = True
        except Exception:
            self.pynvml_available = False

        # 检测psutil可用性 (用于系统内存监控)
        try:
            import psutil
            self.psutil_available = True
        except ImportError:
            self.psutil_available = False

    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """初始化时检测设备类型"""
        if torch.cuda.is_available():
            self.device_type = 'cuda'
        else:
            try:
                import torch_npu
                if torch.npu.is_available():
                    self.device_type = 'npu'
            except ImportError:
                pass

        if self.device_type is None:
            self.device_type = 'cpu'

        logger.info(f'DeviceResourceMonitor initialized. Device type: {self.device_type}')

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """每步结束后收集资源信息"""
        self._collect_resource_stats(state)

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """每轮结束后收集资源信息"""
        self._collect_resource_stats(state)

    def _collect_resource_stats(self, state: TrainerState):
        """收集设备资源统计信息"""
        stats = {}

        # CUDA/GPU 监控
        if self.device_type == 'cuda' and torch.cuda.is_available():
            stats.update(self._collect_cuda_stats(state))

        # NPU 监控
        elif self.device_type == 'npu':
            stats.update(self._collect_npu_stats(state))

        # 系统内存监控
        if self.psutil_available:
            stats.update(self._collect_system_memory_stats())

        self.resource_stats = stats

    def _collect_cuda_stats(self, state: TrainerState) -> dict:
        """收集CUDA设备统计信息"""
        stats = {}
        try:
            device = torch.cuda.current_device()
            
            # 显存信息
            memory_allocated = torch.cuda.memory_allocated(device) / 1024**3  # GB
            memory_reserved = torch.cuda.memory_reserved(device) / 1024**3  # GB
            max_memory_allocated = torch.cuda.max_memory_allocated(device) / 1024**3  # GB
            
            # 获取总显存
            total_memory = torch.cuda.get_device_properties(device).total_memory / 1024**3  # GB
            memory_usage_ratio = memory_allocated / total_memory if total_memory > 0 else 0

            stats['gpu/memory_allocated_gb'] = memory_allocated
            stats['gpu/memory_reserved_gb'] = memory_reserved
            stats['gpu/memory_max_allocated_gb'] = max_memory_allocated
            stats['gpu/memory_total_gb'] = total_memory
            stats['gpu/memory_usage_ratio'] = memory_usage_ratio

            # 显存使用率警告
            if memory_usage_ratio > self.memory_warning_threshold:
                logger.warning(
                    f'[Step {state.global_step}] 显存使用率警告: '
                    f'{memory_usage_ratio * 100:.1f}% 超过阈值 {self.memory_warning_threshold * 100:.1f}%'
                )

            # GPU温度和功耗 (需要pynvml)
            if self.pynvml_available:
                stats.update(self._collect_nvidia_smi_stats(device, state))

        except Exception as e:
            logger.debug(f'Failed to collect CUDA stats: {e}')

        return stats

    def _collect_nvidia_smi_stats(self, device: int, state: TrainerState) -> dict:
        """通过pynvml收集GPU温度和功耗信息"""
        stats = {}
        try:
            import pynvml
            handle = pynvml.nvmlDeviceGetHandleByIndex(device)
            
            # GPU温度
            temperature = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
            stats['gpu/temperature_c'] = temperature
            
            if temperature > self.temperature_warning_threshold:
                logger.warning(
                    f'[Step {state.global_step}] GPU温度警告: '
                    f'{temperature}°C 超过阈值 {self.temperature_warning_threshold}°C'
                )

            # GPU功耗
            power_usage = pynvml.nvmlDeviceGetPowerUsage(handle) / 1000  # mW -> W
            stats['gpu/power_w'] = power_usage

            # GPU利用率
            utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
            stats['gpu/utilization_percent'] = utilization.gpu
            stats['gpu/memory_utilization_percent'] = utilization.memory

        except Exception as e:
            logger.debug(f'Failed to collect NVIDIA SMI stats: {e}')

        return stats

    def _collect_npu_stats(self, state: TrainerState) -> dict:
        """收集NPU设备统计信息"""
        stats = {}
        try:
            import torch_npu
            device = torch.npu.current_device()
            
            memory_allocated = torch.npu.memory_allocated(device) / 1024**3  # GB
            memory_reserved = torch.npu.memory_reserved(device) / 1024**3  # GB
            max_memory_allocated = torch.npu.max_memory_allocated(device) / 1024**3  # GB

            stats['npu/memory_allocated_gb'] = memory_allocated
            stats['npu/memory_reserved_gb'] = memory_reserved
            stats['npu/memory_max_allocated_gb'] = max_memory_allocated

        except Exception as e:
            logger.debug(f'Failed to collect NPU stats: {e}')

        return stats

    def _collect_system_memory_stats(self) -> dict:
        """收集系统内存统计信息"""
        stats = {}
        try:
            import psutil
            mem = psutil.virtual_memory()
            stats['system/memory_used_gb'] = mem.used / 1024**3
            stats['system/memory_total_gb'] = mem.total / 1024**3
            stats['system/memory_percent'] = mem.percent

            # CPU使用率
            stats['system/cpu_percent'] = psutil.cpu_percent(interval=None)

        except Exception as e:
            logger.debug(f'Failed to collect system memory stats: {e}')

        return stats

    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs=None, **kwargs):
        """将设备资源信息集成到日志中推送到可视化工具"""
        if logs is None or not self.resource_stats:
            return

        for key, value in self.resource_stats.items():
            if isinstance(value, float):
                logs[key] = round(value, 4)
            else:
                logs[key] = value

        logger.debug(
            f'Resource Stats - '
            f'GPU Memory: {self.resource_stats.get("gpu/memory_allocated_gb", 0):.2f}GB, '
            f'GPU Temp: {self.resource_stats.get("gpu/temperature_c", "N/A")}°C, '
            f'GPU Power: {self.resource_stats.get("gpu/power_w", "N/A")}W'
        )

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the callback.py integration and the on_log path described in the issue, then reproduce the custom resource-monitor callback with TensorBoard and SwanLab enabled. Check how logging.jsonl and console metrics differ from metrics forwarded to the visualization backends. Done means the custom resource statistics and a custom metric such as perplexity are persisted and visible in both tools.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, observability
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.