利用 ms-swift 监控 PID 资源占用预防 GPU 内存泄漏

在大模型训练日益成为 AI 工程核心的今天，GPU 集群的稳定性早已不再只是“能不能跑起来”的问题，而是“能不能长期稳定运行”的挑战。尤其当多个团队共享一套算力资源时，一个微小的显存泄漏可能就会像雪崩前的第一片雪花——起初毫无察觉，最终却导致整个节点宕机、任务中断、成本飙升。

更令人头疼的是，这类问题往往具有极强的隐蔽性：训练仍在继续，loss 曲线看似正常，但 GPU 显存却以每小时几十 MB 的速度缓慢爬升。等到 OOM（Out of Memory）报错出现时，通常已经错过了最佳干预时机，checkpoint 丢失，日志残缺，根因难寻。

有没有一种方法，能在泄漏发生的早期阶段就识别出异常？答案是肯定的。借助 ms-swift 框架出色的工程封装能力，结合系统级监控工具，我们完全可以构建一套轻量、高效、低侵入的 GPU 内存泄漏预警机制，实现从“被动救火”到“主动防御”的转变。

---

ms-swift：不只是训练框架，更是可观测性的基石

很多人把 ms-swift 当作一个“能快速启动 SFT 或 DPO 任务”的工具链，这没错，但它真正的价值在于其全链路标准化的设计理念。无论是 Qwen3 还是 Llama4，不管是 LoRA 微调还是多模态对齐，ms-swift 都通过统一的接口完成了环境初始化、分布式启动和资源分配。这种一致性，恰恰为精细化监控提供了前提条件。

更重要的是，ms-swift 在执行训练任务时会生成独立的子进程，并返回主训练进程的 PID。这一点至关重要——它让我们能够将抽象的“某个训练任务”映射到操作系统层面的具体资源消耗者。换句话说，每一个 PID 不再只是一个数字，而是承载了模型类型、任务配置、用户身份等上下文信息的可观测单元。

比如你运行：

swift sft --model_type qwen3-7b --peft_type lora ...

后台会拉起一个 Python 进程，假设它的 PID 是 12345。只要我们能追踪这个 PID 的生命周期，就能实时掌握这块 GPU 显存到底被谁用了、怎么用的、有没有异常增长。

---

显存泄漏的本质与检测逻辑

先明确一点：不是所有显存上升都是泄漏。合理的场景包括：

• 动态 batch size 扩展；
• 缓存机制渐进填充（如 KV Cache）；
• 序列长度逐步增加（packing 策略下常见）；
• 分布式梯度累积过程中的临时驻留。

真正危险的是那种无业务逻辑支撑的持续增长。例如，在固定 batch 和 sequence length 下，显存本应趋于平稳，但如果每过一小时就上涨 50MB，且连续数小时不停止，那基本可以判定存在引用未释放或缓存堆积的问题。

检测的关键就在于建立“趋势感知”能力。我们需要做的不是看某一时刻的绝对值，而是分析一段时间内的变化斜率。这里不需要复杂的机器学习模型，简单的滑动窗口增长率判断就足够有效。

而实现这一目标的技术组合非常朴素：nvidia-smi + psutil + Python 脚本。它们共同构成了一个外挂式监控探针，无需修改 ms-swift 源码，也不影响原有训练流程，真正做到“即插即用”。

---

如何动手实现一个监控脚本？

下面这段代码虽然简短，但在实际生产中已经帮助我们捕获过多起潜在泄漏事件：

import subprocess
import time
import re
import psutil
from typing import Dict, List

