監控大模型訓練

一蓑烟雨度平生發表於2024-03-13

監控大模型訓練

大模型訓練時間久,而且過程中容易出現各種各樣的問題而中斷,中斷之後不及時續練的話對GPU資源是很大的浪費,但是我們又不能一直盯著程式。所以本文將介紹如何編寫一個監控程式來監控大模型的訓練,以方便我們在大模型訓練出現異常時及時通知給我們。

監控的方式有很多,這裡介紹兩個方式。

根據log檔案大小變化監控訓練是否進行

linux nohup指令詳解中,我們提到了使用Linux的nohup命令來執行訓練指令碼,該命令會建立一個日誌檔案,大模型在訓練的過程中會不斷輸出內容,因此該日誌檔案的大小是隨時在變化的。因此,我們可以透過隔一段時間判斷該日誌檔案的大小是否變化來判斷大模型的訓練是否出現異常。如果大模型訓練出現異常,我們需要透過一種常用的通訊方式來告知自己,這裡選擇使用郵箱(也可以使用簡訊、QQ/微信通知等方式)。

核心程式碼如下:

import os
import time
import smtplib
from email.mime.text import MIMEText

# 配置郵箱資訊
SMTP_SERVER = 'smtp.qq.com'
SMTP_PORT = 587
EMAIL_USERNAME = '作為SMTP伺服器的QQ郵箱'
EMAIL_PASSWORD = '你的QQ郵箱SMTP服務的金鑰'
EMAIL_FROM = '作為SMTP伺服器的QQ郵箱'
EMAIL_TO = '你接收通知郵件的QQ郵箱'
EMAIL_SUBJECT = '大模型訓練終止提醒'

# 監測的檔案路徑
FILE_PATH = '/...../nohup.out'


def send_email(message):
    msg = MIMEText(message)
    msg['From'] = EMAIL_FROM
    msg['To'] = EMAIL_TO
    msg['Subject'] = EMAIL_SUBJECT

    with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as server:
        server.starttls()
        server.login(EMAIL_USERNAME, EMAIL_PASSWORD)
        server.sendmail(from_addr=EMAIL_FROM, to_addrs=[EMAIL_TO], msg=msg.as_string())


def monitor_file():
    # 獲取初始檔案大小
    initial_size = os.path.getsize(FILE_PATH)

    while True:
        # 等待10分鐘
        time.sleep(600)

        # 獲取當前檔案大小
        current_size = os.path.getsize(FILE_PATH)

        if current_size == initial_size:
            # 檔案大小沒有變化,傳送警告郵件
            message = '日誌檔案在十分鐘內沒有發生變化,大模型訓練可能已終止!'
            send_email(message)
        else:
            # 檔案大小發生變化,更新初始檔案大小
            print(f'log changed: {current_size-initial_size}')
            initial_size = current_size
            

if __name__ == '__main__':
    monitor_file()

根據GPU視訊記憶體佔用率監控訓練是否進行

訓練大模型時,GPU的視訊記憶體佔用率一般都比較高,所以我們也可以透過GPU視訊記憶體的佔用率來判斷大模型的訓練是否出現異常。這裡我們同樣使用郵箱來通知自己。

核心程式碼如下:

import subprocess
import smtplib
from email.mime.text import MIMEText
import time

def get_gpu_memory_usage():
    output = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader'])
    memory_used = [int(x) for x in output.decode().strip().split('\n')]
    return memory_used

# 配置郵箱資訊
SMTP_SERVER = 'smtp.qq.com'
SMTP_PORT = 587
EMAIL_USERNAME = '作為SMTP伺服器的QQ郵箱'
EMAIL_PASSWORD = '你的QQ郵箱SMTP服務的金鑰'
EMAIL_FROM = '作為SMTP伺服器的QQ郵箱'
EMAIL_TO = '你接收通知郵件的QQ郵箱'
EMAIL_SUBJECT = '大模型訓練終止提醒'


def send_email(message):
    msg = MIMEText(message)
    msg['From'] = EMAIL_FROM
    msg['To'] = EMAIL_TO
    msg['Subject'] = EMAIL_SUBJECT

    with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as server:
        server.starttls()
        server.login(EMAIL_USERNAME, EMAIL_PASSWORD)
        server.sendmail(from_addr=EMAIL_FROM, to_addrs=[EMAIL_TO], msg=msg.as_string())

def main():
    while True:
        memory_used_list = get_gpu_memory_usage()
        # memory_total是你伺服器總的視訊記憶體量,此處使用的伺服器有8張40G的A100,因此總視訊記憶體量為40960*8
        memory_total = 40960 * 8
        memory_used = 0
        for memory in memory_used_list:
            memory_used += memory
        memory_usage_percent = (memory_used / memory_total) * 100

        if memory_usage_percent < 10:
            subject = '伺服器視訊記憶體佔用率過低警告'
            body = f'視訊記憶體佔用率為 {memory_usage_percent}%,低於10%。請檢查伺服器。'
            send_email(subject, body)
        else:
            print(f'GPU Memory %: {memory_usage_percent}')

        time.sleep(600)  # 等待10分鐘(600秒)

if __name__ == '__main__':
    main()

我們可以使用linux nohup指令詳解中介紹的nohup命令執行這兩個程式中的一個,以監控大模型的訓練是否正常進行。

相關文章