modelscope / modelscope/ms-swift
如何将自定义指标保存至tensorboard和swanlab
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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