def get_gpu_memory_by_pid(pid: int) -> List[Dict[str, int]]:
    """
    查询指定 PID 占用的 GPU 显存（单位: MB）
    使用 nvidia-smi 查询所有 MIG 实例和 GPU 实例中的进程信息
    """
    try:
        result = subprocess.run([
            "nvidia-smi", "--query-compute-apps=pid,used_memory",
            "--format=csv,noheader,nounits"
        ], capture_output=True, text=True, check=True)

        lines = result.stdout.strip().split('\n')
        records = []
        for line in lines:
            if not line:
                continue
            parts = line.split(', ')
            p = int(parts[0])
            mem = int(parts[1])
            if p == pid:
                records.append({"pid": p, "gpu_memory_mb": mem})
        return records
    except Exception as e:
        print(f"[ERROR] Failed to query GPU memory: {e}")
        return []

def monitor_process_resource(pid: int, interval_sec: int = 60, max_duration_hours: float = 24):
    """
    监控指定 PID 的资源占用情况，检测显存泄漏趋势
    """
    start_time = time.time()
    max_duration_sec = max_duration_hours * 3600
    history = []

    print(f"Starting monitoring for PID {pid}, every {interval_sec} seconds...")

    while True:
        current_time = time.time()
        if current_time - start_time > max_duration_sec:
            print("Monitoring duration exceeded. Exiting.")
            break

        # 获取 GPU 显存
        gpu_mems = get_gpu_memory_by_pid(pid)
        total_gpu_mem = sum(item["gpu_memory_mb"] for item in gpu_mems)

        # 获取 CPU 和 RAM 使用率
        try:
            process = psutil.Process(pid)
            cpu_percent = process.cpu_percent()
            ram_mb = process.memory_info().rss / 1024 / 1024  # 转换为 MB
        except psutil.NoSuchProcess:
            print(f"Process {pid} no longer exists.")
            break

        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        history.append({
            "timestamp": timestamp,
            "pid": pid,
            "cpu_percent": cpu_percent,
            "ram_mb": ram_mb,
            "gpu_memory_mb": total_gpu_mem
        })

        print(f"[{timestamp}] PID={pid}, CPU={cpu_percent:.1f}%, RAM={ram_mb:.1f}MB, GPU={total_gpu_mem}MB")

        # 简单泄漏检测：过去3次平均增长率 > 50MB/hour 视为可疑
        if len(history) >= 4:
            recent = history[-3:]
            time_diff_h = (time.mktime(time.strptime(recent[-1]["timestamp"], "%Y-%m-%d %H:%M:%S")) -
                           time.mktime(time.strptime(recent[0]["timestamp"], "%Y-%m-%d %H:%M:%S"))) / 3600
            mem_diff = recent[-1]["gpu_memory_mb"] - recent[0]["gpu_memory_mb"]
            if time_diff_h > 0:
                growth_rate_per_hour = mem_diff / time_diff_h
                if growth_rate_per_hour > 50:  # 超过每小时增长 50MB 视为警告
                    print(f"[ALERT] High GPU memory growth rate: {growth_rate_per_hour:.1f} MB/hour!")

        time.sleep(interval_sec)

# 示例调用
if __name__ == "__main__":
    TRAINING_PID = 12345  # 替换为实际 ms-swift 启动的训练进程 PID
    monitor_process_resource(TRAINING_PID, interval_sec=60, max_duration_hours=12)

几点关键说明：

• nvidia-smi 是最稳定的显存查询方式，兼容性强，即使容器内也能使用；
• psutil 提供了跨平台的进程信息访问能力，适合部署在异构环境中；
• 建议采样间隔设为 30~60 秒，过于频繁不仅增加系统负载，还容易误判波动；
• 增长率阈值可根据具体模型调整，例如 7B 模型可设为 50MB/h，而 70B 模型由于本身缓存较大，可放宽至 100~150MB/h；
• 可扩展动作包括：写入日志文件、触发钉钉/邮件告警、自动保存 checkpoint 并终止进程。

---

架构设计与落地实践

整个监控体系可以分为四层，结构清晰，职责分明：

+------------------+       +----------------------------+
|   ms-swift CLI   | ----> | Training Job (PID: 12345)  |
+------------------+       +-------------+--------------+
                                         |
                                         v
                             +---------------------------+
                             | NVIDIA Driver + CUDA      |
                             | - GPU Memory Allocation   |
                             +-------------+-------------+
                                           |
                                           v
                            +------------------------------+
                            | Monitoring Script (Python)   |
                            | - Poll nvidia-smi & psutil   |
                            | - Detect memory trend        |
                            +-------------+----------------+
                                          |
                                          v
                           +------------------------------+
                           | Alert System / Auto Recovery |
                           | - Log, Notify, Kill Process  |
                           +------------------------------+

实际工作流示例

1. 用户通过 Web UI 启动一个 Qwen-VL 的多模态微调任务；
2. ms-swift 返回训练主进程 PID（如 12345），并记录该 PID 对应的任务 ID、用户名、模型名称等元数据；
3. 监控服务监听新任务事件，自动绑定对该 PID 的轮询采集；
4. 训练进行两小时后，显存从 9GB 缓慢上升至 12.8GB，远超理论峰值；
5. 监控脚本检测到过去三小时平均增长率达 76MB/h，触发一级告警；
6. 系统向负责人发送通知：“任务 [task-20250405-001] 存在显存泄漏风险，请检查是否启用冗余缓存”；
7. 若未及时响应，一小时后系统尝试调用 torch.cuda.empty_cache() 并安全退出；
8. 所有监控数据持久化存储，用于后续归因分析和模型健康评分。

经验性优化建议

• 监控进程独立部署：不要将监控脚本和训练任务放在同一容器内，避免资源争抢导致数据失真；
• 引入白名单机制：某些合法的大显存行为（如 LongLoRA 中的动态 context 扩展）应被排除在告警之外；
• 结合 loss 曲线交叉验证：如果 loss 已收敛但仍持续增长显存，大概率是泄漏；若 loss 仍在下降，则可能是正常训练动态；
• 权限控制要到位：确保监控服务有足够的权限读取目标进程状态，推荐使用相同用户运行或配置 sudo 免密；
• 支持远程节点采集：在 Slurm/PBS 集群中，可通过 scontrol show job 获取远程节点上的 PID，再通过 SSH 执行远程查询。

---

为什么这个方案特别适合企业级场景？

很多公司一开始靠人工巡检 nvidia-smi 输出，或者依赖 Prometheus + Node Exporter 做通用监控。但这些方式在面对大模型训练时都有明显短板：

• 通用监控看不到 进程级细粒度 数据；
• 无法关联 任务上下文（是谁、跑什么模型、用了哪种策略）；
• 缺乏对 缓慢增长型泄漏 的敏感度；
• 报警动作单一，难以做到自动化恢复。

而基于 ms-swift 的这套方案，天然具备以下优势：

• 精准定位：每个 PID 都能回溯到具体的 YAML 配置和用户操作；
• 早发现：在 OOM 前数小时即可预警，留足处置时间；
• 低成本整合：无需改造现有系统，只需加一个守护脚本；
• 可扩展性强：未来可接入 Grafana 展示趋势图，或集成进 CI/CD 流程作为训练稳定性测试环节。

更重要的是，它推动了一种新的运维范式：把资源稳定性当作软件质量的一部分来管理。就像单元测试保证代码正确性一样，资源监控也应该成为模型上线前的必检项。

---

展望：迈向自愈型训练平台

当前这套方案仍是“外挂式”的，下一步自然是要将其能力内建到 ms-swift 生态中。设想一下未来的 Web UI 中：

• 每个训练任务面板都自带“显存趋势图”；
• 系统自动标注“当前增长率：正常 / 警告 / 危险”；
• 点击“诊断”按钮可一键生成内存快照分析报告；
• 支持设置“自动熔断”策略，超过阈值则暂停训练并通知维护者。

这才是真正意义上的“自愈型大模型训练平台”。而今天我们所做的，正是为这样的未来打下第一根桩基。

技术的进步不一定要惊天动地。有时候，一段小小的监控脚本，反而比一个庞大的系统更能体现工程智慧——因为它解决了真实世界中最痛的那个点。